From 52b715e7a273f549e5ad8eee332144676460a555 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Tue, 15 Sep 2026 01:56:20 -0700 Subject: [PATCH 01/16] [None][perf] Fuse the Kimi K3 routed RMSNorm into the MNNVL all-reduce Signed-off-by: Zhenhuan Chen Signed-off-by: Mingyang Hao --- .../allReduceFusionKernels.h | 6 + .../mnnvlAllreduceKernels.cu | 114 +++++++++++++++++- .../mnnvlAllreduceKernels.h | 1 + .../kernels/customAllReduceKernels.h | 3 + cpp/tensorrt_llm/thop/allreduceOp.cpp | 36 +++++- tensorrt_llm/_torch/distributed/ops.py | 1 + .../_torch/models/modeling_kimi_linear.py | 16 +-- tensorrt_llm/functional.py | 8 +- 8 files changed, 166 insertions(+), 19 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h index 769776273ef6..fa206a3cf98a 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h @@ -68,6 +68,9 @@ enum class AllReduceFusionPattern : int // AllReduce + RMSNorm without residual addition. Useful when there is no // external residual to add (e.g. models where residual is handled externally). kARRMSNorm = 6, + // AllReduce + RMSNorm + FP8 1x128 block-scale quant (UE8M0, deep_gemm's + // packed MN-major layout). No residual, no bf16 norm output. + kARRMSNormFP8BlockScale = 7, }; enum class QuantType : int @@ -75,6 +78,7 @@ enum class QuantType : int kNone = 0, kFP8 = 1, kFP4 = 2, + kFP8BlockScale = 3, }; template @@ -105,6 +109,8 @@ DEFINE_FUSION_PATTERN_TRAITS( DEFINE_FUSION_PATTERN_TRAITS( AllReduceFusionPattern::kARResidualRMSNormOutFP4Quant, false, true, true, true, true, QuantType::kFP4); DEFINE_FUSION_PATTERN_TRAITS(AllReduceFusionPattern::kARRMSNorm, false, false, false, true, true, QuantType::kNone); +DEFINE_FUSION_PATTERN_TRAITS( + AllReduceFusionPattern::kARRMSNormFP8BlockScale, false, false, false, true, false, QuantType::kFP8BlockScale); #undef DEFINE_FUSION_PATTERN_TRAITS template diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu index 0e25a25eb46c..3126b3505d15 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu @@ -336,6 +336,7 @@ struct MnnvlAllReduceKernelParams uint32_t* bufferFlags; bool waitForResults; QuantizationSFLayout layout; + int sfLeadingDim = 0; }; template @@ -538,6 +539,80 @@ inline __device__ void quantizeEpilogue(PackedVec const& value, } reinterpret_cast(params.quantOutPtr)[packedAccessIdx] = quantized; } + else if constexpr (ar_fusion::GetQuantType == ar_fusion::QuantType::kFP8BlockScale) + { + // 1x128 UE8M0 block scale in deep_gemm's packed MN-major layout. + // 16 consecutive threads cover one 128-element block, which holds + // because loadStride == blockDim.x and blockChunkSize % 128 == 0. + static_assert(kELTS_PER_THREAD == 8, "FP8 block-scale quant expects eight elements per 16-byte access."); + constexpr int kQuantBlock = 128; + constexpr int kThreadsPerQuantBlock = kQuantBlock / kELTS_PER_THREAD; + float amax = 0.f; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; ++i) + { + amax = fmaxf(amax, fabsf(cuda_cast(value.elements[i]))); + } +#pragma unroll + for (int off = kThreadsPerQuantBlock / 2; off > 0; off >>= 1) + { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, off, kThreadsPerQuantBlock)); + } + amax = fmaxf(amax, 1e-10f); + + float rcp448; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(rcp448) : "f"(448.0f)); + __nv_fp8_e8m0 ue8m0Scale; + ue8m0Scale.__x = __nv_cvt_float_to_e8m0(amax * rcp448, __NV_SATFINITE, cudaRoundPosInf); + constexpr float kFp32ExponentBias = 127.f; + float const quantScale + = (ue8m0Scale.__x == 0) ? 1.0f : exp2f(kFp32ExponentBias - static_cast(ue8m0Scale.__x)); + + union StoreTrick + { + uint2 pack; + __nv_fp8_e4m3 v[kELTS_PER_THREAD]; + }; + + StoreTrick storeTrick; +#pragma unroll + for (int i = 0; i < kELTS_PER_THREAD; ++i) + { + storeTrick.v[i] = __nv_fp8_e4m3(cuda_cast(value.elements[i]) * quantScale); + } + reinterpret_cast(params.quantOutPtr)[packedAccessIdx] = storeTrick.pack; + + auto* sfBytes = reinterpret_cast(params.scaleOutPtr); + int const numSfK = params.tokenDim / kQuantBlock; + if (accessIdInToken % kThreadsPerQuantBlock == 0) + { + int const sfK = accessIdInToken / kThreadsPerQuantBlock; + sfBytes[(static_cast(sfK / 4) * params.sfLeadingDim + token) * 4 + (sfK % 4)] + = static_cast(ue8m0Scale.__x); + } + // Bytes no real store covers but the GEMM still reads: this row's tail + // of the last packed uint32, and the rows padded up to sfLeadingDim. + if (accessIdInToken == 0) + { + int const numPackedSfK = (numSfK + 3) / 4; + for (int sfK = numSfK; sfK < numPackedSfK * 4; ++sfK) + { + sfBytes[(static_cast(sfK / 4) * params.sfLeadingDim + token) * 4 + (sfK % 4)] = 0; + } + if (token == 0) + { + for (int row = params.numTokens; row < params.sfLeadingDim; ++row) + { + for (int p = 0; p < numPackedSfK; ++p) + { + reinterpret_cast( + params.scaleOutPtr)[static_cast(p) * params.sfLeadingDim + row] + = 0u; + } + } + } + } + } else { static_assert(ar_fusion::GetQuantType == ar_fusion::QuantType::kNone, "Invalid quant type"); @@ -743,6 +818,11 @@ using detail::requiresTwoAccessScaleGroup; void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) { + if (params.pattern == ar_fusion::AllReduceFusionPattern::kARRMSNormFP8BlockScale) + { + TLLM_CHECK_WITH_INFO(params.tokenDim % 128 == 0 && params.sfLeadingDim > 0, + "[MNNVL] FP8 block-scale quant needs tokenDim %% 128 == 0 and a scale leading dim"); + } static int const kSMVersion = tensorrt_llm::common::getSMVersion(); TLLM_CHECK_WITH_INFO(kSMVersion >= 90, "[MNNVL AllReduceOneShot] requires SM 90 or newer."); @@ -825,6 +905,17 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) { \ LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARRMSNorm); \ } \ + else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARRMSNormFP8BlockScale) \ + { \ + if constexpr (!std::is_same_v) \ + { \ + LAUNCH_ALLREDUCE_KERNEL(WORLD_SIZE, T, ar_fusion::AllReduceFusionPattern::kARRMSNormFP8BlockScale); \ + } \ + else \ + { \ + TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceOneShot] FP8 block-scale quant needs FP16/BF16 input."); \ + } \ + } \ else \ { \ TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceOneShot] Unsupported fusion pattern."); \ @@ -843,7 +934,7 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) MnnvlAllReduceKernelParams kernelParams{output, residualOut, input, residualIn, gamma, ucPtrs, reinterpret_cast(params.bufferPtrLocal), mcPtr, params.quantOut, params.scaleOut, params.scaleFactor, numTokens, tokenDim, params.nRanks, params.rank, static_cast(params.epsilon), params.bufferFlags, - false, params.layout}; + false, params.layout, params.sfLeadingDim}; switch (params.nRanks) { @@ -1195,6 +1286,11 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) { + if (params.pattern == ar_fusion::AllReduceFusionPattern::kARRMSNormFP8BlockScale) + { + TLLM_CHECK_WITH_INFO(params.tokenDim % 128 == 0 && params.sfLeadingDim > 0, + "[MNNVL] FP8 block-scale quant needs tokenDim %% 128 == 0 and a scale leading dim"); + } static int const kSMVersion = tensorrt_llm::common::getSMVersion(); TLLM_CHECK_WITH_INFO(kSMVersion >= 90, "[MNNVL AllReduceTwoShot] requires SM 90 or newer."); int const numTokens = params.numTokens; @@ -1237,7 +1333,7 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) reinterpret_cast(params.residualIn), reinterpret_cast(params.gamma), ucPtrs, reinterpret_cast(params.bufferPtrLocal), mcastPtr, params.quantOut, params.scaleOut, params.scaleFactor, numTokens, tokenDim, params.nRanks, params.rank, static_cast(params.epsilon), params.bufferFlags, - !params.rmsNormFusion, params.layout}; + !params.rmsNormFusion, params.layout, params.sfLeadingDim}; switch (params.nRanks) { case 2: LAUNCH_ALLREDUCE_KERNEL(2, T); return true; @@ -1356,6 +1452,18 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) { \ RUN_RMSNORM_KERNEL(T, ar_fusion::AllReduceFusionPattern::kARRMSNorm, USE_CGA, LOADS_PER_THREAD); \ } \ + else if (params.pattern == ar_fusion::AllReduceFusionPattern::kARRMSNormFP8BlockScale) \ + { \ + if constexpr (!std::is_same_v) \ + { \ + RUN_RMSNORM_KERNEL( \ + T, ar_fusion::AllReduceFusionPattern::kARRMSNormFP8BlockScale, USE_CGA, LOADS_PER_THREAD); \ + } \ + else \ + { \ + TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] FP8 block-scale quant needs FP16/BF16 input."); \ + } \ + } \ else \ { \ TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] Unsupported RMSNorm fusion pattern."); \ @@ -1371,7 +1479,7 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) reinterpret_cast(params.bufferPtrsDev), reinterpret_cast(params.bufferPtrLocal), reinterpret_cast(params.multicastPtr), params.quantOut, params.scaleOut, params.scaleFactor, numTokens, tokenDim, params.nRanks, params.rank, static_cast(params.epsilon), params.bufferFlags, - false, params.layout}; + false, params.layout, params.sfLeadingDim}; if (rnUseCGA) { DISPATCH_RMSNORM_PATTERN(T, true, 1); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h index f2006f52240c..fb5cb7c49876 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h @@ -67,6 +67,7 @@ struct AllReduceFusionParams void* quantOut = nullptr; //!< Quantized RMSNorm output (used by quantized fusion patterns) void* scaleOut = nullptr; //!< NVFP4 scale-factor output (used by NVFP4 fusion patterns) float const* scaleFactor = nullptr; //!< Quantization scale factor + int sfLeadingDim = 0; //!< uint32 elements per packed-sf_k row (FP8 block-scale quant only) QuantizationSFLayout layout = QuantizationSFLayout::SWIZZLED; //!< NVFP4 scale-factor layout cudaStream_t stream; //!< CUDA stream for asynchronous kernel execution diff --git a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h index 93f67ffdd911..5e31343e6a78 100644 --- a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h +++ b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h @@ -81,6 +81,8 @@ enum class AllReduceFusionOp : int8_t RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7, MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 8, RMS_NORM = 9, + //! AllReduce + RMSNorm + FP8 1x128 block-scale quant (UE8M0, packed MN-major). + RMS_NORM_QUANT_FP8_BLOCKSCALE = 10, }; inline std::ostream& operator<<(std::ostream& os, AllReduceFusionOp op) @@ -99,6 +101,7 @@ inline std::ostream& operator<<(std::ostream& os, AllReduceFusionOp op) os << "MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM"; break; case AllReduceFusionOp::RMS_NORM: os << "RMS_NORM"; break; + case AllReduceFusionOp::RMS_NORM_QUANT_FP8_BLOCKSCALE: os << "RMS_NORM_QUANT_FP8_BLOCKSCALE"; break; default: os << "UNKNOWN"; break; } return os; diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 85fbd59793d5..4ccd26b10ab9 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -2185,6 +2185,7 @@ MnnvlFusionPattern getMnnvlFusionPattern(AllReduceFusionOp fusionOp) { case AllReduceFusionOp::NONE: return MnnvlFusionPattern::kAllReduce; case AllReduceFusionOp::RMS_NORM: return MnnvlFusionPattern::kARRMSNorm; + case AllReduceFusionOp::RMS_NORM_QUANT_FP8_BLOCKSCALE: return MnnvlFusionPattern::kARRMSNormFP8BlockScale; case AllReduceFusionOp::RESIDUAL_RMS_NORM: return MnnvlFusionPattern::kARResidualRMSNorm; case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8: return MnnvlFusionPattern::kARResidualRMSNormFP8Quant; case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4: return MnnvlFusionPattern::kARResidualRMSNormFP4Quant; @@ -2259,7 +2260,8 @@ std::vector mnnvlFusionAllReduce(torch::Tensor& input, torch::opt { // kARRMSNorm has no residual operand (kHasResidual/kHasResidualOut // are both false), so residual_in is optional for it alone. - bool const needsResidual = fusionOp != AllReduceFusionOp::RMS_NORM; + bool const needsResidual + = fusionOp != AllReduceFusionOp::RMS_NORM && fusionOp != AllReduceFusionOp::RMS_NORM_QUANT_FP8_BLOCKSCALE; TORCH_CHECK(gamma.has_value() && epsilon.has_value(), "[mnnvlFusionAllReduce] gamma and epsilon must be provided for RMSNorm fusion"); TORCH_CHECK(!needsResidual || residual_in.has_value(), @@ -2317,6 +2319,25 @@ std::vector mnnvlFusionAllReduce(torch::Tensor& input, torch::opt } } + // FP8 1x128 block-scale quant: shapes replicate fp8_quantize_1x128_packed_ue8m0 + // (value buffer padded to 4 rows, scale buffer a [num_packed_sf_k, + // scaleLeadingDim] int32 block viewed MN-major), so the consumer GEMM cannot + // tell this from a standalone-quantized activation. + int64_t scaleLeadingDim = 0; + int64_t numPackedSfK = 0; + at::Tensor sfBuf; + if (fusionOp == AllReduceFusionOp::RMS_NORM_QUANT_FP8_BLOCKSCALE) + { + TORCH_CHECK(hiddenDim % 128 == 0, "[mnnvlFusionAllReduce] FP8 block-scale quant needs hidden % 128 == 0, got ", + hiddenDim); + numPackedSfK = ((hiddenDim / 128) + 3) / 4; + int64_t const tPadded = (numTokens + 3) / 4 * 4; + constexpr int64_t kTmaAlignedUint32Elems = 4; + scaleLeadingDim = (tPadded + kTmaAlignedUint32Elems - 1) / kTmaAlignedUint32Elems * kTmaAlignedUint32Elems; + quantOut = at::detail::empty_cuda({tPadded, hiddenDim}, torch::kFloat8_e4m3fn, input.device(), std::nullopt); + sfBuf = at::detail::empty_cuda({numPackedSfK, scaleLeadingDim}, torch::kInt32, input.device(), std::nullopt); + } + auto allreduce_params = tensorrt_llm::kernels::mnnvl::AllReduceFusionParams(); allreduce_params.nRanks = mcast_mem->getWorldSize(); allreduce_params.rank = mcast_mem->getRank(); @@ -2353,6 +2374,12 @@ std::vector mnnvlFusionAllReduce(torch::Tensor& input, torch::opt allreduce_params.scaleOut = scaleOut.defined() ? scaleOut.mutable_data_ptr() : nullptr; allreduce_params.scaleFactor = scale.value().data_ptr(); } + else if (fusionOp == AllReduceFusionOp::RMS_NORM_QUANT_FP8_BLOCKSCALE) + { + allreduce_params.quantOut = quantOut.mutable_data_ptr(); + allreduce_params.scaleOut = sfBuf.mutable_data_ptr(); + allreduce_params.sfLeadingDim = static_cast(scaleLeadingDim); + } allreduce_params.rmsNormFusion = hasRmsNormFusion; allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); @@ -2392,6 +2419,13 @@ std::vector mnnvlFusionAllReduce(torch::Tensor& input, torch::opt { case AllReduceFusionOp::NONE: return {reduceOut}; case AllReduceFusionOp::RMS_NORM: return {normOut}; + case AllReduceFusionOp::RMS_NORM_QUANT_FP8_BLOCKSCALE: + // Scales go back MN-major -- [T, num_packed_sf_k] with stride + // (1, scaleLeadingDim) -- which is the layout the deep_gemm runner wants. + return {quantOut.slice(0, 0, numTokens), + at::from_blob( + sfBuf.mutable_data_ptr(), {numTokens, numPackedSfK}, {1, scaleLeadingDim}, + [keep = sfBuf](void*) mutable {}, sfBuf.options())}; case AllReduceFusionOp::RESIDUAL_RMS_NORM: return {normOut, residualOut}; case AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8: return {quantOut, residualOut}; case AllReduceFusionOp::RESIDUAL_RMS_NORM_OUT_QUANT_FP8: return {normOut, quantOut, residualOut}; diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 1ea2b401e9a9..7b42db517de9 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -874,6 +874,7 @@ class MNNVLAllReduce(nn.Module): SUPPORTED_FUSION_OPS: frozenset[AllReduceFusionOp] = frozenset({ AllReduceFusionOp.RMS_NORM, + AllReduceFusionOp.RMS_NORM_QUANT_FP8_BLOCKSCALE, AllReduceFusionOp.RESIDUAL_RMS_NORM, AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_FP8, AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 5a1b6c41c544..8ad22f849c4b 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -110,16 +110,9 @@ from ...logger import logger from ...mapping import Mapping from ...models.modeling_utils import QuantAlgo, QuantConfig -from ..attention.attention import _helix_cp_pad -from ..attention.backends import AttentionMetadata -from ..distributed import ( - AllReduce, - AllReduceFusionOp, - AllReduceParams, - AllReduceStrategy, - allgather, -) -from ..distributed.symm_mem_allgather import SymmetricMemoryAllGather +from ..attention_backend import AttentionMetadata +from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from ..kimi_k3_cache_policy import KIMI_K3_BF16_KV_LAYERS_ENV, get_kimi_k3_bf16_kv_layer_ids from ..model_config import ModelConfig from ..modules.gated_mlp import GatedMLP from ..modules.kimi_kda import KimiKDALinearAttention @@ -1820,10 +1813,9 @@ def _shared_output(): if self._defer_shared_allreduce and not overlap_shared_allreduce: shared_out = self.shared_experts.down_proj.all_reduce(shared_out) if self._reduce_routed_output: - # RMSNorm folded into the routed all-reduce (RMS_NORM fusion): the + # RMSNorm folded into the routed all-reduce (kARRMSNorm): the # MNNVL kernel normalizes straight out of the lamport buffer, so # the reduced row never round-trips through HBM for a norm kernel. - # Strategies without the fusion fall back inside AllReduce. routed_latent = moe_all_reduce( routed_out, all_reduce_params=AllReduceParams( diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py index 4258696f2c83..17f6f3f7a81f 100644 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -129,6 +129,7 @@ class AllReduceFusionOp(IntEnum): RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7 MOE_FINALIZE_ALLREDUCE_RESIDUAL_RMS_NORM = 8 RMS_NORM = 9 + RMS_NORM_QUANT_FP8_BLOCKSCALE = 10 class AllReduceParams: @@ -157,9 +158,10 @@ def __init__( # For torch path only, has no effect on TRT path self.enable_allreduce = enable_allreduce self.trigger_completion_at_end = trigger_completion_at_end - assert fusion_op in (AllReduceFusionOp.NONE.value, - AllReduceFusionOp.RMS_NORM.value) or (residual - is not None) + assert fusion_op in ( + AllReduceFusionOp.NONE.value, AllReduceFusionOp.RMS_NORM.value, + AllReduceFusionOp.RMS_NORM_QUANT_FP8_BLOCKSCALE.value) or ( + residual is not None) def has_affine(self): return 1 if self.norm_weight is not None else 0 From c4051d7cab728e0d5b83daa6be55eb3c4c597332 Mon Sep 17 00:00:00 2001 From: Sunny Gai Date: Thu, 17 Sep 2026 15:11:53 -0700 Subject: [PATCH 02/16] [None][fix] Join x_ready_event in FC12/B12x CuteDSL MoE backends Signed-off-by: Sunny Gai Signed-off-by: Mingyang Hao --- .../fused_moe/fused_moe_cute_dsl_b12x.py | 289 ++++++++++++++++++ .../fused_moe/fused_moe_cute_dsl_fc12.py | 264 ++++++++++++++++ 2 files changed, 553 insertions(+) create mode 100644 tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py create mode 100644 tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py new file mode 100644 index 000000000000..110457fa1641 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py @@ -0,0 +1,289 @@ +# 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 dataclasses import replace +from typing import Optional, Tuple, Union + +import torch + +from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.models.modeling_utils import QuantAlgo + +from ...utils import ActivationType, Fp4QuantizedTensor +from .fused_moe_cute_dsl import CuteDslFusedMoE +from .fused_moe_cutlass import CutlassFusedMoE +from .impl_contract import ( + MoEDeployment, + MoEEligibility, + MoEProblem, + MoERejectReason, + MoERunContext, + MoEStaticCapability, + require_comm_plan, +) +from .impl_environment import MoEDep +from .interface import _reject + +# Shared MoE output buffer pool, keyed by (max_num_tokens, hidden_size, dtype, +# device). ``B12xMoEWrapper.__init__`` allocates a private +# ``(max_num_tokens, hidden_size)`` output tensor per instance; with one +# wrapper per MoE layer that is ``num_layers * max_num_tokens * hidden_size`` +# bytes of GPU memory holding identical-shape buffers that are written +# sequentially. We fold them into a single shared buffer because MoE layers +# run sequentially on the same CUDA stream, and the wrapper consumes its +# previous output before the next layer is dispatched. +_SHARED_MOE_OUTPUT_BUF: dict = {} + +# ActivationType -> b12x activation string. b12x currently exposes "relu2" +# (Nemotron-style x * relu(x)) and "silu" (SwiGLU-style x * silu(gate)). +_ACTIVATION_MAP = { + ActivationType.Relu2: "relu2", + ActivationType.Swiglu: "silu", +} + + +class CuteDslB12xFusedMoE(CuteDslFusedMoE): + """B12x NVFP4 fused-MoE backend for SM120 / SM121. + + Large prefill chunks use CUTLASS; decode uses FlashInfer's b12x kernel. + """ + + # Restated rather than inherited: the LoRA gate this replaces compared the + # exact class and answered False here, while the DWDP gate used isinstance + # and answered True through CuteDslFusedMoE. + capabilities = MoEStaticCapability(supports_moe_lora=False, supports_dwdp=True) + + # SM versions on which the FlashInfer b12x NVFP4 MoE kernel is available. + # SM120 = desktop Blackwell (RTX 5090 / GB202); SM121 = GB10 / DGX Spark. + _SUPPORTED_SM_VERSIONS = frozenset({120, 121}) + + # Prefill chunks (``x.shape[0] >= threshold``) route via CUTLASS NVFP4 + # GroupGEMM; decode (``x.shape[0] < threshold``) uses b12x. 64 cleanly + # separates conc=1 prefill (m=2048 with ``max_num_tokens=2048``) from + # decode (m=1) and stays robust against future chunked-prefill splits + # that might shrink prefill chunk size. + _PREFILL_VIA_CUTLASS_THRESHOLD = 64 + + @classmethod + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + sm_version = d.env.sm + if sm_version not in cls._SUPPORTED_SM_VERSIONS: + sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) + return _reject( + MoERejectReason.SM_UNSUPPORTED, + f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}", + ) + if p.quant_algo not in {QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4}: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + f"CuteDslB12xFusedMoE only supports NVFP4 or W4A16_NVFP4 quantization " + f"(got quant_algo={p.quant_algo})", + ) + if p.dtype_act not in {torch.float16, torch.bfloat16}: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, + f"CuteDslB12xFusedMoE NVFP4 requires float16 or bfloat16 " + f"activation dtype (got {p.dtype_act})", + ) + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "CuteDslB12xFusedMoE does not support swiglu_gptoss_style", + ) + if p.activation_type not in _ACTIVATION_MAP: + supported = ", ".join(a.name for a in _ACTIVATION_MAP) + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"CuteDslB12xFusedMoE does not support activation " + f"{p.activation}; supported: {supported}", + ) + # The decode kernel ships in the FlashInfer wheel. + if not d.env.has_dep(MoEDep.FLASHINFER): + return _reject( + MoERejectReason.DEP_MISSING, + "CuteDslB12xFusedMoE requires the flashinfer package", + ) + # No expert-parallel dispatch/combine kernel: EP must stay at 1. + if d.ep_size != 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"CuteDslB12xFusedMoE requires ep_size == 1 (got {d.ep_size})", + ) + # Attention-DP is a separate axis from EP: with moe_tp == tp the layer + # can have ep_size == 1 and still sit behind a DP allgather / + # reducescatter that the b12x wrapper has never been exercised under. + # ``use_dp and parallel_size > 1`` is exactly ``mapping.dp_size > 1`` + # (``Mapping.dp_size`` is ``tp_size`` when attention-DP is on). + if d.use_dp and d.parallel_size > 1: + return _reject( + MoERejectReason.TOPOLOGY_UNSUPPORTED, + f"CuteDslB12xFusedMoE does not support attention-DP " + f"(parallel_size={d.parallel_size})", + ) + return MoEEligibility.ok() + + def __init__(self, *args, **kwargs): + # ``ModelConfig`` is consumed by the inherited ``__init__`` for cache + # / mapping setup but isn't kept on ``self``. b12x's wrapper needs the + # ``use_cuda_graph`` flag at construction time, so capture it here + # before delegating. + model_config = kwargs.get("model_config", None) + self._b12x_use_cuda_graph = bool(getattr(model_config, "use_cuda_graph", False)) + + super().__init__(*args, **kwargs) + + # Eligibility (SM / quant / activation / EP) is owned by + # ``can_implement``. ``enable_alltoall`` is a run-time flag not yet on + # ``MoEDeployment``, so keep the construction guard here. + if self.enable_alltoall: + raise ValueError("CuteDslB12xFusedMoE does not support MoE alltoall communication.") + + self._b12x_weights: Optional[dict] = None + self.b12x_wrapper = None + + def _get_quant_method(self): + # Route NVFP4 to the b12x-aware quant method so weight prep + # (SF un-normalization, ``convert_sf_to_mma_layout``, + # ``B12xMoEWrapper`` instantiation) lives next to the rest of the + # NVFP4 quant-method family, while every other quant algo (and the + # unquantized fallback) continues to resolve via the parent. + if ( + self.quant_config is not None + and self.quant_config.layer_quant_mode.has_any_quant(exclude_kv_cache=True) + and self.quant_config.layer_quant_mode.has_nvfp4() + ): + from .quantization import NVFP4CuteDslB12xFusedMoEMethod + + return NVFP4CuteDslB12xFusedMoEMethod() + return super()._get_quant_method() + + def _route_to_cutlass(self, x) -> bool: + """Return ``True`` iff this call should fall back to the inherited + CUTLASS path (NVFP4 prefill chunk). ``Fp4QuantizedTensor`` inputs + always stay on the b12x path (which rejects them) so the existing + error message is preserved.""" + quant_config = getattr(self, "quant_config", None) + if quant_config is not None and quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: + return False + return isinstance(x, torch.Tensor) and x.shape[0] >= self._PREFILL_VIA_CUTLASS_THRESHOLD + + # ``post_load_weights`` is inherited from ``CutlassFusedMoE`` and + # dispatches to ``self.quant_method.transform_weights(self)`` — for this + # backend ``self.quant_method`` is ``NVFP4CuteDslB12xFusedMoEMethod`` + # (see ``_get_quant_method`` override), which performs the SF un-normalization, + # ``convert_sf_to_mma_layout`` reshape, ``B12xMoEWrapper`` instantiation, + # and the cross-layer shared output buffer dance. The wrapper and the + # bundled weight dict are attached to this module as ``self.b12x_wrapper`` + # / ``self._b12x_weights``, which the decode path below consumes. + + @nvtx_range("[b12x] quantize_input") + def quantize_input( + self, + x: Union[torch.Tensor, Fp4QuantizedTensor], + post_quant_comm: bool = True, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Hybrid dispatch entrypoint for activation handling. + + NVFP4 prefill chunks take the inherited + :meth:`CutlassFusedMoE.quantize_input` path so the downstream + ``run_moe`` can call CUTLASS NVFP4 GroupGEMM. Decode chunks and + W4A16_NVFP4 chunks pass through unchanged because b12x quantizes + activations internally (consumes a bf16 / fp16 ``x`` and produces its + own scale factors). + """ + if self._route_to_cutlass(x): + return CutlassFusedMoE.quantize_input( + self, x, post_quant_comm=post_quant_comm, **kwargs + ) + if isinstance(x, Fp4QuantizedTensor): + raise ValueError( + "CuteDslB12xFusedMoE does not accept Fp4QuantizedTensor input " + "on the b12x decode path; b12x performs its own input quantization." + ) + return x, None + + @nvtx_range("[b12x] run_moe") + def run_moe( + self, + ctx: MoERunContext, + *, + workspace: Optional[dict] = None, + ) -> torch.Tensor: + plan = require_comm_plan(self, ctx) + x = ctx.x + # Join the async input quantization: the scheduler forks the inherited + # ``quantize_input_async`` (which runs the polymorphic + # ``self.quantize_input`` on the MoE aux stream). On the + # prefill-via-CUTLASS route that is a real FP4 quant kernel and + # ``CutlassFusedMoE.run_moe`` below never joins the event; on the b12x + # decode route ``quantize_input`` is a pass-through, so the wait is a + # no-op. + if ctx.x_ready_event is not None: + ctx.x_ready_event.wait() + if self._route_to_cutlass(x): + # ``CutlassFusedMoE.run_moe`` forwards ``output_dtype`` straight + # into the C++ ``trtllm::fused_moe`` op, which requires a concrete + # high-precision ``ScalarType`` (uint8 / FP4-packed activations are + # rejected at the kernel epilogue with "Invalid output type Byte"). + # ``ConfigurableMoE.forward`` always fills ``output_dtype``, so this + # only narrows the type for anything driving ``run_moe`` without it. + _HIGH_PRECISION = {torch.float16, torch.bfloat16, torch.float32} + cutlass_output_dtype = ctx.output_dtype + if cutlass_output_dtype is None: + cutlass_output_dtype = ( + x.dtype + if isinstance(x, torch.Tensor) and x.dtype in _HIGH_PRECISION + else torch.bfloat16 + ) + return CutlassFusedMoE.run_moe( + self, + replace(ctx, output_dtype=cutlass_output_dtype), + workspace=workspace, + ) + token_selected_experts = ctx.token_selected_experts + token_final_scales = ctx.token_final_scales + x_sf = ctx.x_sf + moe_output = plan.moe_output + if self.b12x_wrapper is None or self._b12x_weights is None: + raise RuntimeError( + "CuteDslB12xFusedMoE.run_moe called before process_weights_after_loading completed." + ) + if x_sf is not None: + raise ValueError( + "CuteDslB12xFusedMoE expects unquantized input (x_sf=None) " + "on the b12x decode path; got a precomputed scale factor." + ) + + # Annotate the kwargs spread + wrapper entry separately so we can + # attribute the per-layer Python dispatch cost vs. the kernel cost. + with nvtx_range("[b12x] wrapper.run"): + out = self.b12x_wrapper.run( + x=x, + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + **self._b12x_weights, + ) + + # B12xMoEWrapper allocates its own output buffer for CUDA-graph + # compatibility. If the caller provided ``moe_output`` (e.g. an + # alltoall workspace tensor), copy into it; CuteDslB12xFusedMoE + # currently rejects alltoall in __init__, so this is a defensive + # path for future workspace-driven uses. + if moe_output is not None: + with nvtx_range("[b12x] out_copy"): + moe_output.copy_(out) + return moe_output + return out diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py new file mode 100644 index 000000000000..b0f75ee6f637 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py @@ -0,0 +1,264 @@ +# 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. +"""CuteDslFc12FusedMoE: FC1+FC2-fused CuteDSL NVFP4 MoE backend for Rubin (SM107). + +Drives ``cute_dsl_nvfp4_fc12_fused_rubin``: one persistent kernel does gather + +FC1 GEMM + gated epilogue + requant + FC2 GEMM + finalize. The epilogue is a +trace-time specialization of SwiGLU or SiTU (Kimi K3); Relu2 is not gated and +is not supported. uGPU is not enabled for this backend. + +``quantize_input``, ``run_moe``, the weight lifecycle, capabilities, and EPLB +are inherited from :class:`CuteDslFusedMoE`. The SiTU constants live on the +parent (``situ_beta`` / ``situ_linear_beta``) and are forwarded into the fused +op; they are part of the compiled-kernel cache key. + +See ``CuteDslFc12FusedMoE-integration-plan.txt`` for the original integration +plan. +""" + +from typing import Optional + +import torch + +from tensorrt_llm.models.modeling_utils import QuantAlgo + +from ...autotuner import AutoTuner +from ...cute_dsl_utils import IS_CUTLASS_DSL_INTERNAL_AVAILABLE +from ...utils import ActivationType +from .fused_moe_cute_dsl import (CuteDslFusedMoE, CuteDslFusedMoENvfp4Runner, + NvFp4WeightView) +from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, + MoERejectReason) +from .interface import _reject + + +class CuteDslFc12FusedMoENvfp4Runner(CuteDslFusedMoENvfp4Runner): + """Outer autotune runner for the fused FC12 backend. + + Identical to the parent except the routing-tile candidate set: the fused + kernel supports the 128-wide tile (1-CTA) and the 256-wide tile (2-CTA, + cluster (2,1)), so restrict ``_tile_sizes`` to ``[128, 256]`` (the parent + also offers 512, which the fused v1 kernel does not support). The inner + fused runner derives mma_tiler_m == tile_size and cluster M == tile_size // + 128 from the selected tile, mirroring the CuteDSL grouped-GEMM runners. + """ + + @staticmethod + def _tile_sizes(): + return [128, 256] + + +class CuteDslFc12FusedMoE(CuteDslFusedMoE): + """FC1+FC2-fused CuteDSL NVFP4 MoE backend (Rubin/SM107). + + SwiGLU and SiTU share the gated FC1 geometry; the fused kernel specializes + the epilogue at trace time. uGPU is not enabled for this backend. + """ + + @classmethod + def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: + """FC12 fused CuteDSL: NVFP4 on Rubin (SM107) only. + + Pure gate -- reads only ``p`` and ``d`` (no ``get_sm_version()``, no + import probes, no ``os.environ``); the frozen environment lives in + ``d.env``. Narrower than the parent (SM107 + NVFP4 only) because the FC12 + fused kernel targets Rubin NVFP4; broaden when other shapes are ported. + """ + sm_version = d.env.sm + + # Output is hardcoded to bfloat16 (inherited), so activation must match. + if p.dtype_act != torch.bfloat16: + return _reject( + MoERejectReason.DTYPE_UNSUPPORTED, + "CuteDslFc12FusedMoE only supports bfloat16 activation, " + f"got {p.dtype_act}") + + if p.activation_type not in (ActivationType.Swiglu, ActivationType.SiTu): + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "CuteDslFc12FusedMoE supports only SwiGLU and SiTU, " + f"got {p.activation_type.name}") + + if p.swiglu_gptoss_style: + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + "CuteDslFc12FusedMoE does not support swiglu_gptoss_style") + + if p.quant_algo != QuantAlgo.NVFP4: + return _reject( + MoERejectReason.QUANT_UNSUPPORTED, + "CuteDslFc12FusedMoE only supports NVFP4, " + f"got quant_algo={p.quant_algo}") + + if sm_version != 107: + return _reject( + MoERejectReason.SM_UNSUPPORTED, + "CuteDslFc12FusedMoE targets Rubin (SM107), " + f"got SM{sm_version}") + + if not IS_CUTLASS_DSL_INTERNAL_AVAILABLE: + return _reject( + MoERejectReason.DEP_MISSING, + "CuteDslFc12FusedMoE (SM107 NVFP4) requires CuTE DSL internal") + + return MoEEligibility.ok() + + # ------------------------------------------------------------------ + # Fused-kernel dispatch. Mirrors the parent's non-uGPU NVFP4 path but + # (1) uses a distinct autotuner key + outer runner so FC12 tactics do + # not collide with the parent CuteDSL backend, and (2) drives the fused + # single-op path in ``run_moe_nvfp4_impl``. uGPU is never used here. + # ------------------------------------------------------------------ + def run_moe_nvfp4( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + x_sf: Optional[torch.Tensor] = None, + moe_output: Optional[torch.Tensor] = None, + enable_alltoall: bool = False, + weight_view: Optional[NvFp4WeightView] = None, + x_ready_event: Optional[torch.cuda.Event] = None, + ) -> torch.Tensor: + assert self.has_nvfp4 + assert weight_view is not None + output_dtype = torch.bfloat16 + + if moe_output is None: + moe_output = torch.empty( + (token_final_scales.size(0), self.hidden_size), + dtype=output_dtype, + device=x.device) + else: + assert moe_output.size() == (token_final_scales.size(0), + self.hidden_size) + assert moe_output.dtype == output_dtype + + # Empty micro-batches: skip autotuning (synthetic grouped-GEMM inputs + # require at least one output row). + if token_selected_experts.size(0) == 0: + return moe_output + + # Deferred join for ``quantize_input_async`` (see the parent class): + # the fused kernel is the first consumer of ``x``/``x_sf``, so + # ``run_moe_nvfp4_impl`` joins right after ``moe_sort``. + self._pending_x_ready_event = x_ready_event + + effective_top_k = token_selected_experts.size(-1) + tuner = AutoTuner.get() + runner = CuteDslFc12FusedMoENvfp4Runner( + forward_impl=self.run_moe_nvfp4_impl, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=weight_view.expert_size_per_partition, + local_expert_offset=weight_view.slot_start, + enable_finalize_fusion=self.use_fused_finalize, + enable_alltoall=enable_alltoall, + # Inner kernel_cache already keys on these; the outer AutoTuner + # unique_id must too, or SwiGLU/SiTU share a tile-size tactic. + workload_identity=( + int(self.activation_type), + self.situ_beta, + self.situ_linear_beta, + ), + ) + inputs = [ + x, token_selected_experts, token_final_scales, x_sf, moe_output, + weight_view + ] + _, best_tactic = tuner.choose_one( + "CuteDslFc12FusedMoE::run_moe_nvfp4", + [runner], + runner.get_tuning_config(), + inputs, + ) + return runner(inputs, tactic=best_tactic) + + def run_moe_nvfp4_impl( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + x_sf: torch.Tensor, + moe_output: torch.Tensor, + weight_view: NvFp4WeightView, + enable_alltoall: bool = False, + tile_size: int = 128, + ) -> torch.Tensor: + """Single fused FC1+FC2 op (replaces the parent's two-op sequence).""" + effective_top_k = token_selected_experts.size(1) + esp = weight_view.expert_size_per_partition + slot_start = weight_view.slot_start + + (tile_idx_to_expert_idx, tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=self.num_slots, + top_k=effective_top_k, + local_expert_offset=slot_start, + local_num_experts=esp, + tile_tokens_dim=tile_size, + ) + + # Join the async input quantization (``quantize_input_async``): the + # fused op below is the first consumer of ``x``/``x_sf``; ``moe_sort`` + # above only reads routing tensors, so it overlaps the aux-stream + # quantization. + self._join_pending_x_ready() + + # One fused op: gather + FC1 GEMM + gated act (SwiGLU or SiTU) + requant + # + FC2 GEMM + finalize (scatter-add into moe_output). + # fc1_alpha/fc2_alpha map 1:1 to the two-op path's per-expert global + # scales (the kernel takes split alphas). The three atomic counters are + # allocated + memset inside the op runner. The moe_output zeroing memset + # is issued inside the op too (right before the fused kernel, so it is + # the kernel's immediate stream predecessor and the PDL prologue can + # overlap it); expanded_idx_to_permuted_idx / ep_size / enable_alltoall + # are passed through for that memset. + situ = ActivationType(self.activation_type) == ActivationType.SiTu + torch.ops.trtllm.cute_dsl_nvfp4_fc12_fused_rubin( + input=x.view(torch.float4_e2m1fn_x2), + fc1_weight=weight_view.w3_w1_weight.view(torch.float4_e2m1fn_x2), + input_scale=x_sf.view(torch.uint8), + fc1_weight_scale=weight_view.fc1_weight_scale.view(torch.uint8), + fc1_alpha=weight_view.fc1_global_scale, + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + global_sf=self.fc2_input_scale, + fc2_weight=weight_view.w2_weight.view(torch.float4_e2m1fn_x2), + fc2_weight_scale=weight_view.fc2_weight_scale.view(torch.uint8), + fc2_alpha=weight_view.fc2_global_scale, + output=moe_output, + token_final_scales=token_final_scales, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=esp, + local_expert_offset=slot_start, + tile_size=tile_size, + swiglu_limit=self.swiglu_limit_scalar, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + scaling_vector_size=16, + activation_type=int(self.activation_type), + situ_beta=self.situ_beta if situ else -1.0, + situ_linear_beta=self.situ_linear_beta if situ else -1.0, + ) + return moe_output From 87359f85fd9f08451b0407cc26157c1dcf578ec2 Mon Sep 17 00:00:00 2001 From: Xin Guan Date: Wed, 14 Oct 2026 19:27:25 -0700 Subject: [PATCH 03/16] [None][fix] DFlash/DSpark: size slot-indexed rejection-sampling buffers by the seq-slot pool and share them with CUDA-graph metadata Signed-off-by: Xin Guan Signed-off-by: Mingyang Hao --- .../_torch/pyexecutor/model_engine.py | 106 +++++++++++++++++- tensorrt_llm/_torch/speculative/interface.py | 22 +++- tensorrt_llm/_torch/speculative/utils.py | 2 + 3 files changed, 121 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 40789830cb7d..2372d00d16ea 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1012,11 +1012,107 @@ def get_mm_encoder_item_cache_keys( return None return self._mm_item_scheduler.item_cache_keys(request) - def invalidate_multimodal_encoder_cache(self) -> None: - """Clear cached MM encoder outputs when no request is using them.""" - encoder_cache = self.mm_encoder_cache - if encoder_cache is not None: - encoder_cache.clear() + def _resolve_bytes_per_mm_encoder_embedding(self) -> int: + """Bytes occupied by one multimodal encoder output embedding. + + Prefers the mixin's explicit `embedding_dim`/`embedding_dtype` + contract, then the text embedding layer's weight, then the + pretrained config's hidden size with the loaded weights' dtype — + both mixin properties are optional and most VLMs implement + neither. + """ + model = self.model + try: + return (model.embedding_dim * torch.empty( + (), dtype=model.embedding_dtype).element_size()) + except NotImplementedError: + pass + try: + weight = model.text_embedding_layer.weight + return weight.shape[-1] * weight.element_size() + except NotImplementedError: + pass + pretrained = model.model_config.pretrained_config + hidden_size = getattr(pretrained, "hidden_size", None) + if hidden_size is None: + hidden_size = getattr(getattr(pretrained, "text_config", None), + "hidden_size", None) + if hidden_size is None: + raise ValueError( + "Cannot derive the MM embedding row size: the model " + "implements neither embedding_dim/embedding_dtype nor " + "text_embedding_layer, and its pretrained config exposes " + "no (text_config.)hidden_size") + element_size = next(model.parameters()).dtype.itemsize + return hidden_size * element_size + + def _set_up_multimodal_encoder_attn_metadata(self) -> None: + """Construct AttentionMetadata for any multimodal encoders inside the + loaded model, using the engine's encoder token budget + (`encoder_max_num_tokens`, falling back to the LLM-side + `max_num_tokens`). + + Mirrors `_set_up_attn_metadata` for the LLM backbone: encoders opt in + by inheriting `MultimodalEncoderMixin`, and the engine drives the construction + so the sizes match the resolved encoder token budget rather than being + hardcoded inside each encoder's `__init__`. The optional + per-segment capacity combines the encoder token budget with the input + processor's largest supported item. + """ + max_seq_len = self.encoder_max_num_tokens + if isinstance(self.input_processor, BaseMultimodalDummyInputsBuilder): + max_tokens_per_item = ( + self.input_processor.get_mm_max_tokens_per_item()) + max_seq_len = max(max_seq_len, + max(max_tokens_per_item.values(), default=0)) + + for module in self.model.modules(): + if isinstance(module, MultimodalEncoderMixin): + setup_kwargs: Dict[str, Any] = dict( + max_num_tokens=self.encoder_max_num_tokens) + if self.mm_encoder_attention_metadata_capacity is not None: + setup_kwargs["attention_metadata_capacity"] = ( + self.mm_encoder_attention_metadata_capacity) + module.setup_attn_metadata(**setup_kwargs) + module.set_attn_max_seq_len(max_seq_len) + + def _set_up_spec_metadata( + self, + spec_resource_manager: Optional[BaseResourceManager], + no_cache=False): + spec_config = self.spec_config if self.enable_spec_decode else None + # The disaggregated attention-DP overlap path opts into larger metadata + # buffers. Passing None preserves the established max_num_requests + # fallback for other configurations, including PP. + # Size slot-indexed spec buffers (e.g. rejection-sampling draft_probs) + # by the same pool that sizes the SeqSlotManager + # (py_executor_creator passes model_engine.max_num_seq_slots), so every + # py_seq_slot has a row. With max_num_requests only, DSpark + rejection + # sampling on disaggregated attention-DP gen indexed past the buffer. + num_seq_slots = self.max_num_seq_slots + if no_cache: + return get_spec_metadata( + spec_config, + self.model.config, + self.batch_size, + max_num_tokens=self.max_num_tokens, + spec_resource_manager=spec_resource_manager, + is_draft_model=self.is_draft_model, + max_seq_len=self.max_seq_len, + num_seq_slots=num_seq_slots) + + if self.spec_metadata is not None: + return self.spec_metadata + self.spec_metadata = get_spec_metadata( + spec_config, + self.model.config, + self.batch_size, + max_num_tokens=self.max_num_tokens, + spec_resource_manager=spec_resource_manager, + is_draft_model=self.is_draft_model, + max_seq_len=self.max_seq_len, + num_seq_slots=num_seq_slots) + return self.spec_metadata def cleanup(self) -> None: """Release resources owned by this model engine. diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 407e8ee100be..f41d241ce25d 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1126,6 +1126,17 @@ def write_padding_onehot_draft_probs(self, padding_slot_ids, draft_len): return onehot_vocab = (self.draft_probs_last_dim if self.draft_probs_last_dim > 0 else self.draft_probs_vocab_size) + bad = [s for s in padding_slot_ids + if s is None or s < 0 or s >= self.draft_probs.shape[0]] + if bad: + # Writing a different row would leave the real slot's distribution + # stale and break lossless acceptance, so fail loudly instead. + raise RuntimeError( + f"rejection sampling: padding slot ids {bad[:8]} outside " + f"draft_probs rows [0, {self.draft_probs.shape[0]}) " + f"(num_seq_slots={self.num_seq_slots}, " + f"max_num_requests={self.max_num_requests}, " + f"is_cuda_graph={self.is_cuda_graph})") slots = torch.tensor(padding_slot_ids, dtype=torch.long, device=self.draft_probs.device) @@ -1164,11 +1175,14 @@ def create_cuda_graph_metadata(self, max_batch_size: int): if self.is_cuda_graph: return self - # Slot-indexed state persists across steps that may run eagerly or under - # different graphs, so every view must share one allocation: allocate it - # here before the shallow copy rather than lazily in each view. + # Allocate the slot-indexed rejection-sampling buffers on the base + # metadata BEFORE the shallow copy so every CUDA-graph view shares + # them. A copy taken while draft_probs is still None allocates its own + # buffer sized by the graph batch size: the drafter then scatters into + # one buffer while the captured rejection kernel reads another, and + # slot ids >= the graph batch size index out of bounds (seen with + # DSpark + attention DP, write_padding_onehot_draft_probs). self.prepare_rejection_sampling_buffers() - self.prepare_penalty_buffers() cuda_graph_metadata = copy.copy(self) cuda_graph_metadata.is_cuda_graph = True cuda_graph_metadata.max_num_requests = max_batch_size diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index f8e1a7340145..b2c63bc05f7b 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -456,6 +456,7 @@ def _build_spec_metadata(spec_config, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, + num_seq_slots=num_seq_slots, ) if spec_config.spec_dec_mode.is_dspark(): target_layer_ids = getattr(spec_config, 'target_layer_ids', None) @@ -476,6 +477,7 @@ def _build_spec_metadata(spec_config, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, + num_seq_slots=num_seq_slots, ) if spec_config.spec_dec_mode.is_draft_target_one_model(): return DraftTargetOneModelSpecMetadata( From 272f7bf274b786c9101ae3e37523569c934373fb Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 15 Oct 2026 22:08:55 -0700 Subject: [PATCH 04/16] [None][feat] Transfer the standalone DSpark draft KV from CTX to GEN on rubin-advance (ports !10774 + !10802) Signed-off-by: Zhenhuan Chen Signed-off-by: Mingyang Hao --- .../kv_cache_manager_v2/common.h | 7 + .../kv_cache_manager_v2/config.h | 6 +- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 6 +- .../kv_cache_manager_v2/lifeCycleRegistry.h | 44 +- .../kv_cache_manager_v2/storage/config.cpp | 22 +- .../batch_manager/kvCacheManagerV2.cpp | 23 +- .../unit_tests/batch_manager/CMakeLists.txt | 85 +- .../kvCacheManagerV2ConfigTest.cpp | 158 +- .../backends/fp4_mla/cache_manager.py | 47 +- .../sparse/deepseek_v4/cache_manager.py | 191 +- .../backends/sparse/dsa/cache_manager.py | 50 +- .../kernels/blackwell/dspark/attention.py | 49 +- .../blackwell/dspark/attention_kernel.py | 346 +- .../kernels/blackwell/dspark_rmsnorm_rope.py | 288 +- .../_torch/attention_backend/dspark.py | 70 + .../sparse/minimax_m3/cache_manager.py | 473 ++ .../custom_ops/dspark_attention_custom_op.py | 1113 +++- .../dspark_rmsnorm_rope_custom_op.py | 566 +- .../_torch/disaggregation/native/auxiliary.py | 137 +- .../disaggregation/native/bounce/impl.py | 12 +- .../native/mixers/attention/peer.py | 39 +- .../_torch/disaggregation/native/peer.py | 77 +- .../_torch/disaggregation/native/transfer.py | 3724 +++++------- .../disaggregation/resource/cache_reuse.py | 20 +- .../disaggregation/resource/kv_extractor.py | 100 +- .../_torch/disaggregation/resource/page.py | 11 +- .../_torch/disaggregation/transceiver.py | 288 +- tensorrt_llm/_torch/models/modeling_dflash.py | 139 +- tensorrt_llm/_torch/models/modeling_dspark.py | 808 +-- tensorrt_llm/_torch/pyexecutor/_util.py | 282 +- .../_torch/pyexecutor/kv_cache_manager_v2.py | 5326 +++++++++++++++++ .../_torch/pyexecutor/mamba_cache_manager.py | 4723 +++++++++++++++ .../_torch/pyexecutor/model_engine.py | 793 ++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 14 +- .../_torch/pyexecutor/sampler/sampler.py | 19 +- .../pyexecutor/standalone_draft_cache.py | 66 + tensorrt_llm/_torch/speculative/dflash.py | 602 +- tensorrt_llm/_torch/speculative/dspark.py | 550 +- tensorrt_llm/_torch/speculative/interface.py | 15 + tensorrt_llm/_torch/speculative/utils.py | 6 +- .../runtime/kv_cache_manager_v2/__init__.py | 14 + .../runtime/kv_cache_manager_v2/__init__.pyi | 20 +- .../executor/test_kv_cache_estimation.py | 1010 ++++ .../hw_agnostic/test_dflash_worker.py | 190 +- .../hw_agnostic/test_dspark_attention.py | 1062 +++- .../hw_agnostic/test_dspark_cuda_graph.py | 235 +- .../hw_agnostic/test_dspark_worker.py | 278 +- .../test_kimi_k3_dspark_semantics.py | 87 + .../test_dspark_cute_dsl_attention.py | 1556 +++-- .../unittest/disaggregated/region/test_aux.py | 22 + .../unittest/disaggregated/test_extractor.py | 1 + 51 files changed, 19690 insertions(+), 6080 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/dspark.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py create mode 100644 tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py create mode 100644 tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py create mode 100644 tests/unittest/_torch/executor/test_kv_cache_estimation.py diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h index 4d867eb906a4..79ae6c456226 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h @@ -51,6 +51,13 @@ enum class PageStatus : int DROPPABLE = 2, // Allow eviction and dropping. }; +// Standalone draft pools precede target pools in storage layout. +enum class CacheDomain : int +{ + kStandaloneDraft = 0, + kTarget = 1, +}; + enum class CacheTier : int { GPU_MEM = 0, diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index 19600e0fd52a..e071b04e7b6a 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -204,10 +204,10 @@ struct AttentionLayerConfig // nullopt or 0 = no sink tokens. std::optional numSinkTokens; - // Private working pages never constrain prefix matching or enter the tree. - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + // Layers in different ownership domains must not share lifecycle or storage pools. + CacheDomain cacheDomain = CacheDomain::kTarget; - [[nodiscard]] std::optional windowSize() const noexcept + std::optional windowSize() const noexcept { return slidingWindowSize; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index a4c57cacea07..891210d92a9c 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -39,11 +39,7 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) return SsmLifeCycle{}; } else - { - bool const isSparse = !cfg.buffers.empty() && cfg.buffers.front().isSparse; - return AttnLifeCycle::make( - cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, cfg.reusePolicy); - } + return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, cfg.cacheDomain); }, layer); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index bdf550e6fcc1..a5ae27b78d89 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -40,23 +41,7 @@ struct AttnLifeCycle { std::optional windowSize; // nullopt = no sliding window int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) - bool isSparse = false; - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; - - bool isPrivate() const noexcept - { - return reusePolicy == AttentionReusePolicy::PRIVATE; - } - - bool requiresPrefixCoverage() const noexcept - { - return reusePolicy == AttentionReusePolicy::REQUIRED; - } - - bool allowsCheckpoint() const noexcept - { - return reusePolicy == AttentionReusePolicy::OPTIONAL; - } + CacheDomain cacheDomain = CacheDomain::kTarget; HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const { @@ -74,36 +59,24 @@ struct AttnLifeCycle bool operator==(AttnLifeCycle const& o) const noexcept { - return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse - && reusePolicy == o.reusePolicy; + return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && cacheDomain == o.cacheDomain; } bool operator<(AttnLifeCycle const& o) const noexcept { - if (windowSize != o.windowSize) - { - return windowSize < o.windowSize; - } - if (numSinkBlocks != o.numSinkBlocks) - { - return numSinkBlocks < o.numSinkBlocks; - } - if (isSparse != o.isSparse) - { - return isSparse < o.isSparse; - } - return reusePolicy < o.reusePolicy; + return std::tie(windowSize, numSinkBlocks, cacheDomain) + < std::tie(o.windowSize, o.numSinkBlocks, o.cacheDomain); } static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock, - bool isSparse = false, AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED) + CacheDomain cacheDomain = CacheDomain::kTarget) { TLLM_CHECK_DEBUG(tokensPerBlock > 0); TLLM_CHECK_DEBUG(!ws.has_value() || *ws > 0); TLLM_CHECK_DEBUG(!numSinkTokens.has_value() || *numSinkTokens >= 0); TLLM_CHECK_DEBUG((!numSinkTokens.has_value() || *numSinkTokens == 0) || ws.has_value()); int sinkBlocks = divUp(numSinkTokens.value_or(0), tokensPerBlock); - return AttnLifeCycle{ws, sinkBlocks, isSparse, reusePolicy}; + return AttnLifeCycle{ws, sinkBlocks, cacheDomain}; } }; @@ -163,7 +136,8 @@ inline HalfOpenRange computeScratchRange( LifeCycle const& lc, int historyLength, int capacity, int tokensPerBlock, int maxRewindLen) { auto const* attn = std::get_if(&lc); - if (!attn || !attn->windowSize.has_value()) + // Draft history must retain private pages for later drafting or transfer. + if (!attn || !attn->windowSize.has_value() || attn->cacheDomain == CacheDomain::kStandaloneDraft) { return {BlockOrdinal{0}, BlockOrdinal{0}}; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index a5d17a7817e1..c19e63250ffd 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -206,28 +206,14 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) slotGroups.push_back(std::move(var)); } - // Distinct sparsity and reuse policies need separate GPU pools and eviction queues. - struct PoolGroupKey - { - AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; - std::vector slotSizes; - bool isSparse = false; - - bool operator<(PoolGroupKey const& other) const - { - return std::tie(reusePolicy, slotSizes, isSparse) - < std::tie(other.reusePolicy, other.slotSizes, other.isSparse); - } - }; - - std::map> poolGroups; + // Only matching ownership domains and slot layouts may share pools. + std::map>, std::vector> poolGroupsBySizes; for (auto& sg : slotGroups) { auto const sizes = sg.slotSizeList(); auto const* attn = std::get_if(®istry.getLifeCycle(sg.lifeCycleId)); - bool const isSparse = attn != nullptr && attn->isSparse; - auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; - poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .isSparse = isSparse}].push_back(std::move(sg)); + CacheDomain const domain = attn ? attn->cacheDomain : CacheDomain::kTarget; + poolGroupsBySizes[{domain, sizes.raw()}].push_back(std::move(sg)); } StorageConfig out; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index c503e0250648..42a585fab383 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -901,6 +901,10 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("DROPPABLE", kv::PageStatus::DROPPABLE) .export_values(); + nb::enum_(m, "CacheDomain") + .value("STANDALONE_DRAFT", kv::CacheDomain::kStandaloneDraft) + .value("TARGET", kv::CacheDomain::kTarget); + nb::enum_(m, "CacheTier") .value("GPU_MEM", kv::CacheTier::GPU_MEM) .value("HOST_MEM", kv::CacheTier::HOST_MEM) @@ -1407,18 +1411,13 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("OPTIONAL", kv::AttentionReusePolicy::OPTIONAL); nb::class_(m, "AttnLifeCycle") - .def(nb::init, int, bool, kv::AttentionReusePolicy>(), nb::arg("window_size").none(), - nb::arg("num_sink_blocks"), nb::arg("is_sparse") = false, - nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) - // Sink tokens round up to whole blocks. Bound rather than repeated in Python so the - // connector's view of a life cycle is built by the same code as the allocator's. + .def(nb::init, int, kv::CacheDomain>(), nb::arg("window_size").none(), + nb::arg("num_sink_blocks"), nb::arg("cache_domain") = kv::CacheDomain::kTarget) .def_static("make", &kv::AttnLifeCycle::make, nb::arg("window_size").none(), nb::arg("num_sink_tokens").none(), - nb::arg("tokens_per_block"), nb::arg("is_sparse") = false, - nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) + nb::arg("tokens_per_block"), nb::arg("cache_domain") = kv::CacheDomain::kTarget) .def_prop_ro("window_size", [](kv::AttnLifeCycle const& self) { return self.windowSize; }) .def_ro("num_sink_blocks", &kv::AttnLifeCycle::numSinkBlocks) - .def_ro("is_sparse", &kv::AttnLifeCycle::isSparse) - .def_ro("reuse_policy", &kv::AttnLifeCycle::reusePolicy) + .def_ro("cache_domain", &kv::AttnLifeCycle::cacheDomain) .def("get_stale_range", &kv::AttnLifeCycle::getStaleRange, nb::arg("history_length"), nb::arg("tokens_per_block")) .def("__eq__", &kv::AttnLifeCycle::operator==); @@ -1627,14 +1626,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, - kv::AttentionReusePolicy>(), + kv::CacheDomain>(), nb::arg("layer_id"), nb::arg("buffers"), nb::arg("sliding_window_size") = std::nullopt, - nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) + nb::arg("num_sink_tokens") = std::nullopt, nb::arg("cache_domain") = kv::CacheDomain::kTarget) .def_rw("layer_id", &kv::AttentionLayerConfig::layerId) .def_rw("buffers", &kv::AttentionLayerConfig::buffers) .def_rw("sliding_window_size", &kv::AttentionLayerConfig::slidingWindowSize) .def_rw("num_sink_tokens", &kv::AttentionLayerConfig::numSinkTokens) - .def_rw("reuse_policy", &kv::AttentionLayerConfig::reusePolicy) + .def_rw("cache_domain", &kv::AttentionLayerConfig::cacheDomain) .def_prop_ro("window_size", &kv::AttentionLayerConfig::windowSize) DEF_COPY(kv::AttentionLayerConfig); nb::enum_(m, "LayerType") diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 0f2879b1dd77..30b6ea318e39 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -290,73 +290,18 @@ add_tllm_gtest( add_tllm_gtest( kvCacheManagerV2DigestPoolTest - OWNER - batch-manager - SOURCES - kvCacheManagerV2DigestPoolTest.cpp - LINK_LIBRARIES - tllm::kv_cache - $ - LABELS - "unit;batch-manager;host") - -add_tllm_gtest( - kvCacheUtilsTest - OWNER - batch-manager - SOURCES - kvCacheUtilsTest.cpp - LINK_LIBRARIES - tllm::kv_cache_legacy - LABELS - "unit;batch-manager;gpu") - -add_tllm_gtest( - llmRequestTest - OWNER - batch-manager - SOURCES - llmRequestTest.cpp - LINK_LIBRARIES - tllm::kv_cache_legacy - LABELS - "unit;batch-manager;gpu") - -add_tllm_gtest( - microBatchSchedulerTest - OWNER - batch-manager - SOURCES - microBatchSchedulerTest.cpp - LINK_LIBRARIES - tllm::batch_scheduler - LABELS - "unit;batch-manager;gpu") - -add_tllm_gtest(peftCacheManagerTest SOURCES peftCacheManagerTest.cpp FULL_STACK - LABELS "unit;full-stack;gpu") - -add_tllm_gtest(rnnCacheFormatterTest SOURCES rnnCacheFormatterTest.cpp - FULL_STACK LABELS "unit;full-stack;host") - -add_tllm_gtest( - agentTreeTest - OWNER - batch-manager - SOURCES - agentTreeTest.cpp - LINK_LIBRARIES - tllm::batch_scheduler - LABELS - "unit;batch-manager;host") - -add_tllm_gtest( - truncateBlocksTest - OWNER - batch-manager - SOURCES - truncateBlocksTest.cpp - LINK_LIBRARIES - tllm::kv_cache_legacy - LABELS - "unit;batch-manager;gpu") + PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager + ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256) +add_gtest(kvCacheUtilsTest kvCacheUtilsTest.cpp) +add_gtest(llmRequestTest llmRequestTest.cpp) +add_gtest(microBatchSchedulerTest microBatchSchedulerTest.cpp) +add_gtest(peftCacheManagerTest peftCacheManagerTest.cpp) +add_gtest(staticThreadPoolTest staticThreadPoolTest.cpp) +add_gtest(rnnCacheFormatterTest rnnCacheFormatterTest.cpp) +add_gtest(agentTreeTest agentTreeTest.cpp) +add_gtest(truncateBlocksTest truncateBlocksTest.cpp) + +add_gtest(kvCacheManagerV2ConfigTest kvCacheManagerV2ConfigTest.cpp) +target_include_directories( + kvCacheManagerV2ConfigTest + PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index 99756c07befc..ef9b0d7ab3df 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp @@ -15,151 +15,63 @@ * limitations under the License. */ -#include "kvCacheManagerV2TestUtils.h" -#include "kv_cache_manager_v2/lifeCycleRegistry.h" -#include "kv_cache_manager_v2/storage/config.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h" #include -#include -#include -#include - namespace { - using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; -using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeConfig; -using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeTieredConfig; - -AttentionLayerConfig makeAttentionLayer(LayerId layerId, bool isSparse, size_t bufferSize = 4096) -{ - return AttentionLayerConfig{ - .layerId = layerId, .buffers = {BufferConfig{.role = "key", .size = bufferSize, .isSparse = isSparse}}}; -} -TEST(KvCacheManagerV2ConfigTest, ExistingBufferConfigurationDefaultsToDense) +TEST(KvCacheManagerV2ConfigTest, AttentionLifecycleFactoryPreservesCacheDomain) { - BufferConfig const buffer{"key", 4096, 2}; - EXPECT_FALSE(buffer.isSparse); - EXPECT_EQ(buffer.tokensPerBlockOverride, 2); - EXPECT_NO_THROW(makeConfig().validate()); -} - -TEST(KvCacheManagerV2ConfigTest, SparseBuffersRequireHostMemoryAtLevelOne) -{ - auto config = makeConfig(); - config.layers = {makeAttentionLayer(0, true)}; - EXPECT_THROW(config.validate(), std::invalid_argument); - - config.cacheTiers.emplace_back(DiskCacheTierConfig{4 << 20, "/tmp"}); - EXPECT_THROW(config.validate(), std::invalid_argument); - - config.cacheTiers[1] = GpuCacheTierConfig{4 << 20}; - config.cacheTiers.emplace_back(HostCacheTierConfig{4 << 20}); - EXPECT_THROW(config.validate(), std::invalid_argument); - - config.cacheTiers[1] = HostCacheTierConfig{4 << 20}; - EXPECT_NO_THROW(config.validate()); + auto const target = AttnLifeCycle::make(128, 5, 4); + auto const draft = AttnLifeCycle::make(128, 5, 4, CacheDomain::kStandaloneDraft); + EXPECT_EQ(target.numSinkBlocks, 2); + EXPECT_EQ(target.cacheDomain, CacheDomain::kTarget); + EXPECT_EQ(draft.cacheDomain, CacheDomain::kStandaloneDraft); + EXPECT_FALSE(target == draft); + EXPECT_TRUE(target < draft || draft < target); + EXPECT_TRUE(draft == AttnLifeCycle::make(128, 8, 4, CacheDomain::kStandaloneDraft)); } -TEST(KvCacheManagerV2ConfigTest, MixedBuffersInOneLayerAreRejected) +TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsOnlyWithinEachCacheDomain) { - auto config = makeTieredConfig(); - for (bool const firstIsSparse : {false, true}) - { - auto layer = makeAttentionLayer(0, firstIsSparse); - layer.buffers.push_back(BufferConfig{.role = "value", .size = 4096, .isSparse = !firstIsSparse}); - config.layers = {layer}; - EXPECT_THROW(config.validate(), std::invalid_argument); - EXPECT_THROW(createStorageConfig(config), std::invalid_argument); - } -} - -TEST(KvCacheManagerV2ConfigTest, SparseSsmBuffersAreRejected) -{ - auto config = makeTieredConfig(); - config.commitMinSnapshot = true; - config.layers - = {SsmLayerConfig{.layerId = 0, .buffers = {BufferConfig{.role = "state", .size = 4096, .isSparse = true}}}}; - EXPECT_THROW(config.validate(), std::invalid_argument); - EXPECT_THROW(createStorageConfig(config), std::invalid_argument); -} - -TEST(KvCacheManagerV2ConfigTest, SparseAndDenseBuffersHaveSeparateLifecyclesAndPools) -{ - auto config = makeTieredConfig(); - config.layers = {makeAttentionLayer(0, false), makeAttentionLayer(1, true), makeAttentionLayer(2, false), - makeAttentionLayer(3, true)}; - ASSERT_NO_THROW(config.validate()); - - LifeCycleRegistry const registry(config); - auto const storage = createStorageConfig(config); - auto const attributes = storage.bufferAttributes(); - auto const grouping = storage.lifeCycleGrouping(); - auto const& dense = attributes.at(BufferId{0, "key"}); - auto const& sparse = attributes.at(BufferId{1, "key"}); - - EXPECT_EQ(registry.size(), LifeCycleId{2}); - EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{2}); - EXPECT_NE(dense.lifeCycleId, sparse.lifeCycleId); - EXPECT_NE(grouping[dense.lifeCycleId], grouping[sparse.lifeCycleId]); - EXPECT_FALSE(std::get(registry[dense.lifeCycleId]).isSparse); - EXPECT_TRUE(std::get(registry[sparse.lifeCycleId]).isSparse); - EXPECT_FALSE(registry[dense.lifeCycleId] == registry[sparse.lifeCycleId]); - - auto const& otherDense = attributes.at(BufferId{2, "key"}); - auto const& otherSparse = attributes.at(BufferId{3, "key"}); - EXPECT_EQ(dense.lifeCycleId, otherDense.lifeCycleId); - EXPECT_EQ(sparse.lifeCycleId, otherSparse.lifeCycleId); - EXPECT_EQ(sparse.poolIndex, otherSparse.poolIndex); - EXPECT_EQ(otherSparse.offset, sparse.offset + sparse.size); - EXPECT_EQ(otherDense.offset, dense.offset + dense.size); -} - -TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsWithinEachSparsity) -{ - auto config = makeTieredConfig(); - config.layers.clear(); + KVCacheManagerConfig config; + config.tokensPerBlock = 32; + config.cacheTiers + = {GpuCacheTierConfig{4 << 20}, HostCacheTierConfig{4 << 20}, DiskCacheTierConfig{4 << 20, "/tmp"}}; for (LayerId layerId = 0; layerId < 4; ++layerId) { - auto layer = makeAttentionLayer(layerId, layerId >= 2); + AttentionLayerConfig layer; + layer.layerId = layerId; + layer.buffers = {BufferConfig{"key", 4096, std::nullopt}}; layer.slidingWindowSize = layerId % 2 == 0 ? 128 : 256; + layer.cacheDomain = layerId < 2 ? CacheDomain::kTarget : CacheDomain::kStandaloneDraft; config.layers.push_back(layer); } ASSERT_NO_THROW(config.validate()); + LifeCycleRegistry const registry(config); auto const storage = createStorageConfig(config); auto const attributes = storage.bufferAttributes(); auto const grouping = storage.lifeCycleGrouping(); - auto poolGroup = [&](LayerId layerId) { return grouping[attributes.at(BufferId{layerId, "key"}).lifeCycleId]; }; - EXPECT_EQ(storage.numLifeCycles(), LifeCycleId{4}); + auto lifeCycle = [&](LayerId layerId) { return attributes.at(BufferId{layerId, "key"}).lifeCycleId; }; + auto poolGroup = [&](LayerId layerId) { return grouping[lifeCycle(layerId)]; }; + EXPECT_EQ(registry.size(), LifeCycleId{4}); EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{2}); - EXPECT_EQ(poolGroup(0), poolGroup(1)); - EXPECT_EQ(poolGroup(2), poolGroup(3)); - EXPECT_NE(poolGroup(0), poolGroup(2)); -} - -TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceByExpandedSize) -{ - auto config = makeTieredConfig(); - auto firstLayer = makeAttentionLayer(0, true, 2048); - firstLayer.buffers.front().tokensPerBlockOverride = 2; - config.layers = {firstLayer, makeAttentionLayer(1, true, 4096)}; - ASSERT_NO_THROW(config.validate()); - - auto const storage = createStorageConfig(config); - auto const attributes = storage.bufferAttributes(); - auto const& first = attributes.at(BufferId{0, "key"}); - auto const& second = attributes.at(BufferId{1, "key"}); - EXPECT_EQ(storage.numLifeCycles(), LifeCycleId{1}); - EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{1}); - EXPECT_EQ(first.lifeCycleId, second.lifeCycleId); - EXPECT_EQ(first.poolIndex, second.poolIndex); - EXPECT_EQ(first.size, second.size); - EXPECT_EQ(second.offset, first.offset + first.size); - EXPECT_EQ(first.expansion, 2); - EXPECT_EQ(second.expansion, 1); + for (LayerId first = 0; first < 4; ++first) + { + auto const& attention = std::get(registry[lifeCycle(first)]); + EXPECT_EQ(attention.cacheDomain, first < 2 ? CacheDomain::kTarget : CacheDomain::kStandaloneDraft); + for (LayerId second = first + 1; second < 4; ++second) + { + EXPECT_NE(lifeCycle(first), lifeCycle(second)); + EXPECT_FALSE(registry[lifeCycle(first)] == registry[lifeCycle(second)]); + EXPECT_EQ(poolGroup(first) == poolGroup(second), first / 2 == second / 2); + } + } } } // namespace diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index b6dcc2d963c4..e19dae50c301 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -288,6 +288,8 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. dtype and logical shape as an all-BF16 cache manager. """ local_layer = self.layer_offsets[layer_idx] + if self._is_standalone_draft_layer(local_layer): + return self.get_draft_buffers(self.draft_layer_ids.index(layer_idx), kv_layout) is_linear_layer = getattr(self, "_is_local_mamba_layer", None) if callable(is_linear_layer) and is_linear_layer(local_layer): return None @@ -322,8 +324,10 @@ def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: if layer_idx is None: return self.dtype local_layer = self.layer_offsets[layer_idx] - if local_layer in self._bf16_mla_local_layer_indices(): - return DataType.BF16 + if self._is_standalone_draft_layer(local_layer): + return self.dtype_per_layer[local_layer] + if local_layer in self._fp8_mla_local_layer_indices(): + return DataType.FP8 return self.dtype def get_kv_cache_num_blocks(self, layer_idx: int) -> int: @@ -969,10 +973,15 @@ def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: int(layer_id): self._fp4_mla_compact_to_local[compact_layer] for compact_layer, layer_id in enumerate(self._hp_manager_layer_ids) } + draft_global = { + self.layer_offsets[layer_id]: 2 * layer_id for layer_id in self.draft_layer_ids + } result = [] for layer_id in local_layer_ids: internal_layer = int(layer_id) - if internal_layer in cache_local: + if internal_layer in draft_global: + result.append(draft_global[internal_layer]) + elif internal_layer in cache_local: result.append(2 * int(self.pp_layers[cache_local[internal_layer]])) elif internal_layer in hp_local: result.append(2 * int(self.pp_layers[hp_local[internal_layer]]) + 1) @@ -998,6 +1007,8 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[list[int], list[Opti for local_layer in fp4_local_layers ) windows.extend([self._fp4_mla_hp_pool_size] * len(fp4_local_layers)) + sizes.extend(self._draft_layer_sizes) + windows.extend([self.draft_retention_window_size] * len(self._draft_layer_sizes)) return sizes, windows def _get_generation_request_capacity(self) -> int: @@ -1007,7 +1018,7 @@ def _get_generation_request_capacity(self) -> int: def get_cache_bytes_per_token(self) -> int: sizes, _ = self._get_runtime_cache_size_layer_components() - return sum(sizes) + return sum(sizes) + (self._helix_cp_size - 1) * sum(self._draft_layer_sizes) @staticmethod def get_cache_size_per_token( @@ -1061,7 +1072,33 @@ def get_cache_size_per_token( * torch.empty((), dtype=torch.bfloat16).element_size() * num_bf16_layers ) - return cache_bytes, hp_bytes * max_batch_size * mapping.pp_size + fixed_bytes = hp_bytes * max_batch_size * mapping.pp_size + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + _estimate_draft_cache_size_components, + _get_draft_layer_sizes, + _get_generation_kv_capacity, + ) + + draft_sizes = _get_draft_layer_sizes(**kwargs) + if draft_sizes: + scratch = kwargs.get("draft_scratch_tokens", 0) + _, headroom = _get_generation_kv_capacity(spec_config, is_draft=False) + draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( + draft_sizes, + tokens_per_block, + scratch_tokens=scratch, + window_size=kwargs.get("draft_window_size"), + generation_capacity_headroom=headroom + scratch, + helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, + ) + # The shared allocation envelope reserves space in target pools too. + draft_fixed += scratch * (cache_bytes + hp_bytes / (HP_BLOCK_SIZE + rewind)) + fixed_bytes += int(draft_fixed * max_batch_size * mapping.pp_size) + fixed_bytes += (draft_context - draft_generation) * int( + kwargs.get("max_num_tokens") or 0 + ) + cache_bytes += draft_generation + return cache_bytes, fixed_bytes class Fp4MlaKVCacheManagerV2(Fp4MlaV2CacheLayoutPolicy, KVCacheManagerV2): diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index 1d97a1f2f1cd..492e965e80c4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -21,16 +21,13 @@ import torch from tensorrt_llm._torch.pyexecutor import llm_request -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - _GUARD_PAGE_REQUEST_ID, - _RESERVED_REQUEST_IDS, - CUDA_GRAPH_DUMMY_REQUEST_ID, +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( GPU_LEVEL, KVCacheManagerV2, - _fill_kv_pages, + _estimate_draft_cache_size_components, + _get_draft_layer_sizes, + _get_generation_kv_capacity, ) -from tensorrt_llm._torch.pyexecutor.resource_manager import scan_kv_buffer_for_invalid_values -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import ( TensorWrapper, convert_to_torch_tensor, @@ -949,22 +946,14 @@ def _get_extra_quota_padding(self) -> int: """Ensure each attention type has minimal space when max_tokens is small.""" return len(DeepseekV4AttentionType) * (2 << 20) - def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + def _get_cache_cost_components(self) -> tuple[int, int, int, int]: + """Return context/generation bytes per token, then fixed bytes per request.""" compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] has_fp8_kv_cache = self.dtype == DataType.FP8 - non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, - ) + non_sliding_attn_size_per_token = self.get_cache_bytes_per_token() ( context_swa_size_per_token, - _, + context_swa_size_per_request, ) = _estimate_swa_cache_size( self.head_dim, self.index_head_dim, @@ -992,65 +981,60 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: indexer_k_dtype=self._indexer_k_dtype, use_fp8_ds_mla=self.use_fp8_ds_mla, ) + context_size_per_token = non_sliding_attn_size_per_token + context_swa_size_per_token + generation_size_per_token = non_sliding_attn_size_per_token + generation_swa_size_per_token + if self._draft_layer_sizes: + draft_context, draft_generation, draft_per_request = ( + _estimate_draft_cache_size_components( + self._draft_layer_sizes, + self.tokens_per_block, + scratch_tokens=self._standalone_draft_reserve, + window_size=self.draft_window_size, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + ) + ) + draft_per_request += self._standalone_draft_reserve * context_size_per_token + context_size_per_token += draft_context + generation_size_per_token += draft_generation + context_swa_size_per_request += draft_per_request + generation_swa_size_per_request += draft_per_request + return ( + context_size_per_token, + generation_size_per_token, + context_swa_size_per_request, + generation_swa_size_per_request, + ) + + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + ( + context_size_per_token, + generation_size_per_token, + _, + generation_size_per_request, + ) = self._get_cache_cost_components() max_context_tokens = ( self._max_num_tokens if self._max_num_tokens is not None else max_tokens ) context_tokens = min(max_tokens, max_context_tokens) generation_tokens = max_tokens - context_tokens - generation_quota = ( - max_tokens * non_sliding_attn_size_per_token - + generation_tokens * generation_swa_size_per_token - + self.max_batch_size * generation_swa_size_per_request + return int( + context_tokens * context_size_per_token + + generation_tokens * generation_size_per_token + + self.max_batch_size * generation_size_per_request + + self._get_extra_quota_padding() ) - context_extra_quota = context_tokens * context_swa_size_per_token - padding = self._get_extra_quota_padding() - return int(generation_quota + context_extra_quota + padding) def _get_max_tokens_from_quota(self, quota: int) -> float: - compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] - has_fp8_kv_cache = self.dtype == DataType.FP8 - non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, - ) - context_swa_size_per_token, _ = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=True, - scratch=self.enable_swa_scratch_reuse, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - ) ( - generation_swa_size_per_token, - generation_swa_size_per_request, - ) = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=False, - scratch=False, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - ) + context_size_per_token, + generation_size_per_token, + _, + generation_size_per_request, + ) = self._get_cache_cost_components() padding = self._get_extra_quota_padding() - size_per_batch = self.max_batch_size * generation_swa_size_per_request + padding + size_per_batch = self.max_batch_size * generation_size_per_request + padding if quota < size_per_batch: return 0 - context_size_per_token = non_sliding_attn_size_per_token + context_swa_size_per_token if self._max_num_tokens is None: return (quota - size_per_batch) / context_size_per_token @@ -1058,7 +1042,6 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: if quota <= context_limit_quota: return (quota - size_per_batch) / context_size_per_token - generation_size_per_token = non_sliding_attn_size_per_token + generation_swa_size_per_token if generation_size_per_token <= 0: return float("inf") return self._max_num_tokens + (quota - context_limit_quota) / generation_size_per_token @@ -1249,6 +1232,14 @@ def _add_layer( constraints=constraints, ) + def _append_standalone_draft_layers( + self, config: KVCacheManagerConfigPy, **kwargs + ) -> KVCacheManagerConfigPy: + # Preserve target virtual-layer indices when registering draft layers. + return super()._append_standalone_draft_layers( + config, register_model_layers=False, **kwargs + ) + def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None: # Indexer compressor cache layout. Two modes are supported: # - "fp8" (FP8 blockwise): 1 byte per value + 1 fp32 scale per 128 @@ -1433,34 +1424,15 @@ def _get_generation_bytes(self, request: llm_request.LlmRequest) -> int: return self._get_cache_bytes_for_tokens(total_tokens, context=False) def _get_cache_bytes_for_tokens(self, total_tokens: int, *, context: bool) -> int: - has_fp8_kv_cache = self.dtype == DataType.FP8 - compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] - non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - has_nvfp4_compress=self._use_nvfp4_compress, - nvfp4_residual_dim=self.nvfp4_residual_dim, - ) - swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( - self.head_dim, - self.index_head_dim, - compress_ratios, - has_fp8_kv_cache, - self.tokens_per_block, - self._swa_window_size, - context=context, - scratch=self.enable_swa_scratch_reuse, - indexer_k_dtype=self._indexer_k_dtype, - use_fp8_ds_mla=self.use_fp8_ds_mla, - ) - return int( - total_tokens * (non_sliding_attn_size_per_token + swa_size_per_token) - + swa_size_per_request - ) + ( + context_size_per_token, + generation_size_per_token, + context_size_per_request, + generation_size_per_request, + ) = self._get_cache_cost_components() + if context: + return int(total_tokens * context_size_per_token + context_size_per_request) + return int(total_tokens * generation_size_per_token + generation_size_per_request) def get_needed_resource_to_completion(self, request: llm_request.LlmRequest) -> int: if self._is_generation_request(request): @@ -1474,6 +1446,8 @@ def get_layer_bytes_per_token( local_layer_idx: int, data_role: DataRole, ) -> int: + if local_layer_idx in self.cp_as_tp_layer_ids: + return super().get_layer_bytes_per_token(local_layer_idx, data_role) # The generic layers in the base config are replaced by # _build_cache_config, so their buffer sizes are only placeholders. return 1 @@ -2058,6 +2032,33 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwar use_fp8_ds_mla=use_fp8_ds_mla, ) max_batch_size = int(kwargs.get("max_batch_size") or 0) + draft_layer_sizes = _get_draft_layer_sizes(**kwargs) + draft_scratch_tokens = kwargs.get("draft_scratch_tokens", 0) + if draft_layer_sizes: + _, headroom = _get_generation_kv_capacity(kwargs.get("spec_config"), is_draft=False) + target_context_swa_bytes, _ = _estimate_swa_cache_size( + head_dim, + index_head_dim, + compress_ratios, + has_fp8_kv_cache, + kwargs["tokens_per_block"], + model_config.sparse_attention_config.window_size, + context=True, + scratch=False, + indexer_k_dtype=indexer_k_dtype, + ) + _, draft_generation, draft_per_request = _estimate_draft_cache_size_components( + draft_layer_sizes, + kwargs["tokens_per_block"], + scratch_tokens=draft_scratch_tokens, + window_size=kwargs.get("draft_window_size"), + generation_capacity_headroom=headroom + draft_scratch_tokens, + ) + draft_per_request += draft_scratch_tokens * ( + non_sliding_attn_size_per_token + target_context_swa_bytes + ) + swa_size_per_token += draft_generation + swa_size_per_request += draft_per_request return ( non_sliding_attn_size_per_token + swa_size_per_token, swa_size_per_request * max_batch_size, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py index dc653921e1eb..4dc3a34a7c8e 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py @@ -612,9 +612,15 @@ def get_unique_primary_pool(self) -> torch.Tensor: first_head_dim = self.head_dim_per_layer[0] first_num_heads = self.num_kv_heads_per_layer[0] - if any(head_dim != first_head_dim for head_dim in self.head_dim_per_layer): + if any( + head_dim != first_head_dim + for head_dim in self.head_dim_per_layer[: self.num_local_layers] + ): raise ValueError("DSA requires a uniform KV head dimension") - if any(num_heads != first_num_heads for num_heads in self.num_kv_heads_per_layer): + if any( + num_heads != first_num_heads + for num_heads in self.num_kv_heads_per_layer[: self.num_local_layers] + ): raise ValueError("DSA requires a uniform KV head count") first_converter = self.impl.get_page_index_converter(0, Role.KEY) @@ -657,13 +663,22 @@ def get_unique_primary_pool(self) -> torch.Tensor: ) return self._unique_primary_pool + def _append_standalone_draft_layers(self, config, **kwargs): + # DSA masks and fused primary views describe target layers only. + return super()._append_standalone_draft_layers( + config, register_model_layers=False, **kwargs + ) + def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) -> int: if data_role == Role.INDEX_KEY: return _get_indexer_k_cache_bytes_per_token( self.index_head_dim, self.quant_block_size, self.use_fp4 ) cache_bytes = super().get_layer_bytes_per_token(local_layer_idx, data_role) - if self.dtype == DataType.NVFP4 and self.mla_kv_cache_residual_dim > 0: + if ( + self.dtype_per_layer[local_layer_idx] == DataType.NVFP4 + and self.mla_kv_cache_residual_dim > 0 + ): if data_role == Role.KEY: cache_bytes += self.mla_kv_cache_residual_dim // 2 elif data_role == Role.KEY_BLOCK_SCALE: @@ -673,12 +688,15 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) - cache_bytes += self.mla_kv_cache_residual_dim // 16 if ( data_role == Role.ALL - and self._indexer_k_cache_layer_mask[self.pp_layers[local_layer_idx]] + and local_layer_idx < len(self.indexer_k_cache_local_layer_mask) + and self.indexer_k_cache_local_layer_mask[local_layer_idx] ): cache_bytes += self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY) return cache_bytes def get_cache_bytes_per_token(self) -> int: + if self._draft_layer_sizes: + return super().get_cache_bytes_per_token() return sum( self.get_layer_bytes_per_token(local_layer_idx, Role.ALL) for local_layer_idx in range(self.num_local_layers) @@ -709,6 +727,30 @@ def get_cache_size_per_token( model_config, mapping, num_layers ) cache_bytes += num_attention_layers * (residual_dim // 2 + math.ceil(residual_dim / 16)) + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + _estimate_draft_cache_size_components, + _get_draft_layer_sizes, + _get_generation_kv_capacity, + ) + + draft_sizes = _get_draft_layer_sizes(**kwargs) + if draft_sizes: + scratch = kwargs.get("draft_scratch_tokens", 0) + _, headroom = _get_generation_kv_capacity(kwargs.get("spec_config"), is_draft=False) + context, generation, fixed = _estimate_draft_cache_size_components( + draft_sizes, + kwargs["tokens_per_block"], + scratch_tokens=scratch, + window_size=kwargs.get("draft_window_size"), + generation_capacity_headroom=headroom + scratch, + helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, + ) + fixed += scratch * cache_bytes + return ( + cache_bytes + generation, + fixed * int(kwargs.get("max_batch_size") or 0) + + (context - generation) * int(kwargs.get("max_num_tokens") or 0), + ) return cache_bytes def shutdown(self) -> None: diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py index 458d2877f12d..15902ef1cb4e 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py @@ -3,11 +3,12 @@ """Blackwell DSpark attention specialization and persistent scheduler.""" -from typing import Tuple, Type +from typing import Optional, Tuple, Type import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute +from cuda.bindings import driver as cuda from .attention_kernel import DSparkAttentionKernel @@ -195,7 +196,7 @@ def __new_from_mlir_values__(self, values): class DSparkAttention(DSparkAttentionKernel): - """Attention over a 128-token rolling window and one 5/6-token draft block.""" + """Attention over a 128-token paged context window and one 5/6-token draft block.""" window_size = 128 block_size = 6 @@ -244,7 +245,13 @@ def __init__( seq_len_q: int = block_size, mma_qk_tiler_k: int = qk_tiler_k, inverse_rope_dim: int = 0, + output_groups: int | None = None, + history_page_size: int = 32, + num_heads: int = num_heads, ): + if num_heads not in (64, 128): + raise ValueError(f"DSpark attention must have 64 or 128 heads, got {num_heads}") + self.num_heads = num_heads expected_config = { "acc_dtype": (acc_dtype, cutlass.Float32), "mma_qk_tiler_mn": (mma_qk_tiler_mn, self.qk_tiler_mn), @@ -262,6 +269,8 @@ def __init__( raise ValueError("Unsupported DSpark kernel configuration: " + ", ".join(mismatches)) if seq_len_q not in (5, 6): raise ValueError(f"DSpark block size must be 5 or 6, got {seq_len_q}") + if history_page_size not in (16, 32, 64, 128, 256): + raise ValueError("DSpark history page size must be 16, 32, 64, 128, or 256") if inverse_rope_dim not in (0, 64): raise ValueError(f"DSpark inverse_rope_dim must be 0 or 64, got {inverse_rope_dim}") @@ -283,22 +292,30 @@ def __init__( self.tma_page_size_draft = self.qk_tiler_mn[1] self.fixed_cache_seq_len = self.window_size + seq_len_q self.inverse_rope_dim = inverse_rope_dim + self.output_groups = output_groups + self.history_page_size = history_page_size + # Subdivide the 64-row loads at physical-page and cyclic-window boundaries. + self.history_tile_sizes = tuple(sorted({1, 4, 16, min(history_page_size, 64)})) @cute.jit def wrapper( self, batch_size: cutlass.Int32, - num_window_pages: cutlass.Int32, - window_page_stride: cutlass.Int64, + num_pages: cutlass.Int32, + page_stride: cutlass.Int64, + table_width: cutlass.Int32, + table_stride: cutlass.Int64, q_ptr: cute.Pointer, - window_ptr: cute.Pointer, + pages_ptr: cute.Pointer, draft_ptr: cute.Pointer, page_table_ptr: cute.Pointer, cache_seqs_ptr: cute.Pointer, valid_len_ptr: cute.Pointer, + capacities_ptr: cute.Pointer, attn_sink_ptr: cute.Pointer, inverse_rope_freqs_ptr: cute.Pointer, o_tensor: cute.Tensor, + o_scale_tensor: Optional[cute.Tensor], softmax_scale: cutlass.Float32, stream: cuda.CUstream, ): @@ -318,11 +335,11 @@ def wrapper( ), ) q = cute.make_tensor(q_ptr, layout=q_layout) - window = cute.make_tensor( - window_ptr, + pages = cute.make_tensor( + pages_ptr, layout=cute.make_layout( - (self.window_size, self.head_dim, num_window_pages), - stride=(self.head_dim, 1, window_page_stride), + (self.history_page_size, self.head_dim, num_pages), + stride=(self.head_dim, 1, page_stride), ), ) draft = cute.make_tensor( @@ -338,9 +355,12 @@ def wrapper( ) page_table = cute.make_tensor( page_table_ptr, - layout=cute.make_layout((1, batch_size), stride=(1, 1)), + layout=cute.make_layout((table_width, batch_size), stride=(1, table_stride)), ) - output = cute.make_tensor(o_tensor.iterator, layout=q_layout) + if cutlass.const_expr(self.output_groups is None): + output = cute.make_tensor(o_tensor.iterator, layout=q_layout) + else: + output = o_tensor cache_seqs = cute.make_tensor( cache_seqs_ptr, layout=cute.make_layout((batch_size,), stride=(1,)), @@ -349,6 +369,9 @@ def wrapper( valid_len_ptr, layout=cute.make_layout((batch_size,), stride=(1,)), ) + capacities = cute.make_tensor( + capacities_ptr, layout=cute.make_layout((batch_size,), stride=(1,)) + ) attn_sink = cute.make_tensor( attn_sink_ptr, layout=cute.make_layout((self.num_heads,), stride=(1,)), @@ -363,12 +386,14 @@ def wrapper( self( q, - window, + pages, draft, page_table, output, + o_scale_tensor, cache_seqs, valid_len, + capacities, softmax_scale, cutlass.Float32(1.0), attn_sink, diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py index d75f16040bf9..4e0edae9bbf3 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py @@ -50,14 +50,14 @@ cutlass_memory = utils -"""DSpark rolling-window MQA attention for NVIDIA Blackwell-family GPUs. +"""DSpark paged MQA attention for NVIDIA Blackwell-family GPUs. A TMA + tcgen05 tensor-core warp-specialized persistent kernel that fuses the Q*K^T matmul, attention-sink online softmax, P*V matmul, and the inverse-RoPE epilogue into a single launch. KV arrives as two paged streams (the 128-row -rolling window and the draft block's own rows); see +context window and the draft block's own rows); see ``attention.py`` for the supported DSpark contract and -``tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py`` for +``tests/unittest/_torch/speculative/test_dspark_paged_attention.py`` for validation coverage. The split-KV/LSE/variable-sequence machinery of the generic attention kernel and its @@ -116,7 +116,6 @@ def __init__( # normally match them, but a specialization may use a wider logical # span so one OOB-zero-filling TMA replaces many tiny page copies. self.tma_page_size_draft = page_size_draft - self.tma_page_size_win = page_size_win self.fixed_cache_seq_len = None self.attn_sink_is_scaled = True # Fixed two-stream specializations can receive the initialized-row @@ -225,6 +224,7 @@ def __call__( o: cute.Tensor, cache_seqs: Optional[cute.Tensor], window_valid_lens: Optional[cute.Tensor], + capacities: cute.Tensor, softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, attn_sink_unscaled: cute.Tensor, @@ -236,7 +236,7 @@ def __call__( KV is split into two streams concatenated along the seq_len_k dimension: - Sliding-window: first `mma_qk_tiler[1]` rows (default 128). - Draft-block: remaining rows. - The rolling-window stream uses a runtime page-table slot. The draft + The context stream uses the request's physical page table. The draft block is stored contiguously, one physical block per request. """ @@ -325,7 +325,6 @@ def __call__( ) cta_kv_m = qk_tiled_mma.op.shape_mnk[0] // qk_tiled_mma.thr_id.shape kc_page_tile_size_draft = min(self.tma_page_size_draft, cta_kv_m) - kc_page_tile_size_win = min(self.tma_page_size_win, cta_kv_m) kc_smem_layout_for_tma_base = sm100_utils.make_smem_layout( OperandMajorMode.K, @@ -336,8 +335,9 @@ def __call__( kc_smem_layout_for_tma_draft = cute.tiled_divide( kc_smem_layout_for_tma_base, (kc_page_tile_size_draft, self.mma_qk_tiler[2]) ) - kc_smem_layout_for_tma_win = cute.tiled_divide( - kc_smem_layout_for_tma_base, (kc_page_tile_size_win, self.mma_qk_tiler[2]) + kc_smem_layout_for_tma_win = tuple( + cute.tiled_divide(kc_smem_layout_for_tma_base, (rows, self.mma_qk_tiler[2])) + for rows in self.history_tile_sizes ) p_smem_layout_staged = sm100_utils.make_smem_layout_a( @@ -357,7 +357,6 @@ def __call__( self.load_kv_stage, ) vc_page_tile_size_draft = min(self.tma_page_size_draft, self.mma_pv_tiler[2]) - vc_page_tile_size_win = min(self.tma_page_size_win, self.mma_pv_tiler[2]) vc_smem_layout_for_tma_base = sm100_utils.make_smem_layout( OperandMajorMode.MN, (self.mma_pv_tiler[1] // pv_tiled_mma.thr_id.shape, self.mma_pv_tiler[2]), @@ -369,9 +368,9 @@ def __call__( vc_smem_layout_for_tma_base, (vc_pv_n, vc_page_tile_size_draft), ) - vc_smem_layout_for_tma_win = cute.tiled_divide( - vc_smem_layout_for_tma_base, - (vc_pv_n, vc_page_tile_size_win), + vc_smem_layout_for_tma_win = tuple( + cute.tiled_divide(vc_smem_layout_for_tma_base, (vc_pv_n, rows)) + for rows in self.history_tile_sizes ) # TMA load for Q latent tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) @@ -396,16 +395,19 @@ def __call__( is_k_load=True, page_size=self.tma_page_size_draft, ) - kc_smem_layout_win = cute.select(kc_smem_layout_for_tma_win, mode=[0]) - tma_atom_c_latent_win, tma_tensor_c_latent_win = self.make_paged_tiled_tma_atom( - tma_load_op, - c_latent_win, - kc_smem_layout_win, - (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), - qk_tiled_mma, - is_k_load=True, - page_size=self.tma_page_size_win, - ) + tma_atom_c_latent_win, tma_tensor_c_latent_win = (), () + for tile in cutlass.range_constexpr(len(self.history_tile_sizes)): + atom, tensor = self.make_paged_tiled_tma_atom( + tma_load_op, + c_latent_win, + cute.select(kc_smem_layout_for_tma_win[tile], mode=[0]), + (self.mma_qk_tiler[1], self.mma_qk_tiler[2]), + qk_tiled_mma, + is_k_load=True, + page_size=self.history_tile_sizes[tile], + ) + tma_atom_c_latent_win += (atom,) + tma_tensor_c_latent_win += (tensor,) # TMA load for window and draft-block V vc_smem_layout_draft = cute.select(vc_smem_layout_for_tma_draft, mode=[0]) tma_atom_c_latent_transpose_draft, tma_tensor_c_latent_transpose_draft = ( @@ -419,18 +421,19 @@ def __call__( page_size=self.tma_page_size_draft, ) ) - vc_smem_layout_win = cute.select(vc_smem_layout_for_tma_win, mode=[0]) - tma_atom_c_latent_transpose_win, tma_tensor_c_latent_transpose_win = ( - self.make_paged_tiled_tma_atom( + tma_atom_c_latent_transpose_win, tma_tensor_c_latent_transpose_win = (), () + for tile in cutlass.range_constexpr(len(self.history_tile_sizes)): + atom, tensor = self.make_paged_tiled_tma_atom( tma_load_op, c_latent_win_transpose, - vc_smem_layout_win, + cute.select(vc_smem_layout_for_tma_win[tile], mode=[0]), (self.mma_pv_tiler[1], self.mma_pv_tiler[2]), pv_tiled_mma, is_k_load=False, - page_size=self.tma_page_size_win, + page_size=self.history_tile_sizes[tile], ) - ) + tma_atom_c_latent_transpose_win += (atom,) + tma_tensor_c_latent_transpose_win += (tensor,) q_latent_copy_size = ( cute.size_in_bytes(self.q_dtype, q_latent_smem_layout) @@ -487,7 +490,7 @@ class AttentionKernelSharedStorage: cute.struct.MemRange[self.q_dtype, cute.cosize(p_smem_layout_staged)], 1024, ] - smem_page_table: cute.struct.MemRange[cutlass.Int32, self.load_pt_stage] + smem_page_table: cute.struct.MemRange[cutlass.Int32, 128 * self.load_pt_stage] required_smem_bytes = AttentionKernelSharedStorage.__sizeof__() available_smem_bytes = cutlass_memory.get_smem_capacity_in_bytes(self.arch_str) @@ -515,6 +518,7 @@ class AttentionKernelSharedStorage: o, cache_seqs, window_valid_lens, + capacities, softmax_scale_log2, output_scale, attn_sink_unscaled, @@ -578,18 +582,19 @@ def kernel( tiled_mma_pv: cute.TiledMma, tma_atom_q_latent: Optional[cute.CopyAtom], mQL: cute.Tensor, - tma_atom_c_latent_win: Optional[cute.CopyAtom], - mCL_win: cute.Tensor, + tma_atom_c_latent_win: tuple, + mCL_win: tuple, tma_atom_c_latent_draft: Optional[cute.CopyAtom], mCL_draft: cute.Tensor, - tma_atom_c_latent_transpose_win: Optional[cute.CopyAtom], - mCLT_win: cute.Tensor, + tma_atom_c_latent_transpose_win: tuple, + mCLT_win: tuple, tma_atom_c_latent_transpose_draft: Optional[cute.CopyAtom], mCLT_draft: cute.Tensor, mPT_win: cute.Tensor, mO: Optional[cute.Tensor], cache_seqs: cute.Tensor, window_valid_lens: cute.Tensor, + capacities: cute.Tensor, softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, attn_sink_unscaled: cute.Tensor, @@ -598,9 +603,9 @@ def kernel( kc_smem_layout_staged: cute.ComposedLayout, p_smem_layout_staged: cute.ComposedLayout, vc_smem_layout_staged: cute.ComposedLayout, - kc_smem_layout_for_tma_win: cute.ComposedLayout, + kc_smem_layout_for_tma_win: tuple, kc_smem_layout_for_tma_draft: cute.ComposedLayout, - vc_smem_layout_for_tma_win: cute.ComposedLayout, + vc_smem_layout_for_tma_win: tuple, vc_smem_layout_for_tma_draft: cute.ComposedLayout, cta_layout_vmnk: cute.Layout, tile_sched_params, @@ -624,9 +629,11 @@ def kernel( # Prefetch tma descriptor if warp_idx == self.mma_warp_id: cpasync.prefetch_descriptor(tma_atom_q_latent) - cpasync.prefetch_descriptor(tma_atom_c_latent_win) + for tile in cutlass.range_constexpr(len(self.history_tile_sizes)): + cpasync.prefetch_descriptor(tma_atom_c_latent_win[tile]) cpasync.prefetch_descriptor(tma_atom_c_latent_draft) - cpasync.prefetch_descriptor(tma_atom_c_latent_transpose_win) + for tile in cutlass.range_constexpr(len(self.history_tile_sizes)): + cpasync.prefetch_descriptor(tma_atom_c_latent_transpose_win[tile]) cpasync.prefetch_descriptor(tma_atom_c_latent_transpose_draft) # Alloc @@ -679,10 +686,10 @@ def kernel( sKC = storage.smem_kc.get_tensor( kc_smem_layout_staged.outer, swizzle=kc_smem_layout_staged.inner ) - # Two TMA-views over the same physical SMEM, one per page-size - sKC_for_tma_win = storage.smem_kc.get_tensor( - kc_smem_layout_for_tma_win.outer, - swizzle=kc_smem_layout_for_tma_win.inner, + # TMA tile views share the unchanged physical K/V SMEM layout. + sKC_for_tma_win = tuple( + storage.smem_kc.get_tensor(layout.outer, swizzle=layout.inner) + for layout in kc_smem_layout_for_tma_win ) sKC_for_tma_draft = storage.smem_kc.get_tensor( kc_smem_layout_for_tma_draft.outer, @@ -692,13 +699,15 @@ def kernel( # reuse smem sVC_ptr = cute.recast_ptr(sKC.iterator, vc_smem_layout_staged.inner) sVC = cute.make_tensor(sVC_ptr, vc_smem_layout_staged.outer) - sVC_for_tma_win = cute.make_tensor(sVC_ptr, vc_smem_layout_for_tma_win.outer) + sVC_for_tma_win = tuple( + cute.make_tensor(sVC_ptr, layout.outer) for layout in vc_smem_layout_for_tma_win + ) sVC_for_tma_draft = cute.make_tensor(sVC_ptr, vc_smem_layout_for_tma_draft.outer) # (MMA, MMA_H, MMA_K) sP = storage.smem_p.get_tensor( p_smem_layout_staged.outer, swizzle=p_smem_layout_staged.inner ) - sPT = storage.smem_page_table.get_tensor(cute.make_layout((1, self.load_pt_stage))) + sPT = storage.smem_page_table.get_tensor(cute.make_layout((128, self.load_pt_stage))) # (compute_threads,) softmax_smem_exchange = storage.softmax_smem_exchange.get_tensor( cute.make_layout(self.num_compute_warps * self.threads_per_warp) @@ -733,6 +742,9 @@ def kernel( blk_coord=blk_coord, load_pt_pipeline=load_pt_pipeline, mPT_win=mPT_win, + positions=cache_seqs, + valid_lengths=window_valid_lens, + capacities=capacities, sPT=sPT, tidx=tidx, ) @@ -765,6 +777,8 @@ def kernel( # Construct fixed common/tma_qk/tma_pv params for load_tma tma_common_params = SimpleNamespace( blk_coord=blk_coord, + position=cutlass.Int32(cache_seqs[blk_coord[2]]), + window_end=cutlass.Int32(cache_seqs[blk_coord[2]]) % 128, load_q_pipeline=load_q_pipeline, load_kv_pipeline=load_kv_pipeline, sPT=sPT, @@ -925,7 +939,9 @@ def kernel( smem_exchange=softmax_smem_exchange, mO=mO, K=self.get_cache_seq_len(cache_seqs, blk_coord[2]), - window_valid_len=self.get_window_valid_len(window_valid_lens, blk_coord[2]), + window_valid_len=self.get_window_valid_len( + window_valid_lens, capacities, cache_seqs, blk_coord[2] + ), window_end_pos=self.get_window_end_pos(cache_seqs, blk_coord[2]), L=mCL_draft.shape[1], tmem_ptr=tmem_ptr, @@ -1013,15 +1029,18 @@ def get_cache_seq_len( @cute.jit def get_window_valid_len( - self, window_valid_lens: Optional[cute.Tensor], batch_idx: cutlass.Int32 + self, + window_valid_lens: cute.Tensor, + capacities: cute.Tensor, + positions: cute.Tensor, + batch_idx: cutlass.Int32, ) -> cutlass.Int32: - """Return the number of initialized rows in the window stream.""" - if cutlass.const_expr(self.window_valid_len_from_tensor): - return min( - max(cutlass.Int32(window_valid_lens[batch_idx]), cutlass.Int32(0)), - cutlass.Int32(self.tma_page_size_win), - ) - return cutlass.Int32(self.tma_page_size_win) + """Return the initialized context length, excluding synthetic requests.""" + length = min( + max(cutlass.Int32(window_valid_lens[batch_idx]), 0), + min(128, cutlass.Int32(positions[batch_idx])), + ) + return length if capacities[batch_idx] > 0 else cutlass.Int32(0) @cute.jit def get_window_end_pos( @@ -1048,8 +1067,8 @@ def is_score_valid( ) if cutlass.const_expr(self.window_valid_len_from_tensor): age = ( - window_end_pos - local_col + cutlass.Int32(self.tma_page_size_win) - ) % cutlass.Int32(self.tma_page_size_win) + window_end_pos - 1 - local_col + cutlass.Int32(self.page_size_win) + ) % cutlass.Int32(self.page_size_win) valid = valid and (k_index != 0 or cute.elem_less(age, window_valid_len)) return valid @@ -1059,34 +1078,29 @@ def load_page_table( common_params: SimpleNamespace, load_pt_producer_state: pipeline.PipelineState, ) -> pipeline.PipelineState: - """Stage the rolling-window slot followed by the implicit draft-block slot.""" - mPT_win = common_params.mPT_win[None, common_params.blk_coord[2]] - tidx = common_params.tidx % self.threads_per_warp + """Stage physical pages for each context row, then the draft block.""" + batch = common_params.blk_coord[2] + lane = common_params.tidx % self.threads_per_warp + end = cutlass.Int32(common_params.positions[batch]) + length = cutlass.Int32(common_params.valid_lengths[batch]) + capacity = cutlass.Int32(common_params.capacities[batch]) load_pt_pipeline = common_params.load_pt_pipeline - atom_async_copy = cute.make_copy_atom( - cpasync.CopyG2SOp(cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS), - cutlass.Int32, - num_bits_per_copy=cutlass.Int32.width, - ) - mPT_win_for_copy = cute.flat_divide(mPT_win, (1,)) - sPT_for_copy = cute.flat_divide(common_params.sPT, (1,)) - load_pt_pipeline.producer_acquire(load_pt_producer_state) - if tidx == 0: - cute.copy( - atom_async_copy, - mPT_win_for_copy[None, 0], - sPT_for_copy[None, 0, load_pt_producer_state.index], - ) + for i in cutlass.range_constexpr(4): + row = lane + i * 32 + position = end - 1 - (end - 1 - row + 128) % 128 + page = cutlass.Int32(-1) + if position >= max(0, end - length) and position < capacity: + page = common_params.mPT_win[position // self.history_page_size, batch] + common_params.sPT[row, load_pt_producer_state.index] = page load_pt_pipeline.producer_commit(load_pt_producer_state) load_pt_producer_state.advance() load_pt_pipeline.producer_acquire(load_pt_producer_state) - if tidx == 0: - sPT_for_copy[None, 0, load_pt_producer_state.index].fill(common_params.blk_coord[2]) + if lane == 0: + common_params.sPT[0, load_pt_producer_state.index] = batch load_pt_pipeline.producer_commit(load_pt_producer_state) load_pt_producer_state.advance() - return load_pt_producer_state @cute.jit @@ -1117,21 +1131,6 @@ def load_tma( cta_kv_m = qk_params.tiled_mma_qk.op.shape_mnk[0] // qk_params.tiled_mma_qk.thr_id.shape - # === K partition for window stream === - cta_m_win = min(cta_kv_m, self.tma_page_size_win) - page_tile_size_k_win = min(self.tma_page_size_win, cta_m_win) - gCL_win = cute.tiled_divide(qk_params.mCL_win, (page_tile_size_k_win, self.mma_qk_tiler[2])) - tSgCL_win = ( - gCL_win[ - None, - common_params.blk_coord[0] % qk_params.tiled_mma_qk.thr_id.shape, - None, - None, - ] - if cta_m_win < self.tma_page_size_win - else gCL_win[None, 0, None, None] - ) - # === K partition for draft-block stream === cta_m_draft = min(cta_kv_m, self.tma_page_size_draft) page_tile_size_k_draft = min(self.tma_page_size_draft, cta_m_draft) @@ -1158,13 +1157,20 @@ def load_tma( cute.group_modes(tSgQL, 0, 3), ) - tKCsKC_win, tCLgCL_win = cpasync.tma_partition( - qk_params.tma_atom_c_latent_win, - 0, - cute.make_layout(1), - qk_params.sKC_win, - tSgCL_win, - ) + tKCsKC_win, tCLgCL_win = (), () + for tile in cutlass.range_constexpr(len(self.history_tile_sizes)): + gCL_win = cute.tiled_divide( + qk_params.mCL_win[tile], (self.history_tile_sizes[tile], self.mma_qk_tiler[2]) + ) + smem, gmem = cpasync.tma_partition( + qk_params.tma_atom_c_latent_win[tile], + 0, + cute.make_layout(1), + qk_params.sKC_win[tile], + gCL_win, + ) + tKCsKC_win += (smem,) + tCLgCL_win += (gmem,) tKCsKC_draft, tCLgCL_draft = cpasync.tma_partition( qk_params.tma_atom_c_latent_draft, 0, @@ -1178,14 +1184,25 @@ def load_tma( ] # === V partition for window stream === - page_tile_size_v_win = min(self.tma_page_size_win, self.mma_pv_tiler[2]) - gCLT_win = cute.flat_divide(v_params.mCLT_win, (self.mma_pv_tiler[1], page_tile_size_v_win)) cta_n = self.mma_pv_tiler[1] // v_params.tiled_mma_pv.thr_id.shape - gCLT_win = cute.logical_divide(gCLT_win, (cta_n,))[ - (None, common_params.blk_coord[0]), None, None, None, None - ] - tOgCLT_win = cute.tiled_divide(gCLT_win, (cta_n, page_tile_size_v_win)) - tOgCLT_win = tOgCLT_win[None, 0, 0, None, None, None] + tVCsVC_win, tCLTgCLT_win = (), () + for tile in cutlass.range_constexpr(len(self.history_tile_sizes)): + rows = self.history_tile_sizes[tile] + gCLT_win = cute.flat_divide(v_params.mCLT_win[tile], (self.mma_pv_tiler[1], rows)) + gCLT_win = cute.logical_divide(gCLT_win, (cta_n,))[ + (None, common_params.blk_coord[0]), None, None, None, None + ] + tOgCLT_win = cute.tiled_divide(gCLT_win, (cta_n, rows)) + tOgCLT_win = tOgCLT_win[None, 0, 0, None, None, None] + smem, gmem = cpasync.tma_partition( + v_params.tma_atom_c_latent_transpose_win[tile], + 0, + cute.make_layout(1), + v_params.sVC_win[tile], + tOgCLT_win, + ) + tVCsVC_win += (smem,) + tCLTgCLT_win += (gmem,) # === V partition for draft-block stream === page_tile_size_v_draft = min(self.tma_page_size_draft, self.mma_pv_tiler[2]) @@ -1198,13 +1215,6 @@ def load_tma( tOgCLT_draft = cute.tiled_divide(gCLT_draft, (cta_n, page_tile_size_v_draft)) tOgCLT_draft = tOgCLT_draft[None, 0, 0, None, None, None] - tVCsVC_win, tCLTgCLT_win = cpasync.tma_partition( - v_params.tma_atom_c_latent_transpose_win, - 0, - cute.make_layout(1), - v_params.sVC_win, - tOgCLT_win, - ) tVCsVC_draft, tCLTgCLT_draft = cpasync.tma_partition( v_params.tma_atom_c_latent_transpose_draft, 0, @@ -1269,6 +1279,75 @@ def load_tma( load_pt_release_state, ) + @cute.jit + def load_history( + self, + common_params: SimpleNamespace, + page_table_stage: cutlass.Int32, + atoms: tuple, + gmem: tuple, + smem: tuple, + row_base: cutlass.Int32, + head_tile: cutlass.Int32, + kv_stage: cutlass.Int32, + tma_bar_ptr: cute.Pointer, + is_k_load: cutlass.Constexpr, + tile: cutlass.Constexpr = None, + start_row: cutlass.Int32 = 0, + ) -> None: + """Fill a history tile, subdividing only at page or validity boundaries.""" + if cutlass.const_expr(tile is None): + tile = len(self.history_tile_sizes) - 1 + extent = 64 + else: + extent = self.history_tile_sizes[tile + 1] + rows = self.history_tile_sizes[tile] + for row in cutlass.range(start_row, start_row + extent, rows): + context_row = row_base + row + page = common_params.sPT[context_row, page_table_stage] + legal = cutlass.Boolean(True) + if cutlass.const_expr(rows > 1): + legal = ( + (page >= 0) + & ( + (common_params.window_end <= context_row) + | (common_params.window_end >= context_row + rows) + ) + & (page == common_params.sPT[context_row + rows - 1, page_table_stage]) + ) + if legal: + # Recover the absolute position represented by this cyclic row. + # A 256-token page contains two distinct 128-token window cycles. + position = ( + common_params.position + - 1 + - (common_params.position - 1 - context_row + 128) % 128 + ) + offset = position % self.history_page_size // rows + if cutlass.const_expr(is_k_load): + src = gmem[tile][None, offset, head_tile, page] + dst = smem[tile][None, row // rows, 0, kv_stage] + else: + src = gmem[tile][None, head_tile, offset, page] + dst = smem[tile][None, 0, row // rows, kv_stage] + cute.copy(atoms[tile], src, dst, tma_bar_ptr=tma_bar_ptr) + else: + if cutlass.const_expr(tile > 0): + self.load_history( + common_params, + page_table_stage, + atoms, + gmem, + smem, + row_base, + head_tile, + kv_stage, + tma_bar_ptr, + is_k_load, + tile=tile - 1, + start_row=row, + ) + @cute.jit def load_tma_qk_one_k_tile( self, @@ -1307,11 +1386,17 @@ def load_tma_qk_one_k_tile( tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) load_kv_pipeline.producer_acquire(load_kv_producer_state) if cutlass.const_expr(is_window): - cute.copy( + self.load_history( + common_params, + page_table_stage, qk_params.tma_atom_c_latent_win, - qk_params.tCLgCL_win[None, i, k_idx], - qk_params.tKCsKC_win[None, 0, 0, load_kv_producer_state.index], - tma_bar_ptr=tma_bar_ptr, + qk_params.tCLgCL_win, + qk_params.tKCsKC_win, + common_params.blk_coord[0] % 2 * 64, + i, + load_kv_producer_state.index, + tma_bar_ptr, + is_k_load=True, ) else: cute.copy( @@ -1336,8 +1421,6 @@ def load_tma_v_one_k_tile( """Issue V TMA copies for one compile-time-selected stream.""" page_table_stage = load_pt_release_state.index k_idx = common_params.sPT[0, page_table_stage] - common_params.load_pt_pipeline.consumer_release(load_pt_release_state) - load_pt_release_state.advance() load_kv_pipeline = common_params.load_kv_pipeline # Pre-init tma_bar_ptr so its type is established before the loop @@ -1348,11 +1431,17 @@ def load_tma_v_one_k_tile( tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) load_kv_pipeline.producer_acquire(load_kv_producer_state) if cutlass.const_expr(is_window): - cute.copy( + self.load_history( + common_params, + page_table_stage, v_params.tma_atom_c_latent_transpose_win, - v_params.tCLTgCLT_win[None, j, i, k_idx], - v_params.tVCsVC_win[None, 0, 0, load_kv_producer_state.index], - tma_bar_ptr=tma_bar_ptr, + v_params.tCLTgCLT_win, + v_params.tVCsVC_win, + i * self.mma_pv_tiler[2], + j, + load_kv_producer_state.index, + tma_bar_ptr, + is_k_load=False, ) else: cute.copy( @@ -1363,6 +1452,8 @@ def load_tma_v_one_k_tile( ) load_kv_producer_state.advance() + common_params.load_pt_pipeline.consumer_release(load_pt_release_state) + load_pt_release_state.advance() return load_kv_producer_state, load_pt_release_state @cute.jit @@ -1637,7 +1728,14 @@ def compute( tTR_tS_for_sink = tmem_thr_copy.partition_D(cS) local_row = tTR_tS_for_sink[0][0] head_idx = common_params.blk_coord[0] * cta_qk_tiler[0] + local_row - my_sink = common_params.attn_sink_unscaled[head_idx] + my_sink = self.acc_dtype(0) + # TMA zero-fills query rows outside the head dimension. The matching + # sink loads must also be masked for a partially occupied head tile. + if cutlass.const_expr(cute.size(common_params.attn_sink_unscaled) < self.mma_qk_tiler[0]): + if head_idx < cute.size(common_params.attn_sink_unscaled): + my_sink = common_params.attn_sink_unscaled[head_idx] + else: + my_sink = common_params.attn_sink_unscaled[head_idx] if cutlass.const_expr(self.attn_sink_is_scaled): my_sink = my_sink * LOG2_E / softmax_params.softmax_scale_log2 @@ -2096,8 +2194,12 @@ def _tmem_load_partition( common_params.blk_coord[2], ), ) + # Keep both CTA head tiles in the coordinate layout, including when + # only the first 64 heads are stored. A unit head-tile mode has zero + # stride and would otherwise alias the second CTA to the first. + output_coord_shape = (self.mma_pv_tiler[0], *common_params.mO.shape[1:]) cO = cute.local_tile( - cute.make_identity_tensor(common_params.mO.shape), + cute.make_identity_tensor(output_coord_shape), cta_pv_tiler_mn, ( common_params.blk_coord[0], diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py index bfdaa933119a..b35be0314cc1 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py @@ -171,12 +171,198 @@ def kernel( output[row, out_imag] = (imag * cos + real * sin).to(output.element_type) -class DSparkRMSNormRoPECacheWriteKernel(DSparkRMSNormRoPEKernel): - """Apply RMSNorm/RoPE and scatter each row into the rolling KV cache.""" +class DSparkRMSNormFP8QuantKernel: + """Apply RMSNorm and emit packed 1x128 MXFP8 activations.""" - def __init__(self, hidden_dim: int, rope_dim: int, eps: float, window_size: int): + num_threads = cute.arch.WARP_SIZE + quant_block = 128 + + def __init__(self, hidden_dim: int, eps: float): + if hidden_dim % self.quant_block != 0: + raise ValueError( + f"hidden_dim must be divisible by {self.quant_block}; got {hidden_dim}" + ) + self.hidden_dim = hidden_dim + self.num_quant_blocks = hidden_dim // self.quant_block + self.values_per_thread = self.quant_block // self.num_threads + self.eps = eps + + @cute.jit + def __call__( + self, + x: cute.Tensor, + weight: cute.Tensor, + output: cute.Tensor, + output_scale: cute.Tensor, + stream: cuda.CUstream, + ): + self.kernel(x, weight, output, output_scale).launch( + grid=[x.shape[0], 1, 1], + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x: cute.Tensor, + weight: cute.Tensor, + output: cute.Tensor, + output_scale: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + + if cutlass.dynamic_expr(row < x.shape[0]): + sum_sq = cutlass.Float32(0.0) + for quant_block in cutlass.range_constexpr(self.num_quant_blocks): + block_base = quant_block * self.quant_block + for item in cutlass.range_constexpr(self.values_per_thread): + dim = block_base + tidx * self.values_per_thread + item + value = cutlass.Float32(x[row, dim]) + sum_sq += value * value + sum_sq = cute.arch.warp_reduction_sum(sum_sq) + inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) + + for quant_block in cutlass.range_constexpr(self.num_quant_blocks): + block_base = quant_block * self.quant_block + values = cute.make_rmem_tensor((self.values_per_thread,), cutlass.BFloat16) + local_absmax = cutlass.Float32(0.0) + for item in cutlass.range_constexpr(self.values_per_thread): + dim = block_base + tidx * self.values_per_thread + item + value = ( + cutlass.Float32(x[row, dim]) * inverse_rms * cutlass.Float32(weight[dim]) + ).to(cutlass.BFloat16) + values[item] = value + value_f32 = cutlass.Float32(value) + local_absmax = cute.arch.fmax( + local_absmax, + cute.arch.fmax(value_f32, -value_f32), + ) + + block_absmax = cute.arch.warp_reduction_max(local_absmax) + block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) + scale = (block_absmax * cutlass.Float32(1.0 / 448.0)).to(cutlass.Float8E8M0FNU) + quant_scale = cute.arch.rcp_approx(cutlass.Float32(scale)) + + for item in cutlass.range_constexpr(self.values_per_thread): + dim = block_base + tidx * self.values_per_thread + item + output[row, dim] = (cutlass.Float32(values[item]) * quant_scale).to( + output.element_type + ) + + if tidx == 0: + scale_offset = ( + (row // 128) * self.num_quant_blocks * 512 + + quant_block * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + ) + for replica in cutlass.range_constexpr(4): + output_scale[scale_offset + replica] = scale + + +class DSparkConcatFP8QuantKernel: + """Concatenate two row ranges while producing packed 1x128 MXFP8.""" + + num_threads = 4 * cute.arch.WARP_SIZE + num_warps = 4 + quant_block = 128 + values_per_thread = quant_block // cute.arch.WARP_SIZE + + def __init__(self, hidden_dim: int): + if hidden_dim % self.quant_block != 0: + raise ValueError( + f"hidden_dim must be divisible by {self.quant_block}; got {hidden_dim}" + ) + self.hidden_dim = hidden_dim + self.num_quant_blocks = hidden_dim // self.quant_block + + @cute.jit + def __call__( + self, + first: cute.Tensor, + second: cute.Tensor, + output: cute.Tensor, + output_scale: cute.Tensor, + stream: cuda.CUstream, + ): + total_rows = first.shape[0] + second.shape[0] + self.kernel(first, second, output, output_scale).launch( + grid=[self.num_quant_blocks, cute.ceil_div(total_rows, self.num_warps), 1], + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + first: cute.Tensor, + second: cute.Tensor, + output: cute.Tensor, + output_scale: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + quant_block, row_block, _ = cute.arch.block_idx() + warp = tidx // cute.arch.WARP_SIZE + lane = tidx % cute.arch.WARP_SIZE + row = row_block * self.num_warps + warp + total_rows = first.shape[0] + second.shape[0] + scale = cutlass.Float8E8M0FNU(0.0) + + if cutlass.dynamic_expr(row < total_rows): + values = cute.make_rmem_tensor((self.values_per_thread,), cutlass.BFloat16) + local_absmax = cutlass.Float32(0.0) + block_base = quant_block * self.quant_block + if cutlass.dynamic_expr(row < first.shape[0]): + for item in cutlass.range_constexpr(self.values_per_thread): + dim = block_base + lane * self.values_per_thread + item + value = first[row, dim] + values[item] = value + value_f32 = cutlass.Float32(value) + local_absmax = cute.arch.fmax( + local_absmax, + cute.arch.fmax(value_f32, -value_f32), + ) + else: + second_row = row - first.shape[0] + for item in cutlass.range_constexpr(self.values_per_thread): + dim = block_base + lane * self.values_per_thread + item + value = second[second_row, dim] + values[item] = value + value_f32 = cutlass.Float32(value) + local_absmax = cute.arch.fmax( + local_absmax, + cute.arch.fmax(value_f32, -value_f32), + ) + + block_absmax = cute.arch.warp_reduction_max(local_absmax) + block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) + scale = (block_absmax * cutlass.Float32(1.0 / 448.0)).to(cutlass.Float8E8M0FNU) + quant_scale = cute.arch.rcp_approx(cutlass.Float32(scale)) + for item in cutlass.range_constexpr(self.values_per_thread): + dim = block_base + lane * self.values_per_thread + item + output[row, dim] = (cutlass.Float32(values[item]) * quant_scale).to( + output.element_type + ) + + if lane == 0: + scale_offset = ( + (row // 128) * self.num_quant_blocks * 512 + + quant_block * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + ) + for replica in cutlass.range_constexpr(4): + output_scale[scale_offset + replica] = scale + + +class DSparkRMSNormRoPEPageWriteKernel(DSparkRMSNormRoPEKernel): + """Apply RMSNorm/RoPE and write each valid row directly into its managed page.""" + + def __init__(self, hidden_dim: int, rope_dim: int, eps: float, page_size: int): super().__init__(hidden_dim, rope_dim, 1, eps, True, True, False) - self.window_size = window_size + self.page_size = page_size @cute.jit def __call__( @@ -185,10 +371,9 @@ def __call__( weight: cute.Tensor, freqs: cute.Tensor, kv_cache: cute.Tensor, - slots: cute.Tensor, + block_tables: cute.Tensor, start_pos: cute.Tensor, - slots_i32: cute.Tensor, - cache_seqs: cute.Tensor, + capacities: cute.Tensor, stream: cuda.CUstream, ): self.kernel( @@ -196,10 +381,9 @@ def __call__( weight, freqs, kv_cache, - slots, + block_tables, start_pos, - slots_i32, - cache_seqs, + capacities, ).launch( grid=[x.shape[0], 1, 1], block=[self.num_threads, 1, 1], @@ -213,55 +397,55 @@ def kernel( weight: cute.Tensor, freqs: cute.Tensor, kv_cache: cute.Tensor, - slots: cute.Tensor, + block_tables: cute.Tensor, start_pos: cute.Tensor, - slots_i32: cute.Tensor, - cache_seqs: cute.Tensor, + capacities: cute.Tensor, ): tidx, _, _ = cute.arch.thread_idx() row, _, _ = cute.arch.block_idx() - slot = cutlass.Int32(slots[row]) - cache_seq = cutlass.Int32(start_pos[row]) - cache_pos = cache_seq % self.window_size - - if tidx == 0: - slots_i32[row] = slot - cache_seqs[row] = cache_seq - - sum_sq = cutlass.Float32(0.0) - for item in cutlass.range_constexpr(self.elements_per_thread): - dim = tidx + item * self.num_threads - value = cutlass.Float32(x[row, dim]) - sum_sq += value * value - sum_sq = cute.arch.warp_reduction_sum(sum_sq) - inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) - - for item in cutlass.range_constexpr(self.nope_elements_per_thread): - dim = tidx + item * self.num_threads - value = cutlass.Float32(x[row, dim]) * inverse_rms - value *= cutlass.Float32(weight[dim]) - kv_cache[slot, cache_pos, dim] = value.to(kv_cache.element_type) - - for item in cutlass.range_constexpr(self.pairs_per_thread): - pair = tidx + item * self.num_threads - real_dim = self.nope_dim + pair * 2 - imag_dim = real_dim + 1 - real = cutlass.Float32(x[row, real_dim]) * inverse_rms - imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms - real *= cutlass.Float32(weight[real_dim]) - imag *= cutlass.Float32(weight[imag_dim]) - cos = cutlass.Float32(freqs[row, pair, 0]) - sin = cutlass.Float32(freqs[row, pair, 1]) - kv_cache[slot, cache_pos, real_dim] = (real * cos - imag * sin).to( - kv_cache.element_type - ) - kv_cache[slot, cache_pos, imag_dim] = (imag * cos + real * sin).to( - kv_cache.element_type - ) + position = start_pos[row] - 1 + page = cutlass.Int32(-1) + if position >= 0 and position < capacities[row]: + logical_page = position // self.page_size + if logical_page < block_tables.shape[1]: + page = block_tables[row, logical_page] + if page >= 0 and page < kv_cache.shape[0]: + cache_pos = position % self.page_size + + sum_sq = cutlass.Float32(0.0) + for item in cutlass.range_constexpr(self.elements_per_thread): + dim = tidx + item * self.num_threads + value = cutlass.Float32(x[row, dim]) + sum_sq += value * value + sum_sq = cute.arch.warp_reduction_sum(sum_sq) + inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) + + for item in cutlass.range_constexpr(self.nope_elements_per_thread): + dim = tidx + item * self.num_threads + value = cutlass.Float32(x[row, dim]) * inverse_rms + value *= cutlass.Float32(weight[dim]) + kv_cache[cutlass.Int64(page), cache_pos, dim] = value.to(kv_cache.element_type) + + for item in cutlass.range_constexpr(self.pairs_per_thread): + pair = tidx + item * self.num_threads + real_dim = self.nope_dim + pair * 2 + imag_dim = real_dim + 1 + real = cutlass.Float32(x[row, real_dim]) * inverse_rms + imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms + real *= cutlass.Float32(weight[real_dim]) + imag *= cutlass.Float32(weight[imag_dim]) + cos = cutlass.Float32(freqs[row, pair, 0]) + sin = cutlass.Float32(freqs[row, pair, 1]) + kv_cache[cutlass.Int64(page), cache_pos, real_dim] = (real * cos - imag * sin).to( + kv_cache.element_type + ) + kv_cache[cutlass.Int64(page), cache_pos, imag_dim] = (imag * cos + real * sin).to( + kv_cache.element_type + ) class DSparkRMSNormRoPEDraftBlockKernel(DSparkRMSNormRoPEKernel): - """Apply RMSNorm/RoPE and materialize a fixed-size zero-padded draft block.""" + """Apply RMSNorm/RoPE and materialize a zero-padded physical KV page.""" def __init__( self, diff --git a/tensorrt_llm/_torch/attention_backend/dspark.py b/tensorrt_llm/_torch/attention_backend/dspark.py new file mode 100644 index 000000000000..2c0c9a1a26fd --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/dspark.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Direct paged captured-context writes for embedded DSpark.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _write_context_kernel( + values, + positions, + mask, + pool, + tables, + capacities, + M: tl.constexpr, + D: tl.constexpr, + TABLE_STRIDE: tl.constexpr, + PAGE_STRIDE: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + token = tl.program_id(1) + d = tl.arange(0, BLOCK) + # RoPE uses frame p + 1; the managed cache stores logical token p. + position = tl.load(positions + row * M + token) - 1 + capacity = tl.load(capacities + row) + valid = tl.load(mask + row * M + token) & (position >= 0) & (position < capacity) + page = tl.load(tables + row * TABLE_STRIDE + position // PAGE_SIZE, valid, other=-1) + valid = valid & (page >= 0) + value = tl.load(values + (row * M + token) * D + d, (d < D) & valid, other=0) + tl.store( + pool + page.to(tl.int64) * PAGE_STRIDE + (position % PAGE_SIZE) * TOKEN_STRIDE + d, + value, + (d < D) & valid, + ) + + +def write_dspark_context( + values: torch.Tensor, + positions: torch.Tensor, + mask: torch.Tensor, + pool: torch.Tensor, + block_tables: torch.Tensor, + capacities: torch.Tensor, +) -> None: + """Write [batch, tokens, head_dim] KV, masking invalid and dummy rows before stores.""" + batch, tokens, head_dim = values.shape + if batch == 0 or tokens == 0: + return + _write_context_kernel[(batch, tokens)]( + values.contiguous(), + positions.contiguous(), + mask.contiguous(), + pool, + block_tables, + capacities, + tokens, + head_dim, + block_tables.stride(0), + pool.stride(0), + pool.stride(1), + pool.shape[1], + triton.next_power_of_2(head_dim), + ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py new file mode 100644 index 000000000000..b9d2d9adc58f --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -0,0 +1,473 @@ +# 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 management for MiniMax-M3 sparse attention. + +Provides: + * :class:`MiniMaxM3SparseIndexCache` — plain-tensor side cache used by + algorithm-only unit tests (no pyexecutor dependency required for + construction). + * :class:`MiniMaxM3KVCacheManagerV2` — :class:`KVCacheManagerV2` + subclass that registers a per-sparse-layer ``Role.INDEX_KEY`` paged + buffer alongside the standard K/V buffers. +""" + +from __future__ import annotations + +from typing import List, Optional, Sequence + +import torch + +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind +from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode +from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX +from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole + +from ....pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role + + +class MiniMaxM3SparseIndexCache: + """Plain-tensor side cache for the M3 sparse index branch. + + Slot layout matches the main ``KVCacheManagerV2`` paged buffer + geometry so the same ``req_to_token`` mapping addresses both + caches. One ``[num_slots, 1, sparse_index_dim]`` index-K buffer is + allocated per sparse layer; index-V is allocated only for layers + not listed in ``disable_index_value_layer_ids``. + """ + + def __init__( + self, + *, + num_layers: int, + sparse_layer_ids: List[int], + disable_index_value_layer_ids: List[int], + num_slots: int, + sparse_index_dim: int, + dtype: torch.dtype, + device: torch.device, + ): + if num_slots <= 0: + raise ValueError(f"num_slots must be > 0, got {num_slots}") + if sparse_index_dim <= 0: + raise ValueError(f"sparse_index_dim must be > 0, got {sparse_index_dim}") + self.num_layers = int(num_layers) + self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids) + self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids) + self.num_slots = int(num_slots) + self.sparse_index_dim = int(sparse_index_dim) + self.dtype = dtype + self.device = device + + self._index_k: dict[int, torch.Tensor] = {} + self._index_v: dict[int, torch.Tensor] = {} + for layer_idx in self.sparse_layer_ids: + if not (0 <= layer_idx < self.num_layers): + raise ValueError(f"sparse layer_idx {layer_idx} outside [0, {self.num_layers})") + self._index_k[layer_idx] = torch.zeros( + (self.num_slots, 1, self.sparse_index_dim), + dtype=dtype, + device=device, + ) + if layer_idx not in self.disable_index_value_layer_ids: + self._index_v[layer_idx] = torch.zeros( + (self.num_slots, 1, self.sparse_index_dim), + dtype=dtype, + device=device, + ) + + def has_index_value(self, layer_idx: int) -> bool: + return layer_idx in self._index_v + + def get_index_k_buffer(self, layer_idx: int) -> torch.Tensor: + if layer_idx not in self._index_k: + raise KeyError( + f"layer_idx {layer_idx} is not a sparse layer; " + f"sparse layers: {self.sparse_layer_ids}" + ) + return self._index_k[layer_idx] + + def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._index_v.get(layer_idx) + + def set_index_k(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_k: torch.Tensor) -> None: + """Write ``idx_k`` into the index-K cache at ``out_cache_loc``.""" + buf = self.get_index_k_buffer(layer_idx) + if idx_k.shape[1] != 1: + raise ValueError( + f"index K is replicated single-head; expected shape " + f"[N, 1, {self.sparse_index_dim}], got {tuple(idx_k.shape)}" + ) + buf.index_copy_(0, out_cache_loc.to(torch.long), idx_k.to(buf.dtype)) + + def set_index_v(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_v: torch.Tensor) -> None: + """Write ``idx_v`` into the index-V cache (only when allocated).""" + buf = self.get_index_v_buffer(layer_idx) + if buf is None: + raise RuntimeError( + f"layer {layer_idx} has disable_index_value=True; index V is not allocated" + ) + if idx_v.shape[1] != 1: + raise ValueError( + f"index V is replicated single-head; expected shape " + f"[N, 1, {self.sparse_index_dim}], got {tuple(idx_v.shape)}" + ) + buf.index_copy_(0, out_cache_loc.to(torch.long), idx_v.to(buf.dtype)) + + +class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): + """KVCacheManagerV2 subclass with a V2-managed paged index-K cache + per sparse layer. + + Each sparse local layer registers a ``Role.INDEX_KEY`` + :class:`BufferConfig` via :meth:`_extra_buffers_per_layer` so the + index-K cache participates in the V2 paged lifecycle (allocation, + free, prefix reuse) and shares ``life_cycle_id`` with main K/V. + + The optional index-V branch is kept as a plain CUDA tensor for the + rare ``disable_index_value=False`` test path. M3 production sets + ``disable_index_value=True`` everywhere, so ``_index_v_buffers`` + stays empty in production. + + Constructor extras (forwarded kwargs go to :class:`KVCacheManagerV2`): + * ``sparse_layer_ids`` — layer indices using sparse attention. + * ``disable_index_value_layer_ids`` — subset whose index-V is + omitted. + * ``sparse_index_dim`` — width of the index-K/V vectors. + """ + + def __init__( + self, + *args, + sparse_layer_ids=None, + disable_index_value_layer_ids=None, + sparse_index_dim: Optional[int] = None, + **kwargs, + ): + # Resolve M3 sparse-layer metadata from explicit kwargs first, + # then from ``sparse_attn_config``, then from the M3 checkpoint + # convention (layers 0..2 dense, 3..N-1 sparse, + # disable_index_value=True, sparse_index_dim=128). + sparse_attn_config = kwargs.get("sparse_attn_config") + num_layers = kwargs.get("num_layers") + + if sparse_index_dim is None: + sparse_index_dim = int(getattr(sparse_attn_config, "sparse_index_dim", 0) or 0) or 128 + if sparse_layer_ids is None: + if num_layers is not None: + sparse_layer_ids = list(range(3, int(num_layers))) + else: + sparse_layer_ids = [] + if disable_index_value_layer_ids is None: + disable_index_value_layer_ids = list(sparse_layer_ids) + + # Must be set BEFORE super().__init__ — the base + # ``_build_base_config`` invokes ``_extra_buffers_per_layer`` + # which reads these attributes. + self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids) + self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids) + self.sparse_index_dim = int(sparse_index_dim) + + super().__init__(*args, **kwargs) + + index_v_layer_ids = set(self.sparse_layer_ids) - self.disable_index_value_layer_ids + if self.is_disagg and index_v_layer_ids: + raise ValueError( + "MiniMax M3 disaggregated serving requires disable_index_value=True " + "for every sparse layer because the optional test-only index-V cache " + "is not managed or transferred by KVCacheManagerV2; enabled layers=" + f"{sorted(index_v_layer_ids)}" + ) + + # Optional plain-tensor index-V cache for non-disabled sparse + # layers (test-only; production has disable_index_value=True + # on every sparse layer). + num_total_slots = self._compute_num_total_slots() + torch_dtype = self._torch_dtype_for_index_cache() + device = torch.device("cuda") + self._index_v_buffers: dict[int, torch.Tensor] = {} + for layer_idx in self.sparse_layer_ids: + if layer_idx not in self.layer_offsets: + continue + if layer_idx not in self.disable_index_value_layer_ids: + self._index_v_buffers[layer_idx] = torch.zeros( + (num_total_slots, 1, self.sparse_index_dim), + dtype=torch_dtype, + device=device, + ) + + def _extra_buffers_per_layer(self, *, tokens_per_block): + """Register a per-sparse-layer ``Role.INDEX_KEY`` :class:`BufferConfig`. + + ``size`` is bytes per **block**: ``1 * sparse_index_dim * + elem_bytes * tokens_per_block``. Keyed by **local** layer id — + the base ``_build_base_config`` iterates local ids, so keying + by global ids would silently skip registration on non-trivial + PP ranks. + """ + torch_dtype = self._torch_dtype_for_index_cache() + elem_bytes = torch.tensor([], dtype=torch_dtype).element_size() + bytes_per_token = 1 * self.sparse_index_dim * elem_bytes + size_per_block = bytes_per_token * tokens_per_block + return { + self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=size_per_block)] + for layer_id in self.sparse_layer_ids + if layer_id in self.layer_offsets + } + + def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: + """Declare MiniMax M3's token-major K/V and replicated index-K.""" + return { + Role.ALL: MapperKind.NHD, + Role.INDEX_KEY: MapperKind.REPLICATED, + } + + def _compute_num_total_slots(self) -> int: + """Total token slots across all blocks in the main K pool. + + Sizes the plain-tensor index-V cache only; the V2-managed + index-K cache pulls its slot count from ``Role.INDEX_KEY``'s + page upper bound. + """ + if not self.layer_offsets: + return int(self.max_batch_size * self.max_seq_len) + any_layer_offset = next(iter(self.layer_offsets.values())) + page_upper = self.impl.get_page_index_upper_bound(any_layer_offset, Role.KEY) + kv_factor = 1 if self.kv_cache_type == CacheTypeCpp.SELFKONLY else 2 + return int((page_upper // kv_factor) * self.tokens_per_block) + + def _torch_dtype_for_index_cache(self) -> torch.dtype: + """Match the main cache dtype where possible, fall back to bf16.""" + if self.dtype == DataType.HALF: + return torch.float16 + if self.dtype == DataType.FLOAT: + return torch.float32 + return torch.bfloat16 + + def get_index_k_buffer(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + """Return the V2-managed paged index-K view for ``layer_idx``. + + NHD shape is ``[num_pages, tokens_per_block, 1, sparse_index_dim]``; + HND shape is ``[num_pages, 1, tokens_per_block, sparse_index_dim]``. + """ + return super().get_index_k_buffer( + layer_idx, + num_heads=1, + head_dim=self.sparse_index_dim, + dtype=self._torch_dtype_for_index_cache(), + kv_layout=kv_layout, + ) + + def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: + """Plain-tensor index-V cache for non-disabled sparse layers.""" + return self._index_v_buffers.get(layer_idx) + + def has_index_value(self, layer_idx: int) -> bool: + return layer_idx in self._index_v_buffers + + def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + """Return a paged K+V view with strides spanning the coalesced pool. + + The base :meth:`KVCacheManagerV2.get_buffers` produces a + ``[num_pages, kv_factor, ...]`` view with contiguous strides + that assume the slot holds exactly one layer's K+V. In M3's + pool the slot packs K+V for *all* layers of the group + (``scale >= 2 * num_layers_in_group``), so the base view's + dim-0 stride does not reach the next slot's K for this layer. + (When INDEX_KEY's per-block size coincides with K/V's, it is + coalesced into the same pool and contributes to ``scale`` too.) + + The override builds a ``[num_slots, scale, ...]`` view rooted + at K's base, then slices ``[:, :2]`` to extract K+V. The slice + preserves the dim-0 stride (``scale * page_stride``), so + ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``. + """ + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported kv_layout: {kv_layout}") + if self.kv_cache_type == CacheTypeCpp.SELFKONLY: + raise NotImplementedError( + "MiniMaxM3KVCacheManagerV2.get_buffers does not support SELFKONLY cache type" + ) + + layer_offset = self.layer_offsets[layer_idx] + addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY) + addr_value = self.impl.get_mem_pool_base_address(layer_offset, Role.VALUE) + page_stride_key = self.impl.get_page_stride(layer_offset, Role.KEY) + page_stride_value = self.impl.get_page_stride(layer_offset, Role.VALUE) + # V2 always lays V immediately after K within the per-layer + # contribution to a slot. The slice ``[:, :2]`` depends on this. + assert addr_key + page_stride_value == addr_value, ( + f"MiniMaxM3 get_buffers requires addr_K + page_stride " + f"== addr_V (V immediately after K in slot); got " + f"addr_K={addr_key} page_stride_V={page_stride_value} " + f"addr_V={addr_value} for layer {layer_idx}." + ) + assert page_stride_key == page_stride_value, ( + f"MiniMaxM3 get_buffers requires equal K and V page " + f"strides; got K={page_stride_key} V=" + f"{page_stride_value}." + ) + + converter = self.impl.get_page_index_converter(layer_offset, Role.KEY) + scale = int(converter.scale) + layer_offset_pages = int(converter.layer_offset) + page_upper_K = self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) + num_slots_total = page_upper_K + layer_offset_pages + assert num_slots_total % scale == 0, ( + f"V2 storage inconsistency: page_upper_K + " + f"layer_offset_pages = {num_slots_total} is not " + f"divisible by scale = {scale}." + ) + num_slots = num_slots_total // scale + + element_per_container = 1 + dtype = self.dtype + if dtype == DataType.NVFP4: + element_per_container = 2 + torch_dtype = torch.int8 + else: + torch_dtype = binding_to_torch_dtype(dtype) + + layer_head_dim = self.head_dim_per_layer[layer_offset] + num_kv_heads = self.num_kv_heads_per_layer[layer_offset] + + if kv_layout == "NHD": + full_slot_shape = [ + num_slots, + scale, + self.tokens_per_block, + num_kv_heads, + layer_head_dim // element_per_container, + ] + else: + full_slot_shape = [ + num_slots, + scale, + num_kv_heads, + self.tokens_per_block, + layer_head_dim // element_per_container, + ] + + full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) + return full_view[:, :2] + + def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> int: + """Pool-mapping offset from the layer's physical position in its pool. + + The base formula ``exact_div(addr_offset, key_bytes * kv_factor * + tokens_per_block)`` assumes each layer contributes exactly K+V to + its pool slot. When index-K coalesces into the K/V pool the layer + stride is non-uniform (sparse layers add an INDEX_KEY sub-page), + so no uniform-stride offset exists. The M3 forward path uses + :meth:`get_buffers` / :meth:`get_index_k_buffer` rather than this + mapping, so the offset just needs to be a consistent per-layer + position. Rank the group's layers by their K base address instead + of by ``layer_grouping`` iteration order: the ordering of + ``layer_grouping`` is not a V2 API contract, while the address + rank always reflects the physical slot layout (and keeps the + NVFP4 ``block_scale_offset == offset`` cross-check in the base + pool-mapping loop meaningful). + """ + layers_by_addr = sorted( + self.impl.layer_grouping[int(layer_group_id)], + key=lambda lid: self.impl.get_mem_pool_base_address( + lid, Role.KEY, PageIndexMode.SHARED + ), + ) + return layers_by_addr.index(int(layer_id)) + + def _get_batch_cache_indices_by_pool_id( + self, + request_ids, + *, + pool_id: int = 0, + is_kv_aggregate: bool = True, + num_blocks_per_seq: Optional[Sequence[int]] = None, + index_scale: Optional[int] = None, + kv_factor: Optional[int] = None, + ): + """Return page indices; padded entries remain ``BAD_PAGE_INDEX`` (-1). + + The base method converts slot ids to V1-style block ids via + ``base_idx * index_scales[pool_id] // kv_factor``, which is + only correct when each layer contributes exactly K+V. M3's slot + packs K+V for all layers of the group, so the scale breaks the + V1 conversion and produces out-of-bounds block ids during V2 + warmup. + + Bypass the conversion: the M3 forward path indexes paged + views (built by :meth:`get_buffers` / + :meth:`get_index_k_buffer`) directly by slot id. + ``BAD_PAGE_INDEX`` slots remain ``-1`` here because disaggregation's + :class:`KVRegionExtractorV1` filters ``region_ids >= 0``. + :meth:`get_block_ids_per_seq` maps them to zero for the attention + metadata's padded tensor. + + Args: + request_ids: Request IDs whose page-index rows are returned. + pool_id: V2 pool whose page indices are requested. + is_kv_aggregate: Kept for compatibility with the base virtual method. + num_blocks_per_seq: Optional per-request truncation limits. When + omitted, preserve the full padded width required by MiniMax + CUDA-graph metadata initialization. + index_scale: Kept for compatibility with the base virtual method; + M3 bypasses the V1 block-id conversion entirely, so any + caller-supplied scale is ignored alongside ``index_scales``. + """ + res = [] + for req_idx, req_id in enumerate(request_ids): + kv_cache = self.kv_cache_map[req_id] + base_page_indices = kv_cache.get_base_page_indices(pool_id) + if num_blocks_per_seq is not None: + num_blocks = min(kv_cache.num_blocks, num_blocks_per_seq[req_idx]) + base_page_indices = base_page_indices[:num_blocks] + res.append(list(base_page_indices)) + return res + + def get_block_ids_per_seq(self, request_ids): + """Return per-request slot ids matching the per-layer paged view's dim-0. + + Drops the base's final ``i // num_local_layers`` step (paired + with the base ``index_scales`` multiplication that's also + bypassed here). Pads with ``0`` to preserve shape. + """ + block_ids_per_seq = self.get_batch_cache_indices(request_ids) + block_ids_per_seq_tensors = [ + torch.tensor( + [i if i != BAD_PAGE_INDEX else 0 for i in sublist], + dtype=torch.int, + ) + for sublist in block_ids_per_seq + ] + padded_tensor = torch.nn.utils.rnn.pad_sequence( + block_ids_per_seq_tensors, + batch_first=True, + padding_value=0, + ) + return padded_tensor + + +def get_minimax_m3_kv_cache_manager_cls(): + """Backward-compatible accessor; prefer importing the class directly.""" + return MiniMaxM3KVCacheManagerV2 + + +__all__ = [ + "MiniMaxM3KVCacheManagerV2", + "MiniMaxM3SparseIndexCache", + "get_minimax_m3_kv_cache_manager_cls", +] diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index 2b12d6e0008b..1a0614d0c101 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -1,16 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Torch custom op for fused DSV4 DSpark rolling-window attention. - -The op runs the tcgen05 MMA kernel -(:class:`DSparkAttention`) on the supported DSV4 DSpark -geometry: 128 heads, head_dim 512, draft block 5 or 6, and a 128-row rolling -window. Other shapes fall back to the pure-PyTorch reference path in -``models/modeling_dspark.py``. -""" +"""Paged DSpark attention and Rubin FP8 projection operators.""" +import functools from collections.abc import Callable +from typing import List, Optional, Tuple import cutlass import cutlass.cute as cute @@ -20,184 +15,115 @@ from ..._utils import get_sm_version from ...logger import logger -from ..attention.kernels.blackwell.dspark.attention import DSparkAttention -from ..kernels.blackwell.utils import make_ptr -from .dspark_rmsnorm_rope_custom_op import ( - _get_dspark_arch_str, - cute_dsl_dspark_rmsnorm_rope_cache_write, - cute_dsl_dspark_rmsnorm_rope_draft_block, +from ...math_utils import ceil_div, pad_up +from ..autotuner import ( + AutoTuner, + ConstraintSpec, + DistributedTuningStrategy, + DynamicTensorSpec, + TunableRunner, + TuningConfig, ) +from ..cute_dsl_kernels.blackwell.dspark.attention import DSparkAttention +from ..cute_dsl_kernels.blackwell.utils import make_ptr +from ..cute_dsl_utils import IS_CUTLASS_DSL_INTERNAL_AVAILABLE +from ..utils import fp8_scale_infer_shape, mxfp8_scale_infer_shape, next_positive_power_of_2 +from .dspark_rmsnorm_rope_custom_op import _get_dspark_arch_str + +_DSPARK_NUM_HEADS = 128 +_DSPARK_HEAD_DIM = 512 +_DSPARK_WINDOW_SIZE = 128 +_DSPARK_PAGE_SIZE = 8 +_DSPARK_BLOCK_SIZES = (5, 6) +_DSPARK_ROPE_DIM = 64 + +_DSPARK_CONTEXT_WKV_MIN_M_BUCKET = 128 +_DSPARK_GENERATION_M_SPEC = DynamicTensorSpec(0, 0) +# These are request-capacity buckets for tactic selection. The Rubin GEMM +# artifact has runtime-dynamic M/N/K, so this does not specialize a kernel for +# each batch value. +_DSPARK_GENERATION_BATCH_BUCKETS = (1, 2, 4, 8, 16, 32, 64, 128, 256) + + +@functools.cache +def _get_dspark_generation_m_spec( + rows_per_request: int, + max_batch: int, + input_idx: int = 0, + dim_idx: int = 0, +) -> DynamicTensorSpec: + """Build the DSpark operator-owned generation-M bucket contract.""" + if rows_per_request < 1 or max_batch < 1: + raise ValueError( + "DSpark generation tuning requires positive rows_per_request and max_batch, " + f"got {rows_per_request=} and {max_batch=}." + ) -_DSV4_DSPARK_NUM_HEADS = 128 -_DSV4_DSPARK_HEAD_DIM = 512 -_DSV4_DSPARK_WINDOW_SIZE = 128 -_DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) -_DSV4_DSPARK_ROPE_DIM = 64 + batch_buckets = tuple( + batch for batch in _DSPARK_GENERATION_BATCH_BUCKETS if batch < max_batch + ) + (max_batch,) + m_buckets = tuple(batch * rows_per_request for batch in batch_buckets) + + def map_to_bucket(m: int) -> int: + batch = ceil_div(m, rows_per_request) + for bucket in batch_buckets: + if batch <= bucket: + return bucket * rows_per_request + # A shape beyond the configured worker capacity is not expected, but + # keeping it exact gives the autotuner its normal cache-miss fallback. + return m + + return DynamicTensorSpec( + input_idx, + dim_idx, + m_buckets, + map_to_bucket, + ) -def _log_unsupported(reason: str, key: str) -> bool: - logger.debug_once( - f"Falling back from fused DSV4 DSpark attention: {reason}", - key=("fused_dsv4_dspark_attention_unsupported", key), +@functools.cache +def _get_dspark_generation_mxfp8_tuning_config( + rows_per_request: int, + max_batch: int, +) -> TuningConfig: + """Return generation buckets for a DSpark MXFP8 projection.""" + return TuningConfig( + dynamic_tensor_specs=(_get_dspark_generation_m_spec(rows_per_request, max_batch),), + constraint_specs=(ConstraintSpec(2, 0, mxfp8_scale_infer_shape),), + use_cold_l2_cache=True, + distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + use_cuda_graph=True, ) - return False -def is_dsv4_dspark_attention_config_supported( - block_size: int, - num_heads: int, - head_dim: int, - window_size: int, -) -> bool: - """Return whether static model geometry matches the DSV4 specialization.""" - return ( - block_size in _DSV4_DSPARK_BLOCK_SIZES - and num_heads == _DSV4_DSPARK_NUM_HEADS - and head_dim == _DSV4_DSPARK_HEAD_DIM - and window_size == _DSV4_DSPARK_WINDOW_SIZE - ) +def _map_dspark_context_wkv_m_profile(m: int) -> int: + """Bucket CTX seed and batched backfill rows without changing kernel shapes.""" + return max(_DSPARK_CONTEXT_WKV_MIN_M_BUCKET, next_positive_power_of_2(m)) -def is_fused_dsv4_dspark_attention_supported( - q: torch.Tensor, - kv_cache: torch.Tensor, - valid_len: torch.Tensor, - attn_sink: torch.Tensor, - inverse_rope_freqs: torch.Tensor, -) -> bool: - """Return whether tensors can use the fused DSV4 DSpark attention kernel.""" - if _get_dspark_arch_str() is None: - return _log_unsupported(f"SM {get_sm_version()} is not SM100 or SM103", "sm_version") - if not all(t.is_cuda for t in (q, kv_cache, valid_len, attn_sink, inverse_rope_freqs)): - return _log_unsupported("all inputs must be CUDA tensors", "device") - if q.dtype != torch.bfloat16: - return _log_unsupported(f"q dtype must be BF16, got {q.dtype}", "q_dtype") - if kv_cache.dtype != q.dtype: - return _log_unsupported( - f"kv_cache dtype must match q, got {kv_cache.dtype}", "kv_cache_dtype" - ) - if valid_len.dtype != torch.int64: - return _log_unsupported( - f"valid_len must use INT64, got {valid_len.dtype}", "valid_len_dtype" - ) - if attn_sink.dtype != torch.float32 or inverse_rope_freqs.dtype != torch.float32: - return _log_unsupported( - "attn_sink and inverse_rope_freqs must be FP32; " - f"got {attn_sink.dtype} and {inverse_rope_freqs.dtype}", - "aux_dtype", - ) - if q.ndim != 4 or kv_cache.ndim != 3: - return _log_unsupported( - f"expected q/kv_cache ranks 4/3, got {q.ndim}/{kv_cache.ndim}", "tensor_ranks" - ) - if q.shape[1] not in _DSV4_DSPARK_BLOCK_SIZES: - return _log_unsupported(f"draft block size must be 5 or 6, got {q.shape[1]}", "block_size") - if q.shape[2:] != (_DSV4_DSPARK_NUM_HEADS, _DSV4_DSPARK_HEAD_DIM): - return _log_unsupported( - "q must have 128 heads and head_dim 512; " - f"got {q.shape[2]} heads and head_dim {q.shape[3]}", - "q_shape", - ) - if kv_cache.shape[1:] != (_DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM): - return _log_unsupported( - f"kv_cache must have trailing shape (128, 512), got {tuple(kv_cache.shape[1:])}", - "kv_cache_shape", - ) - if ( - kv_cache.stride(0) <= 0 - or kv_cache.stride(0) % 8 != 0 - or kv_cache.stride(1) != _DSV4_DSPARK_HEAD_DIM - or kv_cache.stride(2) != 1 - ): - return _log_unsupported( - "kv_cache strides must have a positive 16-byte-aligned page stride " - f"and trailing strides (512, 1), got {kv_cache.stride()}", - "kv_cache_layout", - ) - if valid_len.shape != (q.shape[0],): - return _log_unsupported( - f"valid_len shape must be {(q.shape[0],)}, got {tuple(valid_len.shape)}", - "valid_len_shape", - ) - if attn_sink.shape != (_DSV4_DSPARK_NUM_HEADS,): - return _log_unsupported( - f"attn_sink shape must be {(_DSV4_DSPARK_NUM_HEADS,)}, got {tuple(attn_sink.shape)}", - "attn_sink_shape", - ) - expected_freqs_shape = (q.shape[0], q.shape[1], _DSV4_DSPARK_ROPE_DIM // 2, 2) - if inverse_rope_freqs.shape != expected_freqs_shape: - return _log_unsupported( - f"inverse_rope_freqs shape must be {expected_freqs_shape}, " - f"got {tuple(inverse_rope_freqs.shape)}", - "inverse_rope_freqs_shape", - ) - for name, tensor in ( - ("q", q), - ("valid_len", valid_len), - ("attn_sink", attn_sink), - ("inverse_rope_freqs", inverse_rope_freqs), - ): - if not tensor.is_contiguous(): - return _log_unsupported(f"{name} must be contiguous", f"{name}_layout") - return True +def _gen_dspark_context_wkv_m_profiles(max_m: int) -> Tuple[int, ...]: + max_bucket = _map_dspark_context_wkv_m_profile(max_m) + buckets = [] + bucket = _DSPARK_CONTEXT_WKV_MIN_M_BUCKET + while bucket <= max_bucket: + buckets.append(bucket) + bucket *= 2 + return tuple(buckets) -def _is_fused_dsv4_dspark_attention_input_supported( - q: torch.Tensor, - draft_block: torch.Tensor, - kv_cache: torch.Tensor, - slots_i32: torch.Tensor, - cache_seqs: torch.Tensor, - valid_len: torch.Tensor, - attn_sink: torch.Tensor, - inverse_rope_freqs: torch.Tensor, -) -> bool: - if not is_fused_dsv4_dspark_attention_supported( - q, kv_cache, valid_len, attn_sink, inverse_rope_freqs - ): - return False - if not all(t.is_cuda for t in (draft_block, slots_i32, cache_seqs)): - return _log_unsupported("all fused-op inputs must be CUDA tensors", "input_device") - if draft_block.dtype != q.dtype: - return _log_unsupported( - f"draft_block dtype must match q, got {draft_block.dtype}", "draft_block_dtype" - ) - if slots_i32.dtype != torch.int32 or cache_seqs.dtype != torch.int32: - return _log_unsupported( - "slots_i32 and cache_seqs must use INT32; " - f"got {slots_i32.dtype} and {cache_seqs.dtype}", - "index_dtype", - ) - expected_draft_block_shape = ( - q.shape[0], - _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, - _DSV4_DSPARK_HEAD_DIM, - ) - if draft_block.shape != expected_draft_block_shape: - return _log_unsupported( - f"draft_block shape must be {expected_draft_block_shape}, " - f"got {tuple(draft_block.shape)}", - "draft_block_shape", - ) - expected_batch_shape = (q.shape[0],) - if slots_i32.shape != expected_batch_shape or cache_seqs.shape != expected_batch_shape: - return _log_unsupported( - f"slots_i32 and cache_seqs shapes must be {expected_batch_shape}; " - f"got {tuple(slots_i32.shape)} and {tuple(cache_seqs.shape)}", - "index_shape", - ) - for name, tensor in ( - ("draft_block", draft_block), - ("slots_i32", slots_i32), - ("cache_seqs", cache_seqs), - ): - if not tensor.is_contiguous(): - return _log_unsupported(f"{name} must be contiguous", f"{name}_layout") - return True +# Worker initialization is process-local and single-threaded. These transient +# values let the memoized builder consume the real cache only during precompile. +_DSV4_DSPARK_NUM_HEADS = (64, 128) +_DSV4_DSPARK_HEAD_DIM = 512 +_DSV4_DSPARK_WINDOW_SIZE = 128 +_DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 +_DSV4_DSPARK_BLOCK_SIZES = (5, 6) +_DSV4_DSPARK_ROPE_DIM = 64 -_dspark_attention_kernel_cache: dict[tuple[int, str], Callable[..., None]] = {} +_dspark_attention_kernel_cache: dict[ + tuple[int, int, int, str, int | None], Callable[..., None] +] = {} def _make_compile_gmem_pointer(dtype: Type[Numeric], assumed_align: int) -> Pointer: @@ -213,9 +139,12 @@ def _make_compile_gmem_pointer(dtype: Type[Numeric], assumed_align: int) -> Poin def _compile_dspark_attention( block_size: int, arch_str: str, + page_size: int, + output_groups: int | None, + num_heads: int, ) -> Callable[..., None]: """Compile the pointer host wrapper without runtime tensor specimens.""" - num_heads, head_dim = _DSV4_DSPARK_NUM_HEADS, _DSV4_DSPARK_HEAD_DIM + head_dim = _DSV4_DSPARK_HEAD_DIM batch = cute.sym_int() output_fake = cute.runtime.make_fake_compact_tensor( cutlass.BFloat16, @@ -223,6 +152,22 @@ def _compile_dspark_attention( stride_order=(3, 2, 1, 0), assumed_align=16, ) + output_scale_fake = None + if output_groups is not None: + group_dim = num_heads * head_dim // output_groups + rows = batch * block_size + output_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (output_groups, rows, group_dim), + stride_order=(2, 1, 0), + assumed_align=16, + ) + output_scale_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (output_groups, group_dim // 128, (rows + 3) // 4 * 4), + stride_order=(2, 1, 0), + assumed_align=16, + ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) # Real tensor specimens used to make PyTorch's primary CUDA context current @@ -243,6 +188,9 @@ def _compile_dspark_attention( mma_qk_tiler_k=128, inverse_rope_dim=_DSV4_DSPARK_ROPE_DIM, arch_str=arch_str, + history_page_size=page_size, + output_groups=output_groups, + num_heads=num_heads, ) # Scalars and typed, aligned compile-only pointers define the runtime ABI. # ``output_fake`` keeps one tensor argument for TVM-FFI environment-stream @@ -251,16 +199,20 @@ def _compile_dspark_attention( kernel.wrapper, 1, 1, - _DSV4_DSPARK_WINDOW_SIZE * head_dim, + page_size * head_dim, + 1, + 1, _make_compile_gmem_pointer(cutlass.BFloat16, 16), _make_compile_gmem_pointer(cutlass.BFloat16, 16), _make_compile_gmem_pointer(cutlass.BFloat16, 16), _make_compile_gmem_pointer(cutlass.Int32, 4), - _make_compile_gmem_pointer(cutlass.Int32, 4), + _make_compile_gmem_pointer(cutlass.Int64, 8), + _make_compile_gmem_pointer(cutlass.Int64, 8), _make_compile_gmem_pointer(cutlass.Int64, 8), _make_compile_gmem_pointer(cutlass.Float32, 4), _make_compile_gmem_pointer(cutlass.Float32, 4), output_fake, + output_scale_fake, cutlass.Float32(1.0), stream_fake, options="--opt-level 2 --enable-tvm-ffi", @@ -272,188 +224,753 @@ def _compile_dspark_attention( return compiled -def _run_dspark_attention( +def _run_paged_dspark_attention( q: torch.Tensor, draft_block: torch.Tensor, - kv_cache: torch.Tensor, - slots_i32: torch.Tensor, - cache_seqs: torch.Tensor, - valid_len: torch.Tensor, + kv_pages: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + valid_lengths: torch.Tensor, + capacities: torch.Tensor, attn_sink: torch.Tensor, inverse_rope_freqs: torch.Tensor, softmax_scale: float, -) -> torch.Tensor: - block_size = q.shape[1] + output_groups: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Attend to a 128-token context window and fuse the inverse-RoPE epilogue.""" arch_str = _get_dspark_arch_str() if arch_str is None: - raise RuntimeError("fused DSV4 DSpark attention requires SM100 or SM103") - cache_key = (block_size, arch_str) + raise RuntimeError("Embedded DSpark attention requires SM100, SM103, or SM107") + batch, block, heads, dim = q.shape + page_size = kv_pages.shape[1] + if ( + block not in _DSV4_DSPARK_BLOCK_SIZES + or heads not in _DSV4_DSPARK_NUM_HEADS + or dim != _DSV4_DSPARK_HEAD_DIM + or page_size not in (16, 32, 64, 128, 256) + ): + raise ValueError( + "DSpark requires block 5/6, 64/128 heads, head_dim 512 and 16/32/64/128/256-token pages" + ) + if q.dtype != torch.bfloat16 or draft_block.dtype != q.dtype or kv_pages.dtype != q.dtype: + raise ValueError("DSpark Q, draft KV and pages must use BF16") + if draft_block.shape != (batch, 8, dim) or kv_pages.shape[2] != dim: + raise ValueError("DSpark draft KV must have shape [batch, 8, 512]") + if ( + block_tables.dtype != torch.int32 + or block_tables.shape[0] != batch + or block_tables.stride(1) != 1 + ): + raise ValueError("DSpark page tables must have contiguous INT32 rows") + for value in (positions, valid_lengths, capacities): + if value.dtype != torch.int64 or value.shape != (batch,) or not value.is_contiguous(): + raise ValueError( + "DSpark positions, lengths and capacities must be contiguous INT64 batch vectors" + ) + if attn_sink.dtype != torch.float32 or attn_sink.shape != (heads,): + raise ValueError("DSpark attention sink must be FP32 with one entry per head") + if inverse_rope_freqs.dtype != torch.float32 or inverse_rope_freqs.shape != ( + batch, + block, + 32, + 2, + ): + raise ValueError("DSpark inverse RoPE frequencies must be FP32 [batch, block, 32, 2]") + if kv_pages.stride(1) != dim or kv_pages.stride(2) != 1 or kv_pages.stride(0) % 8: + raise ValueError("DSpark pages require contiguous tokens and an aligned page stride") + for value in (q, draft_block, attn_sink, inverse_rope_freqs): + if not value.is_contiguous(): + raise ValueError("DSpark query, draft, sink and frequencies must be contiguous") + if not all( + value.is_cuda and value.device == q.device + for value in ( + q, + draft_block, + kv_pages, + block_tables, + positions, + valid_lengths, + capacities, + attn_sink, + inverse_rope_freqs, + ) + ): + raise ValueError("All DSpark attention inputs must be on the same CUDA device") + cache_key = (block, heads, page_size, arch_str, output_groups) compiled = _dspark_attention_kernel_cache.get(cache_key) if compiled is None: if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "fused DSV4 DSpark attention must be warmed up before CUDA graph capture" - ) - compiled = _compile_dspark_attention(block_size, arch_str) + raise RuntimeError("DSpark paged attention must run eagerly before graph capture") + compiled = _compile_dspark_attention(block, arch_str, page_size, output_groups, heads) _dspark_attention_kernel_cache[cache_key] = compiled - output = torch.empty_like(q) + output_scale = None + if output_groups is not None: + rows = batch * block + group_dim = heads * dim // output_groups + output = torch.empty( + (output_groups, rows, group_dim), dtype=torch.float8_e4m3fn, device=q.device + ) + output_scale = torch.empty( + (output_groups, group_dim // 128, (rows + 3) // 4 * 4), + dtype=torch.float32, + device=q.device, + ) compiled( - q.shape[0], - kv_cache.shape[0], - kv_cache.stride(0), + batch, + kv_pages.shape[0], + kv_pages.stride(0), + block_tables.shape[1], + block_tables.stride(0), q.data_ptr(), - kv_cache.data_ptr(), + kv_pages.data_ptr(), draft_block.data_ptr(), - slots_i32.data_ptr(), - cache_seqs.data_ptr(), - valid_len.data_ptr(), + block_tables.data_ptr(), + positions.data_ptr(), + valid_lengths.data_ptr(), + capacities.data_ptr(), attn_sink.data_ptr(), inverse_rope_freqs.data_ptr(), output, + output_scale, softmax_scale, ) - return output + return output, output_scale @torch.library.custom_op( - "trtllm::fused_dsv4_dspark_attention", - mutates_args=(), - device_types="cuda", + "trtllm::fused_dsv4_dspark_attention", mutates_args=(), device_types="cuda" ) def fused_dsv4_dspark_attention( q: torch.Tensor, draft_block: torch.Tensor, - kv_cache: torch.Tensor, - slots_i32: torch.Tensor, - cache_seqs: torch.Tensor, - valid_len: torch.Tensor, + kv_pages: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + valid_lengths: torch.Tensor, + capacities: torch.Tensor, attn_sink: torch.Tensor, inverse_rope_freqs: torch.Tensor, softmax_scale: float, ) -> torch.Tensor: - """Run fused DSV4 DSpark attention after validating its tensor contract.""" - - if not _is_fused_dsv4_dspark_attention_input_supported( + return _run_paged_dspark_attention( q, draft_block, - kv_cache, - slots_i32, - cache_seqs, - valid_len, + kv_pages, + block_tables, + positions, + valid_lengths, + capacities, attn_sink, inverse_rope_freqs, - ): - raise ValueError( - "fused_dsv4_dspark_attention requires contiguous supported DSV4 DSpark tensors " - "([B, 5|6, 128, 512] queries, [B, 8, 512] draft blocks, INT32 indices, " - "and a 128-row BF16 window) on SM100 or SM103" - ) + softmax_scale, + )[0] + - return _run_dspark_attention( +@torch.library.register_fake("trtllm::fused_dsv4_dspark_attention") +def _( + q: torch.Tensor, + draft_block: torch.Tensor, + kv_pages: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + valid_lengths: torch.Tensor, + capacities: torch.Tensor, + attn_sink: torch.Tensor, + inverse_rope_freqs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + return torch.empty_like(q) + + +@torch.library.custom_op( + "trtllm::fused_dsv4_dspark_attention_fp8", mutates_args=(), device_types="cuda" +) +def fused_dsv4_dspark_attention_fp8( + q: torch.Tensor, + draft_block: torch.Tensor, + kv_pages: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + valid_lengths: torch.Tensor, + capacities: torch.Tensor, + attn_sink: torch.Tensor, + inverse_rope_freqs: torch.Tensor, + softmax_scale: float, + output_groups: int, +) -> tuple[torch.Tensor, torch.Tensor]: + output, scales = _run_paged_dspark_attention( q, draft_block, - kv_cache, - slots_i32, - cache_seqs, - valid_len, + kv_pages, + block_tables, + positions, + valid_lengths, + capacities, attn_sink, inverse_rope_freqs, softmax_scale, + output_groups, ) + assert scales is not None + return output, scales -@torch.library.register_fake("trtllm::fused_dsv4_dspark_attention") +@torch.library.register_fake("trtllm::fused_dsv4_dspark_attention_fp8") def _( q: torch.Tensor, draft_block: torch.Tensor, - kv_cache: torch.Tensor, - slots_i32: torch.Tensor, - cache_seqs: torch.Tensor, - valid_len: torch.Tensor, + kv_pages: torch.Tensor, + block_tables: torch.Tensor, + positions: torch.Tensor, + valid_lengths: torch.Tensor, + capacities: torch.Tensor, attn_sink: torch.Tensor, inverse_rope_freqs: torch.Tensor, softmax_scale: float, -) -> torch.Tensor: - del draft_block, kv_cache, slots_i32, cache_seqs, valid_len, attn_sink, inverse_rope_freqs - del softmax_scale - return torch.empty_like(q) + output_groups: int, +) -> tuple[torch.Tensor, torch.Tensor]: + rows = q.shape[0] * q.shape[1] + group_dim = q.shape[2] * q.shape[3] // output_groups + return ( + torch.empty((output_groups, rows, group_dim), dtype=torch.float8_e4m3fn, device=q.device), + torch.empty( + (output_groups, group_dim // 128, (rows + 3) // 4 * 4), + dtype=torch.float32, + device=q.device, + ), + ) -def warmup_fused_dsv4_dspark_attention(block_size: int, eps: float) -> None: - """Best-effort prewarm of the production DSV4 DSpark fused-op path.""" - if _get_dspark_arch_str() is None or block_size not in _DSV4_DSPARK_BLOCK_SIZES: - return +if IS_CUTLASS_DSL_INTERNAL_AVAILABLE: + from .cute_dsl_custom_ops import CuteDSLFp8RubinBmmQuantRunner, CuteDSLMXFP8RubinLinear - batch = 1 - device = torch.device("cuda") - try: - with torch.inference_mode(): - q = torch.zeros( - (batch, block_size, _DSV4_DSPARK_NUM_HEADS, _DSV4_DSPARK_HEAD_DIM), - dtype=torch.bfloat16, - device=device, - ) - main_x = torch.zeros( - (batch, 1, _DSV4_DSPARK_HEAD_DIM), dtype=torch.bfloat16, device=device + @functools.cache + def _get_dspark_generation_fp8_bmm_quant_tuning_config( + rows_per_request: int, + max_batch: int, + ) -> TuningConfig: + # One-engine DSpark runs target verification (K+1 rows/request) + # before the K-row draft. Tune the shared OA BMM for both real + # generation families when the draft first publishes its capacity. + draft_spec = _get_dspark_generation_m_spec( + rows_per_request, max_batch, input_idx=0, dim_idx=1 + ) + target_spec = _get_dspark_generation_m_spec( + rows_per_request + 1, max_batch, input_idx=0, dim_idx=1 + ) + m_buckets = tuple( + sorted(set(draft_spec.gen_tuning_buckets + target_spec.gen_tuning_buckets)) + ) + + def map_to_bucket(m: int) -> int: + for bucket in m_buckets: + if m <= bucket: + return bucket + return m + + return TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=0, + dim_idx=1, + gen_tuning_buckets=m_buckets, + map_to_tuning_buckets=map_to_bucket, + ), + ), + constraint_specs=( + ConstraintSpec(2, 2, fp8_scale_infer_shape), + ConstraintSpec(4, 1, lambda input_shapes: input_shapes[0][1]), + ConstraintSpec( + 5, + 0, + lambda input_shapes: pad_up(input_shapes[0][1], 128) + * pad_up( + ceil_div(input_shapes[0][0] * input_shapes[1][1], 32), + 4, + ), + ), + ), + distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + ) + + class CuteDSLDSparkFp8RubinBmmQuantRunner(CuteDSLFp8RubinBmmQuantRunner): + """DSpark OA runner with operator-owned M buckets and launch plans.""" + + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec(0, 1),), + constraint_specs=( + ConstraintSpec(2, 2, fp8_scale_infer_shape), + ConstraintSpec(4, 1, lambda input_shapes: input_shapes[0][1]), + ConstraintSpec( + 5, + 0, + lambda input_shapes: pad_up(input_shapes[0][1], 128) + * pad_up( + ceil_div(input_shapes[0][0] * input_shapes[1][1], 32), + 4, + ), + ), + ), + distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + ) + + def _forward_tvm_ffi_prepared( + self, + inputs: List[torch.Tensor], + compiled_gemm, + ) -> None: + if self.fp8_quantize_1x128: + ( + a_tensor, + b_tensor, + a_sf_tensor, + b_sf_tensor, + c_tensor, + scale_tensor, + ) = inputs + else: + a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, c_tensor = inputs + scale_tensor = None + + batch_size, m, k = a_tensor.shape + n = b_tensor.shape[1] + runtime_args = ( + m, + n, + k, + pad_up(m, 4), + ceil_div(n, 128), + ceil_div(k, 128), + batch_size, + a_tensor.data_ptr(), + b_tensor.data_ptr(), + a_sf_tensor.data_ptr(), + b_sf_tensor.data_ptr(), + c_tensor.permute(1, 2, 0), ) - block_x = torch.zeros( - (batch, block_size, _DSV4_DSPARK_HEAD_DIM), - dtype=torch.bfloat16, - device=device, + if self.fp8_quantize_1x128: + compiled_gemm(*runtime_args, scale_tensor) + else: + compiled_gemm(*runtime_args) + + def __init__(self, use_tvm_ffi: bool = True): + super().__init__(use_tvm_ffi=use_tvm_ffi) + self._tvm_ffi_launch_plans = {} + + def unique_id(self): + return (self.output_dtype, self.use_tvm_ffi) + + @functools.cache + def _get_cute_dsl_dspark_fp8_rubin_bmm_quant_runner( + use_tvm_ffi: bool, + ) -> CuteDSLDSparkFp8RubinBmmQuantRunner: + return CuteDSLDSparkFp8RubinBmmQuantRunner(use_tvm_ffi=use_tvm_ffi) + + class CuteDSLDSparkMXFP8RubinLinear(CuteDSLMXFP8RubinLinear): + """DSpark MXFP8 runner with process-local tuning and launch plans.""" + + tuning_config = TuningConfig( + # DSpark warmup runs the real generation shapes. Keep those M + # values exact so runtime lookup uses the tactic measured for + # the captured graph instead of remapping it to a power of two. + dynamic_tensor_specs=(_DSPARK_GENERATION_M_SPEC,), + constraint_specs=(ConstraintSpec(2, 0, mxfp8_scale_infer_shape),), + use_cold_l2_cache=True, + distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + use_cuda_graph=True, + ) + + def _forward_tvm_ffi_prepared( + self, + inputs: List[torch.Tensor], + launch_plan, + bias: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + compiled_gemm, swap_ab = launch_plan + if len(inputs) == 5: + a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, alpha_tensor = inputs + output_tensor = None + else: + ( + a_tensor, + b_tensor, + a_sf_tensor, + b_sf_tensor, + alpha_tensor, + output_tensor, + ) = inputs + + m, k = a_tensor.shape + n = b_tensor.shape[0] + partition_id = kwargs.get("partition_id", -1) + if output_tensor is not None: + c_tensor = output_tensor[:, partition_id * n : (partition_id + 1) * n] + elif self.to_userbuffers: + c_tensor = torch.ops.trtllm.create_userbuffers_tensor([m, n], self.output_dtype) + else: + c_tensor = torch.empty( + (m, n), + dtype=self.output_dtype, + device=a_tensor.device, + ) + + if swap_ab: + c_tensor = c_tensor.permute(1, 0) + + real_k = k * self.packed_k_multiplier + sf_m = pad_up(m, 128) + sf_k = pad_up(real_k // self.sf_vec_size, 4) + sf_n = pad_up(n, 128) + if swap_ab: + kernel_m = n + kernel_n = m + kernel_sf_m = sf_n + kernel_sf_n = sf_m + kernel_a_tensor = b_tensor + kernel_a_sf_tensor = b_sf_tensor + kernel_b_tensor = a_tensor + kernel_b_sf_tensor = a_sf_tensor + else: + kernel_m = m + kernel_n = n + kernel_sf_m = sf_m + kernel_sf_n = sf_n + kernel_a_tensor = a_tensor + kernel_a_sf_tensor = a_sf_tensor + kernel_b_tensor = b_tensor + kernel_b_sf_tensor = b_sf_tensor + + c_ld = output_tensor.shape[1] if output_tensor is not None else 0 + compiled_gemm( + kernel_m, + kernel_n, + real_k, + kernel_sf_m // 128, + kernel_sf_n // 128, + sf_k // 4, + kernel_a_tensor.data_ptr(), + kernel_b_tensor.data_ptr(), + kernel_a_sf_tensor.data_ptr(), + kernel_b_sf_tensor.data_ptr(), + c_tensor.data_ptr(), + alpha_tensor, + c_ld, ) - weight = torch.ones((_DSV4_DSPARK_HEAD_DIM,), dtype=torch.bfloat16, device=device) - main_freqs = torch.zeros( - (batch, _DSV4_DSPARK_ROPE_DIM // 2, 2), - dtype=torch.float32, - device=device, + + if swap_ab: + c_tensor = c_tensor.permute(1, 0) + if output_tensor is None: + if bias is not None: + if bias.ndim != 1 or bias.shape[0] != c_tensor.shape[-1]: + raise ValueError( + f"bias must be a 1-D tensor of shape [N]={c_tensor.shape[-1]}, " + f"got shape {tuple(bias.shape)}" + ) + c_tensor = c_tensor + bias + return c_tensor + return + + def __init__( + self, + output_dtype: torch.dtype, + to_userbuffers: bool = False, + use_tvm_ffi: bool = True, + ): + super().__init__( + output_dtype=output_dtype, + to_userbuffers=to_userbuffers, + use_tvm_ffi=use_tvm_ffi, ) - block_freqs = torch.zeros( - (batch * block_size, _DSV4_DSPARK_ROPE_DIM // 2, 2), - dtype=torch.float32, - device=device, + self._tvm_ffi_launch_plans = {} + + @functools.cache + def _get_cute_dsl_mxfp8_rubin_alpha(device: torch.device) -> torch.Tensor: + """Keep the scalar-one GEMM input alive for each CUDA device.""" + return torch.ones((), dtype=torch.float32, device=device) + + @functools.cache + def _get_cute_dsl_dspark_mxfp8_rubin_runner( + output_dtype: torch.dtype, + use_tvm_ffi: bool, + ) -> CuteDSLDSparkMXFP8RubinLinear: + return CuteDSLDSparkMXFP8RubinLinear( + output_dtype=output_dtype, + use_tvm_ffi=use_tvm_ffi, + ) + + _mxfp8_prevalidated_launch_plans = {} + + def _materialize_dspark_tuning_profile_winners( + custom_op: str, + runner: TunableRunner, + tuning_config: TuningConfig, + inputs: List[torch.Tensor], + ) -> None: + """Compile every cached DSpark bucket winner in this process.""" + tuner = AutoTuner.get() + if not tuner.is_tuning_mode: + return + + profiles = tuner._optimization_profiles(tuning_config, inputs) + key = ( + custom_op, + inputs[0].device, + runner.unique_id(), + tuple(profile.get_opt_shapes() for profile in profiles), + ) + materialized = getattr(runner, "_dspark_materialized_profiles", None) + if materialized is None: + materialized = set() + runner._dspark_materialized_profiles = materialized + if key in materialized: + return + + complete = True + for profile in profiles: + found, _, tactic, _ = tuner.profiling_cache.search_cache( + custom_op, + [runner], + profile.get_opt_shapes(), + tuning_config, + apply_map_to_tuning_buckets=False, ) - kv_cache = torch.zeros( - (batch, _DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM), - dtype=torch.bfloat16, - device=device, + if not found: + complete = False + continue + profile_inputs = tuner._prepare_input_tensors(profile, inputs) + runner(profile_inputs, tactic=tactic) + if complete: + materialized.add(key) + + def _run_cute_dsl_dspark_mxfp8_gemm_rubin_prevalidated( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + use_tvm_ffi: bool = True, + materialize_launch: bool = False, + tuning_config: Optional[TuningConfig] = None, + ) -> torch.Tensor: + """Run DSpark MXFP8 after bind-time contract validation.""" + alpha = _get_cute_dsl_mxfp8_rubin_alpha(input.device) + runner = _get_cute_dsl_dspark_mxfp8_rubin_runner(output_dtype, use_tvm_ffi) + inputs = [input, weight, input_scale, weight_scale, alpha] + tuner = AutoTuner.get() + materialized_key = None + if materialize_launch and use_tvm_ffi and not tuner.is_tuning_mode: + m, k = input.shape + n = weight.shape[0] + materialized_key = ( + input.device, + output_dtype, + m, + n, + k, + input_scale.numel(), + weight_scale.numel(), ) - slots = torch.zeros((batch,), dtype=torch.int64, device=device) - start_pos = torch.zeros((batch,), dtype=torch.int64, device=device) - slots_i32, cache_seqs = cute_dsl_dspark_rmsnorm_rope_cache_write( - main_x, - weight, - main_freqs, - kv_cache, - slots, - start_pos, - eps, + launch_plan = _mxfp8_prevalidated_launch_plans.get(materialized_key) + if launch_plan is not None: + return runner._forward_tvm_ffi_prepared(inputs, launch_plan) + + active_tuning_config = tuning_config or runner.__class__.tuning_config + custom_op = "trtllm::dspark_cute_dsl_mxfp8_gemm_rubin" + _, best_tactic = tuner.choose_one( + custom_op, + [runner], + active_tuning_config, + inputs, + ) + output = runner(inputs, tactic=best_tactic) + if tuning_config is not None and tuner.is_tuning_mode: + _materialize_dspark_tuning_profile_winners( + custom_op, runner, active_tuning_config, inputs ) - draft_block = cute_dsl_dspark_rmsnorm_rope_draft_block( - block_x, - weight, - block_freqs, - eps, + if materialized_key is not None: + launch_plan_key = ( + best_tactic, + m, + n, + k, + input_scale.numel(), + weight_scale.numel(), + 0, + -1, ) - fused_dsv4_dspark_attention( - q, - draft_block, - kv_cache, - slots_i32, - cache_seqs, - torch.zeros((batch,), dtype=torch.int64, device=device), - torch.zeros((_DSV4_DSPARK_NUM_HEADS,), dtype=torch.float32, device=device), - block_freqs.view( - batch, - block_size, - _DSV4_DSPARK_ROPE_DIM // 2, - 2, + launch_plan = runner._tvm_ffi_launch_plans.get(launch_plan_key) + if launch_plan is not None: + _mxfp8_prevalidated_launch_plans[materialized_key] = launch_plan + return output + + class CuteDSLDSparkContextWkvRubinLinear(CuteDSLDSparkMXFP8RubinLinear): + """MXFP8 runner with operator-owned buckets for DSpark CTX M.""" + + tuning_config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + 0, + 0, + _gen_dspark_context_wkv_m_profiles, + _map_dspark_context_wkv_m_profile, ), - 1.0, + ), + constraint_specs=(ConstraintSpec(2, 0, mxfp8_scale_infer_shape),), + use_cold_l2_cache=True, + distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + use_cuda_graph=True, + ) + + @functools.cache + def _get_cute_dsl_dspark_context_wkv_rubin_runner( + output_dtype: torch.dtype, + use_tvm_ffi: bool, + ) -> CuteDSLDSparkContextWkvRubinLinear: + return CuteDSLDSparkContextWkvRubinLinear( + output_dtype=output_dtype, + use_tvm_ffi=use_tvm_ffi, + ) + + @torch.library.custom_op( + "trtllm::cute_dsl_dspark_fp8_bmm_quantize_rubin_out", + mutates_args=("output", "packed_scale"), + device_types="cuda", + ) + def cute_dsl_dspark_fp8_bmm_quantize_rubin_out( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output: torch.Tensor, + packed_scale: torch.Tensor, + use_tvm_ffi: bool = True, + dspark_max_batch: int = 0, + dspark_rows_per_request: int = 0, + ) -> None: + if get_sm_version() != 107: + raise ValueError("CuTe DSL FP8 BMM quantization requires SM107.") + batch_size, m, _ = input.shape + n = weight.shape[1] + if n % 128 != 0: + raise ValueError(f"Fused FP8 BMM quantization requires N divisible by 128, got {n}.") + expected_scale_numel = pad_up(m, 128) * pad_up(ceil_div(batch_size * n, 32), 4) + if output.shape != (batch_size, m, n): + raise ValueError( + "Fused FP8 BMM output must have shape " + f"{(batch_size, m, n)}, got {tuple(output.shape)}." + ) + if ( + output.dtype != torch.float8_e4m3fn + or packed_scale.dtype != torch.uint8 + or not packed_scale.is_contiguous() + or packed_scale.shape != (expected_scale_numel,) + ): + raise ValueError( + "Fused FP8 BMM requires E4M3 output and a flat contiguous " + f"uint8 R128c4 scale tensor with {expected_scale_numel} elements." + ) + + runner = _get_cute_dsl_dspark_fp8_rubin_bmm_quant_runner(use_tvm_ffi) + inputs = [ + input, + weight, + input_scale, + weight_scale, + output, + packed_scale, + ] + if dspark_max_batch > 0 and dspark_rows_per_request > 0: + tuning_config = _get_dspark_generation_fp8_bmm_quant_tuning_config( + dspark_rows_per_request, + dspark_max_batch, + ) + runner._dspark_generation_tuning_config = tuning_config + else: + tuning_config = getattr( + runner, "_dspark_generation_tuning_config", runner.__class__.tuning_config + ) + custom_op = "trtllm::cute_dsl_dspark_fp8_bmm_quantize_rubin::gemm" + tuner = AutoTuner.get() + _, best_tactic = tuner.choose_one( + custom_op, + [runner], + tuning_config, + inputs, + ) + runner(inputs, tactic=best_tactic) + if tuning_config is not runner.__class__.tuning_config and tuner.is_tuning_mode: + _materialize_dspark_tuning_profile_winners( + custom_op, + runner, + tuning_config, + inputs, + ) + + @torch.library.register_fake("trtllm::cute_dsl_dspark_fp8_bmm_quantize_rubin_out") + def _( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output: torch.Tensor, + packed_scale: torch.Tensor, + use_tvm_ffi: bool = True, + dspark_max_batch: int = 0, + dspark_rows_per_request: int = 0, + ) -> None: + del input_scale, weight_scale, use_tvm_ffi + del dspark_max_batch, dspark_rows_per_request + batch_size, m, _ = input.shape + n = weight.shape[1] + assert output.dtype == torch.float8_e4m3fn + assert output.shape == (batch_size, m, n) + expected_scale_numel = pad_up(m, 128) * pad_up(ceil_div(batch_size * n, 32), 4) + assert packed_scale.dtype == torch.uint8 + assert packed_scale.shape == (expected_scale_numel,) + assert packed_scale.is_contiguous() + + @torch.library.custom_op( + "trtllm::cute_dsl_dspark_context_wkv_gemm_rubin", + mutates_args=(), + device_types="cuda", + ) + def cute_dsl_dspark_context_wkv_gemm_rubin( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + use_tvm_ffi: bool = True, + ) -> torch.Tensor: + """Run stacked DSpark CTX WKV with operator-owned M buckets.""" + if output_dtype != torch.bfloat16: + raise ValueError( + f"DSpark context WKV GEMM only supports bfloat16 output, got {output_dtype}" ) - torch.cuda.synchronize() - except RuntimeError as e: - logger.warning( - "DSV4 DSpark CuTe DSL attention prewarm failed; the op will " - f"self-JIT on first use. {type(e).__name__}: {e}" + if (sm_version := get_sm_version()) != 107: + raise ValueError(f"DSpark context WKV GEMM requires SM107, got SM{sm_version}.") + alpha = _get_cute_dsl_mxfp8_rubin_alpha(input.device) + runner = _get_cute_dsl_dspark_context_wkv_rubin_runner(output_dtype, use_tvm_ffi) + inputs = [input, weight, input_scale, weight_scale, alpha] + _, best_tactic = AutoTuner.get().choose_one( + "trtllm::cute_dsl_dspark_context_wkv_gemm_rubin", + [runner], + runner.__class__.tuning_config, + inputs, ) + return runner(inputs, tactic=best_tactic) + + @torch.library.register_fake("trtllm::cute_dsl_dspark_context_wkv_gemm_rubin") + def _( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype = torch.bfloat16, + use_tvm_ffi: bool = True, + ) -> torch.Tensor: + shape = list(mat_a.shape) + shape[-1] = mat_b.shape[-2] + return mat_a.new_empty(shape, dtype=torch.bfloat16) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 9fcb8ecfa12f..31e8ee1f7275 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -10,25 +10,22 @@ import torch from ..._utils import get_sm_version -from ..attention.kernels.blackwell.dspark_rmsnorm_rope import ( - DSparkRMSNormRoPECacheWriteKernel, +from ..cute_dsl_kernels.blackwell.dspark_rmsnorm_rope import ( + DSparkConcatFP8QuantKernel, + DSparkRMSNormFP8QuantKernel, DSparkRMSNormRoPEDraftBlockKernel, DSparkRMSNormRoPEKernel, + DSparkRMSNormRoPEPageWriteKernel, ) _DSV4_DSPARK_HEAD_DIM = 512 _DSV4_DSPARK_ROPE_DIM = 64 -_DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 -_DSV4_DSPARK_BLOCK_SIZES = (5, 6) + _DSV4_DSPARK_ARCH_BY_SM = { 100: "sm_100", 103: "sm_103", } -# The RMSNorm/RoPE kernels are plain SIMT CuTe kernels with no architecture-specific -# instructions, so they run on more GPUs than the DSV4 DSpark attention kernel gated -# by _DSV4_DSPARK_ARCH_BY_SM. -_DSPARK_RMSNORM_ROPE_SM_VERSIONS = frozenset((100, 103, 107)) @functools.cache @@ -39,14 +36,6 @@ def _get_dspark_arch_str(sm_version: int | None = None) -> str | None: return _DSV4_DSPARK_ARCH_BY_SM.get(sm_version) -@functools.cache -def _is_dspark_rmsnorm_rope_sm_supported(sm_version: int | None = None) -> bool: - """Return whether the fused RMSNorm/RoPE kernels support this GPU.""" - if sm_version is None: - sm_version = get_sm_version() - return sm_version in _DSPARK_RMSNORM_ROPE_SM_VERSIONS - - def _has_regular_row_stride(x: torch.Tensor) -> bool: """Return whether leading dimensions flatten to non-overlapping strided rows.""" if x.is_contiguous(): @@ -117,70 +106,6 @@ def is_fused_dspark_rmsnorm_rope_supported( ) -def _is_dspark_cache_write_supported( - x: torch.Tensor, - weight: torch.Tensor, - freqs: torch.Tensor, - kv_cache: torch.Tensor, - slots: torch.Tensor, - start_pos: torch.Tensor, -) -> bool: - batch = x.shape[0] if x.ndim == 3 else -1 - return ( - x.shape == (batch, 1, _DSV4_DSPARK_HEAD_DIM) - and is_fused_dspark_rmsnorm_rope_supported( - x, weight, freqs, num_heads=1, rope_dim=_DSV4_DSPARK_ROPE_DIM - ) - and kv_cache.is_cuda - and kv_cache.dtype == x.dtype - and kv_cache.ndim == 3 - and kv_cache.shape[1:] == (_DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM) - and kv_cache.stride(1) == _DSV4_DSPARK_HEAD_DIM - and kv_cache.stride(2) == 1 - and slots.is_cuda - and slots.dtype == torch.int64 - and slots.shape == (batch,) - and slots.is_contiguous() - and start_pos.is_cuda - and start_pos.dtype == torch.int64 - and start_pos.shape == (batch,) - and start_pos.is_contiguous() - ) - - -def _is_dspark_draft_block_supported( - x: torch.Tensor, - weight: torch.Tensor, - freqs: torch.Tensor, -) -> bool: - block_size = x.shape[1] if x.ndim == 3 else -1 - return ( - block_size in _DSV4_DSPARK_BLOCK_SIZES - and x.shape[-1:] == (_DSV4_DSPARK_HEAD_DIM,) - and is_fused_dspark_rmsnorm_rope_supported( - x, weight, freqs, num_heads=1, rope_dim=_DSV4_DSPARK_ROPE_DIM - ) - ) - - -def is_fused_dspark_attention_preparation_supported( - main_x: torch.Tensor, - block_x: torch.Tensor, - weight: torch.Tensor, - main_freqs: torch.Tensor, - block_freqs: torch.Tensor, - kv_cache: torch.Tensor, - slots: torch.Tensor, - start_pos: torch.Tensor, -) -> bool: - """Return whether RMSNorm/RoPE can directly prepare the attention inputs.""" - return ( - main_x.shape[0] == block_x.shape[0] - and _is_dspark_cache_write_supported(main_x, weight, main_freqs, kv_cache, slots, start_pos) - and _is_dspark_draft_block_supported(block_x, weight, block_freqs) - ) - - @functools.cache def _compile_fused_dspark_rmsnorm_rope( hidden_dim: int, @@ -293,105 +218,6 @@ def _compile_fused_dspark_rope_into( ) -@functools.cache -def _compile_dspark_rmsnorm_rope_cache_write(eps: float): - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "DSpark RMSNorm/RoPE cache-write must be warmed up before CUDA graph capture" - ) - rows = cute.sym_int() - pages = cute.sym_int() - x_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (rows, _DSV4_DSPARK_HEAD_DIM), stride_order=(1, 0) - ) - weight_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (_DSV4_DSPARK_HEAD_DIM,), stride_order=(0,) - ) - freqs_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (rows, _DSV4_DSPARK_ROPE_DIM // 2, 2), - stride_order=(2, 1, 0), - ) - cache_fake = cute.runtime.make_fake_tensor( - cutlass.BFloat16, - (pages, _DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM), - stride=(cute.sym_int64(), _DSV4_DSPARK_HEAD_DIM, 1), - ) - slots_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int64, (rows,), stride_order=(0,)) - start_pos_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int64, (rows,), stride_order=(0,) - ) - slots_i32_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (rows,), stride_order=(0,) - ) - cache_seqs_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (rows,), stride_order=(0,) - ) - stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - kernel = DSparkRMSNormRoPECacheWriteKernel( - _DSV4_DSPARK_HEAD_DIM, - _DSV4_DSPARK_ROPE_DIM, - eps, - _DSV4_DSPARK_WINDOW_SIZE, - ) - return cute.compile( - kernel, - x_fake, - weight_fake, - freqs_fake, - cache_fake, - slots_fake, - start_pos_fake, - slots_i32_fake, - cache_seqs_fake, - stream_fake, - options="--opt-level 2 --enable-tvm-ffi", - ) - - -@functools.cache -def _compile_dspark_rmsnorm_rope_draft_block(block_size: int, eps: float): - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "DSpark RMSNorm/RoPE draft-block must be warmed up before CUDA graph capture" - ) - rows = cute.sym_int() - batch = cute.sym_int() - x_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (rows, _DSV4_DSPARK_HEAD_DIM), stride_order=(1, 0) - ) - weight_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (_DSV4_DSPARK_HEAD_DIM,), stride_order=(0,) - ) - freqs_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (rows, _DSV4_DSPARK_ROPE_DIM // 2, 2), - stride_order=(2, 1, 0), - ) - output_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, - (batch, _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM), - stride_order=(2, 1, 0), - ) - stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - kernel = DSparkRMSNormRoPEDraftBlockKernel( - _DSV4_DSPARK_HEAD_DIM, - _DSV4_DSPARK_ROPE_DIM, - eps, - block_size, - _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, - ) - return cute.compile( - kernel, - x_fake, - weight_fake, - freqs_fake, - output_fake, - stream_fake, - options="--opt-level 2 --enable-tvm-ffi", - ) - - @torch.library.custom_op( "trtllm::cute_dsl_dspark_rmsnorm_rope", mutates_args=(), @@ -531,51 +357,192 @@ def _( return None +def _validate_dspark_preparation( + x: torch.Tensor, weight: torch.Tensor, freqs: torch.Tensor +) -> None: + if ( + x.ndim != 3 + or x.shape[-1] != _DSV4_DSPARK_HEAD_DIM + or x.dtype != torch.bfloat16 + or weight.dtype != x.dtype + or weight.shape != (_DSV4_DSPARK_HEAD_DIM,) + or freqs.dtype != torch.float32 + or freqs.shape != (x.shape[0] * x.shape[1], _DSV4_DSPARK_ROPE_DIM // 2, 2) + or not all( + t.is_cuda and t.device == x.device and t.is_contiguous() for t in (x, weight, freqs) + ) + ): + raise ValueError( + "DSpark preparation requires contiguous BF16 [batch, tokens, 512] inputs, " + "BF16 [512] weights and FP32 [batch * tokens, 32, 2] frequencies" + ) + + +@functools.cache +def _compile_dspark_rmsnorm_rope_page_write(page_size: int, eps: float, device: int): + # Device identity scopes compilation to the active CUDA target. + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "DSpark RMSNorm/RoPE page-write must be warmed up before CUDA graph capture" + ) + rows = cute.sym_int() + pages = cute.sym_int() + x_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (rows, _DSV4_DSPARK_HEAD_DIM), stride_order=(1, 0) + ) + weight_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (_DSV4_DSPARK_HEAD_DIM,), stride_order=(0,) + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (rows, _DSV4_DSPARK_ROPE_DIM // 2, 2), + stride_order=(2, 1, 0), + ) + cache_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (pages, page_size, _DSV4_DSPARK_HEAD_DIM), + stride=(cute.sym_int64(), _DSV4_DSPARK_HEAD_DIM, 1), + ) + tables_fake = cute.runtime.make_fake_tensor( + cutlass.Int32, (rows, cute.sym_int()), stride=(cute.sym_int64(), 1) + ) + start_pos_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, (rows,), stride_order=(0,) + ) + capacities_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, (rows,), stride_order=(0,) + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkRMSNormRoPEPageWriteKernel( + _DSV4_DSPARK_HEAD_DIM, + _DSV4_DSPARK_ROPE_DIM, + eps, + page_size, + ) + return cute.compile( + kernel, + x_fake, + weight_fake, + freqs_fake, + cache_fake, + tables_fake, + start_pos_fake, + capacities_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + @torch.library.custom_op( - "trtllm::cute_dsl_dspark_rmsnorm_rope_cache_write", + "trtllm::cute_dsl_dspark_rmsnorm_rope_page_write", mutates_args=("kv_cache",), device_types="cuda", ) -def cute_dsl_dspark_rmsnorm_rope_cache_write( +def cute_dsl_dspark_rmsnorm_rope_page_write( x: torch.Tensor, weight: torch.Tensor, freqs: torch.Tensor, kv_cache: torch.Tensor, - slots: torch.Tensor, + block_tables: torch.Tensor, start_pos: torch.Tensor, + capacities: torch.Tensor, eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - """Prepare attention metadata after the caller validates the fused contract.""" - compiled = _compile_dspark_rmsnorm_rope_cache_write(eps) - slots_i32 = torch.empty_like(slots, dtype=torch.int32) - cache_seqs = torch.empty_like(start_pos, dtype=torch.int32) - compiled( - x.view(-1, _DSV4_DSPARK_HEAD_DIM), - weight, - freqs, - kv_cache, - slots, - start_pos, - slots_i32, - cache_seqs, - ) - return slots_i32, cache_seqs +) -> None: + """Normalize/rotate [batch, 1, 512] KV into pages; zero capacities mask dummy rows.""" + _validate_dspark_preparation(x, weight, freqs) + batch = x.shape[0] + if ( + x.shape[1] != 1 + or kv_cache.ndim != 3 + or kv_cache.shape[1] not in (16, 32, 64, 128, 256) + or kv_cache.shape[2] != _DSV4_DSPARK_HEAD_DIM + or kv_cache.dtype != x.dtype + or kv_cache.stride(1) != _DSV4_DSPARK_HEAD_DIM + or kv_cache.stride(2) != 1 + or block_tables.ndim != 2 + or block_tables.shape[0] != batch + or block_tables.dtype != torch.int32 + or block_tables.stride(1) != 1 + or not all( + t.is_cuda and t.device == x.device + for t in (kv_cache, block_tables, start_pos, capacities) + ) + or not all( + t.shape == (batch,) and t.dtype == torch.int64 and t.is_contiguous() + for t in (start_pos, capacities) + ) + ): + raise ValueError( + "DSpark page write requires BF16 pages with contiguous tokens, INT32 " + "page-table rows and contiguous INT64 positions/capacities on x's device" + ) + with torch.cuda.device(x.device): + compiled = _compile_dspark_rmsnorm_rope_page_write(kv_cache.shape[1], eps, x.device.index) + compiled( + x.view(batch, _DSV4_DSPARK_HEAD_DIM), + weight, + freqs, + kv_cache, + block_tables, + start_pos, + capacities, + ) -@torch.library.register_fake("trtllm::cute_dsl_dspark_rmsnorm_rope_cache_write") +@torch.library.register_fake("trtllm::cute_dsl_dspark_rmsnorm_rope_page_write") def _( x: torch.Tensor, weight: torch.Tensor, freqs: torch.Tensor, kv_cache: torch.Tensor, - slots: torch.Tensor, + block_tables: torch.Tensor, start_pos: torch.Tensor, + capacities: torch.Tensor, eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - del x, weight, freqs, kv_cache, eps - return ( - torch.empty_like(slots, dtype=torch.int32), - torch.empty_like(start_pos, dtype=torch.int32), +) -> None: + return None + + +@functools.cache +def _compile_dspark_rmsnorm_rope_draft_block(block_size: int, eps: float, device: int): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "DSpark RMSNorm/RoPE draft-block must be warmed up before CUDA graph capture" + ) + rows = cute.sym_int() + batch = cute.sym_int() + x_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (rows, _DSV4_DSPARK_HEAD_DIM), stride_order=(1, 0) + ) + weight_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (_DSV4_DSPARK_HEAD_DIM,), stride_order=(0,) + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (rows, _DSV4_DSPARK_ROPE_DIM // 2, 2), + stride_order=(2, 1, 0), + ) + output_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, + (batch, _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM), + stride_order=(2, 1, 0), + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkRMSNormRoPEDraftBlockKernel( + _DSV4_DSPARK_HEAD_DIM, + _DSV4_DSPARK_ROPE_DIM, + eps, + block_size, + _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, + ) + return cute.compile( + kernel, + x_fake, + weight_fake, + freqs_fake, + output_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", ) @@ -591,8 +558,12 @@ def cute_dsl_dspark_rmsnorm_rope_draft_block( eps: float, ) -> torch.Tensor: """Prepare a zero-padded draft block after validating the fused contract.""" + _validate_dspark_preparation(x, weight, freqs) block_size = x.shape[1] - compiled = _compile_dspark_rmsnorm_rope_draft_block(block_size, eps) + if block_size not in (5, 6): + raise ValueError("DSpark draft block size must be 5 or 6") + with torch.cuda.device(x.device): + compiled = _compile_dspark_rmsnorm_rope_draft_block(block_size, eps, x.device.index) output = x.new_empty((x.shape[0], _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM)) compiled( x.view(-1, _DSV4_DSPARK_HEAD_DIM), @@ -612,3 +583,186 @@ def _( ) -> torch.Tensor: del weight, freqs, eps return x.new_empty((x.shape[0], _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM)) + + +_dspark_fp8_input_compile_allowed = False + + +def _r128c4_scale_numel(rows: int, hidden_dim: int) -> int: + return ((rows + 127) // 128 * 128) * (hidden_dim // 32) + + +def precompile_dspark_fp8_input_fusions( + qkv_hidden_dim: int | None, + q_lora_rank: int | None, + eps: float, +) -> None: + """Precompile dynamic-row DSpark FP8 input producers during model init.""" + global _dspark_fp8_input_compile_allowed + if get_sm_version() != 107: + return + if _dspark_fp8_input_compile_allowed: + raise RuntimeError("DSpark FP8 input precompile is already in progress") + + _dspark_fp8_input_compile_allowed = True + try: + if qkv_hidden_dim is not None: + _compile_dspark_concat_fp8_quant(qkv_hidden_dim) + if q_lora_rank is not None: + _compile_dspark_rmsnorm_fp8_quant(q_lora_rank, eps) + finally: + _dspark_fp8_input_compile_allowed = False + + +@functools.cache +def _compile_dspark_rmsnorm_fp8_quant(hidden_dim: int, eps: float): + if not _dspark_fp8_input_compile_allowed: + raise RuntimeError( + "DSpark RMSNorm/FP8 quant kernel was not precompiled during worker initialization" + ) + + rows = cute.sym_int() + scale_elems = cute.sym_int() + x_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (rows, hidden_dim), + stride=(cute.sym_int64(), 1), + ) + weight_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (hidden_dim,), stride_order=(0,) + ) + output_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float8E4M3FN, (rows, hidden_dim), stride_order=(1, 0) + ) + scale_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float8E8M0FNU, (scale_elems,), stride_order=(0,) + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkRMSNormFP8QuantKernel(hidden_dim, eps) + return cute.compile( + kernel, + x_fake, + weight_fake, + output_fake, + scale_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + +@functools.cache +def _compile_dspark_concat_fp8_quant(hidden_dim: int): + if not _dspark_fp8_input_compile_allowed: + raise RuntimeError( + "DSpark concat/FP8 quant kernel was not precompiled during worker initialization" + ) + + first_rows = cute.sym_int() + second_rows = cute.sym_int() + output_rows = cute.sym_int() + scale_elems = cute.sym_int() + first_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (first_rows, hidden_dim), + stride=(cute.sym_int64(), 1), + ) + second_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (second_rows, hidden_dim), + stride=(cute.sym_int64(), 1), + ) + output_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float8E4M3FN, (output_rows, hidden_dim), stride_order=(1, 0) + ) + scale_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float8E8M0FNU, (scale_elems,), stride_order=(0,) + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkConcatFP8QuantKernel(hidden_dim) + return cute.compile( + kernel, + first_fake, + second_fake, + output_fake, + scale_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + +@torch.library.custom_op( + "trtllm::cute_dsl_dspark_rmsnorm_fp8_quant", + mutates_args=(), + device_types="cuda", +) +def cute_dsl_dspark_rmsnorm_fp8_quant( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Produce a prequantized MXFP8 Linear input after model-init precompile.""" + hidden_dim = x.shape[-1] + x_flat = x.view(-1, hidden_dim) + output = x.new_empty(x_flat.shape, dtype=torch.float8_e4m3fn) + output_scale = x.new_empty( + (_r128c4_scale_numel(x_flat.shape[0], hidden_dim),), dtype=torch.uint8 + ) + compiled = _compile_dspark_rmsnorm_fp8_quant(hidden_dim, eps) + compiled( + x_flat, + weight, + output, + output_scale.view(torch.float8_e8m0fnu), + ) + return output, output_scale + + +@torch.library.register_fake("trtllm::cute_dsl_dspark_rmsnorm_fp8_quant") +def _( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del weight, eps + rows = x.numel() // x.shape[-1] + return ( + x.new_empty((rows, x.shape[-1]), dtype=torch.float8_e4m3fn), + x.new_empty((_r128c4_scale_numel(rows, x.shape[-1]),), dtype=torch.uint8), + ) + + +@torch.library.custom_op( + "trtllm::cute_dsl_dspark_concat_fp8_quant", + mutates_args=(), + device_types="cuda", +) +def cute_dsl_dspark_concat_fp8_quant( + first: torch.Tensor, + second: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Concatenate row ranges into one prequantized MXFP8 Linear input.""" + rows = first.shape[0] + second.shape[0] + hidden_dim = first.shape[1] + output = first.new_empty((rows, hidden_dim), dtype=torch.float8_e4m3fn) + output_scale = first.new_empty((_r128c4_scale_numel(rows, hidden_dim),), dtype=torch.uint8) + compiled = _compile_dspark_concat_fp8_quant(hidden_dim) + compiled( + first, + second, + output, + output_scale.view(torch.float8_e8m0fnu), + ) + return output, output_scale + + +@torch.library.register_fake("trtllm::cute_dsl_dspark_concat_fp8_quant") +def _( + first: torch.Tensor, + second: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + rows = first.shape[0] + second.shape[0] + hidden_dim = first.shape[1] + return ( + first.new_empty((rows, hidden_dim), dtype=torch.float8_e4m3fn), + first.new_empty((_r128c4_scale_numel(rows, hidden_dim),), dtype=torch.uint8), + ) diff --git a/tensorrt_llm/_torch/disaggregation/native/auxiliary.py b/tensorrt_llm/_torch/disaggregation/native/auxiliary.py index e7a3c411a7a2..04da6a4c56c2 100644 --- a/tensorrt_llm/_torch/disaggregation/native/auxiliary.py +++ b/tensorrt_llm/_torch/disaggregation/native/auxiliary.py @@ -114,6 +114,63 @@ def build_aux_transfer_layout( AuxSlot = namedtuple("AuxSlot", ["id", "buffer"]) +_DRAFT_HISTORY_VERSION = 2 +_DRAFT_HISTORY_FIELDS = 10 +_DRAFT_DTYPE_CODES = {"torch.float16": 1, "torch.bfloat16": 2, "torch.float8_e4m3fn": 3} +_DRAFT_BACKEND_CODES = {"VANILLA": 1, "TRTLLM": 2, "DSv4": 3} + + +def _encode_draft_history(history: dict[str, Any]) -> list[int]: + """Serialize the manager's validated history and rank-local storage identity.""" + if history is None: + raise ValueError("Standalone draft transfer requires draft history metadata") + layout = history["layout"] + return [ + _DRAFT_HISTORY_VERSION, + history["valid_length"], + history["position"], + layout["num_layers"], + layout["num_kv_heads"], + layout["head_dim"], + _DRAFT_DTYPE_CODES[layout["dtype"]], + _DRAFT_BACKEND_CODES[layout["attention_backend"]], + layout["kv_factor"], + layout["window_size"] or 0, + ] + + +def _decode_draft_history(values: list[int]) -> dict[str, Any]: + if values[0] != _DRAFT_HISTORY_VERSION: + raise ValueError("Missing or unsupported standalone draft history metadata version") + ( + _, + valid_length, + position, + num_layers, + num_kv_heads, + head_dim, + dtype_code, + backend_code, + kv_factor, + window_size, + ) = values + dtypes = {code: name for name, code in _DRAFT_DTYPE_CODES.items()} + backends = {code: name for name, code in _DRAFT_BACKEND_CODES.items()} + layout = { + "num_layers": num_layers, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + "dtype": dtypes.get(dtype_code), + "attention_backend": backends.get(backend_code), + "kv_factor": kv_factor, + "window_size": window_size or None, + } + return { + "valid_length": valid_length, + "position": position, + "layout": layout, + } + class AuxBufferBase(ABC): """ @@ -169,8 +226,15 @@ def get_slot_data(self, slot: int) -> tuple[list[int], list[int], tuple[int, int class AuxBuffer(AuxBufferBase): - def __init__(self, max_slot_num: int, beam_width: int, max_draft_len: int, device: str = "cpu"): - # public constructor args remain the same, internals are private + def __init__( + self, + max_slot_num: int, + beam_width: int, + max_draft_len: int, + device: str = "cpu", + *, + draft_history: bool = False, + ) -> None: self._max_slot_num = int(max_slot_num) self._beam_width = int(beam_width) self._max_draft_len = int(max_draft_len) @@ -178,9 +242,6 @@ def __init__(self, max_slot_num: int, beam_width: int, max_draft_len: int, devic self._free_slots = deque(list(range(self._max_slot_num))) self._occupied_slots: set[int] = set() - self._slot_token_counts: dict[ - int, tuple[int, int] - ] = {} # slot -> (first_tokens_len, draft_tokens_len) data_type = torch.int32 self._first_tokens_buffer = torch.empty( @@ -199,35 +260,32 @@ def __init__(self, max_slot_num: int, beam_width: int, max_draft_len: int, devic self._prompt_token_counts_buffer = torch.zeros( self._max_slot_num, 2, dtype=data_type, device=self._device ) + # Version zero denotes an unfilled or recycled slot. + self._draft_history_buffer = ( + torch.zeros( + self._max_slot_num, _DRAFT_HISTORY_FIELDS, dtype=torch.int64, device=self._device + ) + if draft_history + else None + ) + + buffers = [ + self._first_tokens_buffer, + self._draft_tokens_buffer, + self._token_counts_buffer, + self._prompt_token_counts_buffer, + ] + if self._draft_history_buffer is not None: + buffers.append(self._draft_history_buffer) self._meta = AuxBufferMeta( - ptrs=np.array( - [ - self._first_tokens_buffer.data_ptr(), - self._draft_tokens_buffer.data_ptr(), - self._token_counts_buffer.data_ptr(), - self._prompt_token_counts_buffer.data_ptr(), - ], - dtype=np.int64, - ), + ptrs=np.array([buffer.data_ptr() for buffer in buffers], dtype=np.int64), size=np.array( - [ - self._first_tokens_buffer.numel() * self._first_tokens_buffer.element_size(), - self._draft_tokens_buffer.numel() * self._draft_tokens_buffer.element_size(), - self._token_counts_buffer.numel() * self._token_counts_buffer.element_size(), - self._prompt_token_counts_buffer.numel() - * self._prompt_token_counts_buffer.element_size(), - ], + [buffer.numel() * buffer.element_size() for buffer in buffers], dtype=np.int64, ), item_sizes=np.array( - [ - self._first_tokens_buffer[0].numel() * self._first_tokens_buffer.element_size(), - self._draft_tokens_buffer[0].numel() * self._draft_tokens_buffer.element_size(), - self._token_counts_buffer[0].numel() * self._token_counts_buffer.element_size(), - self._prompt_token_counts_buffer[0].numel() - * self._prompt_token_counts_buffer.element_size(), - ], + [buffer[0].numel() * buffer.element_size() for buffer in buffers], dtype=np.int64, ), device=self._device, @@ -247,7 +305,8 @@ def alloc_slot(self) -> AuxSlot: "This indicates a bug in slot management." ) self._occupied_slots.add(slot_id) - self._slot_token_counts[slot_id] = (0, 0) + if self._draft_history_buffer is not None: + self._draft_history_buffer[slot_id].zero_() return AuxSlot(slot_id, self) def free_slot(self, slot: int) -> None: @@ -261,13 +320,17 @@ def free_slot(self, slot: int) -> None: f"Invalid slot id {slot}. Valid slot indices are in the range 0..{self._max_slot_num - 1}." ) self._occupied_slots.remove(slot) - self._slot_token_counts.pop(slot, None) self._free_slots.append(slot) @property def meta(self) -> AuxBufferMeta: return self._meta + @property + def has_draft_history(self) -> bool: + """Whether this buffer transfers standalone drafter history metadata.""" + return self._draft_history_buffer is not None + def fill_slot(self, slot: int, request: LlmRequest) -> None: if slot not in self._occupied_slots: raise ValueError( @@ -295,7 +358,6 @@ def fill_slot(self, slot: int, request: LlmRequest) -> None: self._draft_tokens_buffer[slot][: len(draft_tokens)].copy_( torch.tensor(draft_tokens, dtype=torch.int32, device=self._device) ) - self._slot_token_counts[slot] = (len(first_gen_tokens), len(draft_tokens)) self._token_counts_buffer[slot].copy_( torch.tensor( [len(first_gen_tokens), len(draft_tokens)], dtype=torch.int32, device=self._device @@ -305,6 +367,11 @@ def fill_slot(self, slot: int, request: LlmRequest) -> None: self._prompt_token_counts_buffer[slot].copy_( torch.tensor([prompt_tokens, cached_tokens], dtype=torch.int32, device=self._device) ) + if self._draft_history_buffer is not None: + values = _encode_draft_history(request.py_draft_transfer_history) + self._draft_history_buffer[slot].copy_( + torch.tensor(values, dtype=torch.int64, device=self._device) + ) @staticmethod def _resolve_prompt_token_counts(request: LlmRequest) -> tuple[int, int]: @@ -335,3 +402,11 @@ def get_slot_data(self, slot: int) -> tuple[list[int], list[int], tuple[int, int first_gen_tokens, draft_tokens = self.get_slot_tokens(slot) prompt_tokens, cached_tokens = self._prompt_token_counts_buffer[slot].tolist() return first_gen_tokens, draft_tokens, (int(prompt_tokens), int(cached_tokens)) + + def get_slot_draft_history(self, slot: int) -> dict[str, Any]: + """Decode transferred history for validation against the receiving request/cache.""" + if slot not in self._occupied_slots: + raise ValueError(f"Cannot read slot {slot}: slot is not currently allocated.") + if self._draft_history_buffer is None: + raise ValueError("Standalone draft history transfer is not enabled for this buffer") + return _decode_draft_history(self._draft_history_buffer[slot].tolist()) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index b6bf6e81addc..bcd58da099c1 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -630,8 +630,11 @@ def block_bytes_per_group(page_table: KVCachePageTable) -> list[int | None]: only once. Non-attention groups retain a ``None`` placeholder so the result remains aligned with receive-request layer-group indices. """ - from tensorrt_llm._torch.disaggregation.resource.page import CacheKind - from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool + from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup + from tensorrt_llm._torch.disaggregation.resource.utils import ( + get_physical_pool, + get_pool_view_num_layers, + ) assert page_table is not None out: list[int | None] = [] @@ -639,6 +642,11 @@ def block_bytes_per_group(page_table: KVCachePageTable) -> list[int | None]: if lg.kind != CacheKind.PAGED: out.append(None) continue + if lg.cp_as_tp: + out.append( + sum(get_pool_view_num_layers(pv) * pv.bytes_per_layer for pv in lg.pool_views) + ) + continue pool_indices = {pool_view.pool_idx for pool_view in lg.pool_views} out.append( sum( diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index f141cce93ff4..58689e5d8a74 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -26,8 +26,12 @@ SpecRegionPair, ) from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind -from tensorrt_llm._torch.disaggregation.resource.utils import find_replicated_role_mismatch +from tensorrt_llm._torch.disaggregation.resource.page import ( + AttentionLayerGroup, + KVCachePageTable, + MapperKind, +) +from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_to_layer_group from tensorrt_llm._utils import nvtx_range @@ -558,21 +562,24 @@ def _tpb_check(self, local: int, peer: int, peer_ri: RankInfo) -> bool: @staticmethod def validate_peer_compatible( - self_page_table: Optional[KVCachePageTable], - peer_page_table: Optional[KVCachePageTable], + self_page_table: KVCachePageTable | None, + peer_page_table: KVCachePageTable | None, ) -> None: - """Reject a peer whose PAGED groups declare different replicated roles. - - Pool matching drops a view with no counterpart silently, so an - index-key side cache the peer never declares would leave the receiver - holding zeroed state instead of raising. Raises ``ValueError``. - """ - differing = find_replicated_role_mismatch(self_page_table, peer_page_table, CacheKind.PAGED) - if differing: - raise ValueError( - "AttentionPolicy.validate_peer_compatible: replicated roles differ on " - f"overlapping layers: {differing}" - ) + """Validate standalone draft storage on overlapping layers.""" + if self_page_table is not None and peer_page_table is not None: + local_layers = get_layer_to_layer_group(self_page_table) + peer_layers = get_layer_to_layer_group(peer_page_table) + for layer in local_layers.keys() & peer_layers.keys(): + local = self_page_table.layer_groups[local_layers[layer]] + peer = peer_page_table.layer_groups[peer_layers[layer]] + if not isinstance(local, AttentionLayerGroup) or not isinstance( + peer, AttentionLayerGroup + ): + continue + if local.cp_as_tp != peer.cp_as_tp: + raise ValueError("Attention peers disagree on full-sequence CP-to-TP storage") + if local.cp_as_tp and local.total_kv_head_num != peer.total_kv_head_num: + raise ValueError("Standalone draft peers have different global KV head counts") def check_peer_compatible(self, peer_ri: RankInfo) -> bool: a = self._ri.attention diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index d8d37658195b..647099dc6fab 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -14,8 +14,8 @@ # limitations under the License. from collections import Counter -from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple, Union +from dataclasses import dataclass, field, replace +from typing import Dict, List, Optional, Tuple import numpy as np @@ -167,6 +167,10 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: if not self._get_policy(CacheKind.PAGED).check_peer_compatible(peer_ri): return False + AttentionPolicy.validate_peer_compatible( + self._self_ext_cache.page_table, peer_ri.page_table + ) + # Recurrent-state (Mamba/KDA) layout gate. Raises ValueError with a # field-level diagnostic instead of returning False, so the precise # mismatch reaches the caller of register(). @@ -400,31 +404,13 @@ def get_kv_map( pool_idx=peer_pi, ) - # Polymorphic dispatch: attention vs mamba policy - policy = self._get_policy(self_lg.kind) - - # For mamba, compute ptr-array layer offsets for partial PP overlap. - # extract_slot returns ptrs for ALL local_layer_ids in sorted order; - # the mapper must slice only the overlapping subset. - extra_kwargs = {} - if self_lg.kind == CacheKind.STATE: - # Positions index the ptr array ``extract_slot`` builds, which is - # ordered by view offset. - self_gid_to_pos = {gid: i for i, gid in enumerate(self_global_ids)} - peer_gid_to_pos = {gid: i for i, gid in enumerate(peer_global_ids)} - self_overlap_positions = [self_gid_to_pos[gid] for gid in overlapping_layers] - peer_overlap_positions = [peer_gid_to_pos[gid] for gid in overlapping_layers] - # Under contiguous PP partitioning, overlapping layers form a - # contiguous block in the ptrs array. - extra_kwargs["src_layer_off"] = ( - self_overlap_positions[0] if self_overlap_positions else 0 - ) - extra_kwargs["dst_layer_off"] = ( - peer_overlap_positions[0] if peer_overlap_positions else 0 - ) - - mapper = policy.build_mapper( - peer_ri=peer_ri, + policy = self._attention_policy + mapper_peer_ri = peer_ri + if self_lg.cp_as_tp: + policy = AttentionPolicy(self._group_rank_info(self._ri, self_lg)) + mapper_peer_ri = self._group_rank_info(peer_ri, peer_lg) + mapper = policy.build_kv_mapper( + peer_ri=mapper_peer_ri, mapper_kind=self_pv.mapper_kind, self_layer_offsets=self_layer_offsets, peer_layer_offsets=peer_layer_offsets, @@ -587,10 +573,39 @@ def should_send_pool( """ layer_group = self._self_ext_cache.page_table.layer_groups[layer_group_id] pool_view = layer_group.pool_views[pool_idx] - # Delegate to the policy's should_send. None means fan-in election. - policy = self._get_policy(layer_group.kind) - result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) - return self._owns_tp_fan_in(peer_rank_info) if result is None else result + if layer_group.cp_as_tp: + peer_lg_id, _ = self.get_pool_mapping(peer_rank_info)[(layer_group_id, pool_idx)] + peer_lg = peer_rank_info.page_table.layer_groups[peer_lg_id] + local = self._group_rank_info(self._ri, layer_group) + peer = self._group_rank_info(peer_rank_info, peer_lg) + if local.tp_size <= peer.tp_size: + if peer.tp_rank // (peer.tp_size // local.tp_size) != local.tp_rank: + return False + elif local.tp_rank // (local.tp_size // peer.tp_size) != peer.tp_rank: + return False + duplicate, _ = AttentionPolicy(local).duplicate_head_factors(peer) + return duplicate <= 1 or peer.dp_rank % duplicate == local.tp_rank % duplicate + if pool_view.mapper_kind == MapperKind.REPLICATED: + return self._owns_tp_fan_in(peer_rank_info) + return self.should_send_kv(peer_overlap, peer_rank_info) + + @staticmethod + def _group_rank_info(rank_info: RankInfo, layer_group: AttentionLayerGroup) -> RankInfo: + tp_size, tp_rank = MambaPolicy._mamba_tp(rank_info) + return replace( + rank_info, + tp_size=tp_size, + tp_rank=tp_rank, + cp_size=1, + cp_rank=0, + dp_size=1, + attention=replace( + rank_info.attention, + kv_heads_per_rank=layer_group.kv_head_num_per_rank, + is_mla=False, + enable_attention_dp=False, + ), + ) def should_send_aux(self, peer_rank_info: RankInfo) -> bool: # to ensure the transfer aux is not duplicated diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 29243078fb1c..807c2417a19c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -15,17 +15,16 @@ from __future__ import annotations -import math import os import queue import struct import threading import time import weakref -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, Optional, Union +from typing import TYPE_CHECKING, List, Optional, Union import msgpack import numpy as np @@ -46,7 +45,9 @@ TransferOp, TransferRequest, ) +from tensorrt_llm._torch.disaggregation.base.region import IndexRange, KVRegionSpec from tensorrt_llm._torch.disaggregation.base.transfer import ( + KVSlice, ReceiverBase, RxSessionBase, SenderBase, @@ -55,49 +56,27 @@ TxSessionBase, WaitResult, ) -from tensorrt_llm._torch.disaggregation.lifecycle.retirement import ( - QuiescenceFatalEvent, - RetirementDeadline, - RetirementWatchdog, -) from tensorrt_llm._torch.disaggregation.native.auxiliary import ( AuxBuffer, build_aux_transfer_layout, get_non_empty_aux_indices, ) -from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message -from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy -from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( - MambaPolicy, - mamba_receiver_payload_bytes, -) -from tensorrt_llm._torch.disaggregation.native.ownership import ( - PhysicalOperationState as _PhysicalOperationState, # noqa: F401 -- native diagnostic alias -) -from tensorrt_llm._torch.disaggregation.native.ownership import ( - ReceiveOperationOwner as _ReceiveOperationOwner, -) -from tensorrt_llm._torch.disaggregation.native.ownership import SendOperationOwner -from tensorrt_llm._torch.disaggregation.native.ownership import ( - TransferNotSubmittedError as _TransferNotSubmittedError, -) +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable, MapperKind from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import CUASSERT, nvtx_range from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle if TYPE_CHECKING: - from tensorrt_llm._torch.disaggregation.base.agent import TransferStatus - from .bounce import Config AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType @@ -115,20 +94,65 @@ _FALLBACK_TX_OVERALL_TIMEOUT_S = 60.0 +class _AgentOperationGate: + """Reader/writer gate for NIXL transfers and remote-metadata updates.""" + + def __init__(self): + self._condition = threading.Condition() + self._active_transfers = 0 + self._metadata_active = False + self._metadata_waiters = 0 + + @contextmanager + def transfer(self): + with self._condition: + # Prefer a waiting metadata update so a continuous transfer stream + # cannot starve first-use peer registration. + while self._metadata_active or self._metadata_waiters > 0: + self._condition.wait() + self._active_transfers += 1 + try: + yield + finally: + with self._condition: + self._active_transfers -= 1 + if self._active_transfers == 0: + self._condition.notify_all() + + @contextmanager + def metadata(self): + with self._condition: + self._metadata_waiters += 1 + try: + while self._metadata_active or self._active_transfers > 0: + self._condition.wait() + self._metadata_active = True + finally: + self._metadata_waiters -= 1 + try: + yield + finally: + with self._condition: + self._metadata_active = False + self._condition.notify_all() + + @dataclass class RecvReqInfo: sender_req_id: int instance_name: str instance_rank: int - # Positional pool slots per layer group (see Chunk): entry i is the - # receiver's slot for block ordinal i, or -1 where it accepts nothing. - block_ids_per_layer_groups: list[np.ndarray] + block_ids_per_layer_groups: list[ + np.ndarray + ] # Block IDs per layer group, each np.ndarray(dtype=np.int64) unique_rid: int - # Optional extra lower bound (block-aligned token offset) below which the - # receiver accepts nothing, on top of the -1 holes already in its table. + # Block-aligned token offset where the receiver's block list starts. + # None means "end-of-range suffix" — sender derives it from len(blocks). dst_start_token: Optional[int] = None aux_slot: Optional[int] = None + mamba_state_index: Optional[int] = None slice_id: Optional[int] = None + need_aux: bool = False bounce_dst_base: Optional[int] = None def to_bytes(self) -> bytes: @@ -143,7 +167,9 @@ def to_bytes(self) -> bytes: "unique_rid": self.unique_rid, "dst_start_token": self.dst_start_token, "aux_slot": self.aux_slot, + "mamba_state_index": self.mamba_state_index, "slice_id": self.slice_id, + "need_aux": self.need_aux, "bounce_dst_base": self.bounce_dst_base, } ) @@ -182,11 +208,12 @@ class WriteMeta: sizes: np.ndarray # dtype=np.int64 dst_device_id: Optional[int] = None slice_id: Optional[int] = None - # Peer task index; slice_id identifies the local sender task. - receiver_slice_id: int = 0 is_last_slice: bool = False meta_type: WriteMetaType = WriteMetaType.KV bounce_dst_base: Optional[int] = None + # Set by Sender._enqueue while the work item owns the Tx session's + # queued-or-active lifetime reservation. Worker completion releases it. + owner_session: Optional["TxSession"] = None class MessageType: @@ -197,7 +224,6 @@ class MessageType: REGISTER_RANK_INFO = b"REGISTER_RANK_INFO" AUX_AGENT_RESULT = b"AUX_AGENT_RESULT" CANCEL_SESSION = b"CANCEL_SESSION" - SESSION_QUIESCED = b"SESSION_QUIESCED" class PeerIncompatibleError(ValueError): @@ -217,149 +243,17 @@ class TaskStatus(Enum): ERROR = "ERROR" -@dataclass(frozen=True) -class _LogicalOutcome: - """A committed delivery decision, independent of physical/report progress.""" - - status: SessionStatus - reason: str = "" - by_peer: bool = False - - -class _LogicalOutcomes: - """Arbitrate task completion against session failure/cancellation at event time. - - A failed or cancelled session ends only pending tasks. Already committed - results, including their original cause, never change. No task, request or - backend handle is retained here: memory ownership remains a separate concern. - """ - - def __init__(self, retirement: Optional[RetirementDeadline] = None) -> None: - self._lock = threading.Lock() if retirement is None else retirement.lock - self._retirement = retirement - self._terminal: Optional[_LogicalOutcome] = None - self._results: list[Optional[_LogicalOutcome]] = [] - self._sealed = False - if retirement is not None: - retirement.bind_timeout_outcome(self._timeout_locked) - - def _timeout_locked(self) -> None: - """Commit timeout metadata without entering session, backend or CUDA code.""" - self._end_locked(_LogicalOutcome(SessionStatus.ERROR, reason="KV transfer timed out")) - - def _end_locked(self, outcome: _LogicalOutcome) -> None: - """End pending results once, preserving every previously committed outcome.""" - if self._terminal is None: - self._terminal = outcome - for index, result in enumerate(self._results): - if result is None: - self._results[index] = outcome - - @property - def terminal(self) -> Optional[_LogicalOutcome]: - """Return the stable terminal decision, including an elapsed request timeout.""" - with self._lock: - if self._retirement is not None: - self._retirement.check() - return self._terminal - - def add_task(self) -> int: - with self._lock: - if self._sealed: - raise RuntimeError("transfer session already registered its final piece") - index = len(self._results) - self._results.append(self._terminal) - return index - - def seal(self) -> None: - """Declare that every expected KV and AUX result has been registered.""" - with self._lock: - self._sealed = True - self._complete_session_locked() - - def _complete_session_locked(self) -> None: - if ( - self._retirement is not None - and self._sealed - and self._results - and all( - result is not None and result.status is SessionStatus.TRANSFERRED - for result in self._results - ) - ): - self._retirement.complete_pieces() - - def get(self, index: int) -> Optional[_LogicalOutcome]: - with self._lock: - if self._retirement is not None: - self._retirement.check() - return self._results[index] - - def complete(self, index: int) -> None: - with self._lock: - if self._retirement is not None: - self._retirement.check() - if self._results[index] is None: - self._results[index] = _LogicalOutcome(SessionStatus.TRANSFERRED) - self._complete_session_locked() - - def fail(self, error: Exception) -> None: - self._end(_LogicalOutcome(SessionStatus.ERROR, reason=str(error))) - - def cancel(self, by_peer: bool) -> None: - self._end(_LogicalOutcome(SessionStatus.CANCELLED, by_peer=by_peer)) - - def _end(self, outcome: _LogicalOutcome) -> None: - with self._lock: - if self._retirement is not None: - self._retirement.check() - if self._terminal is not None: - return - self._end_locked(outcome) - if self._retirement is not None: - self._retirement.request_drain(outcome.reason or outcome.status.value) - - -class _LogicalTask: - def __init__(self) -> None: - self._logical_outcomes = _LogicalOutcomes() - self._logical_index = self._logical_outcomes.add_task() - self._retirement: Optional[RetirementDeadline] = None - - def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: - """Join the session before this task is exposed to a worker or caller.""" - self._logical_outcomes = outcomes - self._logical_index = outcomes.add_task() - self._retirement = outcomes._retirement - - @property - def logical_outcome(self) -> Optional[_LogicalOutcome]: - return self._logical_outcomes.get(self._logical_index) - - class AgentResult(Enum): SUCCESS = "SUCCESS" FAILED = "FAILED" - IN_DOUBT = "IN_DOUBT" - FAILED_QUIESCED = "FAILED_QUIESCED" # KV_AGENT_RESULT prefix in one struct frame (was ascii frames serialized/parsed under the # GIL per slice per writer): instance_rank, unique_rid, slice_id, is_last, status, # transfer_size. The optional bounce tail follows at message[2:]. _KV_RESULT_PREFIX = struct.Struct(" list[bytes]: - return [ - MessageType.AUX_AGENT_RESULT, - str(instance_rank).encode("ascii"), - str(unique_rid).encode("ascii"), - result.value.encode("ascii"), - ] +@dataclass +class _QueuedMessage: + endpoint: str + message: list[bytes] + initial_attempt: bool = True -@dataclass -class _PendingSettlement: - write_meta: WriteMeta - initial_report: Optional[list[bytes]] - send_slot_id: Optional[int] = None - backend_done: bool = False - report_error_logged: bool = False +class _ReliableMessageOutbox: + """Single-threaded, at-least-once delivery for terminal control messages. + All ZMQ DEALER creation/use/destruction stays on the outbox thread. New + messages get a small burst of immediate reconnect attempts; persistent + failures are requeued at the tail so one unavailable endpoint cannot starve + other sessions. A successful ZMQ send is considered reliably queued. + """ + + _IMMEDIATE_ATTEMPTS = 2 + _RETRY_INTERVAL_S = 0.05 + + def __init__(self, owner_name: str, *, start: bool = True): + self._owner_name = owner_name + self._queue: queue.Queue[Optional[_QueuedMessage]] = queue.Queue() + self._stop = threading.Event() + self._state_lock = threading.Lock() + self._dealers: dict[str, ZMQMessenger] = {} + self._thread = threading.Thread(target=self._run, daemon=True) + self._started = False + if start: + self.start() + + def start(self) -> None: + with self._state_lock: + if self._started: + return + if self._stop.is_set(): + raise RuntimeError(f"{self._owner_name} message outbox is stopped") + self._started = True + self._thread.start() + + def submit(self, endpoint: str, message: list[bytes]) -> bool: + if not endpoint: + raise ValueError(f"{self._owner_name}: peer endpoint is empty") + with self._state_lock: + if self._stop.is_set(): + logger.error( + f"Cannot queue reliable message for endpoint {endpoint}: " + f"{self._owner_name} outbox is shutting down" + ) + return False + # Copy the multipart list so callers may safely reuse their object. + self._queue.put(_QueuedMessage(endpoint, list(message))) + return True + + def _get_or_connect_dealer(self, endpoint: str) -> ZMQMessenger: + if endpoint not in self._dealers: + self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) + return self._dealers[endpoint] + + def _drop_dealer(self, endpoint: str) -> None: + dealer = self._dealers.pop(endpoint, None) + if dealer is None: + return + try: + dealer.stop() + except Exception as exc: + logger.warning( + f"{self._owner_name} outbox failed to close dealer for endpoint {endpoint}: {exc}" + ) + + def _try_send(self, pending: _QueuedMessage) -> bool: + attempts = self._IMMEDIATE_ATTEMPTS if pending.initial_attempt else 1 + for attempt in range(1, attempts + 1): + try: + self._get_or_connect_dealer(pending.endpoint).send(pending.message) + return True + except Exception as exc: + self._drop_dealer(pending.endpoint) + logger.warning( + f"{self._owner_name} reliable send failed for endpoint " + f"{pending.endpoint} (attempt {attempt}/{attempts}): {exc}" + ) + return False + + def process_once(self) -> bool: + """Process one queued item; return False for the shutdown sentinel.""" + pending = self._queue.get() + if pending is None: + return False + if not self._try_send(pending): + pending.initial_attempt = False + if not self._stop.is_set(): + self._queue.put(pending) + self._stop.wait(self._RETRY_INTERVAL_S) + return True -class SendTaskBase(_LogicalTask): + def _run(self) -> None: + try: + while self.process_once(): + pass + finally: + for endpoint in list(self._dealers): + self._drop_dealer(endpoint) + + def shutdown(self, timeout: float = 5.0) -> None: + with self._state_lock: + if self._stop.is_set(): + return + self._stop.set() + if not self._started: + return + # FIFO ordering gives every already-queued message one final send + # attempt before the sentinel terminates the worker. + self._queue.put(None) + self._thread.join(timeout=timeout) + if self._thread.is_alive(): + logger.warning( + f"{self._owner_name} reliable-message outbox did not terminate within {timeout:g}s" + ) + + +class SendTaskBase: def __init__(self, params: DisaggregatedParams): - super().__init__() self.status = TaskStatus.INIT self._event = threading.Event() self._exception: Optional[Exception] = None @@ -412,27 +407,13 @@ def __init__(self, params: DisaggregatedParams): self._params = params self._unique_rid: Optional[int] = params.disagg_request_id self._perf_timer = PerfTimer() if perf_log_manager.enabled else None - self._physical_owner = SendOperationOwner() - # Preserve the task's diagnostic view of its owner's exact roots. - self._physical_operations = self._physical_owner._physical_operations - - def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: - """Bind session outcome and deadline before exposing the task to workers. - - Args: - outcomes: The session's stable logical-result arbiter. - """ - super().bind_logical_outcomes(outcomes) - self._physical_owner.bind_retirement(self._retirement) def fail(self, exc: Exception) -> None: - self._logical_outcomes.fail(exc) self._exception = exc self.status = TaskStatus.ERROR self._event.set() def complete(self) -> None: - self._logical_outcomes.complete(self._logical_index) self.status = TaskStatus.TRANSFERRED self._event.set() @@ -444,43 +425,6 @@ def wait(self, timeout: Optional[float] = None) -> bool: def is_done(self) -> bool: return self._event.is_set() - def begin_physical_operation(self, peer_rank: int) -> bool: - """Admit a participant once in the native physical owner.""" - return self._physical_owner.begin_physical_operation(peer_rank) - - def begin_backend_submission(self, peer_rank: int, request: TransferRequest) -> None: - """Root the request before the native backend submission escapes.""" - self._physical_owner.begin_backend_submission(peer_rank, request) - - def record_backend_submission(self, peer_rank: int, status: TransferStatus) -> None: - """Retain the exact status returned by the native backend.""" - self._physical_owner.record_backend_submission(peer_rank, status) - - def mark_physical_operation_in_doubt(self, peer_rank: int) -> None: - """Retain ambiguous access independently of the logical outcome.""" - self._physical_owner.mark_physical_operation_in_doubt(peer_rank) - - def retire_unsubmitted_physical_operation(self, peer_rank: int) -> None: - """Record that a participant never submitted a backend operation.""" - self._physical_owner.retire_unsubmitted_physical_operation(peer_rank) - - def retire_backend_done_physical_operation(self, peer_rank: int) -> bool: - """Settle backend-defined completion through the common deadline arbiter.""" - return self._physical_owner.retire_backend_done_physical_operation(peer_rank) - - def poll_in_doubt_physical_operation(self, peer_rank: int) -> bool: - """Query and settle only the exact retained ambiguous operation.""" - return self._physical_owner.poll_in_doubt_physical_operation(peer_rank) - - def has_started_physical_operation(self, peer_rank: int) -> bool: - """Whether this task has admitted the participant.""" - return self._physical_owner.has_started_physical_operation(peer_rank) - - @property - def resources_drained(self) -> bool: - """Whether operation evidence and the session deadline permit retirement.""" - return self._physical_owner.resources_drained - def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int): if self._perf_timer is None: return @@ -504,35 +448,18 @@ def __init__(self, params: DisaggregatedParams, slot: Optional[int]): class KVSendTask(SendTaskBase): def __init__( self, - chunk: Chunk, + kv_slice: KVSlice, params: DisaggregatedParams, slice_id: int, prompt_len: Optional[int] = None, + beam_width: int = 1, ): super().__init__(params) self.slice_id = slice_id self.transferred_count = 0 - self.expected_transfers = 0 - self._chunk = chunk + self._slice = kv_slice self._prompt_len = prompt_len - - @property - def reports_outstanding(self) -> bool: - """Whether any of this piece's writes has still to finish. - - One peer failing ends the task while the writes to its siblings run on, so the task's own - status cannot answer this. Zero expected means nothing was ever submitted. - """ - return self.transferred_count < self.expected_transfers - - -@dataclass(frozen=True) -class _SessionQuiescence: - """A FIFO marker for one cancelled request's no-future-access acknowledgment.""" - - unique_rid: int - peer_rank: int - endpoint: Optional[str] + self._beam_width = beam_width class Sender(SenderBase): @@ -541,21 +468,20 @@ class Sender(SenderBase): # RecvReqInfo that never gets consumed. Entries older than this # are evicted during periodic sweeps. _STALE_REQ_INFO_TTL_S = 120.0 + # Bound synchronous worker retries before handing delivery to the shared + # reliable outbox. + _AGENT_RESULT_IMMEDIATE_ATTEMPTS = 2 def __init__( self, peer_registrar: PeerRegistrar, agent: BaseTransferAgent, bounce=None, - enforce_physical_ownership: bool = False, - retirement_watchdog: Optional[RetirementWatchdog] = None, - ) -> None: + ): self._registrar = peer_registrar self._device_id = peer_registrar.self_rank_info.device_id self._agent = agent self._bounce = bounce - self._enforce_physical_ownership = enforce_physical_ownership - self._retirement_watchdog = retirement_watchdog self._peer_requests: dict = {} self._peer_requests_timestamps: dict[int, float] = {} # unique_rid -> insert time self._peer_requests_lock = threading.Lock() @@ -563,33 +489,31 @@ def __init__( self._dealers = {} # used by listener thread only (single-threaded path) self._thread_local = threading.local() # per-thread DEALER cache for worker threads self._sessions = {} # unique_rid -> TxSession - self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids - # unique_rid -> whether the peer asked, for cancels that arrive before the session exists. - # Who asked is not recoverable from the id alone, and the two differ: a peer's cancel is a - # transfer error, our own side's is an ordinary end. - # Ownership mode also fences existing sessions: keep one entry per - # cancelled immutable ID until Sender teardown, not merely metadata TTL. - self._pre_cancelled_rids: dict[int, bool] = {} - self._pending_session_quiescence: dict[tuple[int, int], _SessionQuiescence] = {} + self._sessions_lock = threading.Lock() # Protects sessions and cancel tombstones + self._pre_cancelled_rids: set[int] = set() + self._pre_cancelled_timestamps: dict[int, float] = {} + # A pre-cancelled sender may receive duplicate/late REQUEST_DATA. + # Remember the terminal replies already emitted so each receiver-side + # peer reservation is resolved at most once by this path. + self._cancel_failed_result_keys: set[tuple[int, str, int, int, bool]] = set() self._shutdown = False - self._shutdown_requested = False self._instance_rank = self._registrar.self_rank_info.instance_rank # Guards concurrent add() from the listener thread. self._loaded_remote_agents: set[str] = set() self._loaded_remote_agents_lock = threading.Lock() - self._ownership_poisoned: Optional[Exception] = None - self._ownership_poison_lock = ( - threading.Lock() if retirement_watchdog is None else retirement_watchdog.lock - ) + # NIXL remote-metadata mutation is not safe while getXferStatus() is + # polling an active request. The C++ wrapper's shared lock covers + # submit_transfer_requests(), but the returned status outlives that + # lock. Serialize lazy peer loads with the complete submit+wait + # interval so a first request from another peer cannot invalidate an + # unrelated transfer already in flight. + self._agent_operation_gate = _AgentOperationGate() self._num_threads = KV_TRANSFER_NUM_THREADS self._send_task_queues: List[queue.Queue] = [ queue.Queue() for _ in range(self._num_threads) ] - # Each dictionary belongs to its worker, preserving that peer stream's - # result ordering and ZMQ socket affinity even after the first failure. - self._pending_settlements: list[dict[tuple[SendTaskBase, int], _PendingSettlement]] = [ - {} for _ in range(self._num_threads) - ] + # Agent results and CANCEL_SESSION share one reliable, fair outbox. + self._reliable_message_outbox = _ReliableMessageOutbox("Sender") self._worker_threads: List[threading.Thread] = [ threading.Thread(target=self._process_task_queue, args=(i,), daemon=True) for i in range(self._num_threads) @@ -607,6 +531,25 @@ def __init__( def endpoint(self): return self._messenger.endpoint + def _load_remote_agent_serialized(self, name: str, agent_desc: bytes) -> None: + with self._agent_operation_gate.metadata(): + self._agent.load_remote_agent(name, agent_desc) + + def _submit_and_wait_agent_transfer(self, request: TransferRequest): + with self._agent_operation_gate.transfer(): + status = self._agent.submit_transfer_requests(request) + try: + completed = status.wait() + last_status = None + if not completed: + last_status = getattr(status, "last_status_str", lambda: "")() + finally: + # The pybind status destructor calls releaseXferReq(). Drop + # the final local reference before releasing the reader gate, + # so a lazy remote-metadata load cannot race that NIXL call. + del status + return completed, last_status + def _add_req_info(self, unique_rid: int, instance_rank: int, req_info: RecvReqInfo): with self._peer_requests_lock: if unique_rid not in self._peer_requests: @@ -638,69 +581,48 @@ def _remove_req_info(self, unique_rid: int): self._peer_requests_timestamps.pop(unique_rid, None) def sweep_stale_req_infos(self): - """Evict RecvReqInfo entries that have no matching TxSession and exceed the TTL. + """Evict stale request infos and cancel tombstones with no TxSession. Called opportunistically from the listener thread when a new REQUEST_DATA arrives. With gen-first ADP broadcast, non-assigned DP ranks accumulate entries that are never consumed; this sweep prevents unbounded growth. """ now = time.monotonic() - with self._peer_requests_lock: - stale_rids = [ + with self._sessions_lock, self._peer_requests_lock: + stale_req_rids = [ rid for rid, ts in self._peer_requests_timestamps.items() if now - ts > self._STALE_REQ_INFO_TTL_S ] - if not stale_rids: - return - for rid in stale_rids: - with self._sessions_lock, self._peer_requests_lock: + stale_cancel_rids = [ + rid + for rid, ts in self._pre_cancelled_timestamps.items() + if now - ts > self._STALE_REQ_INFO_TTL_S + ] + for rid in stale_req_rids: if rid not in self._sessions and rid in self._peer_requests: self._peer_requests.pop(rid, None) self._peer_requests_timestamps.pop(rid, None) logger.debug(f"Swept stale RecvReqInfo for rid={rid}") + for rid in stale_cancel_rids: + if rid in self._sessions: + continue + self._pre_cancelled_rids.discard(rid) + self._pre_cancelled_timestamps.pop(rid, None) + self._cancel_failed_result_keys = { + key for key in self._cancel_failed_result_keys if key[0] != rid + } + logger.debug(f"Swept stale pre-cancel tombstone for rid={rid}") def setup_session(self, tx_session: "TxSession"): unique_rid = tx_session.disagg_request_id pre_cancel = False - cancelled_by_peer = False with self._sessions_lock: - if self._shutdown_requested: - raise RuntimeError("cannot create a TxSession after Sender shutdown") - if self._enforce_physical_ownership and unique_rid in self._sessions: - raise RuntimeError("cannot replace an owned TxSession") + self._sessions[unique_rid] = weakref.ref(tx_session) if unique_rid in self._pre_cancelled_rids: pre_cancel = True - cancelled_by_peer = self._pre_cancelled_rids[unique_rid] - if not self._enforce_physical_ownership: - self._pre_cancelled_rids.pop(unique_rid) - if not (pre_cancel and self._enforce_physical_ownership): - self._sessions[unique_rid] = ( - tx_session if self._enforce_physical_ownership else weakref.ref(tx_session) - ) if pre_cancel: - if self._enforce_physical_ownership: - tx_session.cancel_local(by_peer=cancelled_by_peer, report_unsubmitted_aux=False) - # Make the cancelled session visible only after its terminal - # state and the previously saved requests are captured. Late - # creation cannot reopen admission while the listener - # acknowledges the generation-first cancellation. - with self._sessions_lock: - if self._shutdown_requested: - raise RuntimeError("cannot create a TxSession after Sender shutdown") - if unique_rid in self._sessions: - raise RuntimeError("cannot replace an owned TxSession") - # Keep the fence for this immutable request ID after - # session cleanup, so delayed creation cannot reopen it. - with self._peer_requests_lock: - req_infos = list(self._peer_requests.get(unique_rid, {}).values()) - self._sessions[unique_rid] = tx_session - for info in req_infos: - self._queue_session_quiescence( - unique_rid, info.instance_name, info.instance_rank - ) - else: - tx_session.cancel(by_peer=cancelled_by_peer) + tx_session.cancel() return req_info = self._get_first_req_info(unique_rid) @@ -716,43 +638,37 @@ def setup_session(self, tx_session: "TxSession"): return def _get_session(self, unique_rid: Optional[int]) -> Optional["TxSession"]: - session_entry = self._sessions.get(unique_rid) - if session_entry is None: - return None - if isinstance(session_entry, weakref.ReferenceType): - session = session_entry() - if session is None: - logger.warning(f"TxSession {unique_rid} has been garbage collected") - return session - return session_entry - - def _begin_task_operation(self, task: SendTaskBase, peer_rank: int) -> Optional[bool]: - """Return True for admission, None for a duplicate, or False for rejection.""" - if not task.begin_physical_operation(peer_rank): + session_ref = self._sessions.get(unique_rid) + if session_ref is None: return None - with self._ownership_poison_lock: - if self._ownership_poisoned is not None: - task.retire_unsubmitted_physical_operation(peer_rank) - return False - with self._sessions_lock: - session = self._get_session(task._unique_rid) + session = session_ref() if session is None: - task.retire_unsubmitted_physical_operation(peer_rank) - return False - with session.lock: - if session._closed or session.has_failed() or task.is_done: - task.retire_unsubmitted_physical_operation(peer_rank) - return False - if task.status == TaskStatus.INIT: - task.status = TaskStatus.TRANSFERRING - return True + logger.warning(f"TxSession {unique_rid} has been garbage collected") + return None + return session def _enqueue(self, write_meta: WriteMeta): # Route by (unique_rid, peer_rank) so that: # - Same peer's slices stay ordered on one thread (is_last_slice correctness) # - Different peers can run on different threads (better load balancing) thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads - self._send_task_queues[thread_idx].put(write_meta) + with self._sessions_lock: + session = self._get_session(write_meta.unique_rid) + if session is not None: + with session.lock: + # A terminal/closed session no longer exposes source buffers. + # Still enqueue the metadata so the worker sends FAILED and + # resolves the receiver's already-dispatched peer reservation. + if session._terminal_status is None and not session._closed: + session._begin_nixl_op_locked() + write_meta.owner_session = session + try: + self._send_task_queues[thread_idx].put(write_meta) + except Exception: + if write_meta.owner_session is not None: + write_meta.owner_session._finish_nixl_op() + write_meta.owner_session = None + raise def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger: """Get or create a per-thread DEALER socket via threading.local(). @@ -768,63 +684,52 @@ def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) return dealers[endpoint] - def _get_result_dealer( - self, endpoint: Optional[str], *, legacy_shared: bool = False - ) -> ZMQMessenger: - """Keep legacy early-failure routing outside the ownership bridge.""" - if legacy_shared and not self._enforce_physical_ownership: - return self._get_or_connect_dealer(endpoint) - return self._get_or_connect_thread_dealer(endpoint) + def _drop_thread_dealer(self, endpoint: str) -> None: + """Discard the current thread's failed socket before reconnecting.""" + dealers = getattr(self._thread_local, "dealers", None) + if not dealers: + return + dealer = dealers.pop(endpoint, None) + if dealer is None: + return + try: + dealer.stop() + except Exception as exc: + logger.warning(f"Failed to close agent-result dealer for endpoint {endpoint}: {exc}") - def _submit_transfer( + def _try_send_agent_result( self, - task: SendTaskBase, - peer_rank: int, - request: TransferRequest, - ) -> tuple[bool, Optional[str]]: - if not self._enforce_physical_ownership: - status = self._agent.submit_transfer_requests(request) - completed = status.wait() - detail = None if completed else getattr(status, "last_status_str", lambda: None)() - return completed, detail - with self._ownership_poison_lock: - if self._ownership_poisoned is not None: - error = self._ownership_poisoned - task.retire_unsubmitted_physical_operation(peer_rank) - raise _TransferNotSubmittedError( - "NIXL sender rejected transfer before backend submission" - ) from error - task.begin_backend_submission(peer_rank, request) - try: - # SUBMITTING already retains the request. Neither backend admission - # nor completion may hold the lock needed by deadline containment. - status = self._agent.submit_transfer_requests(request) - task.record_backend_submission(peer_rank, status) - except Exception as error: - with self._ownership_poison_lock: - self._ownership_poisoned = error - task.mark_physical_operation_in_doubt(peer_rank) - return False, str(error) - try: - if not status.wait(): - # A non-success query is a logical transfer failure, but the - # current binding does not expose a backend-defined proof that - # the request handle has released every memory accessor. Treat - # it as physically unresolved until that contract is available. - detail = getattr(status, "last_status_str", lambda: "")() - error = RuntimeError(f"NIXL transfer outcome is ambiguous: {detail}") - with self._ownership_poison_lock: - self._ownership_poisoned = error - task.mark_physical_operation_in_doubt(peer_rank) - return False, detail - except Exception as error: - with self._ownership_poison_lock: - self._ownership_poisoned = error - task.mark_physical_operation_in_doubt(peer_rank) - return False, str(error) - if not task.retire_backend_done_physical_operation(peer_rank): - return False, "quiescence deadline expired before backend completion" - return True, None + endpoint: str, + message: list[bytes], + *, + attempts: int, + ) -> bool: + """Try an agent-result send, rebuilding the thread-local socket on error.""" + for attempt in range(1, attempts + 1): + try: + self._get_or_connect_thread_dealer(endpoint).send(message) + return True + except Exception as exc: + self._drop_thread_dealer(endpoint) + logger.warning( + f"Agent-result send failed for endpoint {endpoint} " + f"(attempt {attempt}/{attempts}): {exc}" + ) + return False + + def _send_agent_result(self, endpoint: str, message: list[bytes]) -> bool: + """Send on the worker fast path, queueing only transport failures. + + Delivery is at-least-once. Receiver result processing is idempotent, so + retrying after an ambiguous transport failure is safe. + """ + if self._try_send_agent_result( + endpoint, + message, + attempts=self._AGENT_RESULT_IMMEDIATE_ATTEMPTS, + ): + return True + return self._reliable_message_outbox.submit(endpoint, message) def _process_task_queue(self, thread_idx: int): device_id = self._device_id @@ -834,42 +739,13 @@ def _process_task_queue(self, thread_idx: int): task_queue = self._send_task_queues[thread_idx] try: while True: - if self._pending_settlements[thread_idx]: - self._poll_in_doubt_transfers(thread_idx) - if any( - pending.initial_report is not None - for pending in self._pending_settlements[thread_idx].values() - ): - # A later queued FAILED may describe only an unsubmitted - # chunk. It must never overtake the earlier IN_DOUBT. - time.sleep(0.01) - continue - try: - write_meta = task_queue.get(timeout=0.01) - except queue.Empty: - continue - else: - write_meta = task_queue.get() + write_meta = task_queue.get() if write_meta is None: break - if isinstance(write_meta, _SessionQuiescence): - if not self._send_session_quiesced(write_meta): - task_queue.put(write_meta) - time.sleep(0.01) - continue - if isinstance(write_meta, tuple): - endpoint, message = write_meta - try: - self._get_or_connect_thread_dealer(endpoint).send(message) - except Exception as e: - logger.warning(f"failed to send transfer rejection: {e}") - continue try: if write_meta.meta_type == WriteMetaType.AUX: logger.debug( - "_process_task_queue[%d]: delivering aux task to agent: %s", - thread_idx, - write_meta, + f"_process_task_queue[{thread_idx}]: delivering aux task to agent: {write_meta}" ) self._deliver_aux_to_agent(write_meta) else: @@ -880,6 +756,10 @@ def _process_task_queue(self, thread_idx: int): f"unique_rid={write_meta.unique_rid}: {e}" ) write_meta.task.fail(e) + finally: + if write_meta.owner_session is not None: + write_meta.owner_session._finish_nixl_op() + write_meta.owner_session = None finally: # Clean up this thread's DEALER sockets. threading.local storage # is only accessible from the owning thread, so shutdown must @@ -896,69 +776,6 @@ def _process_task_queue(self, thread_idx: int): ) dealers.clear() - def _retain_in_doubt_transfer( - self, - write_meta: WriteMeta, - initial_report: list[bytes], - send_slot_id: Optional[int] = None, - ) -> None: - thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads - key = (write_meta.task, write_meta.peer_rank) - if key in self._pending_settlements[thread_idx]: - return - pending = _PendingSettlement(write_meta, initial_report, send_slot_id) - self._pending_settlements[thread_idx][key] = pending - try: - self._get_result_dealer(write_meta.peer_endpoint).send(initial_report) - except Exception as error: - logger.warning(f"Failed to report ambiguous transfer; retaining evidence: {error}") - pending.report_error_logged = True - else: - pending.initial_report = None - - def _poll_in_doubt_transfers(self, thread_idx: int) -> None: - pending_transfers = self._pending_settlements[thread_idx] - for key, pending in list(pending_transfers.items()): - meta = pending.write_meta - try: - dealer = self._get_result_dealer(meta.peer_endpoint) - if pending.initial_report is not None: - dealer.send(pending.initial_report) - pending.initial_report = None - if not pending.backend_done: - if not meta.task.poll_in_doubt_physical_operation(meta.peer_rank): - continue - pending.backend_done = True - if pending.send_slot_id is not None: - self._bounce.release_send(pending.send_slot_id) - pending.send_slot_id = None - if meta.meta_type == WriteMetaType.AUX: - message = _make_aux_result_msg( - self._instance_rank, meta.unique_rid, AgentResult.FAILED_QUIESCED - ) - else: - message = _make_kv_result_msg( - self._instance_rank, - meta.unique_rid, - meta.receiver_slice_id, - True, - AgentResult.FAILED_QUIESCED, - ) - dealer.send(message) - except Exception as error: - # Sending may have escaped before raising. A duplicate settlement - # is idempotent; losing the only remaining report is not safe. - if not pending.report_error_logged: - logger.warning(f"Failed to report physical settlement; will retry: {error}") - pending.report_error_logged = True - continue - with meta.task.lock: - if meta.meta_type == WriteMetaType.AUX: - meta.task._transfer_count += 1 - else: - meta.task.transferred_count += 1 - del pending_transfers[key] - @staticmethod @nvtx_range("_make_agent_request") def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferRequest": @@ -1004,7 +821,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"WriteMeta ptr/size mismatch for unique_rid={write_meta.unique_rid}" ) - owned = self._enforce_physical_ownership + assert write_meta.slice_id is not None with self._sessions_lock: session = self._get_session(write_meta.unique_rid) if session is None: @@ -1012,11 +829,21 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"_deliver_kv_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" ) logger.error(msg) - if owned: - write_meta.task.retire_unsubmitted_physical_operation(write_meta.peer_rank) write_meta.task.fail(RuntimeError(msg)) + # REQUEST_DATA was already accepted by this sender. Resolve the + # receiver reservation even though cancellation/cleanup removed + # the source session before this queued item reached a worker. + self._send_agent_result( + write_meta.peer_endpoint, + _make_kv_result_msg( + self._instance_rank, + write_meta.unique_rid, + write_meta.slice_id, + write_meta.is_last_slice, + AgentResult.FAILED, + ), + ) return - assert write_meta.slice_id is not None task = session.kv_tasks[write_meta.slice_id] timer = task._perf_timer if timer: @@ -1024,13 +851,13 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): # Hold session.lock to serialize the INIT→TRANSFERRING transition with # cancel(): prevents cancel_request() from freeing KV pages while a # worker is about to write into them. + active_nixl_op = write_meta.src_ptrs.size > 0 with session.lock: status = session.status if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): should_abort = True else: - if not owned: - task.status = TaskStatus.TRANSFERRING + task.status = TaskStatus.TRANSFERRING should_abort = False if should_abort: @@ -1040,19 +867,18 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): ) # Task may have been enqueued after cancel() already iterated kv_tasks, # so its future was never set by cancel(). Set it here as a fallback. - if owned: - task.retire_unsubmitted_physical_operation(write_meta.peer_rank) task.fail( RuntimeError(f"session {write_meta.unique_rid} {status.value}, transfer aborted") ) - self._get_result_dealer(write_meta.peer_endpoint, legacy_shared=True).send( + self._send_agent_result( + write_meta.peer_endpoint, _make_kv_result_msg( self._instance_rank, write_meta.unique_rid, - write_meta.receiver_slice_id, + write_meta.slice_id, True, # is_last_slice — ensures receiver resolves its task future AgentResult.FAILED, - ) + ), ) return @@ -1060,46 +886,18 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): agent_result = AgentResult.SUCCESS send_slot_id = None - submitted_and_completed = False - if write_meta.src_ptrs.size > 0: - try: + try: + if active_nixl_op: request, send_slot_id = build_send_request( self._bounce, write_meta, lambda: Sender._make_agent_request(write_meta, device_id=self._device_id), ) - except Exception as e: - # Don't let a gather fault escape: without a result the receiver would hang and its - # region leak. Tell the receiver it failed and fail the local task instead. - logger.error( - f"_deliver_kv_to_agent: failed to build the KV send request for " - f"{write_meta.unique_rid} slice={write_meta.slice_id}: {e}" - ) - if owned: - task.retire_unsubmitted_physical_operation(write_meta.peer_rank) - task.fail(RuntimeError(f"build_send_request failed: {e}")) - self._get_result_dealer(write_meta.peer_endpoint, legacy_shared=True).send( - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.receiver_slice_id, - True, # is_last_slice — ensures receiver resolves its task future - AgentResult.FAILED, - ) - ) - return - if timer: - timer.record_transfer_start(write_meta.peer_rank) - transfer_finished = False - try: - transfer_finished, last_status = self._submit_transfer( - task, write_meta.peer_rank, request - ) - if transfer_finished: - submitted_and_completed = True - del request - if not transfer_finished: - agent_result = AgentResult.IN_DOUBT if owned else AgentResult.FAILED + if timer: + timer.record_transfer_start(write_meta.peer_rank) + completed, last_status = self._submit_and_wait_agent_transfer(request) + if not completed: + agent_result = AgentResult.FAILED agent_name = getattr(self._agent, "name", "") detail = ( f"KV transfer agent failed: " @@ -1114,53 +912,47 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"nixl_status={last_status} agent={agent_name}" ) logger.error(detail) - error = RuntimeError(detail) - task.fail(error) - except _TransferNotSubmittedError as error: - # Backend admission failed, so local resources are safe to retire. - transfer_finished = True - agent_result = AgentResult.FAILED - task.fail(error) - finally: - if send_slot_id is not None and (not owned or transfer_finished): - self._bounce.release_send(send_slot_id) - elif owned: - task.retire_unsubmitted_physical_operation(write_meta.peer_rank) - if timer: - timer.record_transfer_end(write_meta.peer_rank) + task.fail(RuntimeError(detail)) + except Exception as exc: + agent_result = AgentResult.FAILED + task.fail(exc) + logger.error( + f"KV transfer raised for rid={write_meta.unique_rid}, " + f"peer_rank={write_meta.peer_rank}: {exc}" + ) + finally: + if send_slot_id is not None: + self._bounce.release_send(send_slot_id) + if timer: + try: + timer.record_transfer_end(write_meta.peer_rank) + except Exception as exc: + logger.warning( + f"Failed to record KV transfer end for rid={write_meta.unique_rid}: {exc}" + ) - # Report every chunk so failures reach the receiver immediately. + ## TODO: just last slice need to send task state? tail = ( encode_result_tail(write_meta) if send_slot_id is not None and agent_result == AgentResult.SUCCESS else None ) - # Attested written bytes: only what was actually submitted to the - # agent AND completed. The receiver checks the per-slice sum against - # the byte total of its published destinations before admitting the - # transferred range to KV block reuse (verified-range admission), so - # this must reflect the completed submission, never intent — gate on - # submitted_and_completed, not on agent_result, so a SUCCESS message - # that did not come from a completed submit attests zero bytes. - # Matches the perf timer's transfer-size accounting (kv_sizes.sum() - # recorded in _build_kv_write_meta) but is populated even with perf - # logging off. - transfer_size = 0 - if submitted_and_completed and agent_result == AgentResult.SUCCESS: - transfer_size = int(write_meta.sizes.sum()) + transfer_size = timer.get_transfer_size(write_meta.peer_rank) if timer else 0 result_msg = _make_kv_result_msg( self._instance_rank, write_meta.unique_rid, - write_meta.receiver_slice_id, + write_meta.slice_id, write_meta.is_last_slice, agent_result, transfer_size=transfer_size, tail=tail, ) - if agent_result == AgentResult.IN_DOUBT: - self._retain_in_doubt_transfer(write_meta, result_msg, send_slot_id) - return - self._get_result_dealer(write_meta.peer_endpoint).send(result_msg) + self._send_agent_result(write_meta.peer_endpoint, result_msg) + + if timer: + timer.record_task_end(write_meta.peer_rank) + ri = self._registrar.self_rank_info + task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) with task.lock: task.transferred_count += 1 @@ -1181,17 +973,6 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): if all(t.status == TaskStatus.TRANSFERRED for t in session.kv_tasks): session.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() - # Diagnostics must not delay authoritative completion or turn it into failure. - try: - if timer: - timer.record_task_end(write_meta.peer_rank) - ri = self._registrar.self_rank_info - task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) - except Exception as error: - logger.warning( - f"KV send perf logging failed for request {write_meta.unique_rid}: {error}" - ) - logger.debug( f"deliver_kv_to_agent completed: unique_rid={write_meta.unique_rid}, " f"slice_id={write_meta.slice_id}, agent_result={agent_result}" @@ -1199,14 +980,20 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): @nvtx_range("_deliver_aux_to_agent") def _deliver_aux_to_agent(self, write_meta: WriteMeta): - owned = self._enforce_physical_ownership session = self._get_session(write_meta.unique_rid) if session is None: msg = f"_deliver_aux_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" logger.error(msg) - if owned: - write_meta.task.retire_unsubmitted_physical_operation(write_meta.peer_rank) write_meta.task.fail(RuntimeError(msg)) + self._send_agent_result( + write_meta.peer_endpoint, + [ + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(write_meta.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ], + ) return aux_task = session.aux_task assert aux_task is not None, f"aux_task is None for session {write_meta.unique_rid}" @@ -1214,69 +1001,72 @@ def _deliver_aux_to_agent(self, write_meta: WriteMeta): if timer: timer.record_push_end(write_meta.peer_rank) + active_nixl_op = write_meta.src_ptrs.size > 0 + with session.lock: + status = session.status + if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + should_abort = True + else: + aux_task.status = TaskStatus.TRANSFERRING + should_abort = False + + if should_abort: + aux_task.fail( + RuntimeError( + f"session {write_meta.unique_rid} {status.value}, aux transfer aborted" + ) + ) + self._send_agent_result( + write_meta.peer_endpoint, + [ + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(write_meta.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ], + ) + return + agent_result = AgentResult.SUCCESS - if write_meta.src_ptrs.size > 0: - try: + try: + if active_nixl_op: request = Sender._make_agent_request(write_meta, device_id=self._device_id) - except Exception as error: - logger.error( - "_deliver_aux_to_agent: failed to build the auxiliary send " - f"request for {write_meta.unique_rid}: {error}" - ) - if owned: - aux_task.retire_unsubmitted_physical_operation(write_meta.peer_rank) - aux_task.fail(RuntimeError(f"build aux transfer request failed: {error}")) - # An admitted operation normally makes this worker the only - # eligible claimant. Keep the claim at the send boundary so - # duplicate dispatch cannot publish contradictory evidence. - if not owned or session._claim_aux_terminal_result(write_meta.peer_rank): - self._get_result_dealer(write_meta.peer_endpoint).send( - _make_aux_result_msg( - self._instance_rank, - write_meta.unique_rid, - AgentResult.FAILED, - ) - ) - return - if timer: - timer.record_transfer_start(write_meta.peer_rank) - try: - transfer_finished, last_status = self._submit_transfer( - aux_task, write_meta.peer_rank, request - ) - if transfer_finished: - del request - if not transfer_finished: - agent_result = AgentResult.IN_DOUBT if owned else AgentResult.FAILED - session.set_exception(f"aux transfer agent request failed: {last_status}") - except _TransferNotSubmittedError as error: - agent_result = AgentResult.FAILED - aux_task.fail(error) - if timer: - timer.record_transfer_end(write_meta.peer_rank) - elif owned: - aux_task.retire_unsubmitted_physical_operation(write_meta.peer_rank) - - # An admitted operation normally makes this worker the only eligible - # claimant. Keep the claim at the send boundary so future cleanup-path - # changes cannot publish contradictory evidence. - if not owned or session._claim_aux_terminal_result(write_meta.peer_rank): - if agent_result == AgentResult.IN_DOUBT: - self._retain_in_doubt_transfer( - write_meta, - _make_aux_result_msg(self._instance_rank, write_meta.unique_rid, agent_result), - ) - return - self._get_result_dealer(write_meta.peer_endpoint).send( - _make_aux_result_msg( - self._instance_rank, - write_meta.unique_rid, - agent_result, - ) + if timer: + timer.record_transfer_start(write_meta.peer_rank) + completed, _ = self._submit_and_wait_agent_transfer(request) + if not completed: + agent_result = AgentResult.FAILED + session.set_exception("aux transfer agent request failed") + except Exception as exc: + agent_result = AgentResult.FAILED + aux_task.fail(exc) + logger.error( + f"Aux transfer raised for rid={write_meta.unique_rid}, " + f"peer_rank={write_meta.peer_rank}: {exc}" ) + finally: + if timer: + try: + timer.record_transfer_end(write_meta.peer_rank) + except Exception as exc: + logger.warning( + f"Failed to record aux transfer end for rid={write_meta.unique_rid}: {exc}" + ) - if agent_result == AgentResult.IN_DOUBT: - return + self._send_agent_result( + write_meta.peer_endpoint, + [ + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(write_meta.unique_rid).encode("ascii"), + agent_result.value.encode("ascii"), + ], + ) + + if timer: + timer.record_task_end(write_meta.peer_rank) + ri = self._registrar.self_rank_info + aux_task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) with aux_task.lock: aux_task._transfer_count += 1 @@ -1293,64 +1083,106 @@ def _deliver_aux_to_agent(self, write_meta: WriteMeta): f"aux task received more than {write_meta.expected_transfers} transfers" ) - try: - if timer: - timer.record_task_end(write_meta.peer_rank) - ri = self._registrar.self_rank_info - aux_task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) - except Exception as error: - logger.warning( - f"AUX send perf logging failed for request {write_meta.unique_rid}: {error}" - ) - @staticmethod - def _pair_ordinals( + def _align_kv_blocks( src_block_ids: np.ndarray, dst_block_ids: np.ndarray, - dst_start_block: int = 0, + src_token_start: int, + dst_token_start: int, + tokens_per_block: int, + minimum_token_start: int = 0, ) -> tuple[np.ndarray, np.ndarray]: - """Pair src/dst positional block tables ordinal by ordinal. - - Both tables are indexed by block ordinal (see Chunk); an entry is -1 - where that side has nothing to offer/accept. A block moves only where - both sides hold it, so a sender that has not evicted yet, a receiver - that already reused a prefix, a receiver with a wider speculative - window, or a sender chunk that blanks everything outside its range all - fall out of the same intersection without any per-case arithmetic. - Ordinals below ``dst_start_block`` are excluded as well. + """Align src/dst block arrays using explicit token-start positions. + + Both src_token_start and dst_token_start must be block-aligned + (multiples of tokens_per_block), which is always true for prefix-cache + boundaries in current KV cache managers. minimum_token_start applies a + shared lower bound, such as an SWA stale boundary, and trims both arrays + instead of relabeling their first blocks. + + Returns the (src, dst) sub-arrays that cover the shared token overlap. + Returns a pair of empty arrays when there is no overlap (i.e. this + context slice is entirely within generation's already-cached prefix). + + This handles four cases without special-casing: + 1. No prefix cache on either side → identity (start_token == 0 both) + 2. Context prefix cache (src starts later than 0) → trim dst head + 3. Generation prefix cache (dst starts later than 0) → trim src head + 4. Chunked context (each slice has its own token_range) → correct + overlap even when the slice is entirely before dst_token_start + """ + overlap_start = max(src_token_start, dst_token_start, minimum_token_start) + src_skip = (overlap_start - src_token_start) // tokens_per_block + dst_skip = (overlap_start - dst_token_start) // tokens_per_block + n_transfer = min(src_block_ids.size - src_skip, dst_block_ids.size - dst_skip) + if n_transfer <= 0: + return np.array([], dtype=np.int64), np.array([], dtype=np.int64) + return ( + src_block_ids[src_skip : src_skip + n_transfer], + dst_block_ids[dst_skip : dst_skip + n_transfer], + ) + + @staticmethod + def _beam0_block_count(block_ids: np.ndarray, total_blocks: int, beam_width: int) -> int: + """Return the number of beam-0 blocks in a packed 1-D beam layout.""" + if beam_width <= 1 or block_ids.size <= total_blocks: + return block_ids.size + return max(0, block_ids.size - (beam_width - 1)) + + @staticmethod + def _helix_windowed_block_ids( + src_block_ids: np.ndarray, + total_blocks: int, + beam_width: int, + cp_rank: int, + cp_size: int, + ) -> np.ndarray: + """Select the blocks of a windowed (suffix) list owned by one Helix rank. + + The list covers global ordinals ``[total_blocks - k, total_blocks)``; + rank ``cp_rank`` owns every ordinal congruent to it modulo ``cp_size``. + """ + k = Sender._beam0_block_count(src_block_ids, total_blocks, beam_width) + first_ordinal = total_blocks - k + keep = [i for i in range(k) if (first_ordinal + i) % cp_size == cp_rank] + return src_block_ids[keep] + + @staticmethod + def _trim_receiver_window_head( + src_block_ids: np.ndarray, + dst_block_ids: np.ndarray, + peer_window_size: Optional[int], + beam_width: int, + ) -> np.ndarray: + """Drop the receiver's extra leading blocks for a windowed layer group. + + A windowed receiver keeps a larger window when only it runs speculative + decoding, so its suffix starts earlier and the extra blocks are at the + head. Both starts are derived from list length, so trimming the tail + instead would shift every block one position early. + + Non-windowed lists are both trimmed to ceil(prompt_len / tpb) in + _create_kv_slice, so there dst must not exceed src. A smaller dst + (generation prefix-cache reuse) is handled via dst_start. """ - if src_block_ids.size != dst_block_ids.size: + block_diff = dst_block_ids.size - src_block_ids.size + if block_diff <= 0: + return dst_block_ids + if peer_window_size is None or beam_width > 1: raise ValueError( - f"positional block tables differ in length: src={src_block_ids.size} " - f"dst={dst_block_ids.size}; both sides must span the same prompt" + f"src/dst block count mismatch: {src_block_ids.size} vs " + f"{dst_block_ids.size} (dst must not exceed src)" ) - mask = (src_block_ids >= 0) & (dst_block_ids >= 0) - if dst_start_block > 0: - mask[:dst_start_block] = False - idx = np.flatnonzero(mask) - return src_block_ids[idx], dst_block_ids[idx] + return dst_block_ids[block_diff:] @nvtx_range("_build_kv_write_meta") def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> WriteMeta: peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) - self_ri = self._registrar.self_rank_info - token_range = task._chunk.token_range - # Only a piece short of the whole prompt runs the chunk projection below, and that is the - # part with no answer under CP; a sole piece covering everything takes the untouched path. - is_partial = token_range.start > 0 or token_range.end < task._prompt_len - if is_partial and (self_ri.cp_size > 1 or peer_ri.cp_size > 1): - raise ValueError( - "enable_pipelined_transfer is not supported with context parallelism " - f"(sender cp_size={self_ri.cp_size}, receiver cp_size={peer_ri.cp_size})" - ) timer = task._perf_timer if timer: timer.record_prepare_args_start(peer_ri.instance_rank) targets = self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank) expected_transfers = len(targets.ranks) - # Recorded on the task as well: the count on the write meta settles the task, while the - # handle has to answer how many writes are still out after one of them has already failed. - task.expected_transfers = expected_transfers # Aggregate fragment pointers from all matching pool pairs. # Each pool pair produces one or more region pairs (rp), each containing @@ -1369,55 +1201,109 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write ) pool_mapping = self._registrar.get_pool_mapping(peer_ri) dst_block_ids_per_groups = req_info.block_ids_per_layer_groups - src_block_ids_per_groups = task._chunk.block_ids_per_layer_groups + src_block_ids_per_groups = task._slice.block_ids_per_layer_groups # Aggregate fragments from all matching pools using numpy concatenation. # Send ownership is per pool: replicated pools elect one fan-in # owner, sharded pools keep head-duplication routing. for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): - if (self_lg, self_pi) in task._chunk.excluded_pool_views: - continue if not self._registrar.should_send_pool(targets, peer_ri, self_lg, self_pi): continue - - lg_info = extractor.page_table.layer_groups[self_lg] src_block_ids = src_block_ids_per_groups[self_lg] dst_block_ids = dst_block_ids_per_groups[peer_lg] - - # Extract regions: STATE = per-layer ptrs, PAGED = per-block ptrs - if lg_info.kind == CacheKind.STATE: - if src_block_ids.size == 0 or dst_block_ids.size == 0: - continue - if not MambaPolicy.is_paired(self._registrar.self_rank_info, peer_ri): - continue - src_region = extractor.extract_slot(int(src_block_ids[0]), self_lg, self_pi) - dst_region = peer_extractor.extract_slot(int(dst_block_ids[0]), peer_lg, peer_pi) - else: - src_block_ids = np.asarray(src_block_ids, dtype=np.int64) - dst_block_ids = np.asarray(dst_block_ids, dtype=np.int64) - if peer_ri.cp_size > 1 and self_ri.cp_size == 1: - # Helix: the receiver owns global blocks [cp_rank::cp_size] - # (same protocol as partition_context_for_helix), so its - # table is the strided subset of ours; block reuse is - # rejected under helix. + tpb = extractor.page_table.tokens_per_block + token_range = task._slice.token_range + lg_info = extractor.page_table.layer_groups[self_lg] + window_size = getattr(lg_info, "sliding_window_size", None) + if ( + peer_ri.cp_size > 1 + and self._registrar.self_rank_info.cp_size == 1 + and not getattr(lg_info, "cp_as_tp", False) + ): + # Helix: the receiver owns global blocks [cp_rank::cp_size] + # (same protocol as partition_context_for_helix); block reuse + # is rejected under helix. Windowed lists are a suffix of the + # global ordinals, so route them by ordinal, not by position. + if window_size is None: src_block_ids = src_block_ids[peer_ri.cp_rank :: peer_ri.cp_size] - tpb = extractor.page_table.tokens_per_block - dst_start_block = ( - req_info.dst_start_token // tpb if req_info.dst_start_token is not None else 0 - ) - # Both tables are positional (index == block ordinal, -1 == not - # held), so pairing is a plain intersection: SWA eviction, prefix - # reuse, and pipelined chunk bounds are all already encoded as - # holes by whichever side owns that state. - src_block_ids, dst_block_ids = Sender._pair_ordinals( - src_block_ids, dst_block_ids, dst_start_block=dst_start_block + else: + slice_end = token_range.end if token_range is not None else 0 + src_block_ids = Sender._helix_windowed_block_ids( + src_block_ids, + total_blocks=(slice_end + tpb - 1) // tpb, + beam_width=task._beam_width, + cp_rank=peer_ri.cp_rank, + cp_size=peer_ri.cp_size, + ) + + # Non-windowed lists are both trimmed to ceil(prompt_len / tpb) in + # _create_kv_slice, so dst must not exceed src. A smaller dst + # (generation prefix-cache reuse) is handled via dst_start below. + block_diff = dst_block_ids.size - src_block_ids.size + if block_diff > 0: + # A windowed receiver keeps a larger window when only it runs + # speculative decoding, so its suffix starts earlier. Drop the + # extra head: dst_start is derived from len(dst), so trimming + # the tail instead would shift every block by block_diff. + peer_lg_info = peer_extractor.page_table.layer_groups[peer_lg] + peer_window = getattr(peer_lg_info, "sliding_window_size", None) + if peer_window is not None and task._beam_width <= 1: + dst_block_ids = dst_block_ids[block_diff:] + else: + raise ValueError( + f"src/dst block count mismatch: {src_block_ids.size} vs " + f"{dst_block_ids.size} (dst must not exceed src)" + ) + + # Block lists are the suffix of [..., slice_end); cached prefix + # is implicit in their size. token_start = (total_blocks - n) * tpb. + slice_end = token_range.end if token_range is not None else 0 + total_blocks = (slice_end + tpb - 1) // tpb + src_beam0_blocks = Sender._beam0_block_count( + src_block_ids, total_blocks, task._beam_width + ) + dst_beam0_blocks = Sender._beam0_block_count( + dst_block_ids, total_blocks, task._beam_width + ) + assert src_beam0_blocks <= total_blocks, ( + f"src beam-0 block list ({src_beam0_blocks}) exceeds total slice " + f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" + ) + assert dst_beam0_blocks <= total_blocks, ( + f"dst beam-0 block list ({dst_beam0_blocks}) exceeds total slice " + f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" + ) + src_start = (total_blocks - src_beam0_blocks) * tpb + dst_start = (total_blocks - dst_beam0_blocks) * tpb + minimum_start = req_info.dst_start_token or 0 + if window_size is not None: + # SWA stale_end uses the request prompt_len (not slice_end — + # they differ for non-final slices). prompt_len must be plumbed + # via the session; falling back to slice_end is wrong on + # non-final slices. + assert task._prompt_len is not None, ( + "SWA layer requires session.prompt_len; " + "set TxSession(prompt_len=request.prompt_len)." ) - if src_block_ids.size == 0: - continue - src_region = extractor.extract(src_block_ids, self_lg, self_pi) - dst_region = peer_extractor.extract(dst_block_ids, peer_lg, peer_pi) + stale_end = max(0, (task._prompt_len + 1 - window_size) // tpb) + minimum_start = max(minimum_start, stale_end * tpb) + src_block_ids, dst_block_ids = Sender._align_kv_blocks( + src_block_ids, + dst_block_ids, + src_token_start=src_start, + dst_token_start=dst_start, + tokens_per_block=tpb, + minimum_token_start=minimum_start, + ) - # Map and collect fragments (same for all layer types) + src_region = extractor.extract(src_block_ids, layer_group_id=self_lg, pool_idx=self_pi) + dst_region = peer_extractor.extract( + dst_block_ids, layer_group_id=peer_lg, pool_idx=peer_pi + ) + if token_range is not None: + token_spec = KVRegionSpec(tokens=IndexRange(token_range.start, token_range.end - 1)) + src_region = src_region._replace(spec=token_spec) + dst_region = dst_region._replace(spec=token_spec) mapper = self._registrar.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) region_pair = mapper.map(src_region, dst_region) region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] @@ -1440,6 +1326,20 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write dst_frags = np.array([], dtype=np.int64) kv_sizes = np.array([], dtype=np.int64) + # handle mamba fragments + m_src, m_dst, m_sizes = MambaPolicy.collect_frags( + self_page_table=extractor.page_table, + peer_page_table=peer_extractor.page_table, + src_slot=task._slice.mamba_state_index, + dst_slot=req_info.mamba_state_index, + self_ri=self._registrar.self_rank_info, + peer_ri=peer_ri, + ) + if m_src: + src_frags = np.concatenate([src_frags, np.array(m_src, dtype=np.int64)]) + dst_frags = np.concatenate([dst_frags, np.array(m_dst, dtype=np.int64)]) + kv_sizes = np.concatenate([kv_sizes, np.array(m_sizes, dtype=np.int64)]) + if timer: timer.record_prepare_args_end(peer_ri.instance_rank) timer.record_transfer_sizes(peer_ri.instance_rank, int(kv_sizes.sum()), dst_frags.size) @@ -1456,8 +1356,7 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write peer_endpoint=peer_ri.self_endpoint, unique_rid=task._unique_rid, slice_id=task.slice_id, - receiver_slice_id=req_info.slice_id if req_info.slice_id is not None else 0, - is_last_slice=task._chunk.is_last, + is_last_slice=task._slice.is_last_slice, bounce_dst_base=req_info.bounce_dst_base, ) @@ -1519,44 +1418,33 @@ def dispatch_task( # critical section small. When not provided, we fetch it here (legacy / standalone path). if req_info_snapshot is None: req_info_snapshot = dict(self._get_req_info(task._unique_rid) or {}) - first_error: Optional[Exception] = None - for info in req_info_snapshot.values(): + req_infos = list(req_info_snapshot.values()) + for info_idx, info in enumerate(req_infos): try: - self._dispatch_task_to_peer(task, info) - except Exception as error: - if not self._enforce_physical_ownership: - raise - if first_error is None: - first_error = error - if first_error is not None: - raise first_error - - def _dispatch_task_to_peer(self, task: KVSendTask | AuxSendTask, info: RecvReqInfo) -> None: - owned = False - if self._enforce_physical_ownership: - admission = self._begin_task_operation(task, info.instance_rank) - if admission is None: - return - if not admission: - self._send_failed_task_result_to_receiver(task, info) - return - owned = True - try: - if task._perf_timer is not None: - task._perf_timer.record_task_start(info.instance_rank) - trans_meta = ( - self._build_kv_write_meta(task, info) - if isinstance(task, KVSendTask) - else self._build_aux_write_meta(task, info) - ) - if task._perf_timer is not None: - task._perf_timer.record_push_start(trans_meta.peer_rank) - self._enqueue(trans_meta) - except Exception: - if owned: - task.retire_unsubmitted_physical_operation(info.instance_rank) - self._send_failed_task_result_to_receiver(task, info) - raise + if task._perf_timer is not None: + task._perf_timer.record_task_start(info.instance_rank) + if isinstance(task, KVSendTask): + trans_meta = self._build_kv_write_meta(task, info) + else: + trans_meta = self._build_aux_write_meta(task, info) + if task._perf_timer is not None: + task._perf_timer.record_push_start(trans_meta.peer_rank) + self._enqueue(trans_meta) + except Exception as exc: + task.fail(exc) + # The receiver reserved one peer transfer before publishing + # every RecvReqInfo in this snapshot. Entries before info_idx + # already own an enqueued WriteMeta and will report their own + # terminal result. This entry and every later one have not + # been dispatched, so resolve each reservation exactly once + # before propagating the build/enqueue failure. + for undispatched_info in req_infos[info_idx:]: + self._send_failed_result_to_receiver( + undispatched_info, + slice_id=(task.slice_id if isinstance(task, KVSendTask) else None), + aux=isinstance(task, AuxSendTask), + ) + raise def _start_listener(self): def handle_message(messages: list[bytes]): @@ -1586,24 +1474,19 @@ def handle_message(messages: list[bytes]): self._messenger.start_listener(handle_message) def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): - if self._shutdown_requested: + # Skip late messages so we don't race shutdown's invalidate loop. + if self._shutdown: return torch.cuda.set_device(self._device_id) CUASSERT(cudart.cudaSetDevice(self._device_id)) ri: RankInfo = RankInfo.from_bytes(message[1]) + + self._registrar.register(ri.instance_name, ri.instance_rank, ri) + agent_name = ri.instance_name + str(ri.instance_rank) logger.debug(f"Loading remote transfer agent descriptor for peer '{agent_name}'") - # NIXL serializes descriptor load against submit internally. This lock - # only closes the local load-vs-shutdown window; active transfers do - # not need to drain before a new peer is registered. + self._load_remote_agent_serialized(agent_name, ri.transfer_engine_info) with self._loaded_remote_agents_lock: - if self._shutdown_requested: - return - self._registrar.register(ri.instance_name, ri.instance_rank, ri) - self._agent.load_remote_agent( - agent_name, - ri.transfer_engine_info, - ) self._loaded_remote_agents.add(agent_name) logger.debug( f"Completed handling REGISTER_RANK_INFO for instance='{ri.instance_name}', rank={ri.instance_rank}" @@ -1612,235 +1495,201 @@ def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): def _handle_cancel_session(self, message: list[bytes]): unique_rid = int(message[1]) session = None + pre_cancel_infos: list[RecvReqInfo] = [] with self._sessions_lock: - session = self._get_session(unique_rid) - if session is None or self._enforce_physical_ownership: - # Parked with the peer as who asked: the session built afterwards reads its - # cancellation from here, and an id on its own would read as a local cancel. - if self._enforce_physical_ownership: - self._pre_cancelled_rids.setdefault(unique_rid, True) - else: - self._pre_cancelled_rids[unique_rid] = True - if session is not None: - if self._enforce_physical_ownership: - session.cancel_local(by_peer=True) - else: - session.cancel(by_peer=True) - if self._enforce_physical_ownership and len(message) == 4: - # Routing is carried by CANCEL itself, not by expiring REQUEST_DATA. - self._queue_session_quiescence(unique_rid, message[2].decode(), int(message[3])) - - def _queue_session_quiescence( - self, unique_rid: int, instance_name: str, instance_rank: int - ) -> None: - """Queue proof on the same worker/socket as this peer's piece results.""" - endpoint = self._registrar.get_peer_rank_info(instance_name, instance_rank).self_endpoint - key = (unique_rid, instance_rank) - with self._sessions_lock: - if self._shutdown_requested or key in self._pending_session_quiescence: - return - marker = _SessionQuiescence(unique_rid, instance_rank, endpoint) - self._pending_session_quiescence[key] = marker - self._send_task_queues[hash(key) % self._num_threads].put(marker) - - def _send_session_quiesced(self, marker: _SessionQuiescence) -> bool: - """Acknowledge a fenced, drained session; return False to retry later. - - Called only by the ordered worker, after earlier queued piece reports. - Session absence is safe only with the persistent admission fence. An - existing session must also prove that all KV and AUX accesses drained. - """ - key = (marker.unique_rid, marker.peer_rank) - with self._sessions_lock: - if marker.unique_rid not in self._pre_cancelled_rids: - return False - session = self._get_session(marker.unique_rid) + self._pre_cancelled_rids.add(unique_rid) + self._pre_cancelled_timestamps[unique_rid] = time.monotonic() + session_ref = self._sessions.get(unique_rid) + if session_ref is None: + with self._peer_requests_lock: + req_info_map = self._peer_requests.pop(unique_rid, None) + self._peer_requests_timestamps.pop(unique_rid, None) + if req_info_map: + pre_cancel_infos = list(req_info_map.values()) + else: + session = session_ref() + if session is None: + with self._peer_requests_lock: + req_info_map = self._peer_requests.pop(unique_rid, None) + self._peer_requests_timestamps.pop(unique_rid, None) + if req_info_map: + pre_cancel_infos = list(req_info_map.values()) if session is not None: + session.cancel() with session.lock: - if not session.has_failed() or not session.resources_drained(): - return False - thread_idx = hash(key) % self._num_threads - if any( - pending.write_meta.unique_rid == marker.unique_rid - and pending.write_meta.peer_rank == marker.peer_rank - for pending in self._pending_settlements[thread_idx].values() - ): - return False - try: - self._get_or_connect_thread_dealer(marker.endpoint).send( - [ - MessageType.SESSION_QUIESCED, - str(self._instance_rank).encode("ascii"), - str(marker.unique_rid).encode("ascii"), - ] - ) - except Exception as error: - logger.warning(f"Failed to acknowledge fenced session {marker.unique_rid}: {error}") - return False - with self._sessions_lock: - self._pending_session_quiescence.pop(key, None) - return True + need_kv_result = not session.kv_tasks + need_aux_result = session._need_aux and session.aux_task is None + if need_kv_result or need_aux_result: + # cancel() has made future send()/send_aux() calls fail, so no + # worker can be created for the missing task type. Remove the + # saved requests and resolve those reservations directly. + with self._peer_requests_lock: + req_info_map = self._peer_requests.pop(unique_rid, None) + self._peer_requests_timestamps.pop(unique_rid, None) + direct_infos = list(req_info_map.values()) if req_info_map else [] + for info in direct_infos: + self._send_pre_cancel_failed_results( + info, + send_kv=need_kv_result, + send_aux=need_aux_result, + ) + else: + # These REQUEST_DATA messages were accepted, but no TxSession will + # ever enqueue a worker result on this (possibly non-selected ADP) + # rank. Explicitly resolve the receiver's per-peer reservations. + for info in pre_cancel_infos: + self._send_pre_cancel_failed_results(info) @nvtx_range("_respond_with_kv") def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): # _sessions_lock prevents a race between session lookup and req_info save. - # session.lock atomically saves peer info and snapshots tasks against send(). + # session.lock serializes _enqueue calls from both paths. + self.sweep_stale_req_infos() info: RecvReqInfo = RecvReqInfo.from_bytes(message[1]) + pre_cancelled = False with self._sessions_lock: session = self._get_session(info.unique_rid) - fenced = ( - self._enforce_physical_ownership and info.unique_rid in self._pre_cancelled_rids - ) - if session is None and not fenced: - self._save_peer_req_info(info) - return - if fenced: - self._queue_session_quiescence(info.unique_rid, info.instance_name, info.instance_rank) + if session is None: + pre_cancelled = info.unique_rid in self._pre_cancelled_rids + if not pre_cancelled: + self._save_peer_req_info(info) + return + if pre_cancelled: + # CANCEL_SESSION may race ahead of REQUEST_DATA. Do not retain a + # request that can never acquire a source session; acknowledge it + # exactly once instead. + self._send_pre_cancel_failed_results(info) return - assert session is not None with session.lock: - with self._sessions_lock: - if self._get_session(info.unique_rid) is not session or session._closed: - tasks = None - terminal = True - include_aux = bool(session._claim_unsubmitted_aux_failures_locked((info,))) - else: - self._save_peer_req_info(info) - tasks = list(session.kv_tasks) - terminal = session.has_failed() - include_aux = terminal and bool( - session._claim_unsubmitted_aux_failures_locked((info,)) - ) - if tasks is None: - self._send_failed_result_to_receiver(info, include_aux=include_aux) - return + terminal = session.status in ( + SessionStatus.ERROR, + SessionStatus.CANCELLED, + ) if terminal: - self._send_failed_result_to_receiver(info, include_aux=include_aux) - return + with self._peer_requests_lock: + existing = self._peer_requests.get(info.unique_rid, {}).get(info.instance_rank) + already_known = existing is not None and existing.slice_id == info.slice_id + tasks = list(session.kv_tasks) + aux_task_exists = session.aux_task is not None + else: + self._save_peer_req_info(info) + tasks = list(session.kv_tasks) + if session.aux_task is not None: + tasks.append(session.aux_task) + already_known = False + aux_task_exists = False + if terminal: + # Do not enqueue duplicate work into a cancelled/error session. + # A previously known request with an existing task already has a + # queued/active worker which will report its own terminal result; + # otherwise this late request needs a direct per-peer reply. + self._send_pre_cancel_failed_results( + info, + send_kv=not already_known or not tasks, + send_aux=(info.need_aux and (not already_known or not aux_task_exists)), + ) + return for task in tasks: - self._dispatch_task_to_peer(task, info) - - def _send_failed_result_to_receiver( - self, - info: RecvReqInfo, - *, - include_aux: bool = False, - defer_to_worker: bool = False, - ): - try: - endpoint = self._get_result_receiver_endpoint(info) - slice_id = info.slice_id if info.slice_id is not None else 0 - messages = [ - _make_kv_result_msg( - self._instance_rank, - info.unique_rid, - slice_id, - True, # is_last_slice - AgentResult.FAILED, + try: + if task._perf_timer is not None: + task._perf_timer.record_task_start(info.instance_rank) + trans_meta = ( + self._build_kv_write_meta(task, info) + if isinstance(task, KVSendTask) + else self._build_aux_write_meta(task, info) ) - ] - if include_aux: - messages.append( - _make_aux_result_msg( - self._instance_rank, - info.unique_rid, - AgentResult.FAILED, - ) + if task._perf_timer is not None: + task._perf_timer.record_push_start(trans_meta.peer_rank) + self._enqueue(trans_meta) + except Exception as exc: + task.fail(exc) + session.set_exception(str(exc)) + is_aux = isinstance(task, AuxSendTask) + self._send_failed_result_to_receiver( + info, slice_id=-1 if is_aux else task.slice_id, aux=is_aux ) - self._route_result_messages_to_receiver( - info, - endpoint, - messages, - defer_to_worker=defer_to_worker, - ) - except Exception as e: - logger.warning( - f"_respond_with_kv: failed to abort receiver for rid={info.unique_rid}: {e}" - ) + raise + + def _cancel_failed_result_key( + self, info: RecvReqInfo, *, aux: bool + ) -> tuple[int, str, int, int, bool]: + slice_id = -1 if aux else (info.slice_id if info.slice_id is not None else 0) + return ( + info.unique_rid, + info.instance_name, + info.instance_rank, + slice_id, + aux, + ) + + def _send_pre_cancel_failed_result_once(self, info: RecvReqInfo, *, aux: bool) -> bool: + key = self._cancel_failed_result_key(info, aux=aux) + with self._peer_requests_lock: + if key in self._cancel_failed_result_keys: + return True + self._cancel_failed_result_keys.add(key) + sent = self._send_failed_result_to_receiver( + info, + slice_id=info.slice_id, + aux=aux, + ) + if not sent: + # A later duplicate CANCEL/REQUEST_DATA may retry a transport + # failure; successful sends remain idempotent. + with self._peer_requests_lock: + self._cancel_failed_result_keys.discard(key) + return sent - def _send_aux_result_to_receiver( + def _send_pre_cancel_failed_results( self, info: RecvReqInfo, - result: AgentResult, *, - defer_to_worker: bool = True, - ) -> None: - try: - endpoint = self._get_result_receiver_endpoint(info) - message = _make_aux_result_msg( - self._instance_rank, - info.unique_rid, - result, - ) - self._route_result_messages_to_receiver( - info, - endpoint, - [message], - defer_to_worker=defer_to_worker, - ) - except Exception as error: - logger.warning( - f"failed to report auxiliary transfer result for rid={info.unique_rid}: {error}" - ) + send_kv: bool = True, + send_aux: Optional[bool] = None, + ) -> bool: + if send_aux is None: + send_aux = info.need_aux + sent = True + if send_kv: + sent = self._send_pre_cancel_failed_result_once(info, aux=False) + if send_aux: + sent = self._send_pre_cancel_failed_result_once(info, aux=True) and sent + return sent - def _send_failed_task_result_to_receiver(self, task: SendTaskBase, info: RecvReqInfo) -> None: - if not task.is_done: - task.fail(RuntimeError("transfer rejected before backend submission")) + def _send_failed_result_to_receiver( + self, + info: RecvReqInfo, + slice_id: Optional[int] = None, + aux: bool = False, + ) -> bool: try: - endpoint = self._get_result_receiver_endpoint(info) - if isinstance(task, KVSendTask): - receiver_slice_id = info.slice_id if info.slice_id is not None else 0 - message = _make_kv_result_msg( - self._instance_rank, - info.unique_rid, - receiver_slice_id, - True, - AgentResult.FAILED, - ) + peer_ri = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) + if aux: + message = [ + MessageType.AUX_AGENT_RESULT, + str(self._instance_rank).encode("ascii"), + str(info.unique_rid).encode("ascii"), + AgentResult.FAILED.value.encode("ascii"), + ] else: - with self._sessions_lock: - session = self._get_session(info.unique_rid) - if session is not None and not session._claim_aux_terminal_result( - info.instance_rank - ): - return - message = _make_aux_result_msg( + resolved_slice_id = ( + slice_id + if slice_id is not None + else (info.slice_id if info.slice_id is not None else 0) + ) + message = _make_kv_result_msg( self._instance_rank, info.unique_rid, + resolved_slice_id, + True, # is_last_slice AgentResult.FAILED, ) - self._route_result_messages_to_receiver( - info, - endpoint, - [message], - defer_to_worker=True, + return self._send_agent_result(peer_ri.self_endpoint, message) + except Exception as e: + logger.warning( + f"_respond_with_kv: failed to abort receiver for rid={info.unique_rid}: {e}" ) - except Exception as error: - logger.warning(f"failed to settle rejected transfer {info.unique_rid}: {error}") - - def _get_result_receiver_endpoint(self, info: RecvReqInfo) -> Optional[str]: - return self._registrar.get_peer_rank_info( - info.instance_name, info.instance_rank - ).self_endpoint - - def _route_result_messages_to_receiver( - self, - info: RecvReqInfo, - endpoint: Optional[str], - messages: Iterable[list[bytes]], - *, - defer_to_worker: bool, - ) -> None: - if defer_to_worker or self._enforce_physical_ownership: - # Listener-side rejection must share the worker's ordered stream; - # it cannot overtake an ambiguous write's pending IN_DOUBT report. - thread_idx = hash((info.unique_rid, info.instance_rank)) % self._num_threads - for message in messages: - self._send_task_queues[thread_idx].put((endpoint, message)) - return - dealer = self._get_or_connect_dealer(endpoint) - for message in messages: - dealer.send(message) + return False def _get_or_connect_dealer(self, endpoint: Optional[str]): if endpoint is None: @@ -1888,52 +1737,65 @@ def send_cancel_to_receivers(self, unique_rid: int) -> None: peer_ri = self._registrar.get_peer_rank_info( req_info.instance_name, req_info.instance_rank ) - self._get_or_connect_dealer(peer_ri.self_endpoint).send( - [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] + message = [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] + outbox = getattr(self, "_reliable_message_outbox", None) + accepted = ( + outbox.submit(peer_ri.self_endpoint, message) + if outbox is not None + else self._send_agent_result(peer_ri.self_endpoint, message) ) + if not accepted: + logger.error(f"send_cancel_to_receivers: outbox rejected rid={unique_rid}") except Exception as e: logger.warning(f"send_cancel_to_receivers: failed for rid={unique_rid}: {e}") def shutdown(self): - watchdog = getattr(self, "_retirement_watchdog", None) - if watchdog is not None: - watchdog.request_shutdown() - watchdog.require_retired() - with self._sessions_lock: - if self._shutdown: - return - if self._enforce_physical_ownership and self._sessions: - raise RuntimeError("Sender refuses shutdown while transfer ownership is active") - if any(self._pending_settlements): - raise RuntimeError("Sender refuses shutdown while settlement reports are pending") - self._shutdown_requested = True + if getattr(self, "_shutdown", False): + return + self._shutdown = True - self._messenger.stop() + # Quiesce listener before invalidate to avoid set/map mutation races. + messenger = getattr(self, "_messenger", None) + if messenger is not None: + messenger.stop() - for q in self._send_task_queues: + for q in getattr(self, "_send_task_queues", []): q.put(None) - for t in self._worker_threads: + for t in getattr(self, "_worker_threads", []): t.join(timeout=5) + # Workers cannot enqueue any more terminal results after joining. Let + # the shared result/cancel outbox make one final attempt for every + # queued message, then stop. + reliable_outbox = getattr(self, "_reliable_message_outbox", None) + if reliable_outbox is not None: + reliable_outbox.shutdown() + # Snapshot under lock as defense in depth. - with self._loaded_remote_agents_lock: - loaded_agents = list(self._loaded_remote_agents) - self._loaded_remote_agents.clear() + loaded_agents: list[str] = [] + loaded_remote_agents_lock = getattr(self, "_loaded_remote_agents_lock", None) + if loaded_remote_agents_lock is not None: + with loaded_remote_agents_lock: + loaded_remote_agents = getattr(self, "_loaded_remote_agents", set()) + loaded_agents = list(loaded_remote_agents) + loaded_remote_agents.clear() # Invalidate all loaded remote agents to release fabric/POSIX FD resources. + agent = getattr(self, "_agent", None) for agent_name in loaded_agents: try: - self._agent.invalidate_remote_agent(agent_name) + if agent is not None: + agent.invalidate_remote_agent(agent_name) except Exception as e: logger.warning( f"Failed to invalidate remote agent '{agent_name}' during shutdown: {e}" ) - for dealer in self._dealers.values(): + dealers = getattr(self, "_dealers", {}) + for dealer in dealers.values(): try: dealer.stop() except Exception as e: logger.warning(f"Failed to stop dealer during Sender shutdown: {e}") - self._dealers.clear() - self._shutdown = True + dealers.clear() def __del__(self): try: @@ -1949,9 +1811,6 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): class TxSession(TxSessionBase): - cancelled_by_peer: bool = False - """Only meaningful once the session is CANCELLED.""" - def __init__( self, request_id: int, @@ -1960,41 +1819,34 @@ def __init__( aux_buffer: Optional[AuxBuffer] = None, timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, + beam_width: int = 1, overall_timeout_s: Optional[float] = None, - *, - retirement_watchdog: Optional[RetirementWatchdog] = None, ): - super().__init__(sender, SessionArgsBase(params, prompt_len=prompt_len)) + super().__init__( + sender, + SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), + ) self._timeout_s = timeout_s self._overall_timeout_s = overall_timeout_s self._deadline_monotonic_s: Optional[float] = None - self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST - self._enforce_physical_ownership = getattr(sender, "_enforce_physical_ownership", False) + self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST or ( + aux_buffer is not None and aux_buffer.has_draft_history + ) self._sender: Sender # narrow base class type for Pylance self.request_id = request_id - self._retirement = None - if retirement_watchdog is not None: - if overall_timeout_s is None: - raise ValueError("deadline retirement requires the overall transfer timeout") - self._retirement = retirement_watchdog.create_owner( - request_id, "send", overall_timeout_s - ) self._aux_buffer = aux_buffer - self.aux_slot = None - try: - if aux_buffer is not None: - self.aux_slot = aux_buffer.alloc_slot().id - except Exception: - if self._retirement is not None: - self._retirement.close() - raise + self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None self.receiver_ready: bool = False self.kv_tasks = [] self.aux_task = None - self._reported_aux_peer_ranks: set[int] = set() - self._has_last_slice = False self.lock = threading.Lock() - self._logical_outcomes = _LogicalOutcomes(self._retirement) + # TaskStatus is a logical aggregate status and cannot represent + # concurrent writes to multiple peers: one failed peer may set ERROR + # while another peer still owns the registered source addresses. + # Track queued-or-active peer writes independently under ``self.lock`` + # so cancellation cannot unregister source addresses before a queued + # worker either finishes NIXL or reports FAILED to the receiver. + self._active_nixl_ops = 0 self._exception: Optional[Exception] = None self._closed = False @@ -2003,20 +1855,7 @@ def __init__( self.transfer_end_time = None # Must be last: makes session visible to listener thread, # so all attributes above must be initialized first. - try: - self._sender.setup_session(self) - except Exception: - with self._sender._sessions_lock: - entry = self._sender._sessions.get(self.disagg_request_id) - if entry is self or (isinstance(entry, weakref.ReferenceType) and entry() is self): - self._sender._sessions.pop(self.disagg_request_id, None) - if self._retirement is not None: - self._retirement.close() - if self._aux_buffer is not None and self.aux_slot is not None: - self._aux_buffer.free_slot(self.aux_slot) - self.aux_slot = None - self._closed = True - raise + self._sender.setup_session(self) @property def disagg_request_id(self) -> int: @@ -2032,40 +1871,31 @@ def disagg_request_id(self) -> int: @property def status(self) -> SessionStatus: - if getattr(self, "_retirement", None) is not None: - terminal = self._logical_outcomes.terminal - if terminal is not None: - return terminal.status if self._terminal_status is not None: return self._terminal_status - # A chunk may fail before the final task is created. - if self._exception is not None or any(t.status == TaskStatus.ERROR for t in self.kv_tasks): - return SessionStatus.ERROR - kv_all_transferred = ( - self._has_last_slice - and bool(self.kv_tasks) - and all(t.status == TaskStatus.TRANSFERRED for t in self.kv_tasks) - ) - aux_done = not self._need_aux or ( - self.aux_task is not None and self.aux_task.status == TaskStatus.TRANSFERRED + kv_all_transferred = bool(self.kv_tasks) and all( + t.status == TaskStatus.TRANSFERRED for t in self.kv_tasks ) - if kv_all_transferred and aux_done: - retirement = getattr(self, "_retirement", None) - if retirement is not None and not retirement.is_complete: - return SessionStatus.TRANSFERRING - return SessionStatus.TRANSFERRED - if kv_all_transferred or ( - self.kv_tasks and any(t.status == TaskStatus.TRANSFERRING for t in self.kv_tasks) - ): + if kv_all_transferred: + if self.aux_task is not None and self.aux_task.status == TaskStatus.TRANSFERRED: + return SessionStatus.FULLY_TRANSFERRED + return SessionStatus.KV_TRANSFERRED + if self.kv_tasks and any(t.status == TaskStatus.TRANSFERRING for t in self.kv_tasks): return SessionStatus.TRANSFERRING return SessionStatus.READY if self.receiver_ready else SessionStatus.INIT - def send(self, chunk: Chunk) -> None: + def send(self, slice: KVSlice) -> None: if self.transfer_start_time is None: self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() with self.lock: - if self._retirement is not None and self._has_last_slice: - raise RuntimeError("cannot send another KV piece after the final slice") + if self._terminal_status in ( + SessionStatus.ERROR, + SessionStatus.CANCELLED, + ): + raise RuntimeError( + f"TxSession {self.disagg_request_id} is " + f"{self._terminal_status.value}; refusing new KV work" + ) if not self.kv_tasks: overall_timeout_s = self._overall_timeout_s if overall_timeout_s is None or overall_timeout_s <= 0: @@ -2074,95 +1904,39 @@ def send(self, chunk: Chunk) -> None: params = self._base_args.params slice_id = len(self.kv_tasks) task = KVSendTask( - chunk, + slice, params, slice_id, prompt_len=self._base_args.prompt_len, ) task._unique_rid = self.disagg_request_id - task.bind_logical_outcomes(self._logical_outcomes) self.kv_tasks.append(task) - self._has_last_slice |= chunk.is_last - if self._retirement is not None and self._has_last_slice: - if not self._need_aux or self.aux_task is not None: - self._logical_outcomes.seal() req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) self._sender.dispatch_task(task, req_info_snapshot) def send_aux(self) -> AuxSendTask: - aux_failures: list[RecvReqInfo] = [] - terminal_error: Optional[Exception] = None - task: Optional[AuxSendTask] = None with self.lock: - if self._retirement is not None and self.aux_task is not None: - return self.aux_task - # In the ownership bridge's generation-first no-retry profile, - # prefill waits for every REQUEST_DATA, so this peer snapshot is sealed. - req_info_snapshot = dict(self._sender._get_req_info(self.disagg_request_id) or {}) - if self._closed or self._terminal_status is not None: - terminal_error = self._exception or RuntimeError( - f"TxSession {self.disagg_request_id} became terminal before auxiliary submission" + if self._terminal_status in ( + SessionStatus.ERROR, + SessionStatus.CANCELLED, + ): + raise RuntimeError( + f"TxSession {self.disagg_request_id} is " + f"{self._terminal_status.value}; refusing new aux work" ) - aux_failures = self._claim_unsubmitted_aux_failures_locked( - req_info_snapshot.values() + if not self._need_aux: + raise RuntimeError( + "Aux transfer requires GENERATION_FIRST scheduling so the " + "receiver reserves its destination buffer before dispatch" ) - else: - params = self._base_args.params - task = AuxSendTask(params, self.aux_slot) - task._unique_rid = self.disagg_request_id - task.bind_logical_outcomes(self._logical_outcomes) - self.aux_task = task - if self._retirement is not None and self._has_last_slice: - self._logical_outcomes.seal() - self._report_unsubmitted_aux_failures(aux_failures) - if terminal_error is not None: - raise terminal_error - assert task is not None + params = self._base_args.params + task = AuxSendTask(params, self.aux_slot) + task._unique_rid = self.disagg_request_id + self.aux_task = task + req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) self._sender.dispatch_task(task, req_info_snapshot) return task - def _claim_unsubmitted_aux_failures_locked( - self, req_infos: Iterable[RecvReqInfo] - ) -> list[RecvReqInfo]: - if not self._enforce_physical_ownership or not self._need_aux: - return [] - claimed = [] - for info in req_infos: - peer_rank = info.instance_rank - if self.aux_task is not None and self.aux_task.has_started_physical_operation( - peer_rank - ): - continue - if not self._claim_aux_terminal_result_locked(peer_rank): - continue - claimed.append(info) - return claimed - - def _claim_aux_terminal_result_locked(self, peer_rank: int) -> bool: - """Claim the only terminal AUX result report permitted for one peer. - - Terminal cleanup and worker completion both use this claim. Exactly one - path may publish terminal evidence so the receiver never sees - contradictory results for the same operation. The caller must hold - ``lock``. - """ - if peer_rank in self._reported_aux_peer_ranks: - return False - self._reported_aux_peer_ranks.add(peer_rank) - return True - - def _claim_aux_terminal_result(self, peer_rank: int) -> bool: - with self.lock: - return self._claim_aux_terminal_result_locked(peer_rank) - - def _claim_unsubmitted_aux_failure(self, info: RecvReqInfo) -> bool: - with self.lock: - return bool(self._claim_unsubmitted_aux_failures_locked((info,))) - - def _report_unsubmitted_aux_failures(self, req_infos: list[RecvReqInfo]) -> None: - for info in req_infos: - self._sender._send_aux_result_to_receiver(info, AgentResult.FAILED) - def pack_aux(self, request: LlmRequest) -> None: """Fill the aux buffer slot with token data from the given request.""" assert self._aux_buffer is not None, "No aux_buffer set for this session" @@ -2171,9 +1945,10 @@ def pack_aux(self, request: LlmRequest) -> None: def is_completed(self) -> bool: """Non-blocking check: has the transfer completed successfully?""" - return self.status is SessionStatus.TRANSFERRED and ( - not getattr(self, "_enforce_physical_ownership", False) or self.resources_drained() - ) + status = self.status + if self._need_aux: + return status == SessionStatus.FULLY_TRANSFERRED + return status in (SessionStatus.KV_TRANSFERRED, SessionStatus.FULLY_TRANSFERRED) def has_failed(self) -> bool: """Non-blocking check: has the transfer failed or been cancelled?""" @@ -2183,73 +1958,45 @@ def has_failed(self) -> bool: return True return self.aux_task is not None and self.aux_task.status == TaskStatus.ERROR - def resources_drained(self) -> bool: - if not getattr(self, "_enforce_physical_ownership", False): - return not any(task.status == TaskStatus.TRANSFERRING for task in self.kv_tasks) - retirement = getattr(self, "_retirement", None) - if retirement is not None and not retirement.can_retire(): - return False - tasks = self.kv_tasks + ([self.aux_task] if self.aux_task is not None else []) - return all(task.resources_drained for task in tasks) - - def _failed_wait_result(self) -> Optional[WaitResult]: - return ( - None - if getattr(self, "_enforce_physical_ownership", False) and not self.resources_drained() - else WaitResult.FAILED - ) - - def _completed_wait_result(self) -> Optional[WaitResult]: - if getattr(self, "_retirement", None) is not None and not self.is_completed(): - return None - return WaitResult.COMPLETED - - def cancel(self, by_peer: bool = False) -> bool: - """Cancel the session and notify the remote receiver. ``True`` if the ask went out. + def cancel(self) -> None: + """Cancel the session and notify the remote receiver. Safe to call multiple times. TRANSFERRING tasks keep running (mid-write). Only INIT tasks have their events signalled immediately. The lock serializes with _deliver_kv_to_agent() so has_transferring_tasks() is accurate the moment this returns. """ - if not self.cancel_local(by_peer=by_peer): - return False - self._sender.send_cancel_to_receivers(self.disagg_request_id) - return True - - def cancel_local(self, by_peer: bool = False, *, report_unsubmitted_aux: bool = True) -> bool: - aux_failures: list[RecvReqInfo] = [] - retirement = getattr(self, "_retirement", None) - if retirement is not None: - retirement.request_drain("peer cancellation" if by_peer else "cancellation") with self.lock: - # A later cancellation must still notify peers after logical failure. - # The logical arbiter preserves whichever outcome already committed. if self._terminal_status == SessionStatus.CANCELLED: - return False - self._logical_outcomes.cancel(by_peer) + return self._terminal_status = SessionStatus.CANCELLED - # Who asked is not recoverable later, and the two differ: the peer asking is a transfer - # error, our own side asking is an ordinary end. - self.cancelled_by_peer = by_peer exc = RuntimeError(f"TxSession {self.disagg_request_id} cancelled") for task in self.kv_tasks: if task.status == TaskStatus.INIT: task.fail(exc) if self.aux_task is not None and self.aux_task.status == TaskStatus.INIT: self.aux_task.fail(exc) - req_infos = list((self._sender._get_req_info(self.disagg_request_id) or {}).values()) - if report_unsubmitted_aux: - aux_failures = self._claim_unsubmitted_aux_failures_locked(req_infos) - self._report_unsubmitted_aux_failures(aux_failures) - return True + # Send outside the lock to avoid holding it during I/O. + self._sender.send_cancel_to_receivers(self.disagg_request_id) def has_transferring_tasks(self) -> bool: - """True if any KV or auxiliary task is currently mid-write. + """True while queued/active peer work still owns session buffers. cancel_request() must return False while this is True. """ - return not self.resources_drained() + with self.lock: + return self._active_nixl_ops > 0 + + def _begin_nixl_op_locked(self) -> None: + """Reserve source-address lifetime; caller must hold ``self.lock``.""" + self._active_nixl_ops += 1 + + def _finish_nixl_op(self) -> None: + with self.lock: + if self._active_nixl_ops <= 0: + logger.error(f"TxSession {self.disagg_request_id} NIXL operation counter underflow") + return + self._active_nixl_ops -= 1 def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: """Poll or block until KV (and optionally aux) transfer finishes. @@ -2263,14 +2010,14 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: or aux is not yet done. """ if self.has_failed(): - return self._failed_wait_result() + return WaitResult.FAILED if not self.kv_tasks: return None if not blocking: has_pending = False for task in self.kv_tasks: if task.status == TaskStatus.ERROR: - return self._failed_wait_result() + return WaitResult.FAILED if task.status != TaskStatus.TRANSFERRED: has_pending = True if has_pending: @@ -2279,12 +2026,10 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: if self.aux_task is None: return None if self.aux_task.status == TaskStatus.ERROR: - return self._failed_wait_result() + return WaitResult.FAILED if self.aux_task.status != TaskStatus.TRANSFERRED: return None - return ( - self._failed_wait_result() if self.has_failed() else self._completed_wait_result() - ) + return WaitResult.COMPLETED # send() normally anchors this once, at the first dispatched KV task. # Keep direct/internal TxSession construction bounded as well, and never @@ -2304,14 +2049,14 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: if wait_slice_s is None or wait_slice_s <= 0: wait_slice_s = _FALLBACK_TX_WAIT_SLICE_S - def wait_for_task(task: SendTaskBase) -> Optional[WaitResult]: + def wait_for_task(task: SendTaskBase) -> WaitResult: while True: # A task/session terminal state observed at the deadline # boundary takes precedence over TIMEOUT. if self.has_failed(): - return self._failed_wait_result() + return WaitResult.FAILED if task.status == TaskStatus.TRANSFERRED: - return self._failed_wait_result() if self.has_failed() else WaitResult.COMPLETED + return WaitResult.COMPLETED remaining_s = None if self._deadline_monotonic_s is not None: @@ -2321,13 +2066,9 @@ def wait_for_task(task: SendTaskBase) -> Optional[WaitResult]: # checks above and the clock read. Preserve boundary # precedence before classifying this as a timeout. if self.has_failed(): - return self._failed_wait_result() + return WaitResult.FAILED if task.status == TaskStatus.TRANSFERRED: - return ( - self._failed_wait_result() - if self.has_failed() - else WaitResult.COMPLETED - ) + return WaitResult.COMPLETED return WaitResult.TIMEOUT timeout_s = wait_slice_s if remaining_s is None else min(wait_slice_s, remaining_s) task.wait(timeout=timeout_s) @@ -2340,7 +2081,7 @@ def wait_for_task(task: SendTaskBase) -> Optional[WaitResult]: if self._need_aux: if self.aux_task is None: # _finalize_send() installs the aux task synchronously before - # returning from respond_and_send_async(). Once every KV task is + # publishing the request to _send_reqs. Once every KV task is # terminal, a missing required aux task is an invariant error, # not an asynchronously pending transfer. with self.lock: @@ -2351,59 +2092,44 @@ def wait_for_task(task: SendTaskBase) -> Optional[WaitResult]: self._exception = RuntimeError( "required auxiliary transfer was not dispatched" ) - self._logical_outcomes.fail(self._exception) self._terminal_status = SessionStatus.ERROR return WaitResult.FAILED result = wait_for_task(self.aux_task) if result != WaitResult.COMPLETED: return result return ( - self._failed_wait_result() + WaitResult.FAILED if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) - else self._completed_wait_result() + else WaitResult.COMPLETED ) - def set_exception(self, reason: str = "") -> None: + def set_exception(self, reason: str = ""): msg = f"TxSession {self.disagg_request_id} exception" if reason: msg += f": {reason}" - aux_failures: list[RecvReqInfo] = [] with self.lock: - if self._exception is None: - self._exception = RuntimeError(msg) - self._logical_outcomes.fail(self._exception) - if self._terminal_status is None: - self._terminal_status = SessionStatus.ERROR + self._exception = RuntimeError(msg) + self._terminal_status = SessionStatus.ERROR for task in self.kv_tasks: if not task.is_done: task.fail(self._exception) if self.aux_task is not None and not self.aux_task.is_done: self.aux_task.fail(self._exception) - req_infos = list((self._sender._get_req_info(self.disagg_request_id) or {}).values()) - aux_failures = self._claim_unsubmitted_aux_failures_locked(req_infos) - self._report_unsubmitted_aux_failures(aux_failures) @property def exception(self) -> Optional[Exception]: return self._exception - def close(self) -> bool: - with self.lock: - if self._closed: - return True - if getattr(self, "_enforce_physical_ownership", False) and not self.resources_drained(): - return False - retirement = getattr(self, "_retirement", None) - if retirement is not None and not retirement.close(): - return False - self._closed = True + def close(self): + if getattr(self, "_closed", False): + return + self._closed = True if self._aux_buffer is not None and self.aux_slot is not None: self._aux_buffer.free_slot(self.aux_slot) self.aux_slot = None # Unregister from Sender; keep fields alive for in-flight worker threads. if self._sender is not None: self._sender.clear_session(self.disagg_request_id) - return True def __enter__(self): return self @@ -2418,169 +2144,53 @@ def __del__(self): logger.warning(f"TxSession.__del__: exception during close: {e}") -class KVRecvTask(_LogicalTask): +class KVRecvTask: def __init__( self, unique_rid: Optional[int], - chunk: Chunk, + kv_slice: KVSlice, slice_id: int, params: DisaggregatedParams, aux_slot: Optional[int], - expected_write_bytes: Optional[int] = None, ): - super().__init__() self._event = threading.Event() self.slice_id = slice_id self.status = TaskStatus.INIT self.expected_transfers = 0 - # One terminal result per writer rank, keyed by the rank the result frame carries. - self._writer_reports: dict[int, bool] = {} - # Verified-range admission accounting: the receiver-computed byte - # total the writers must cover for this piece, and the bytes the - # writers attested as actually submitted-and-completed. None means - # "unknown" and the task can never become write-verified. - self.expected_write_bytes = expected_write_bytes - self.verified_write_bytes = 0 + self.last_slice_count = 0 + # Peers for which REQUEST_DATA was successfully dispatched but no + # KV_AGENT_RESULT has been received yet. This deliberately tracks + # actual dispatches rather than expected topology peers so a partial + # dispatch failure cannot leave a phantom active write forever. + self.pending_peer_ranks: set[int] = set() + self.resolved_peer_ranks: set[int] = set() + self.aux_reserved_peer_ranks: set[int] = set() + self.is_broadcast_dispatch = False self._unique_rid = unique_rid - self._chunk = chunk + self._kv_slice = kv_slice self._params = params self._exception: Optional[Exception] = None self._aux_slot = aux_slot self._perf_timer = PerfTimer() if perf_log_manager.enabled else None - self._physical_owner: Optional[_ReceiveOperationOwner] = None - self._ownership_state_lock: Optional[threading.Lock] = None def fail(self, exc: Exception) -> None: - self._logical_outcomes.fail(exc) - if self._ownership_state_lock is None: - self._exception = exc - self.status = TaskStatus.ERROR - self._event.set() - return - with self._ownership_state_lock: - self._exception = exc - self.status = TaskStatus.ERROR - self._event.set() + self._exception = exc + self.status = TaskStatus.ERROR + self._event.set() def complete(self) -> None: - if self._ownership_state_lock is None: - self._logical_outcomes.complete(self._logical_index) - self.status = TaskStatus.TRANSFERRED - self._event.set() - return - with self._ownership_state_lock: - if self.status == TaskStatus.ERROR: - return - self._logical_outcomes.complete(self._logical_index) - self.status = TaskStatus.TRANSFERRED - self._event.set() + self.status = TaskStatus.TRANSFERRED + self._event.set() def wait(self, timeout: Optional[float] = None) -> bool: """Block until terminal state. Returns True if done, False on timeout.""" return self._event.wait(timeout=timeout) - @property - def write_verified(self) -> bool: - """True when every published destination byte has attested write coverage. - - Strict equality: an over-count (more attested bytes than published) - indicates duplicated or misrouted writes and is as disqualifying as a - missing write. - """ - return ( - self.status == TaskStatus.TRANSFERRED - and self.expected_write_bytes is not None - and self.verified_write_bytes == self.expected_write_bytes - ) - @property def is_done(self) -> bool: return self._event.is_set() - def note_writer_report(self, peer_rank: int, succeeded: bool) -> tuple[bool, bool]: - """Record one writer's terminal word: whether this rank had yet to speak, and whether every - expected writer has now succeeded. - - The first answer is an edge, not a level. A rank keeps the first result it reports, and a - repeat is refused outright, so the caller settles the piece once rather than once per frame. - Callers hold ``RxSession.lock``, which serializes the writes. - """ - # TODO: a report from a rank outside the authorized writer set is counted, not rejected. - if peer_rank in self._writer_reports: - previous = self._writer_reports[peer_rank] - if previous != succeeded: - logger.warning( - f"KV writer {peer_rank} reported success={succeeded} for slice {self.slice_id} " - f"after reporting success={previous}; the first result stands" - ) - return False, False - self._writer_reports[peer_rank] = succeeded - all_succeeded = len(self._writer_reports) == self.expected_transfers and all( - self._writer_reports.values() - ) - return True, all_succeeded - - @property - def reports_outstanding(self) -> bool: - """Whether a writer's terminal word about this task is still missing. - - A logical outcome is not an answer to this: one writer's failure settles the task while - its siblings write on. The ownership owner tracks its own cohort and answers for itself. - """ - if self._physical_owner is not None: - return not self._physical_owner.all_writers_reported - return len(self._writer_reports) < self.expected_transfers - - def begin_publication(self) -> None: - if self._physical_owner is None: - self._physical_owner = _ReceiveOperationOwner(self._retirement) - self._ownership_state_lock = threading.Lock() - self._physical_owner.begin_publication() - - def _get_physical_owner(self) -> _ReceiveOperationOwner: - if self._physical_owner is None: - raise RuntimeError("receive operation owner was not initialized") - return self._physical_owner - - def seal_writer_cohort(self, writer_cohort: Optional[set[int]] = None) -> None: - self._get_physical_owner().seal_writer_cohort(self.expected_transfers, writer_cohort) - - def finish_publication(self) -> None: - self._get_physical_owner().finish_publication() - - def abort_publication(self, published_writers: set[int]) -> None: - self._get_physical_owner().abort_publication(published_writers) - - def cancel_unpublished(self) -> bool: - return self._get_physical_owner().cancel_unpublished() - - def record_writer_in_doubt(self, peer_rank: int) -> bool: - return self._get_physical_owner().record_writer_in_doubt(peer_rank) - - def record_writer_settlement(self, peer_rank: int) -> bool: - return self._get_physical_owner().record_writer_settlement(peer_rank) - - def record_writer_result( - self, - peer_rank: int, - succeeded: bool, - *, - wait_for_local_completion: bool, - ) -> tuple[bool, bool]: - return self._get_physical_owner().record_writer_result( - peer_rank, - succeeded, - wait_for_local_completion=wait_for_local_completion, - ) - - def finish_local_completion(self) -> None: - self._get_physical_owner().finish_local_completion() - - @property - def resources_drained(self) -> bool: - return self._get_physical_owner().resources_drained - def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int): if self._perf_timer is None: return @@ -2600,16 +2210,10 @@ def __init__( peer_registrar: PeerRegistrar, agent: BaseTransferAgent, bounce=None, - enforce_physical_ownership: bool = False, - retirement_watchdog: Optional[RetirementWatchdog] = None, ): self._registrar = peer_registrar self._agent = agent self._bounce = bounce - # Internal component gate only. Production wiring is added after the - # complete Phase 1 ownership path is available and qualified. - self._enforce_physical_ownership = enforce_physical_ownership - self._retirement_watchdog = retirement_watchdog self._dealers = {} self._sender_ep_instance_map = {} # info_endpoint -> diagnostic message for peers that failed the @@ -2621,13 +2225,9 @@ def __init__( self._messenger = ZMQMessenger(mode="ROUTER") self._sessions = {} # unique_rid -> RxSession self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids - # unique_rid -> whether the peer asked, for cancels that arrive before the session exists. - # Who asked is not recoverable from the id alone, and the two differ: a peer's cancel is a - # transfer error, our own side's is an ordinary end. - self._pre_cancelled_rids: dict[int, bool] = {} + self._pre_cancelled_rids: set[int] = set() self._shutdown = False - self._ownership_admission_lock = threading.Lock() - self._ownership_poisoned: Optional[Exception] = None + self._reliable_message_outbox = _ReliableMessageOutbox("Receiver") self._start_listener() logger.info(f"Receiver init with endpoint: {self._messenger.endpoint}") @@ -2637,96 +2237,49 @@ def endpoint(self): return self._messenger.endpoint def shutdown(self): - watchdog = getattr(self, "_retirement_watchdog", None) - if watchdog is not None: - watchdog.request_shutdown() - watchdog.require_retired() - with self._sessions_lock: - if self._shutdown: - return - if self._enforce_physical_ownership and self._sessions: - raise RuntimeError("Receiver refuses shutdown while transfer ownership is active") - self._shutdown = True - for dealer in self._dealers.values(): + if getattr(self, "_shutdown", False): + return + self._shutdown = True + # Quiesce listener callbacks before putting the outbox sentinel behind + # all already-accepted CANCEL_SESSION messages. + messenger = getattr(self, "_messenger", None) + if messenger is not None: + messenger.stop() + reliable_outbox = getattr(self, "_reliable_message_outbox", None) + if reliable_outbox is not None: + reliable_outbox.shutdown() + dealers = getattr(self, "_dealers", {}) + for dealer in dealers.values(): try: dealer.stop() except Exception as e: logger.warning(f"Failed to stop dealer during Receiver shutdown: {e}") - self._dealers.clear() - self._messenger.stop() + dealers.clear() def clear_session(self, unique_rid: int): with self._sessions_lock: self._sessions.pop(unique_rid, None) - def _get_ownership_admission_lock(self) -> threading.Lock: - lock = getattr(self, "_ownership_admission_lock", None) - if lock is None: - lock = threading.Lock() - self._ownership_admission_lock = lock - return lock - - def _acquire_ownership_admission(self) -> Callable[[], None]: - lock = self._get_ownership_admission_lock() - lock.acquire() - watchdog = getattr(self, "_retirement_watchdog", None) - if watchdog is not None and watchdog.fatal is not None: - lock.release() - raise RuntimeError("receiver retirement admission is closed") - error = getattr(self, "_ownership_poisoned", None) - if error is not None: - lock.release() - raise RuntimeError("receiver ownership admission is quarantined") from error - return lock.release - - def _poison_ownership_locked(self, error: Exception) -> None: - if getattr(self, "_ownership_poisoned", None) is None: - self._ownership_poisoned = error - def setup_session(self, rx_session: RxSessionBase): - release_admission = None - if getattr(rx_session, "_enforce_physical_ownership", False): - release_admission = self._acquire_ownership_admission() pre_cancel = False - cancelled_by_peer = False - try: - with self._sessions_lock: - if self._shutdown: - raise RuntimeError("cannot create an RxSession after Receiver shutdown") - if getattr(rx_session, "_enforce_physical_ownership", False): - self._sessions[rx_session.disagg_request_id] = rx_session - else: - self._sessions[rx_session.disagg_request_id] = weakref.ref(rx_session) - if rx_session.disagg_request_id in self._pre_cancelled_rids: - pre_cancel = True - cancelled_by_peer = self._pre_cancelled_rids.pop(rx_session.disagg_request_id) - finally: - if release_admission is not None: - release_admission() + with self._sessions_lock: + self._sessions[rx_session.disagg_request_id] = weakref.ref(rx_session) + if rx_session.disagg_request_id in self._pre_cancelled_rids: + pre_cancel = True + self._pre_cancelled_rids.discard(rx_session.disagg_request_id) if pre_cancel: - if getattr(rx_session, "_enforce_physical_ownership", False): - rx_session.cancel_local(by_peer=cancelled_by_peer) - else: - rx_session.cancel(by_peer=cancelled_by_peer) + rx_session.cancel() def _get_session(self, unique_rid: Optional[int]) -> Optional["RxSession"]: with self._sessions_lock: - session_entry = self._sessions.get(unique_rid) - return self._resolve_session_entry(unique_rid, session_entry) - - @staticmethod - def _resolve_session_entry( - unique_rid: Optional[int], - session_entry: Optional[Union["RxSession", weakref.ReferenceType["RxSession"]]], - ) -> Optional["RxSession"]: - if session_entry is None: + session_ref = self._sessions.get(unique_rid) + if session_ref is None: return None - if isinstance(session_entry, weakref.ReferenceType): - session = session_entry() - if session is None: - logger.warning(f"RxSession {unique_rid} has been garbage collected") - return session - return session_entry + session = session_ref() + if session is None: + logger.warning(f"RxSession {unique_rid} has been garbage collected") + return None + return session def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: self_ri = self._registrar.self_rank_info @@ -2749,11 +2302,13 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: sender_req_id=sender_req_id, instance_name=self_ri.instance_name, instance_rank=self_ri.instance_rank, - block_ids_per_layer_groups=task._chunk.block_ids_per_layer_groups, + block_ids_per_layer_groups=task._kv_slice.block_ids_per_layer_groups, unique_rid=task._unique_rid, dst_start_token=None, aux_slot=task._aux_slot, + mamba_state_index=task._kv_slice.mamba_state_index, slice_id=task.slice_id, + need_aux=(task._params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST), ) @staticmethod @@ -2764,8 +2319,8 @@ def _fanin_bounce_safe( ) -> bool: """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. The split assumes every writer contributes the same size, which holds when: - * duplicate_head_factor == 1 -- else some ranks don't send KV yet still count in - expected_transfers, so the live writers overflow their slots; + * duplicate_head_factor == 1 -- else some ranks don't send KV (should_send_kv) yet still + count in expected_transfers, so the live writers overflow their slots; * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; for PP fan-in, every overlapping stage must hold the same number of layers (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage @@ -2785,6 +2340,19 @@ def _fanin_bounce_safe( # a masked PP stage may advertise no replicated view even though another # stage owns one that is visible in the receiver's page table. if len(overlap.ranks) > 1: + if peer_ri.page_table is not None: + for layer_group in peer_ri.page_table.layer_groups: + if not getattr(layer_group, "cp_as_tp", False): + continue + draft_tp, _ = MambaPolicy._mamba_tp(peer_ri) + if ( + layer_group.total_kv_head_num is None + or layer_group.kv_head_num_per_rank * draft_tp + > layer_group.total_kv_head_num + ): + # Draft head replicas may elect fewer writers than the + # target's heads, so their payload shares are unequal. + return False for page_table in (peer_ri.page_table, receiver_page_table): if page_table is None: continue @@ -2794,18 +2362,6 @@ def _fanin_bounce_safe( return False return True - def _get_mamba_slot(self, req_info: RecvReqInfo) -> Optional[int]: - """Extract the mamba slot index from block_ids_per_layer_groups, or None.""" - pt = self._registrar.self_extractor.page_table - if pt is None: - return None - for lg_idx, lg in enumerate(pt.layer_groups): - if lg.kind == CacheKind.STATE: - block_ids = req_info.block_ids_per_layer_groups[lg_idx] - if block_ids.size > 0: - return int(block_ids[0]) - return None - def dispatch_task(self, task: KVRecvTask) -> None: params = task._params logger.debug( @@ -2826,17 +2382,11 @@ def dispatch_task(self, task: KVRecvTask) -> None: # validation and may still run a transfer that is then # discarded. The task is already in the session's _kv_tasks, no # bounce reservation exists yet, and session._sender_endpoints - # is still empty, so the unpublished owner can close locally. + # is still empty, so no cleanup is needed here. logger.error( "dispatch_task: context peer incompatible, failing request " f"unique_rid={task._unique_rid}: {e}" ) - if self._enforce_physical_ownership: - session = self._get_session(task._unique_rid) - if session is None: - task.cancel_unpublished() - else: - session.cancel_unpublished_task(task) task.fail(e) return @@ -2868,6 +2418,7 @@ def dispatch_task(self, task: KVRecvTask) -> None: task.expected_transfers = len(peer_overlap.ranks) else: task.expected_transfers = len(dp0_overlap.ranks) + task.is_broadcast_dispatch = True # TP fan-in splits ONE region equally, so allow it only for a uniform writer set: # _fanin_bounce_safe() (TP-by-head / even-PP), and never under ADP broadcast (sender_dp_rank # None), where the real writer count exceeds expected_transfers and would overflow the slot. @@ -2885,18 +2436,21 @@ def dispatch_task(self, task: KVRecvTask) -> None: self._registrar.self_extractor.page_table, ) ) - # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks, - # so the bounce region must be sized for those bytes too. + # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks (the sender + # appends its MambaPolicy fragments in _build_kv_write_meta), so the bounce region must be + # sized for those bytes too or the write would overrun into the neighboring slot. extra_bytes = 0 - mamba_dst_slot = self._get_mamba_slot(receiver_req) - if mamba_dst_slot is not None: + if receiver_req.mamba_state_index is not None: if peer_infos.page_table is None: allow_bounce = False # cannot size the recurrent-state payload else: - extra_bytes = mamba_receiver_payload_bytes( + # Receiver-local sizing: RankInfoServer runs on ctx rank 0 + # only, so a sender-side simulation under-reserves for every + # receiver whose paired sender is not rank 0. + extra_bytes = MambaPolicy.receiver_payload_bytes( sender_page_table=peer_infos.page_table, receiver_page_table=self._registrar.self_extractor.page_table, - dst_slot=mamba_dst_slot, + dst_slot=receiver_req.mamba_state_index, ) bounced = allow_bounce and self._bounce.reserve( receiver_req, task.expected_transfers, extra_bytes=extra_bytes @@ -2907,65 +2461,39 @@ def dispatch_task(self, task: KVRecvTask) -> None: f"dispatch_task: RxSession {task._unique_rid} not found; " "session may have been closed before dispatch" ) - sender_endpoints = {peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks} # Fan-in: each sender gets its own sub-region base (writers must not overwrite); else serialize once. fanin_bounce = bounced and task.expected_transfers > 1 key = (receiver_req.unique_rid, receiver_req.slice_id) receiver_req_bytes = None if fanin_bounce else receiver_req.to_bytes() - if not self._enforce_physical_ownership: - session.mark_transferring(task.slice_id) - # Cache sender endpoints so cancel() can send CANCEL_SESSION to them. - session._sender_endpoints.update(sender_endpoints) - for i, rank in enumerate(peer_overlap.ranks): - if task._perf_timer is not None: - task._perf_timer.record_task_start(rank) - if fanin_bounce: - receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) - receiver_req_bytes = receiver_req.to_bytes() - self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) - return - - peer_ranks = list(peer_overlap.ranks) - serialized_requests: list[tuple[int, bytes]] = [] - for i, rank in enumerate(peer_ranks): + successful_dispatches = 0 + for i, rank in enumerate(peer_overlap.ranks): + endpoint = peer_infos.sender_endpoints[rank] + if not session.begin_peer_transfer(task.slice_id, rank, endpoint): + exc = RuntimeError( + f"dispatch_task: RxSession {task._unique_rid} was cancelled " + f"before dispatch to peer_rank={rank}" + ) + task.fail(exc) + raise exc + if task._perf_timer is not None: + task._perf_timer.record_task_start(rank) if fanin_bounce: receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) - payload = receiver_req.to_bytes() - else: - assert receiver_req_bytes is not None - payload = receiver_req_bytes - serialized_requests.append((rank, payload)) - - published_writers: set[int] = set() - - def publish_requests() -> None: - for rank, payload in serialized_requests: - if task._perf_timer is not None: - task._perf_timer.record_task_start(rank) - published_writers.add(rank) - try: - self._request_sender_data(peer_infos.sender_endpoints[rank], payload) - except Exception: - published_writers.discard(rank) - raise - - # Gen-first ADP publishes the destination to every eligible DP group, - # but the qualified immutable-request, no-retry/no-reroute profile - # admits exactly one group. Track that group's distinct terminal - # writers by count instead of sealing the full broadcast union. - writer_cohort = ( - set(peer_ranks) if sender_dp_rank is not None or peer_infos.dp_size == 1 else None - ) - if not session.try_begin_transfer( - task.slice_id, - sender_endpoints, - writer_cohort, - writer_candidates=set(peer_ranks), - publish=publish_requests, - published_writers=published_writers, - ): - self._bounce.release_idle_reservation(key) - task.cancel_unpublished() + receiver_req_bytes = receiver_req.to_bytes() + try: + self._request_sender_data(endpoint, receiver_req_bytes) + except Exception as exc: + # The synchronous send did not complete, so this peer cannot + # own the destination. Keep only earlier successful peer + # reservations alive for the async cancellation/drain path. + session.abort_peer_transfer(task.slice_id, rank) + task.fail(exc) + logger.error( + f"Receiver.dispatch_task failed for rid={task._unique_rid}, " + f"peer_rank={rank}, successful_dispatches={successful_dispatches}: {exc}" + ) + raise + successful_dispatches += 1 return @staticmethod @@ -3004,7 +2532,7 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: finally: messenger.stop() - # Peer compatibility gates on the receiver side. + # Recurrent-state (Mamba/KDA) layout gate on the receiver side. # The sender-side check (PeerRegistrar.register) runs in the # sender's listener thread, where exceptions are only logged, so # reject here — before REGISTER_RANK_INFO is even sent, so no @@ -3013,17 +2541,12 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: # (handled in dispatch_task) so only requests targeting this peer # fail, and cached so later requests fail fast. try: - self._registrar.validate_bounded_replay_compatible(sender_info) MambaPolicy.validate_peer_compatible( self._registrar.self_rank_info, sender_info, self._registrar.self_extractor.page_table, sender_info.page_table, ) - AttentionPolicy.validate_peer_compatible( - self._registrar.self_extractor.page_table, - sender_info.page_table, - ) except ValueError as e: msg = ( f"context peer at '{info_endpoint}' is incompatible: {e} " @@ -3046,22 +2569,18 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: else: return self._sender_ep_instance_map[info_endpoint] - def send_cancel_to_senders(self, unique_rid: int, sender_endpoints: set[str]) -> set[str]: - """Notify senders and return the endpoints whose cancel message was sent.""" - message = [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] - if self._enforce_physical_ownership: - rank_info = self._registrar.self_rank_info - message.extend( - [rank_info.instance_name.encode(), str(rank_info.instance_rank).encode("ascii")] - ) - sent = set() + def send_cancel_to_senders(self, unique_rid: int, sender_endpoints: set[str]) -> None: + """Notify all senders involved in this session to cancel.""" for endpoint in sender_endpoints: try: - self._get_or_connect_dealer(endpoint).send(message) - sent.add(endpoint) + accepted = self._reliable_message_outbox.submit( + endpoint, + [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")], + ) + if not accepted: + logger.error(f"send_cancel_to_senders: outbox rejected rid={unique_rid}") except Exception as e: logger.warning(f"send_cancel_to_senders: failed for rid={unique_rid}: {e}") - return sent def _start_listener(self): def handle_message(messages: list[bytes]) -> bool: @@ -3080,11 +2599,6 @@ def handle_message(messages: list[bytes]) -> bool: self._process_aux_agent_result(send_id, msg) except Exception as e: logger.error(f"Receiver: error handling AUX_AGENT_RESULT: {e}") - case MessageType.SESSION_QUIESCED: - try: - self._process_session_quiesced(msg) - except Exception as e: - logger.error(f"Receiver: error handling SESSION_QUIESCED: {e}") case MessageType.CANCEL_SESSION: try: self._handle_cancel_session(msg) @@ -3096,30 +2610,19 @@ def handle_message(messages: list[bytes]) -> bool: self._messenger.start_listener(handle_message) - def _process_session_quiesced(self, message: list[bytes]) -> None: - """Dispatch no-future-access proof only to the matching live request.""" - if not self._enforce_physical_ownership: - return - peer_rank, unique_rid = int(message[1]), int(message[2]) - session = self._get_session(unique_rid) - if session is not None: - session.process_session_quiesced(peer_rank) - def _handle_cancel_session(self, message: list[bytes]): unique_rid = int(message[1]) session = None with self._sessions_lock: - session_entry = self._sessions.get(unique_rid) - session = self._resolve_session_entry(unique_rid, session_entry) - if session is None: - # Parked with the peer as who asked: the session built afterwards reads its - # cancellation from here, and an id on its own would read as a local cancel. - self._pre_cancelled_rids[unique_rid] = True - if session is not None: - if self._enforce_physical_ownership: - session.cancel_local(by_peer=True) + session_ref = self._sessions.get(unique_rid) + if session_ref is None: + self._pre_cancelled_rids.add(unique_rid) else: - session.cancel(by_peer=True) + session = session_ref() + if session is None: + self._pre_cancelled_rids.add(unique_rid) + if session is not None: + session.cancel() def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): if message[0] != MessageType.KV_AGENT_RESULT: @@ -3127,7 +2630,7 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): f"_process_kv_agent_result: unexpected msg_type={message[0]!r}, expected KV_AGENT_RESULT" ) return - peer_rank, unique_rid, receiver_slice_id, is_last_slice, status_code, transfer_size = ( + peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code, transfer_size = ( _KV_RESULT_PREFIX.unpack(message[1]) ) from .bounce import decode_result_tail @@ -3141,7 +2644,7 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): return session.process_kv_agent_result( peer_rank, - receiver_slice_id, + sender_slice_id, is_last_slice, _AGENT_RESULT_BY_CODE[status_code], dst_ptrs=dst_ptrs, @@ -3182,9 +2685,6 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): class RxSession(RxSessionBase): - cancelled_by_peer: bool = False - """Only meaningful once the session is CANCELLED.""" - def __init__( self, request_id: int, @@ -3193,29 +2693,20 @@ def __init__( aux_buffer: Optional[AuxBuffer] = None, timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, - *, - retirement_watchdog: Optional[RetirementWatchdog] = None, + beam_width: int = 1, ): - super().__init__(receiver, SessionArgsBase(params, prompt_len=prompt_len)) + super().__init__( + receiver, + SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), + ) self._timeout_s = timeout_s - self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST - self._enforce_physical_ownership = getattr(receiver, "_enforce_physical_ownership", False) + self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST or ( + aux_buffer is not None and aux_buffer.has_draft_history + ) self._receiver: Receiver # narrow base class type for Pylance self.request_id = request_id - self._retirement = None - if retirement_watchdog is not None: - if timeout_s is None: - raise ValueError("deadline retirement requires the transfer timeout") - self._retirement = retirement_watchdog.create_owner(request_id, "receive", timeout_s) self._aux_buffer = aux_buffer - self.aux_slot = None - try: - if aux_buffer is not None: - self.aux_slot = aux_buffer.alloc_slot().id - except Exception: - if self._retirement is not None: - self._retirement.close() - raise + self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None self._exception: Optional[Exception] = None self._closed = False self._terminal_status: Optional[SessionStatus] = None @@ -3225,45 +2716,18 @@ def __init__( self._kv_tasks: list[KVRecvTask] = [] self._aux_count = 0 self._aux_status: TaskStatus = TaskStatus.INIT - self._publication_may_have_escaped = False - self._aux_physical_owner: Optional[_ReceiveOperationOwner] = None - if self._enforce_physical_ownership and self._need_aux: - self._aux_physical_owner = _ReceiveOperationOwner(self._retirement) - self._aux_physical_owner.begin_publication() + self._pending_aux_peer_ranks: set[int] = set() + self._resolved_aux_peer_ranks: set[int] = set() self._sender_endpoints: set[str] = set() - self._cancelled_sender_endpoints: set[str] = set() - self._cancel_notification_started = False - # Serialize REQUEST_DATA publication with cancellation notification - # without holding the session state lock across a potentially blocking - # network send. The ordering is publication -> cancellation whenever - # publication wins the state transition. - self._publication_lock = threading.Lock() + # Error-triggered ADP cancellation is endpoint-idempotent. Sender + # notifications are claimed under ``lock`` and sent after releasing it. + self._failure_cancelled_sender_endpoints: set[str] = set() self.lock = threading.Lock() - self._logical_outcomes = _LogicalOutcomes(self._retirement) - self._aux_logical_index = ( - self._logical_outcomes.add_task() - if self._retirement is not None and self._need_aux - else None - ) - try: - self._receiver.setup_session(self) - except Exception: - if self._aux_physical_owner is not None: - self._aux_physical_owner.cancel_unpublished() - if self._retirement is not None: - self._retirement.close() - with self._receiver._sessions_lock: - session_entry = self._receiver._sessions.get(self.disagg_request_id) - if ( - self._receiver._resolve_session_entry(self.disagg_request_id, session_entry) - is self - ): - self._receiver._sessions.pop(self.disagg_request_id, None) - if self._aux_buffer is not None and self.aux_slot is not None: - self._aux_buffer.free_slot(self.aux_slot) - self.aux_slot = None - self._closed = True - raise + # See TxSession._active_nixl_ops. On the receiver side each + # successfully dispatched peer request conservatively owns one + # destination operation until its agent-result message arrives. + self._active_nixl_ops = 0 + self._receiver.setup_session(self) @property def disagg_request_id(self) -> int: @@ -3277,239 +2741,126 @@ def disagg_request_id(self) -> int: return params.ctx_request_id return self.request_id - @contextmanager - def _ownership_evidence_guard(self) -> Iterator[None]: - get_lock = getattr(self._receiver, "_get_ownership_admission_lock", None) - if not self._enforce_physical_ownership or get_lock is None: - yield - return - with get_lock(): - yield - - def _poison_receiver_ownership(self, error: Exception) -> None: - poison = getattr(self._receiver, "_poison_ownership_locked", None) - if poison is not None: - poison(error) - - def _record_ownership_evidence_error(self, error: Exception) -> None: - """Record a fatal ownership-evidence error and close receiver admission.""" - self._logical_outcomes.fail(error) - self._exception = error - if self._terminal_status is None: - self._terminal_status = SessionStatus.ERROR - self._poison_receiver_ownership(error) - @property def status(self) -> SessionStatus: - if getattr(self, "_retirement", None) is not None: - terminal = self._logical_outcomes.terminal - if terminal is not None: - return terminal.status if self._terminal_status is not None: return self._terminal_status if self._exception is not None or any(t.status == TaskStatus.ERROR for t in self._kv_tasks): return SessionStatus.ERROR if self._kv_tasks: kv_all_transferred = all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks) - aux_done = not self._need_aux or self._aux_status == TaskStatus.TRANSFERRED - if kv_all_transferred and aux_done: - retirement = getattr(self, "_retirement", None) - if retirement is not None and not retirement.is_complete: - return SessionStatus.TRANSFERRING - return SessionStatus.TRANSFERRED - if kv_all_transferred or any( - t.status == TaskStatus.TRANSFERRING for t in self._kv_tasks - ): + if kv_all_transferred and self._aux_status == TaskStatus.TRANSFERRED: + return SessionStatus.FULLY_TRANSFERRED + if kv_all_transferred: + return SessionStatus.KV_TRANSFERRED + if any(t.status == TaskStatus.TRANSFERRING for t in self._kv_tasks): return SessionStatus.TRANSFERRING return SessionStatus.INIT - def try_begin_transfer( - self, - slice_id: int, - sender_endpoints: set[str], - writer_cohort: Optional[set[int]] = None, - *, - writer_candidates: Optional[set[int]] = None, - publish: Optional[Callable[[], None]] = None, - published_writers: Optional[set[int]] = None, - ) -> bool: - """Seal and publish a writer cohort unless cancellation won first.""" - if publish is not None and published_writers is None: - raise ValueError("publication must track successfully queued writers") - with self._publication_lock: - # Only metadata and publication reservation use the admission gate. - # Network I/O must not hold the arbiter needed by the watchdog. - retirement = getattr(self, "_retirement", None) - with ( - self._ownership_evidence_guard(), - self.lock, - retirement.lock if retirement is not None else nullcontext(), - ): - if self._closed or self._terminal_status is not None: - return False - error = getattr(self._receiver, "_ownership_poisoned", None) - if error is not None: - raise RuntimeError("receiver ownership admission is quarantined") from error - task = self._kv_tasks[slice_id] - owners = [task._get_physical_owner()] - if self._aux_physical_owner is not None: - owners.append(self._aux_physical_owner) - unpublished = [owner for owner in owners if owner._expected_writers is None] - try: - for owner in owners: - owner.seal_writer_cohort( - task.expected_transfers, - writer_cohort, - published_writers=writer_candidates, - ) - if retirement is not None and not retirement.expose(*owners): - raise RuntimeError("destination retirement admission is closed") - except Exception: - # No send has run for this reservation. Roll back only new - # owners; AUX may already belong to an earlier published slice. - for owner in unpublished: - owner.abort_publication(set()) - raise - task.status = TaskStatus.TRANSFERRING - self._sender_endpoints.update(sender_endpoints) - if publish is not None: - assert published_writers is not None - self._publication_may_have_escaped = True - try: - publish() - expected_publication = ( - writer_candidates if writer_candidates is not None else writer_cohort - ) - if ( - expected_publication is not None - and published_writers != expected_publication - ): - raise RuntimeError("publication did not queue the complete writer cohort") - except Exception: - # ZMQ delivers multipart messages atomically. A successful - # send transfers responsibility to that writer; a failed - # send does not. Retain only the successfully queued prefix - # and wait for each of those writers to settle. - self._receiver._bounce.abort_publication( - (self.disagg_request_id, task.slice_id), - published_writers, - ) - task.abort_publication(published_writers) - if self._aux_physical_owner is not None: - self._aux_physical_owner.abort_publication(published_writers) - raise - task.finish_publication() - if self._aux_physical_owner is not None: - self._aux_physical_owner.finish_publication() + def begin_peer_transfer(self, slice_id: int, peer_rank: int, sender_endpoint: str) -> bool: + """Reserve destination lifetime before dispatching REQUEST_DATA. + + Returns False when cancellation/error won the race, in which case the + caller must not contact the sender. + """ + with self.lock: + if self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED): + return False + task = self._kv_tasks[slice_id] + if peer_rank in task.pending_peer_ranks: + raise RuntimeError( + f"duplicate peer dispatch for request {self.request_id}, " + f"slice={slice_id}, peer_rank={peer_rank}" + ) + task.pending_peer_ranks.add(peer_rank) + reserve_aux = ( + self._need_aux + and peer_rank not in self._pending_aux_peer_ranks + and peer_rank not in self._resolved_aux_peer_ranks + ) + if reserve_aux: + self._pending_aux_peer_ranks.add(peer_rank) + task.aux_reserved_peer_ranks.add(peer_rank) + task.status = TaskStatus.TRANSFERRING + self._sender_endpoints.add(sender_endpoint) + self._active_nixl_ops += 1 + int(reserve_aux) return True - def mark_transferring(self, slice_id: int, writer_cohort: Optional[set[int]] = None) -> None: - if not self._enforce_physical_ownership: - with self.lock: - self._kv_tasks[slice_id].status = TaskStatus.TRANSFERRING + def abort_peer_transfer(self, slice_id: int, peer_rank: int) -> None: + """Undo a reservation whose REQUEST_DATA send raised synchronously.""" + with self.lock: + task = self._kv_tasks[slice_id] + if peer_rank not in task.pending_peer_ranks: + return + task.pending_peer_ranks.remove(peer_rank) + op_count = 1 + if peer_rank in task.aux_reserved_peer_ranks: + task.aux_reserved_peer_ranks.remove(peer_rank) + if peer_rank in self._pending_aux_peer_ranks: + self._pending_aux_peer_ranks.remove(peer_rank) + op_count += 1 + if self._active_nixl_ops < op_count: + logger.error(f"RxSession {self.disagg_request_id} NIXL operation counter underflow") + return + self._active_nixl_ops -= op_count + + def _discard_broadcast_kv_phantoms_locked(self, task: KVRecvTask) -> None: + """Drop non-selected ADP broadcast peers after all real responders reply.""" + if ( + not task.is_broadcast_dispatch + or self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + or len(task.resolved_peer_ranks) < task.expected_transfers + ): return - if writer_cohort is None: - raise ValueError("ownership-mode publication requires an explicit writer cohort") - if not self.try_begin_transfer(slice_id, set(), writer_cohort): - raise RuntimeError( - f"RxSession {self.disagg_request_id} became terminal before publication" - ) + phantom_count = len(task.pending_peer_ranks) + if phantom_count == 0: + return + task.pending_peer_ranks.clear() + self._active_nixl_ops = max(0, self._active_nixl_ops - phantom_count) - def receive(self, chunk: Chunk, expected_write_bytes: Optional[int] = None) -> None: + def _discard_broadcast_aux_phantoms_locked(self, task: KVRecvTask) -> None: + if ( + not task.is_broadcast_dispatch + or self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED) + or len(self._resolved_aux_peer_ranks) < task.expected_transfers + ): + return + phantom_count = len(self._pending_aux_peer_ranks) + if phantom_count == 0: + return + phantom_ranks = set(self._pending_aux_peer_ranks) + self._pending_aux_peer_ranks.clear() + for kv_task in self._kv_tasks: + kv_task.aux_reserved_peer_ranks.difference_update(phantom_ranks) + self._active_nixl_ops = max(0, self._active_nixl_ops - phantom_count) + + def _claim_broadcast_failure_cancel_locked(self, task: KVRecvTask) -> set[str]: + """Claim sender endpoints that still need ERROR-path cancellation.""" + if not task.is_broadcast_dispatch or self._terminal_status != SessionStatus.ERROR: + return set() + endpoints = self._sender_endpoints.difference(self._failure_cancelled_sender_endpoints) + self._failure_cancelled_sender_endpoints.update(endpoints) + return endpoints + + def receive(self, slice: KVSlice) -> None: if self.transfer_start_time is None: self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() - if self._enforce_physical_ownership: - task = self.prepare_receive(chunk, expected_write_bytes=expected_write_bytes) - if task is not None: - self.dispatch_prepared_receive(task) - return params = self._base_args.params slice_id = len(self._kv_tasks) task = KVRecvTask( self.disagg_request_id, - chunk, + slice, slice_id, params, aux_slot=self.aux_slot, - expected_write_bytes=expected_write_bytes, ) - task.bind_logical_outcomes(self._logical_outcomes) self._kv_tasks.append(task) self._receiver.dispatch_task(task) - def prepare_receive( - self, chunk: Chunk, expected_write_bytes: Optional[int] = None - ) -> Optional[KVRecvTask]: - """Create an unpublished task under the cancellation linearization lock.""" - with self.lock: - if self._closed or self._terminal_status is not None: - return None - params = self._base_args.params - task = KVRecvTask( - self.disagg_request_id, - chunk, - len(self._kv_tasks), - params, - aux_slot=self.aux_slot, - expected_write_bytes=expected_write_bytes, - ) - task.bind_logical_outcomes(self._logical_outcomes) - task.begin_publication() - self._kv_tasks.append(task) - if self._retirement is not None and chunk.is_last: - self._logical_outcomes.seal() - return task - - def kv_write_verified(self) -> bool: - """True when every KV task's published destination bytes are attested written. - - This is the admission predicate for verified-range KV block reuse: a - session whose writers merely reported SUCCESS without covering the - published byte range (skipped, partial, or misrouted writes) is not - verified, and the received range must not be committed to the reuse - tree. An empty session (nothing received) is never verified. - """ - with self.lock: - if not self._kv_tasks: - return False - return all(task.write_verified for task in self._kv_tasks) - - def dispatch_prepared_receive(self, task: KVRecvTask) -> None: - try: - self._receiver.dispatch_task(task) - except Exception as error: - if self.cancel_unpublished_task(task): - self._receiver._bounce.release_idle_reservation( - (self.disagg_request_id, task.slice_id) - ) - self.fail_admission(error) - raise - - def fail_admission(self, error: Exception) -> None: - """Fail logical admission without releasing possibly published destinations.""" - with self.lock: - self._logical_outcomes.fail(error) - self._exception = error - if self._terminal_status is None: - self._terminal_status = SessionStatus.ERROR - for task in self._kv_tasks: - if task.status == TaskStatus.INIT: - if self._enforce_physical_ownership: - task.cancel_unpublished() - task.fail(error) - if self._aux_physical_owner is not None and not self._publication_may_have_escaped: - self._aux_physical_owner.cancel_unpublished() - - def cancel_unpublished_task(self, task: KVRecvTask) -> bool: - cancelled = task.cancel_unpublished() - if self._aux_physical_owner is not None: - self._aux_physical_owner.cancel_unpublished() - return cancelled - def process_kv_agent_result( self, peer_rank: int, - receiver_slice_id: int, + sender_slice_id: int, is_last_slice: bool, status: AgentResult, dst_ptrs=None, @@ -3517,255 +2868,167 @@ def process_kv_agent_result( src_base=None, transfer_size: int = 0, ): - with self._ownership_evidence_guard(), self.lock: - if self._retirement is not None and self._closed: - return - assert receiver_slice_id < len(self._kv_tasks), ( - f"Receiver got receiver_slice_id={receiver_slice_id} but only has " + cancel_endpoints: set[str] = set() + with self.lock: + self.kv_cache_size_bytes += transfer_size + assert sender_slice_id < len(self._kv_tasks), ( + f"Receiver got slice_id={sender_slice_id} from sender but only has " f"{len(self._kv_tasks)} receive task(s) for request {self.request_id}. " - f"The sender echoes this index back from RecvReqInfo, so it must " - f"address a task this receiver posted." + f"Sender/receiver slice count mismatch." ) - task = self._kv_tasks[receiver_slice_id] - if status not in ( - AgentResult.SUCCESS, - AgentResult.FAILED, - AgentResult.IN_DOUBT, - AgentResult.FAILED_QUIESCED, - ): - raise ValueError( - f"Session {self.request_id} received unknown task status: {status.value}" - ) - if status == AgentResult.IN_DOUBT: - if not self._enforce_physical_ownership: - raise RuntimeError( - "received IN_DOUBT evidence without physical ownership enabled" - ) - try: - accepted = task.record_writer_in_doubt(peer_rank) - except Exception as error: - task.fail(error) - self._record_ownership_evidence_error(error) - raise - if not accepted: - return - error = RuntimeError( - f"KV transfer quiescence is unproven for request {self.request_id} " - f"slice={receiver_slice_id} peer_rank={peer_rank}" + task = self._kv_tasks[sender_slice_id] + if peer_rank not in task.pending_peer_ranks: + kind = "duplicate" if peer_rank in task.resolved_peer_ranks else "unexpected" + logger.warning( + f"RxSession {self.request_id} ignoring {kind} KV result for " + f"slice={sender_slice_id}, peer_rank={peer_rank}" ) - task.fail(error) - self._record_ownership_evidence_error(error) return - if status == AgentResult.FAILED_QUIESCED: - if not self._enforce_physical_ownership: - raise RuntimeError("received physical settlement without ownership enabled") - try: - accepted = task.record_writer_settlement(peer_rank) - except Exception as error: - self._record_ownership_evidence_error(error) - raise - if accepted: - self._receiver._bounce.record_failure( - (self.disagg_request_id, task.slice_id), peer_rank + try: + if status == AgentResult.SUCCESS: + from .bounce import scatter_write_result + + on_done = None + if is_last_slice: + task.last_slice_count += 1 + if task.last_slice_count == task.expected_transfers: + request_id = self.request_id + + def on_done( + success, + task=task, + peer_rank=peer_rank, + sender_slice_id=sender_slice_id, + request_id=request_id, + ): + if not success: + task.fail( + RuntimeError( + f"KV bounce scatter failed for request {request_id} " + f"slice={sender_slice_id}" + ) + ) + return + if task.status == TaskStatus.ERROR: + return + task.complete() + # Transfer end for perf/time-sync is valid only + # after the bounce data has landed in the real KV pool. + if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): + self.transfer_end_time = ( + tensorrt_llm.bindings.global_steady_clock_now() + ) + logger.debug( + f"KV transfer complete for request {request_id} " + f"slice={sender_slice_id}" + ) + if task._perf_timer is not None: + task._perf_timer.record_task_end(peer_rank) + ri = self._receiver._registrar.self_rank_info + task.print_perf_info(peer_rank, ri.instance_name, ri.instance_rank) + + scatter_write_result( + self._receiver._bounce, + (self.disagg_request_id, task.slice_id), + peer_rank, + dst_ptrs, + sizes, + src_base, + on_done, ) - return - if self._enforce_physical_ownership: - if status == AgentResult.FAILED or is_last_slice: - try: - accepted, all_succeeded = task.record_writer_result( - peer_rank, - status == AgentResult.SUCCESS, - wait_for_local_completion=(status == AgentResult.SUCCESS), - ) - except Exception as error: - task.fail(error) - self._record_ownership_evidence_error(error) - raise - if not accepted: - return + elif status == AgentResult.FAILED: + detail = ( + f"KV transfer failed for request {self.request_id} slice={sender_slice_id} " + f"peer_rank={peer_rank} is_last_slice={is_last_slice} " + f"(reported by remote agent; see sender-side log for nixl_status)" + ) + logger.error(detail) + task.fail(RuntimeError(detail)) + if self._terminal_status is None: # Don't overwrite CANCELLED with ERROR + self._terminal_status = SessionStatus.ERROR else: - all_succeeded = False - else: - all_succeeded = False - # A failure and a last-slice success are both one writer's terminal word. - if status == AgentResult.FAILED or is_last_slice: - accepted, all_succeeded = task.note_writer_report( - peer_rank, status == AgentResult.SUCCESS + raise ValueError( + f"Session {self.request_id} received unknown task status: {status.value}" ) - if not accepted: - return - self.kv_cache_size_bytes += transfer_size - if status == AgentResult.SUCCESS: - # Verified-range accounting: transfer_size is the sender's - # attestation of destination bytes actually submitted and - # completed for this chunk (0 unless the write completed). - task.verified_write_bytes += transfer_size - - from .bounce import scatter_write_result - - on_done = None - if is_last_slice and all_succeeded: - # Completing message: defer task.complete()+perf until the scatter has actually - # landed. scatter_write_result fires this inline for the non-bounced path, or on - # the scatter worker (after cudaStreamSynchronize) for the bounced path, so the - # gen consumer never observes completion before the KV is scattered into place. - request_id = self.request_id - ri = self._receiver._registrar.self_rank_info - instance_name, instance_rank = ri.instance_name, ri.instance_rank - - def on_done( - success, - task=task, - peer_rank=peer_rank, - receiver_slice_id=receiver_slice_id, - request_id=request_id, - instance_name=instance_name, - instance_rank=instance_rank, - ): - # Runs on the scatter worker thread for the bounced path. Do not acquire - # RxSession.lock here: the non-bounced path invokes this callback inline - # while already holding it. Task outcome/ownership locks are independent. - # complete() sets status before _event for wait_complete's status-first poll. - if self._enforce_physical_ownership: - task.finish_local_completion() - if not success: - task.fail( - RuntimeError( - f"KV bounce scatter failed for request {request_id} " - f"slice={receiver_slice_id}" - ) - ) - return - if task.status == TaskStatus.ERROR: - return # a concurrent FAILED writer already failed it; don't un-fail - task.complete() - # Record completion before best-effort diagnostics can block the worker. - if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): - self.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() - try: - if task._perf_timer is not None: - task._perf_timer.record_task_end(peer_rank) - task.print_perf_info(peer_rank, instance_name, instance_rank) - except Exception as e: # perf is best-effort; never block completion - logger.warning( - f"KV transfer perf logging failed for request {request_id} " - f"slice={receiver_slice_id}: {e}" - ) - logger.debug( - f"KV transfer complete for request {request_id} " - f"slice={receiver_slice_id}" - ) - - scatter_write_result( - self._receiver._bounce, - (self.disagg_request_id, task.slice_id), - peer_rank, - dst_ptrs, - sizes, - src_base, - on_done, - ) - elif status == AgentResult.FAILED: - detail = ( - f"KV transfer failed for request {self.request_id} " - f"slice={receiver_slice_id} " - f"peer_rank={peer_rank} is_last_slice={is_last_slice} " - f"(reported by remote agent; see sender-side log for nixl_status)" - ) - logger.error(detail) - # Drain-before-release: record this writer FAILED; the owner frees the shared region - # only once every fan-in writer is terminal (freeing now could race a sibling's RMA). - self._receiver._bounce.record_failure( - (self.disagg_request_id, task.slice_id), peer_rank - ) - task.fail(RuntimeError(detail)) - if self._terminal_status is None: # Don't overwrite CANCELLED with ERROR - self._terminal_status = SessionStatus.ERROR + finally: + # A result is sent only after the remote submit/wait path has + # stopped using our destination. Resolve this peer exactly + # once even when result processing itself raises. + task.pending_peer_ranks.remove(peer_rank) + task.resolved_peer_ranks.add(peer_rank) + if self._active_nixl_ops <= 0: + logger.error( + f"RxSession {self.disagg_request_id} NIXL operation counter underflow" + ) + else: + self._active_nixl_ops -= 1 + self._discard_broadcast_kv_phantoms_locked(task) + cancel_endpoints = self._claim_broadcast_failure_cancel_locked(task) + if cancel_endpoints: + # Never perform messaging I/O while holding the session lock. + self._receiver.send_cancel_to_senders(self.disagg_request_id, cancel_endpoints) def process_aux_agent_result(self, peer_rank: int, status: AgentResult): # Aux is session-level (not per-slice); expected_transfers is identical # across all kv_tasks, so any task provides the right count. - with self._ownership_evidence_guard(), self.lock: - if self._retirement is not None and self._closed: + cancel_endpoints: set[str] = set() + with self.lock: + if peer_rank not in self._pending_aux_peer_ranks: + kind = "duplicate" if peer_rank in self._resolved_aux_peer_ranks else "unexpected" + logger.warning( + f"RxSession {self.request_id} ignoring {kind} aux result " + f"from peer_rank={peer_rank}" + ) return if not self._kv_tasks: logger.warning( f"Aux result received before any KV tasks for request {self.request_id}" ) + self._pending_aux_peer_ranks.remove(peer_rank) + self._resolved_aux_peer_ranks.add(peer_rank) + if self._active_nixl_ops > 0: + self._active_nixl_ops -= 1 return task = self._kv_tasks[0] - all_succeeded = False - if status == AgentResult.IN_DOUBT: - if self._aux_physical_owner is None: - raise RuntimeError( - "received auxiliary IN_DOUBT evidence without physical ownership enabled" - ) - try: - accepted = self._aux_physical_owner.record_writer_in_doubt(peer_rank) - except Exception as error: - self._aux_status = TaskStatus.ERROR - self._record_ownership_evidence_error(error) - raise - if not accepted: - return - error = RuntimeError( - f"auxiliary transfer quiescence is unproven for request " - f"{self.request_id} peer_rank={peer_rank}" - ) - self._aux_status = TaskStatus.ERROR - self._record_ownership_evidence_error(error) - return - if status == AgentResult.FAILED_QUIESCED: - if self._aux_physical_owner is None: - raise RuntimeError("received auxiliary settlement without ownership enabled") - try: - self._aux_physical_owner.record_writer_settlement(peer_rank) - except Exception as error: - self._record_ownership_evidence_error(error) - raise - return - if self._aux_physical_owner is not None: - try: - accepted, all_succeeded = self._aux_physical_owner.record_writer_result( - peer_rank, - status == AgentResult.SUCCESS, - wait_for_local_completion=False, - ) - except Exception as error: - self._aux_status = TaskStatus.ERROR - self._record_ownership_evidence_error(error) - raise - if not accepted: - return - if status == AgentResult.SUCCESS: - self._aux_count += 1 - - if (self._aux_physical_owner is not None and all_succeeded) or ( - self._aux_physical_owner is None and self._aux_count == task.expected_transfers - ): - self._aux_status = TaskStatus.TRANSFERRED - aux_logical_index = getattr(self, "_aux_logical_index", None) - if aux_logical_index is not None: - self._logical_outcomes.complete(aux_logical_index) - elif self._aux_count > task.expected_transfers: + try: + if status == AgentResult.SUCCESS: + self._aux_count += 1 + + if self._aux_count == task.expected_transfers: + self._aux_status = TaskStatus.TRANSFERRED + elif self._aux_count > task.expected_transfers: + self._aux_status = TaskStatus.ERROR + self._exception = RuntimeError( + f"Session {self.request_id} received too many aux transfers" + ) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + logger.error(str(self._exception)) + elif status == AgentResult.FAILED: self._aux_status = TaskStatus.ERROR - self._exception = RuntimeError( - f"Session {self.request_id} received too many aux transfers" - ) - self._logical_outcomes.fail(self._exception) + self._exception = RuntimeError(f"Session {self.request_id} aux transfer failed") if self._terminal_status is None: self._terminal_status = SessionStatus.ERROR - logger.error(str(self._exception)) - elif status == AgentResult.FAILED: - self._aux_status = TaskStatus.ERROR - self._exception = RuntimeError(f"Session {self.request_id} aux transfer failed") - self._logical_outcomes.fail(self._exception) - if self._terminal_status is None: - self._terminal_status = SessionStatus.ERROR - else: - raise ValueError( - f"Session {self.request_id} received unknown aux send status: {status}" - ) + else: + raise ValueError( + f"Session {self.request_id} received unknown aux send status: {status}" + ) + finally: + self._pending_aux_peer_ranks.remove(peer_rank) + self._resolved_aux_peer_ranks.add(peer_rank) + for kv_task in self._kv_tasks: + kv_task.aux_reserved_peer_ranks.discard(peer_rank) + if self._active_nixl_ops <= 0: + logger.error( + f"RxSession {self.disagg_request_id} NIXL operation counter underflow" + ) + else: + self._active_nixl_ops -= 1 + self._discard_broadcast_aux_phantoms_locked(task) + cancel_endpoints = self._claim_broadcast_failure_cancel_locked(task) + if cancel_endpoints: + # Never perform messaging I/O while holding the session lock. + self._receiver.send_cancel_to_senders(self.disagg_request_id, cancel_endpoints) @property def exception(self) -> Optional[Exception]: @@ -3775,6 +3038,8 @@ def unpack_aux(self, request: LlmRequest) -> None: """Read token data from the aux buffer slot into the given request.""" assert self._aux_buffer is not None, "No aux_buffer set for this session" assert self.aux_slot is not None, "No aux_slot set for this session" + if self._aux_buffer.has_draft_history: + self.unpack_draft_history(request) first_gen_tokens, draft_tokens, (prompt_tokens, cached_tokens) = ( self._aux_buffer.get_slot_data(self.aux_slot) ) @@ -3790,131 +3055,58 @@ def unpack_aux(self, request: LlmRequest) -> None: }, } - def is_completed(self) -> bool: - """Non-blocking check: has the transfer completed successfully? + def unpack_draft_history(self, request: LlmRequest) -> None: + """Read standalone history without changing context-first token fields.""" + if self._aux_buffer is None or self.aux_slot is None: + raise ValueError("Standalone draft transfer requires an auxiliary buffer slot") + request.py_draft_transfer_history = self._aux_buffer.get_slot_draft_history(self.aux_slot) - One read of ``status`` answers both halves at one instant. Asking the KV handles and the - auxiliary handle separately would compose two instants, and retirement feeds a cross-rank - consensus that must not see a combination which never held. - """ - return self.status is SessionStatus.TRANSFERRED and ( - not self._enforce_physical_ownership or self.resources_drained() - ) + def is_completed(self) -> bool: + """Non-blocking check: has the transfer completed successfully?""" + status = self.status + if self._need_aux: + return status == SessionStatus.FULLY_TRANSFERRED + return status in (SessionStatus.KV_TRANSFERRED, SessionStatus.FULLY_TRANSFERRED) def has_failed(self) -> bool: """Non-blocking check: has the transfer failed or been cancelled?""" return self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) - def resources_drained(self) -> bool: - if not self._enforce_physical_ownership: - return not any(task.status == TaskStatus.TRANSFERRING for task in self._kv_tasks) - retirement = getattr(self, "_retirement", None) - if retirement is not None and not retirement.can_retire(): - return False - kv_drained = all(task.resources_drained for task in self._kv_tasks) - aux_drained = self._aux_physical_owner is None or self._aux_physical_owner.resources_drained - return kv_drained and aux_drained - - def cancel_local(self, by_peer: bool = False) -> bool: - """Commit cancellation under the same lock used for publication.""" - retirement = getattr(self, "_retirement", None) - if retirement is not None: - retirement.request_drain("peer cancellation" if by_peer else "cancellation") - with self.lock: - # A later cancellation must still notify peers after logical failure. - # The logical arbiter preserves whichever outcome already committed. - if self._terminal_status == SessionStatus.CANCELLED: - return False - self._logical_outcomes.cancel(by_peer) - self._terminal_status = SessionStatus.CANCELLED - # Who asked is not recoverable later, and the two differ: the peer asking is a transfer - # error, our own side asking is an ordinary end. - self.cancelled_by_peer = by_peer - exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") - for task in self._kv_tasks: - rid_slice = (self.disagg_request_id, task.slice_id) - if task.status == TaskStatus.INIT: - if self._enforce_physical_ownership: - task.cancel_unpublished() - self._receiver._bounce.release_idle_reservation(rid_slice) - task.fail(exc) - else: - has_active_access = ( - not task.resources_drained - if self._enforce_physical_ownership - else task.status == TaskStatus.TRANSFERRING - ) - if has_active_access: - self._receiver._bounce.orphan_reservation(rid_slice) - if self._aux_physical_owner is not None and not self._publication_may_have_escaped: - self._aux_physical_owner.cancel_unpublished() - return True + def cancel(self) -> None: + """Cancel the session and notify the remote sender. - def capture_cancel_targets(self) -> set[str]: + Safe to call multiple times. TRANSFERRING tasks keep running (mid-write). + Only INIT tasks have their events signalled immediately. + The lock serializes with process_kv_agent_result() / process_aux_agent_result(). + """ with self.lock: - return set(self._sender_endpoints) - - def notify_cancel(self, sender_endpoints: Optional[set[str]] = None) -> None: - with self._publication_lock: - if sender_endpoints is None: - sender_endpoints = self.capture_cancel_targets() - if self._enforce_physical_ownership: - with self.lock: - self._cancel_notification_started = True - sender_endpoints = sender_endpoints - self._cancelled_sender_endpoints - sent = self._receiver.send_cancel_to_senders(self.disagg_request_id, sender_endpoints) - if self._enforce_physical_ownership and sent: - self._cancelled_sender_endpoints.update(sent) - - def process_session_quiesced(self, peer_rank: int) -> None: - """Settle missing pieces only after every published candidate fences writes.""" - with self._ownership_evidence_guard(), self.lock: - if ( - not self._enforce_physical_ownership - or self._closed - or not self._cancel_notification_started - or not self.has_failed() - ): + if self._terminal_status == SessionStatus.CANCELLED: return - owners = [task._get_physical_owner() for task in self._kv_tasks] - if self._aux_physical_owner is not None: - owners.append(self._aux_physical_owner) - try: - for owner in owners: - if owner._expected_writers is not None and not owner._cancelled_unpublished: - owner.record_session_quiesced(peer_rank) - except Exception as error: - self._record_ownership_evidence_error(error) - raise - - def cancel(self, by_peer: bool = False) -> bool: - """Cancel locally, then notify every writer without holding the lock. - - ``True`` if the ask went out; ``False`` means the session had already ended. - """ - if not self.cancel_local(by_peer=by_peer): - return False - self.notify_cancel() - return True + if self._terminal_status == SessionStatus.ERROR: + # Preserve the diagnostically useful ERROR while still + # notifying any endpoints not covered by failure fan-out. + endpoints = self._sender_endpoints.difference( + self._failure_cancelled_sender_endpoints + ) + self._failure_cancelled_sender_endpoints.update(endpoints) + else: + self._terminal_status = SessionStatus.CANCELLED + exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") + for task in self._kv_tasks: + if task.status == TaskStatus.INIT: + task.fail(exc) + endpoints = set(self._sender_endpoints) + # Send outside the lock to avoid holding it during I/O. + if endpoints: + self._receiver.send_cancel_to_senders(self.disagg_request_id, endpoints) def has_transferring_tasks(self) -> bool: - """True while a destination accessor may remain active. + """True while any dispatched peer may still own destination buffers. cancel_request() must return False while this is True. """ - return not self.resources_drained() - - def _failed_wait_result(self) -> Optional[WaitResult]: - return ( - WaitResult.FAILED - if not self._enforce_physical_ownership or self.resources_drained() - else None - ) - - def _completed_wait_result(self) -> Optional[WaitResult]: - if getattr(self, "_retirement", None) is not None and not self.is_completed(): - return None - return WaitResult.COMPLETED + with self.lock: + return self._active_nixl_ops > 0 def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: """Poll or block until transfer completes. @@ -3924,8 +3116,6 @@ def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: With blocking=True: waits up to _timeout_s for each task. Returns WaitResult.COMPLETED on full success, WaitResult.FAILED on error/timeout. """ - if self.has_failed() and self._enforce_physical_ownership: - return self._failed_wait_result() if not blocking: # Use task.status instead of task.wait(timeout=0): task.complete() # sets status before event, so a GIL switch between the two steps @@ -3934,60 +3124,32 @@ def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: if task.status == TaskStatus.TRANSFERRED: continue if task.status == TaskStatus.ERROR: - return self._failed_wait_result() + return WaitResult.FAILED return None # task not yet done; re-poll next cycle else: timeout = self._timeout_s for task in self._kv_tasks: done = task.wait(timeout=timeout) if not done: - return self._failed_wait_result() + return WaitResult.FAILED # timeout if task.status == TaskStatus.ERROR: - return self._failed_wait_result() + return WaitResult.FAILED if self._need_aux: while True: status = self.status - if status == SessionStatus.TRANSFERRED: - return ( - self._failed_wait_result() - if self.has_failed() - else self._completed_wait_result() - ) + if status == SessionStatus.FULLY_TRANSFERRED: + return WaitResult.COMPLETED elif status in (SessionStatus.ERROR, SessionStatus.CANCELLED): - return self._failed_wait_result() + return WaitResult.FAILED if not blocking: return None # KV done, aux still in flight; re-poll next cycle time.sleep(0.001) - return self._failed_wait_result() if self.has_failed() else self._completed_wait_result() + return WaitResult.COMPLETED - def close(self) -> bool: - if self._enforce_physical_ownership: - with self.lock: - if self._closed: - return True - retirement = getattr(self, "_retirement", None) - if retirement is not None and retirement.controller.shutdown_requested: - # Publication also holds this session lock. Only unsealed - # owners can be cancelled here; a reserved send remains live - # even before its network callback has actually started. - cancelled = [task for task in self._kv_tasks if task.cancel_unpublished()] - aux_cancelled = ( - self._aux_physical_owner is not None - and self._aux_physical_owner.cancel_unpublished() - ) - if cancelled or aux_cancelled: - self._logical_outcomes.cancel(by_peer=False) - for task in cancelled: - task.fail(RuntimeError("receiver shutdown before publication")) - if not self.resources_drained(): - return False - if retirement is not None and not retirement.close(): - return False - self._closed = True - else: - if getattr(self, "_closed", False): - return True - self._closed = True + def close(self): + if getattr(self, "_closed", False): + return + self._closed = True if self._aux_buffer is not None and self.aux_slot is not None: self._aux_buffer.free_slot(self.aux_slot) self.aux_slot = None @@ -3998,24 +3160,16 @@ def close(self) -> bool: for task in self._kv_tasks: self._receiver._bounce.orphan_reservation((self.disagg_request_id, task.slice_id)) self._receiver.clear_session(self.disagg_request_id) - return True def __enter__(self): return self def __exit__(self, _exc_type, _exc, _tb): - if not self.close(): - logger.warning( - f"RxSession {self.disagg_request_id} refused close while resources remain active" - ) + self.close() def __del__(self): try: - if not self.close(): - logger.warning( - f"RxSession {self.disagg_request_id} refused destructor close while " - "resources remain active" - ) + self.close() except Exception as e: logger.warning(f"RxSession.__del__: exception during close: {e}") @@ -4076,13 +3230,7 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): self.shutdown() -def _create_nixl_agent( - name: str, - rank: int, - world_size: int, - agent_buffer_size_mb: int = 0, - agent_bounce_params: Optional[Dict[str, str]] = None, -) -> NixlTransferAgent: +def _create_nixl_agent(name: str, rank: int, world_size: int) -> NixlTransferAgent: num_threads = int(os.environ.get("TRTLLM_NIXL_NUM_THREADS", "8")) kwargs = {} if "TRTLLM_NIXL_SPLIT_BATCH_SIZE" in os.environ: @@ -4093,16 +3241,17 @@ def _create_nixl_agent( num_threads=num_threads, rank=rank, world_size=world_size, - agent_buffer_size_mb=agent_buffer_size_mb, - agent_bounce_params=agent_bounce_params, **kwargs, ) def _make_aux_buffer( - kvm: KVCacheManagerV2, max_slots: int, max_draft_len: Optional[int] = None + kvm: KVCacheManager, max_slots: int, max_draft_len: Optional[int] = None ) -> Optional[AuxBuffer]: + draft_history = bool(getattr(kvm, "draft_layer_ids", ())) if max_slots <= 0: + if draft_history: + raise ValueError("Standalone draft transfer requires auxiliary buffer slots") return None if max_draft_len is None: max_draft_len = max(0, int(getattr(kvm, "max_draft_len", 0))) @@ -4111,12 +3260,13 @@ def _make_aux_buffer( beam_width=max(1, int(getattr(kvm, "max_beam_width", 1))), max_draft_len=max_draft_len, device="cpu", + draft_history=draft_history, ) @dataclass class TransferWorkerConfig: - kv_cache_manager: KVCacheManagerV2 + kv_cache_manager: KVCacheManager device_id: int instance_name: str max_concurrent_sessions: int = 0 @@ -4125,33 +3275,11 @@ class TransferWorkerConfig: rx_timeout_s: Optional[float] = None bounce: Optional["Config"] = None tx_overall_timeout_s: Optional[float] = None - enforce_physical_ownership: bool = False - # Internal activation gate: only the qualified containment policy supplies it. - quiescence_fatal_callback: Optional[Callable[[QuiescenceFatalEvent], None]] = None - # Transfer-agent staging (bounce v2) buffer size in MiB; 0 disables the - # fast path. Derived upstream from CacheTransceiverConfig: it equals - # kv_cache_bounce_size_mb when agent_bounce_buffer_enable is set (in which - # case `bounce` above is off) and 0 otherwise. - agent_buffer_size_mb: int = 0 - # Expert bounce knobs (TRTLLM_NIXL_BOUNCE_* names without the prefix and the - # trailing _BYTES, lowercased); dict > env > default. Ignored when - # agent_buffer_size_mb == 0. - agent_bounce_params: Optional[Dict[str, str]] = None class TransferWorker: def __init__(self, config: TransferWorkerConfig): self._config = config - self._retirement_watchdog = None - if config.quiescence_fatal_callback is not None: - if not config.enforce_physical_ownership: - raise ValueError("deadline retirement requires physical ownership") - for timeout in (config.tx_overall_timeout_s, config.rx_timeout_s): - if timeout is None or not math.isfinite(timeout) or timeout <= 0: - raise ValueError( - "deadline retirement requires finite positive transfer timeouts" - ) - self._retirement_watchdog = RetirementWatchdog(self._on_quiescence_fatal) kvm = config.kv_cache_manager self._aux_buffer = _make_aux_buffer( kvm, config.max_concurrent_sessions, config.max_draft_len @@ -4164,24 +3292,6 @@ def __init__(self, config: TransferWorkerConfig): ) self._setup_peer_infrastructure(kvm) self._setup_transfer_engine() - if self._retirement_watchdog is not None: - self._retirement_watchdog.start() - - def _on_quiescence_fatal(self, event: QuiescenceFatalEvent) -> None: - """Forward fatal evidence while rooting all worker-owned resources. - - Args: - event: Immutable deadline decision; never authorizes resource release. - """ - callback = self._config.quiescence_fatal_callback - assert callback is not None - callback(event) - - def request_shutdown(self) -> None: - """Close admission and start drain before any potentially blocking teardown.""" - watchdog = getattr(self, "_retirement_watchdog", None) - if watchdog is not None: - watchdog.request_shutdown() def populate_instance_and_rank_info(self, endpoints: list[str], layer_num_per_pp: list[int]): assert self._rank_info is not None @@ -4198,8 +3308,8 @@ def create_tx_session(self, request: LlmRequest) -> TxSession: aux_buffer=self._aux_buffer, timeout_s=self._config.tx_timeout_s, prompt_len=request.prompt_len, + beam_width=request.py_beam_width, overall_timeout_s=self._config.tx_overall_timeout_s, - retirement_watchdog=getattr(self, "_retirement_watchdog", None), ) def create_rx_session(self, request: LlmRequest) -> RxSession: @@ -4212,7 +3322,7 @@ def create_rx_session(self, request: LlmRequest) -> RxSession: aux_buffer=self._aux_buffer, timeout_s=self._config.rx_timeout_s, prompt_len=request.prompt_len, - retirement_watchdog=getattr(self, "_retirement_watchdog", None), + beam_width=request.py_beam_width, ) def has_all_peer_req_infos_for_send(self, unique_rid: int) -> bool: @@ -4222,7 +3332,7 @@ def sweep_stale_req_infos(self): """Forward to Sender to evict orphaned RecvReqInfo from ADP broadcast.""" self._sender.sweep_stale_req_infos() - def _setup_peer_infrastructure(self, kvm: KVCacheManagerV2): + def _setup_peer_infrastructure(self, kvm: KVCacheManager): self._rank_info_server = RankInfoServer(self._rank_info) if kvm.mapping.rank == 0 else None self._kv_extractor = KVRegionExtractorV1(kvm) self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor) @@ -4235,22 +3345,7 @@ def _setup_transfer_engine(self): self._rank_info.instance_name + str(self._rank_info.instance_rank), rank=mapping.rank, world_size=mapping.world_size, - agent_buffer_size_mb=self._config.agent_buffer_size_mb, - agent_bounce_params=self._config.agent_bounce_params, ) - # The pure-Python agent has no bounce_enabled attribute and already warns on its own. - if ( - self._config.agent_buffer_size_mb > 0 - and hasattr(self._agent, "bounce_enabled") - and not self._agent.bounce_enabled - ): - logger.warning( - f"TransferWorker: the C++ transfer-agent bounce was requested " - f"(agent_bounce_buffer_enable=True, kv_cache_bounce_size_mb=" - f"{self._config.agent_buffer_size_mb}) but is inactive on agent " - f"{self._agent.name} (init failed); standard per-descriptor NIXL " - "transfers are in use." - ) self._registered_mem: list = [] try: self._register_kv_cache() @@ -4267,20 +3362,8 @@ def _setup_transfer_engine(self): # The bounce object owns its own NIXL descriptors (deregistered by Transport.close in # shutdown), so they are NOT added to _registered_mem — single-source to avoid a # double deregister. - self._sender = Sender( - self._peer_registrar, - self._agent, - bounce=self._bounce, - enforce_physical_ownership=self._config.enforce_physical_ownership, - retirement_watchdog=self._retirement_watchdog, - ) - self._receiver = Receiver( - self._peer_registrar, - self._agent, - bounce=self._bounce, - enforce_physical_ownership=self._config.enforce_physical_ownership, - retirement_watchdog=self._retirement_watchdog, - ) + self._sender = Sender(self._peer_registrar, self._agent, bounce=self._bounce) + self._receiver = Receiver(self._peer_registrar, self._agent, bounce=self._bounce) self._rank_info.transfer_engine_info = bytes(self._agent.get_local_agent_desc()) self._rank_info.self_endpoint = self._receiver.endpoint except Exception: @@ -4334,25 +3417,20 @@ def page_table(self): def shutdown(self): if getattr(self, "_shutdown", False): return - self.request_shutdown() - watchdog = getattr(self, "_retirement_watchdog", None) - if watchdog is not None: - watchdog.require_retired() - watchdog.stop() + self._shutdown = True # Use getattr guards: __init__ may have failed partway, leaving some # attributes unset. Without them, __del__ -> shutdown() raises # AttributeError and ZMQ resources from already-created sub-objects # are never cleaned up. rank_info_server = getattr(self, "_rank_info_server", None) + if rank_info_server is not None: + rank_info_server.shutdown() sender = getattr(self, "_sender", None) if sender is not None: sender.shutdown() receiver = getattr(self, "_receiver", None) if receiver is not None: receiver.shutdown() - if rank_info_server is not None: - rank_info_server.shutdown() - self._shutdown = True # Close the bounce transport (stops its scatter thread, deregisters its own descriptors and # frees the VMM buffers) once the receiver listener is stopped so no new scatter is enqueued. # No-op for the non-bounced path (NoBounceTransport.close); without this the daemon thread + fabric @@ -4367,14 +3445,6 @@ def shutdown(self): # (e.g. when the KV cache manager is recreated after profiling). agent = getattr(self, "_agent", None) if agent is not None: - try: - if getattr(agent, "bounce_enabled", False): - logger.info( - f"TransferWorker.shutdown: agent {agent.name} bounce_submit_count=" - f"{agent.bounce_submit_count} bounce_reject_count={agent.bounce_reject_count}" - ) - except Exception as e: # observability only; never skip deregister/shutdown below - logger.debug(f"TransferWorker.shutdown: bounce counters unavailable: {e}") registered = getattr(self, "_registered_mem", []) while registered: desc = registered.pop(0) diff --git a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py index b97f16d52a53..5dea7ba2eb26 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py +++ b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py @@ -19,8 +19,10 @@ import numpy as np -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode from .page import AttentionLayerGroup @@ -143,7 +145,11 @@ def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 # V2 already returns per-cache-level pool slot indices (not logical block # IDs), and active sequences GPU-lock their pages (_UniqPageLock enforces # cache_level==GPU), so the slot_ids yielded here are already the right - # offsets for primary-pool pointer arithmetic. No translation is needed. + # offsets for primary-pool pointer arithmetic. No translation is needed, + # unlike V1 (see _CacheReuseAdapterV1.get_block_ids). + if lg.cp_as_tp: + ordinals = self.get_block_ordinals(req, group_idx, lg) + return ordinals[ordinals >= 0] return np.fromiter( self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices( group_idx, valid_only=True @@ -155,12 +161,20 @@ def get_block_ordinals(self, req, group_idx, lg): # noqa: ARG002 # valid_only=False yields one entry per block ordinal, with -1 # (BAD_PAGE_INDEX) for out-of-window (SWA-evicted) and unbound blocks. # Position is the index; no length arithmetic needed. - return np.fromiter( + ordinals = np.fromiter( self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices( group_idx, valid_only=False ), dtype=np.int64, ) + if lg.cp_as_tp: + converter = self._mgr.impl.get_page_index_converter( + lg.local_layers[0].local_layer_id, Role.KEY + ) + return np.asarray( + converter(ordinals.tolist(), index_mode=PageIndexMode.SHARED), dtype=np.int64 + ) + return ordinals def commit_blocks_for_reuse(self, req: LlmRequest) -> None: self._mgr.try_commit_blocks(req) diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index b1132c10ddc4..1209490b31eb 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -43,10 +43,15 @@ get_layer_byte_ranges, get_physical_pool, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 -from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManager, + MambaHybridCacheManagerV2, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm._utils import get_size_in_bytes, nvtx_range +from tensorrt_llm.bindings import DataType +from tensorrt_llm.runtime.kv_cache_manager_v2 import PoolGroupDesc, SlotDescVariant # Mapper kinds a V2 manager may declare via get_disagg_role_mapper_kinds(). # A physical pool may mix kinds (V2 storage coalesces buffers purely by @@ -201,7 +206,12 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: if not hasattr(manager, "_layer_attn_to_layer_id"): # Standard: local_layer_id is index into pp_layers - return [manager.pp_layers[lid] for lid in local_layer_ids] + local_to_global = dict(enumerate(manager.pp_layers)) + local_to_global.update( + (manager.layer_offsets[layer_id], layer_id) + for layer_id in getattr(manager, "draft_layer_ids", ()) + ) + return [local_to_global[lid] for lid in local_layer_ids] # Virtual layers: build inverse mapping internal_layer_id -> (model_layer, attn_type) # and encode as model_layer * num_attn_types + attn_type_value @@ -214,7 +224,9 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: # matching this ID is required to be consistent for. inverse = {} for (model_layer, attn_type), layer_id in manager._layer_attn_to_layer_id.items(): - inverse.setdefault(layer_id, (model_layer, attn_type.value)) + inverse[layer_id] = (model_layer, attn_type.value) + for draft_layer in getattr(manager, "draft_layer_ids", ()): + inverse[manager.layer_offsets[draft_layer]] = (draft_layer, 0) # Use the full enum range for consistent encoding across all PP ranks. # Different PP ranks may have different subsets of attention types (e.g., @@ -327,6 +339,56 @@ def _build_pool_views_for_variant( return pool_views +def _build_draft_pool_views( + manager: KVCacheManagerV2, variant: SlotDescVariant, pg_desc: PoolGroupDesc +) -> tuple[PhysicalPoolGroup, list[PoolView]]: + """Expose expanded draft buffers as physical-token pages for transfer.""" + if len(variant.coalesced_buffers) != 1: + raise ValueError("Standalone draft layers must share one equal-sized K/V pool") + buffer = variant.coalesced_buffers[0] + first_layer = buffer.buffer_ids[0].layer_id + expansion = manager.impl.get_page_index_converter(first_layer, Role.KEY).expansion + buffer_size = int(buffer.single_buffer_size) + if buffer_size % expansion: + raise ValueError("Draft buffer size is not divisible by its page expansion") + page_bytes = buffer_size // expansion + physical_pool = pg_desc.pools[0] + slot_bytes = int(physical_pool.slot_bytes) + if slot_bytes % page_bytes: + raise ValueError("Draft pool slots must contain whole physical-token pages") + pool_group = PhysicalPoolGroup( + pools=[ + PhysicalPool( + base_address=int(physical_pool.base_address), + slot_bytes=page_bytes, + num_slots=int(pg_desc.num_slots) * (slot_bytes // page_bytes), + slot_stride_bytes=page_bytes, + ) + ] + ) + entries_by_role = defaultdict(list) + for index, buffer_id in enumerate(buffer.buffer_ids): + if buffer_id.role not in (Role.KEY, Role.VALUE): + raise ValueError("Standalone draft transfer expects only K/V buffers") + entries_by_role[buffer_id.role].append( + (int(buffer_id.layer_id), index * buffer_size, page_bytes) + ) + # K and V occupy separate expanded runs. Keeping one role per view lets + # the existing head mapper copy each physical-token page without treating + # the intervening subpages as part of that page's payload. + views = [ + PoolView( + pool_idx=0, + buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE), + pool_role=frozenset({str(role)}), + mapper_kind=MapperKind.INDEXED, + bytes_per_layer=page_bytes, + ) + for role, entries in entries_by_role.items() + ] + return pool_group, views + + def _build_page_table_v2(manager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManagerV2. @@ -428,19 +490,21 @@ def _window_size_for_layer(internal_layer_id: int): for iid, gid in zip(all_internal_layer_ids, all_global_layer_ids) ] - # A life cycle is recurrent when its layers are the manager's - # mamba layers. Recurrent and attention layers never share a life - # cycle: their buffer configs differ, so V2 groups them apart. - is_recurrent = isinstance(manager, MambaHybridCacheManagerV2) and any( - manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids - ) - if is_recurrent and not all( - manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids - ): - raise ValueError( - f"Layer group {layer_group_id} mixes recurrent and attention layers" + layer_id = int(all_internal_layer_ids[0]) + if layer_id in manager.cp_as_tp_layer_ids: + draft_pool_group, pool_views = _build_draft_pool_views(manager, variant, pg_desc) + draft_pool_group_idx = len(pool_groups) + pool_groups.append(draft_pool_group) + layer_groups_by_id[layer_group_id] = AttentionLayerGroup( + pool_group_idx=draft_pool_group_idx, + local_layers=local_layers, + pool_views=pool_views, + kv_head_num_per_rank=manager.num_kv_heads_per_layer[layer_id], + sliding_window_size=_window_size_for_layer(all_internal_layer_ids[0]), + cp_as_tp=True, + total_kv_head_num=manager._draft_total_num_kv_heads, ) - cache_kind = CacheKind.STATE if is_recurrent else CacheKind.PAGED + continue # Bucket buffer entries by (pool, mapper kind). One PoolView is # emitted per bucket and spans every layer of that role class, diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 1254c06001f6..850257d8f5ab 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -293,6 +293,11 @@ class AttentionLayerGroup(LayerGroup): kv_head_num_per_rank: int = 0 sliding_window_size: Optional[int] = None + local_layers: List[LocalLayer] = field(default_factory=list) + pool_views: List[PoolView] = field(default_factory=list) + # Full-sequence KV whose heads use the flattened TP x CP rank grid. + cp_as_tp: bool = False + total_kv_head_num: Optional[int] = None def to_dict(self) -> dict: return { @@ -300,6 +305,8 @@ def to_dict(self) -> dict: "pool_group_idx": int(self.pool_group_idx), "kv_head_num_per_rank": int(self.kv_head_num_per_rank), "sliding_window_size": self.sliding_window_size, + "cp_as_tp": self.cp_as_tp, + "total_kv_head_num": self.total_kv_head_num, "local_layers": [ll.to_dict() for ll in self.local_layers], "pool_views": [pv.to_dict() for pv in self.pool_views], } @@ -310,8 +317,8 @@ def from_dict(cls, data: dict) -> "AttentionLayerGroup": pool_group_idx=int(data["pool_group_idx"]), local_layers=[LocalLayer.from_dict(x) for x in data.get("local_layers", [])], pool_views=[PoolView.from_dict(pv) for pv in data.get("pool_views", [])], - kv_head_num_per_rank=int(data["kv_head_num_per_rank"]), - sliding_window_size=data.get("sliding_window_size"), + cp_as_tp=bool(data.get("cp_as_tp", False)), + total_kv_head_num=data.get("total_kv_head_num"), ) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 759d408fa362..84c3014471a8 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -58,10 +58,12 @@ ) from tensorrt_llm._torch.disaggregation.resource.utils import get_pool_view_slot_bytes from tensorrt_llm._torch.distributed.communicator import Distributed -from tensorrt_llm._torch.pyexecutor import hang_detector -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - BlockReusePolicy, - KVCacheManagerV2, +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManager, + MambaHybridCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -652,16 +654,11 @@ def _describe_local(self, req: LlmRequest) -> Chunk: tpb = adapter.tokens_per_block assert self._page_table is not None layer_groups = self._page_table.layer_groups - remote_tail_start = getattr(req, "py_csa2_remote_tail_start", None) - transfer_len = remote_tail_start if remote_tail_start is not None else req.prompt_len - prompt_blocks = (transfer_len + tpb - 1) // tpb - remote_tail_split = getattr(req, "py_csa2_remote_tail_split", None) - excluded_pool_views = set() - if remote_tail_split is not None: - exclude = getattr(self._kv_cache_manager, "remote_tail_excluded_pool_views", None) - if exclude is None: - raise ValueError("remote-tail replay requires a CSA2 cache manager") - excluded_pool_views = exclude(self._page_table, remote_tail_split) + prompt_len = self._global_prompt_len(req) + prompt_blocks = (prompt_len + tpb - 1) // tpb + manager = self._kv_cache_manager + cp_size = getattr(manager, "_helix_cp_size", 1) + cp_rank = getattr(manager, "_helix_cp_rank", 0) is_gen_only = req.is_generation_only_request cached_per_lg = ( @@ -670,13 +667,98 @@ def _describe_local(self, req: LlmRequest) -> Chunk: else [0] * len(layer_groups) ) - empty = np.array([], dtype=np.int64) + unified_draft = isinstance(manager, KVCacheManagerV2) and bool(manager.draft_layer_ids) + transfer_length = prompt_len if unified_draft else req.prompt_len + token_range = None + if transfer_length > 0: + # end must match the trimmed block list below (ceil(prompt_len / tpb) + # blocks). num_extra_kv_tokens slots (speculative decoding) are not + # transferred. In the previously added support for ctx disabling + # speculative decoding while gen enables it, both sides currently + # use prompt_len as the transfer range, so the ranges stay + # consistent. + # TODO: the accuracy impact of not transferring num_extra_kv_tokens + # on MTP and other speculative decoding paths is currently unclear; + # revisit whether these extra KV slots need to be transferred. + token_range = TokenRange(start=0, end=transfer_length) + groups = [] kinds = [lg.kind for lg in layer_groups] for idx, lg in enumerate(layer_groups): - if lg.kind == CacheKind.STATE: - slot = self._get_mamba_slot_for_request(req) - group = np.array([slot], dtype=np.int64) if slot is not None else empty + if isinstance(lg, MambaLayerGroup): + groups.append(np.array([], dtype=np.int64)) + continue + if unified_draft: + # KVSlice carries a contiguous retained suffix in this protocol. + local_prompt_blocks = ( + prompt_blocks if lg.cp_as_tp else len(range(cp_rank, prompt_blocks, cp_size)) + ) + stale_end = 0 + if lg.sliding_window_size is not None: + stale_end = max(0, (prompt_len + 1 - lg.sliding_window_size) // tpb) + if not lg.cp_as_tp: + stale_end = len(range(cp_rank, stale_end, cp_size)) + ordinals = adapter.get_block_ordinals(req, idx, lg) + pages = ordinals[stale_end:local_prompt_blocks] + if pages.size != local_prompt_blocks - stale_end or np.any(pages < 0): + raise ValueError("Missing allocated prompt pages for unified draft KV transfer") + cache_skip = max(0, cached_per_lg[idx] // tpb - stale_end) + groups.append(pages[cache_skip:]) + continue + block_ids = adapter.get_block_ids(req, idx, lg) + # Limit to prompt_len blocks, matching C++ cacheFormatter behavior. + total_blocks = (req.prompt_len + tpb - 1) // tpb + window_size = lg.sliding_window_size + + if window_size is not None: + draft_reserve_tokens = get_draft_token_length(req) if is_gen_only else 0 + allocated_blocks = ( + req.prompt_len + + draft_reserve_tokens + + self._kv_cache_manager.num_extra_kv_tokens + + tpb + - 1 + ) // tpb + beam0_block_ids, tail_block_ids = self._split_packed_beam_block_ids( + block_ids, + req.py_beam_width, + allocated_blocks, + ) + if beam0_block_ids.size > allocated_blocks: + beam0_block_ids = beam0_block_ids[:allocated_blocks] + block_ids = ( + np.concatenate([beam0_block_ids, tail_block_ids]) + if tail_block_ids.size > 0 + else beam0_block_ids + ) + # Current PyExecutor cache managers disable KV-cache token sinks, + # so SWA block lists contain an evictable prompt prefix followed + # by the speculative scratch tail. If token sinks are enabled, + # this must use block-ordinal metadata to preserve the sink prefix. + # Remove scratch before trimming stale prompt blocks; otherwise a + # boundary-crossing allocation can displace initialized prompt KV. + scratch_blocks = max(0, allocated_blocks - total_blocks) + if scratch_blocks > 0 and getattr(self._kv_cache_manager, "_has_cp_helix", False): + # Helix ledgers expose one windowed page per + # tokens_per_block * cp_size tokens, so the per-page scratch + # arithmetic above would empty a one-entry HP-tail list. + scratch_blocks = 0 + if scratch_blocks > 0: + if req.py_beam_width != 1: + raise ValueError("speculative scratch blocks require beam_width == 1") + block_ids = ( + block_ids[:-scratch_blocks] + if scratch_blocks < block_ids.size + else np.array([], dtype=np.int64) + ) + # Drop stale blocks the manager may still expose (V1 pre-eviction). + stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb) + expected_valid = max(0, total_blocks - stale_end) + # Stale prefix already pruned above; skip reuse-hit blocks that + # land inside the window. Clamp to 0: ctx side has cached_per_lg + # synthetically 0, and a reuse hit may fall entirely inside the + # stale region (those blocks were already pruned, no extra skip). + cache_skip = max(0, cached_per_lg[idx] // tpb - stale_end) else: # Block lists carry beam 0 only (beam-search attention reads # prompt positions through beam 0's block table), so the @@ -709,10 +791,69 @@ def _positional_window( window[: min(cached_blocks, prompt_blocks)] = -1 return window - @staticmethod - def _need_aux_transfer(req: LlmRequest) -> bool: + def _need_aux_transfer(self, req: LlmRequest) -> bool: params = req.py_disaggregated_params - return params is not None and params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST + return bool(getattr(self._kv_cache_manager, "draft_layer_ids", ())) or ( + params is not None and params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST + ) + + def _global_prompt_len(self, req: LlmRequest) -> int: + if getattr(self._kv_cache_manager, "_has_cp_helix", False): + return req.total_input_len_cp + return req.prompt_len + + def _validate_draft_history_range(self, req: LlmRequest, history: dict) -> None: + # Full-attention drafters can skip an oversized context store; rolling + # drafters retain their complete live suffix. Neither includes scratch. + window_size = history["layout"]["window_size"] + max_positions = self._kv_cache_manager.draft_max_position_embeddings + prompt_len = self._global_prompt_len(req) + expected_length = prompt_len + if window_size is not None: + if type(window_size) is not int or window_size <= 0: + raise ValueError("Invalid draft history window size") + expected_length = min(expected_length, window_size) + elif max_positions is not None and prompt_len > max_positions: + expected_length = min(history["valid_length"], max_positions) + if history["valid_length"] != expected_length or history["position"] != prompt_len: + raise ValueError( + "DSpark transfer requires valid draft history and sequence position " + f"covering the complete prompt ({prompt_len} tokens, " + f"{expected_length} retained)." + ) + + def _pack_draft_history(self, req: LlmRequest) -> None: + manager = getattr(self, "_kv_cache_manager", None) + if not getattr(manager, "draft_layer_ids", ()): + return + history = self._kv_cache_manager.export_draft_history(req.py_request_id) + self._validate_draft_history_range(req, history) + req.py_draft_transfer_history = history + + def _restore_draft_history(self, req: LlmRequest) -> None: + history = getattr(req, "py_draft_transfer_history", None) + manager = getattr(self, "_kv_cache_manager", None) + has_draft = bool(getattr(manager, "draft_layer_ids", ())) + if history is None: + if has_draft: + raise ValueError( + "Standalone DSpark generation requires draft history from a prefill worker " + "with matching speculative configuration; draft history metadata is missing." + ) + return + if not has_draft: + raise ValueError( + "Received standalone DSpark draft history without a manager-owned draft cache." + ) + self._validate_draft_history_range(req, history) + # K/V already occupies receiver-local pages; restore only validity and position. + self._kv_cache_manager.restore_draft_history(req.py_request_id, history) + + def _prepare_received_history(self, session: RxSessionBase, req: LlmRequest) -> None: + if self._need_aux_transfer(req): + self._apply_aux(session, req) + self._assert_disagg_history_declared(req) + self._restore_draft_history(req) def _validate_bridge_req(self, req: LlmRequest, synchronous: bool = False) -> bool: if not getattr(self, "_fp4_mla_bridge_enabled", False): @@ -1015,6 +1156,11 @@ def _close_session_or_raise(self, session: object, rid: int, outcome: str) -> No def _apply_aux(self, session, req: LlmRequest): """Unpack aux tokens from session into request's context_phase_params.""" + params = req.py_disaggregated_params + if params is not None and params.schedule_style != DisaggScheduleStyle.GENERATION_FIRST: + # Context-first tokens and usage already arrived in the context response. + session.unpack_draft_history(req) + return session.unpack_aux(req) first_gen_tokens = req.py_first_gen_tokens # type: ignore[attr-defined] draft_tokens = req.py_draft_tokens @@ -1145,20 +1291,28 @@ def respond_and_send_async(self, req: LlmRequest) -> None: rid = get_unique_rid(req) assert rid is not None session = self._get_or_create_send_session(req) - if session is None: - return - bridge_enabled = getattr(self, "_fp4_mla_bridge_enabled", False) - if bridge_enabled: - # Root the request before backend admission so its source pages - # outlive every ownership-enabled NIXL operation. - self._send_reqs[rid] = req - req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS - # A remote cancellation can win before the local TxSession is - # created. Sender.setup_session() records that terminal state and - # acknowledges the fenced session to every known receiver; - # leave retirement to the normal status path without attempting - # to publish KV or auxiliary memory afterward. - if session.has_failed(): + req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + # Register before dispatch: dispatch_task may synchronously fail after + # enqueueing work for an earlier peer. Status polling must always have + # the matching request while that partial dispatch drains. + self._send_reqs[rid] = req + try: + self._pack_draft_history(req) + session.send(self._create_kv_slice(req)) + self._finalize_send(req, session) + except Exception as exc: + req.py_disagg_send_setup_failed = True + session.cancel() + if session.has_transferring_tasks(): + # At least one worker already entered NIXL. Keep the session, + # request, and pinned pages registered; + # normal status polling will classify the safe cancellation + # as a setup failure once every active operation has drained. + req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + logger.error( + "CTX send setup failed after transfer started; " + f"retaining rid={rid} until NIXL leaves mid-write: {exc}" + ) return try: if self.pipeline_transfer_enabled: @@ -1214,10 +1368,9 @@ def request_and_receive_sync(self, req: LlmRequest) -> None: # blocking path has no cross-rank outcome consensus). req.py_kv_transfer_verified = session.kv_write_verified() # KV-transfer timing setters deferred to #15871 (clock-source consistency); size only. - req.set_kv_cache_size(self._request_kv_bytes(req, paged_bytes, state_bytes)) - if self._need_aux_transfer(req): - self._apply_aux(session, req) - self._assert_disagg_history_declared(req) + req.set_kv_cache_size(self._slice_num_bytes(kv_slice) * self._kv_size_rank_factor) + self._prepare_received_history(session, req) + self._rebuild_fp4_mla_disagg_imported_cache(req) req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE else: req.state = LlmRequestState.DISAGG_TRANS_ERROR @@ -1408,7 +1561,6 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]) -> GenT 0 if need_progress else wait_num, block_all, ) - completed, failed, cancelled = [], [], [] for rid in to_process: session = self._recv_sessions[rid] @@ -1475,19 +1627,49 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]) -> GenT for rid in completed: session = self._recv_sessions[rid] req = self._recv_reqs[rid] - # Stamp the write-verification verdict on the request before the - # session is closed and dropped: commit_blocks_for_reuse runs - # later (batch preparation) and gates reuse admission on it. - req.py_kv_transfer_verified = rid in verified - # transfer_end already stamped at completion detection above. - req.set_kv_cache_size(getattr(req, "py_kv_cache_xfer_bytes", 0)) - if self._need_aux_transfer(req): - self._apply_aux(session, req) - self._assert_disagg_history_declared(req) - self._close_session_or_raise(session, rid, "completed") - req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE - del self._recv_reqs[rid] - del self._recv_sessions[rid] + completed_reqs[rid] = req + try: + self._prepare_received_history(session, req) + self._rebuild_fp4_mla_disagg_imported_cache(req) + except Exception as exc: + local_postprocess_failed.append(rid) + req.state = LlmRequestState.DISAGG_TRANS_ERROR + logger.error( + f"Disagg gen postprocess failed rank={self._dist.rank} rid={rid}: {exc}" + ) + finally: + # NIXL is complete at this point, so unregistering is safe even + # if aux application or an FP4 MLA sidecar rebuild fails. + session.close() + del self._recv_reqs[rid] + del self._recv_sessions[rid] + + # ``completed`` is globally consistent, so guarding this collective is + # itself symmetric and avoids another allgather on empty polling cycles. + if completed: + all_postprocess_failed = self._allgather_or_passthrough( + local_postprocess_failed, + self._gen_allgather, + self._gen_need_sync, + ) + global_postprocess_failed = self._union(all_postprocess_failed) + else: + global_postprocess_failed = set() + postprocess_failed = [rid for rid in completed if rid in global_postprocess_failed] + postprocess_failed_set = set(postprocess_failed) + for rid in completed: + req = completed_reqs[rid] + if rid in postprocess_failed_set: + # This also covers remote-only failure: a locally successful + # rank must not publish TRANS_COMPLETE when another + # participating rank failed the same request. + req.state = LlmRequestState.DISAGG_TRANS_ERROR + else: + req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + + completed = [rid for rid in completed if rid not in postprocess_failed_set] + transfer_failed = failed + failed = transfer_failed + postprocess_failed if failed: logger.warning( f"Disagg gen transfer FAILED rank={self._dist.rank} " diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index b50a55327929..8fb53462c6d1 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -1399,16 +1399,21 @@ def _prepare_dflash_trtllm_gen_buffers( head_dim: int, ) -> None: trtllm_gen_ops = self._dflash_trtllm_gen_ops - workspace_bytes = trtllm_gen_ops.get_workspace_size( - dtype=dtype, - num_tokens=max_batch_size * block_size, - num_gen_tokens=max_batch_size * block_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_size=head_dim, - max_num_requests=max_batch_size, - rotary_embedding_dim=0, - fp8_context_fmha=False, + workspace_bytes = max( + trtllm_gen_ops.get_workspace_size( + dtype=dtype, + num_tokens=max_batch_size * block_size, + num_gen_tokens=max_batch_size * block_size, + num_heads=heads, + num_kv_heads=num_kv_heads, + head_size=head_dim, + max_num_requests=max_batch_size, + rotary_embedding_dim=0, + fp8_context_fmha=False, + ) + # A non-causal block runs as one decode request whose query heads + # are the block's queries times the attention heads. + for heads in (num_heads, num_heads * block_size) ) device = torch.device(device) is_capturing = torch.cuda.is_current_stream_capturing() @@ -1441,7 +1446,7 @@ def _prepare_dflash_trtllm_gen_buffers( sm_count = self._dflash_trtllm_gen_sm_count counter_bytes = trtllm_gen_ops.get_multi_ctas_kv_counter_size( - num_heads, max_batch_size, sm_count + num_heads * block_size, max_batch_size, sm_count ) counters = self._dflash_trtllm_gen_counters counters_need_allocation = ( @@ -1623,20 +1628,14 @@ def dflash_forward( .reshape(-1) .contiguous() ) - else: # VANILLA - # Slot and column that each block K/V row occupies - self._prepare_dflash_index_buffers( - ctx_k_cache.device, - ctx_k_cache.size(0), - block_size, - need_batch_indices=False, - ) - append_batch_indices = ( - cache_batch_idx_i32.view(-1, 1).expand(-1, block_size).reshape(-1).long() - ) - append_positions = ( - (cache_seqlens_i32.view(-1, 1) + self._dflash_block_offsets).reshape(-1).long() - ) + # An FP8 cache takes the block's K/V through direct page writes + # (the same rows and kv positions the append addresses). + fp8_cache = ctx_kv_cache[0].dtype == torch.float8_e4m3fn + if fp8_cache: + append_rows = append_batch_indices.long() + append_pos_long = append_positions.long() + fp8_pages = block_tables[append_rows, append_pos_long // page_size].long() + fp8_offsets = append_pos_long % page_size # Flatten query positions once for the fused QK-norm-RoPE kernel. query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) @@ -1741,23 +1740,41 @@ def dflash_forward( causal, window_size = self._resolve_block_attention(layer_idx) if self.dflash_attention_backend == "TRTLLM": layer_cache = ctx_kv_cache[layer_idx] - trtllm_gen_ops.append_paged_kv_cache( - append_key=k_noise_bshd.reshape( - -1, num_kv_heads_per_rank, head_dim - ).contiguous(), - append_value=v_noise_bshd.reshape( - -1, num_kv_heads_per_rank, head_dim - ).contiguous(), - batch_indices=append_batch_indices, - positions=append_positions, - paged_kv_cache=layer_cache, - kv_indices=kv_indices, - kv_indptr=kv_indptr, - kv_last_page_len=kv_last_page_len, - kv_layout="HND", - ) + if fp8_cache: + cache_bytes = layer_cache.view(torch.uint8) + cache_bytes[fp8_pages, 0, :, fp8_offsets] = ( + k_noise_bshd.reshape(-1, num_kv_heads_per_rank, head_dim) + .to(layer_cache.dtype) + .view(torch.uint8) + ) + cache_bytes[fp8_pages, 1, :, fp8_offsets] = ( + v_noise_bshd.reshape(-1, num_kv_heads_per_rank, head_dim) + .to(layer_cache.dtype) + .view(torch.uint8) + ) + else: + trtllm_gen_ops.append_paged_kv_cache( + append_key=k_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous(), + append_value=v_noise_bshd.reshape( + -1, num_kv_heads_per_rank, head_dim + ).contiguous(), + batch_indices=append_batch_indices, + positions=append_positions, + paged_kv_cache=layer_cache, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=kv_last_page_len, + kv_layout="HND", + ) out = torch.empty_like(Q_bshd) q_flat = Q_bshd.reshape(-1, num_heads_per_rank, head_dim) + if fp8_cache: + # TRTLLM-Gen reads an FP8 cache with FP8 queries (the + # context kernels have no BF16-query variant); scales stay + # 1, matching the K/V writes. + q_flat = q_flat.to(layer_cache.dtype) out_flat = out.reshape(-1, num_heads_per_rank, head_dim) window_left = window_size[0] if causal: @@ -1784,6 +1801,48 @@ def dflash_forward( bmm1_scale_log2=None, multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, ) + elif window_left == -1: + # Every block query attends the same keys (the history and + # the whole block), so the block runs as one decode request + # with block_size * num_heads query heads. Query (j, head + # g * group + r) becomes head (g, j, r): the GQA group + # grows to block_size * group and the decode kernel, which + # also splits a long history across CTAs, reads each + # history once. + group = num_heads_per_rank // num_kv_heads_per_rank + q_folded = q_flat.reshape( + B, block_size, num_kv_heads_per_rank, group, head_dim + ).permute(0, 2, 1, 3, 4) + q_folded = q_folded.reshape(B, -1, head_dim) + out_folded = q_folded.new_empty(q_folded.shape, dtype=out.dtype) + trtllm_gen_ops.batch_decode_with_kv_cache( + query=q_folded, + kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), + workspace_buffer=self._dflash_trtllm_gen_workspace, + block_tables=block_tables, + seq_lens=seq_lens_after, + max_seq_len=pages_per_slot * page_size, + bmm1_scale=head_dim**-0.5, + bmm2_scale=1.0, + window_left=window_left, + out=out_folded, + sinks=None, + enable_pdl=False, + kv_layout="HND", + backend="trtllm-gen", + q_len_per_req=1, + max_q_len=None, + cum_seq_lens_q=None, + kv_cache_sf=None, + uses_shared_paged_kv_idx=True, + bmm1_scale_log2=None, + multi_ctas_kv_counter_buffer=self._dflash_trtllm_gen_counters, + ) + out_flat.view(B, block_size, num_kv_heads_per_rank, group, head_dim).copy_( + out_folded.view( + B, num_kv_heads_per_rank, block_size, group, head_dim + ).permute(0, 2, 1, 3, 4) + ) else: cum_seq_lens_q = torch.arange( 0, diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 81ce6e3aee8d..dc3ee32f041b 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -16,12 +16,8 @@ # DSpark backbone ported from the DeepSeek-V4-Pro-DSpark reference # (`inference/model.py`: DSparkBlock / Transformer.forward_spec). # -# The captured-context attention primitives are ported from DeepSeek's DeepSpec -# reference ``inference/kernel.py`` (``sparse_attn``) and ``inference/model.py`` -# (``get_dspark_topk_idxs``). The reference computes these with a TileLang -# kernel; this is a functional-first pure-PyTorch port with the same math -# (index-gather + online softmax + a learnable attention sink that contributes -# only to the softmax denominator). +# Captured-context attention reads managed KV pages and the noncausal draft +# block, with a learnable attention sink contributing only to the denominator. # # The draft I/O stages are ported from the same reference # (`inference/model.py`: DSparkBlock.forward_embed / forward_head). @@ -59,9 +55,8 @@ class or the edge would become a cycle. ``hc_head`` — the block-draft output head. The heads themselves are shared with the standalone path and live in :mod:`modeling_speculative`. -2. **Captured-context attention primitives** — the dense sliding-window MLA the - draft block attends with (``get_dspark_topk_idxs`` / ``dspark_sparse_attn`` - and the rotary helpers). Hardware-agnostic and unit-testable in isolation. +2. **Captured-context attention** — projections and rotary helpers feeding + paged attention over the retained context and the noncausal draft block. 3. **Block draft I/O** — ``build_draft_input_ids`` (the ``[bonus_token, noise, ...]`` block input) and ``dspark_propose`` (Markov @@ -96,10 +91,9 @@ class or the edge would become a cycle. from tensorrt_llm.quantization.mode import QuantAlgo from ..._utils import get_sm_version, is_sm_100f -from ...models.modeling_utils import QuantConfig -from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams -from ..model_config import ModelConfig +from ..attention_backend.dspark import write_dspark_context +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_INTERNAL_AVAILABLE +from ..distributed import AllReduceParams from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.gated_mlp import GatedMLP @@ -129,15 +123,14 @@ class or the edge would become a cycle. from .modeling_utils import DecoderModel, register_auto_model, register_draft_model if IS_CUTLASS_DSL_AVAILABLE: - from ..custom_ops.dspark_attention_custom_op import ( - fused_dsv4_dspark_attention, - is_fused_dsv4_dspark_attention_supported, - ) + from ..custom_ops import dspark_attention_custom_op # noqa: F401 from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( + _get_dspark_arch_str, + cute_dsl_dspark_concat_fp8_quant, + cute_dsl_dspark_rmsnorm_fp8_quant, cute_dsl_dspark_rmsnorm_rope, - cute_dsl_dspark_rmsnorm_rope_cache_write, cute_dsl_dspark_rmsnorm_rope_draft_block, - is_fused_dspark_attention_preparation_supported, + cute_dsl_dspark_rmsnorm_rope_page_write, is_fused_dspark_rmsnorm_rope_supported, ) @@ -217,148 +210,6 @@ def apply_dspark_rotary_batched( return out.to(orig_dtype) -@lru_cache(maxsize=64) -def _topk_matrix(window_size: int, block_size: int, start_pos: int) -> torch.Tensor: - # [min(window, start_pos+1)] context positions in the rolling KV window, - # followed by [block_size] positions for the current block's own K/V (which - # the caller appends to the window at offset ``window_size``). - ctx = torch.arange(min(window_size, start_pos + 1)) - blk = window_size + torch.arange(block_size) - return torch.cat([ctx, blk]).int() - - -def get_dspark_topk_idxs( - window_size: int, - bsz: int, - block_size: int, - start_pos: int, - device: torch.device | str = "cpu", -) -> torch.Tensor: - """Per-query attended-position indices for the DSpark draft block. - - Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` - query positions attends to the same set — the ``min(window_size, start_pos+1)`` - most-recent context positions in the rolling KV window, then the - ``block_size`` positions of the current block (stored at offset - ``window_size`` in the concatenated KV). Note this is *non-causal* within the - block (every position sees every block position), matching the reference. - - Args: - window_size: sliding-window length of the captured-context KV cache. - bsz: batch size. - block_size: number of draft positions per request. - start_pos: absolute decode position (must be > 0); bounds the context. - device: device for the returned index tensor. - - Returns: - int32 tensor ``[bsz, block_size, topk]`` with - ``topk = min(window_size, start_pos+1) + block_size``. - """ - assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" - matrix = _topk_matrix(int(window_size), int(block_size), int(start_pos)).to(device) - return matrix.view(1, 1, -1).expand(bsz, block_size, -1).contiguous() - - -def get_dspark_topk_idxs_batched( - window_size: int, - block_size: int, - start_pos: torch.Tensor, - valid_len: torch.Tensor | None = None, -) -> torch.Tensor: - """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. - - Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width - ``min(window_size, start_pos+1) + block_size`` depends on the host int - ``start_pos``), this always returns the **fixed** width ``window_size + - block_size`` and masks the unfilled context slots with ``-1``. The masked - slots are excluded by :func:`dspark_sparse_attn` exactly as if they were - absent while the shape remains CUDA-graph safe. - - Every query attends to the actually written circular-window suffix, followed - by the current-block positions. Without ``valid_len`` this preserves the - legacy ``start_pos``-only behavior. - - Args: - window_size: sliding-window length of the captured-context KV cache. - block_size: number of draft positions per request. - start_pos: ``[G]`` int tensor of per-request absolute decode positions. - valid_len: optional ``[G]`` count of actually written rolling-window - entries. When omitted, preserve the legacy ``start_pos`` mask. - - Returns: - int32 tensor ``[G, block_size, window_size + block_size]``. - """ - device = start_pos.device - g = start_pos.shape[0] - ctx_cols = torch.arange(window_size, device=device) # [win] - if valid_len is None: - valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] - else: - # The valid entries are the contiguous logical suffix ending at - # start_pos, but their physical slots wrap modulo window_size. - valid_len = valid_len.clamp(min=0, max=window_size) - age = torch.remainder(start_pos.unsqueeze(1) - ctx_cols.unsqueeze(0), window_size) - valid = age < valid_len.unsqueeze(1) - ctx_idx = torch.where( - valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) - ) - blk_idx = window_size + torch.arange(block_size, device=device) # [block] - blk_idx = blk_idx.unsqueeze(0).expand(g, -1) # [G, block] - row = torch.cat([ctx_idx, blk_idx], dim=1).to(torch.int32) # [G, win+block] - return row.unsqueeze(1).expand(g, block_size, -1).contiguous() - - -def dspark_sparse_attn( - q: torch.Tensor, - kv: torch.Tensor, - attn_sink: torch.Tensor, - topk_idxs: torch.Tensor, - softmax_scale: float, -) -> torch.Tensor: - """Index-gathered multi-query attention with an attention sink. - - Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For - each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by - ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled - dot-product softmax over them, and adds a per-head learnable *sink* logit that - participates only in the softmax denominator (i.e. an "attend-to-nothing" - option with a zero value vector). KV is shared across query heads (MQA). - - Args: - q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). - kv: ``[b, n, d]`` keys/values (shared across heads). - attn_sink: ``[h]`` per-head sink logits (fp32). - topk_idxs: ``[b, m, topk]`` int gather indices into ``kv`` (``-1`` masks). - softmax_scale: scalar applied to the q·k scores (``head_dim ** -0.5``). - - Returns: - ``[b, m, h, d]`` attention output, in ``q.dtype``. - """ - b, m, h, d = q.shape - idx = topk_idxs.long() # [b, m, topk] - valid = idx >= 0 - safe = idx.clamp(min=0) - - # Invalid slots read kv[0, :] (via safe.clamp), but masked_fill below - # zeros their softmax probs, so the einsum nullifies them. - kv_exp = kv.unsqueeze(1).expand(b, m, kv.shape[1], d) - gathered = torch.gather(kv_exp, 2, safe.unsqueeze(-1).expand(b, m, safe.shape[-1], d)).float() - - # Scores [b, m, h, topk]; mask invalid slots to -inf before the softmax. - scores = torch.einsum("bmhd,bmkd->bmhk", q.float(), gathered) * softmax_scale - scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) - - # Online-softmax max is taken over gathered positions only (the sink is added - # to the denominator afterwards), matching the kernel's reduce order. - smax = scores.max(dim=-1, keepdim=True).values # [b, m, h, 1] - smax = torch.where(torch.isinf(smax), torch.zeros_like(smax), smax) - probs = torch.exp(scores - smax) # masked slots -> exp(-inf) = 0 - sink = torch.exp(attn_sink.to(torch.float32).view(1, 1, h) - smax.squeeze(-1)) - denom = probs.sum(dim=-1) + sink # [b, m, h] - out = torch.einsum("bmhk,bmkd->bmhd", probs, gathered) / denom.unsqueeze(-1) - return out.to(q.dtype) - - def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: """RMSNorm matching the DeepSpec reference (fp32 reduce, then * weight).""" dtype = x.dtype @@ -367,19 +218,10 @@ def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: return (weight.float() * xf).to(dtype) -def _rope_last_dims( - t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False -) -> torch.Tensor: - """Apply RoPE to the last ``rope_head_dim`` dims; pass the rest through.""" - nope = t[..., :-rope_head_dim] - rope = apply_dspark_rotary(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) - return torch.cat([nope, rope], dim=-1) - - def _rope_last_dims_batched( t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False ) -> torch.Tensor: - """Per-row variant of :func:`_rope_last_dims` (``freqs_cis`` has a batch axis).""" + """Apply per-request RoPE to the last dimensions, preserving the non-RoPE prefix.""" nope = t[..., :-rope_head_dim] rope = apply_dspark_rotary_batched(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) return torch.cat([nope, rope], dim=-1) @@ -469,104 +311,13 @@ def _rmsnorm_rope_batched( def dspark_attention_forward( - x: torch.Tensor, - main_x: torch.Tensor, - start_pos: int, - kv_cache: torch.Tensor, - *, - wq_a: torch.Tensor, - q_norm_w: torch.Tensor, - wq_b: torch.Tensor, - wkv: torch.Tensor, - kv_norm_w: torch.Tensor, - wo_a: torch.Tensor, - wo_b: torch.Tensor, - attn_sink: torch.Tensor, - n_heads: int, - head_dim: int, - rope_head_dim: int, - n_groups: int, - o_lora_rank: int, - window_size: int, - eps: float, - softmax_scale: float, - freqs_cis: torch.Tensor, - persist: bool = False, - q_b_norm_enabled: bool = True, -) -> torch.Tensor: - """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). - - Functional port of DeepSpec ``DSparkAttention.forward`` for the dense - (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) - with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys - gathered from a rolling captured-context window (``kv_cache``, into which the - projected ``main_x`` context is written at ``start_pos % window_size``) plus the - block's own positions, attention-sink softmax, inverse-RoPE on the output, and a - grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). - - Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / - dequantized projection weights); ``wo_a`` is the raw grouped weight matrix - ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is - ``[b, window_size, head_dim]`` and is updated functionally (cloned). - - Returns: - ``[b, block_size, dim]`` attention output (residual stream contribution). - """ - assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" - b, block, _ = x.shape - rd = rope_head_dim - # Host-int slices run off the end of the table silently, unlike the batched - # path's gathers: a short blk_freqs drops RoPE from the block tail and shows - # up only as lower acceptance. Fail loudly instead. - assert start_pos + block < freqs_cis.shape[0], ( - f"DSpark RoPE table holds {freqs_cis.shape[0]} entries, short of the " - f"{start_pos + block + 1} the block at start_pos {start_pos} needs" - ) - main_freqs = freqs_cis[start_pos : start_pos + 1] - blk_freqs = freqs_cis[start_pos + 1 : start_pos + 1 + block] - - # Captured-context K/V from main_x (MQA, shared across heads). - main_kv = _rmsnorm(F.linear(main_x, wkv), kv_norm_w, eps) # [b, 1, head_dim] - main_kv = _rope_last_dims(main_kv, rd, main_freqs) - - # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm(F.linear(x, wq_a), q_norm_w, eps) - q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [b, block, h, head_dim] - # Per-head RMS in the query dtype (matches the reference inline normalization, - # which is NOT the fp32 RMSNorm path). - if q_b_norm_enabled: - q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) - q = _rope_last_dims(q, rd, blk_freqs) - - # Block K/V. - kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] - kv = _rope_last_dims(kv, rd, blk_freqs) - - # Write the context K/V into the rolling window, then attend over - # [window context | block] with the sink. ``persist=True`` writes through - # to the caller's buffer (cross-step decode, worker-owned window); the - # default clones so single-shot callers (golden / unit tests) stay pure. - cache = kv_cache if persist else kv_cache.clone() - cache[:, start_pos % window_size] = main_kv.squeeze(1) - kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] - topk = get_dspark_topk_idxs(window_size, b, block, start_pos, device=x.device) - o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale) # [b, block, h, head_dim] - o = _rope_last_dims(o, rd, blk_freqs, inverse=True) - - # Grouped low-rank O projection. - o = o.reshape(b, block, n_groups, -1) - wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) - o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) - return F.linear(o.flatten(2), wo_b) - - -def dspark_attention_forward_batched( x: torch.Tensor, main_x: torch.Tensor, start_pos: torch.Tensor, kv_cache: torch.Tensor, - slots: torch.Tensor, - valid_len: torch.Tensor | None = None, + block_tables: torch.Tensor, + capacities: torch.Tensor, + valid_len: torch.Tensor, *, wq_a: torch.Tensor, q_norm_w: torch.Tensor, @@ -585,50 +336,25 @@ def dspark_attention_forward_batched( eps: float, softmax_scale: float, freqs_cis: torch.Tensor, - persist: bool = False, - q_b_norm_enabled: bool = True, + qkv_proj: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + o_a_proj: Optional[torch.Tensor] = None, + o_a_scale: Optional[torch.Tensor] = None, + o_a_dequant: Optional[torch.Tensor] = None, + use_fused_output_projection: Optional[bool] = None, + use_fused_qkv_input_projection: bool = False, + use_fused_q_input_projection: bool = False, ) -> torch.Tensor: - """Batched, CUDA-graph-safe captured-context DSpark draft attention. - - This path belongs to the embedded DeepSeek-V4-Pro draft. Standalone DSpark - drafters (including Kimi-K3-DSpark) use the DFlash paged-KV decode instead. - - Numerically identical, per request, to :func:`dspark_attention_forward`, but - free of host syncs and data-dependent shapes so it can be captured into a CUDA - graph (the one-engine drafter runs inside the target's graph). The differences - from the scalar path are purely mechanical: - - * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen - request) instead of a python int; RoPE phases are *gathered* per request from - the fixed ``freqs_cis`` table rather than sliced. - * the rolling-window context K/V is written/read through the ``slots`` index - into a shared ``kv_cache`` (``persist=True`` writes through to the caller's - worker-owned buffer; otherwise a clone is used), instead of mutating a - per-request cache in place. - * the attended-position list has the fixed width ``window_size + block_size`` - with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). - - Args: - x: ``[G, block, dim]`` block layer input (per gen request). - main_x: ``[G, 1, hidden]`` projected captured context. - start_pos: ``[G]`` int tensor of absolute decode positions (> 0). - kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows - (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). - slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. - valid_len: optional ``[G]`` count of actually written context entries; - masks holes left when absolute positions are bootstrapped without - receiving the corresponding DSpark rolling-window state. - freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; - must satisfy ``maxlen > start_pos.max() + block_size``. + """CUDA-graph-safe DSpark attention over managed context pages and a draft block. - Returns: - ``[G, block, dim]`` attention output (residual stream contribution). + ``x`` is [batch, block, dim], ``main_x`` is [batch, 1, dim], and + ``kv_cache`` is [pages, page_size, head_dim]. The [batch, max_pages] + block table maps absolute token positions into that pool. ``start_pos`` + contains RoPE frames (logical token positions + 1); ``valid_len`` counts + initialized context tokens. Zero capacities disable synthetic page access. """ + if window_size != 128 or rope_head_dim != 64: + raise ValueError("Embedded DSpark attention requires a 128-token window and 64 RoPE dims") g, block, _ = x.shape - if kv_cache.shape[1] != window_size: - raise ValueError( - f"kv_cache window extent {kv_cache.shape[1]} does not match window_size {window_size}" - ) rd = rope_head_dim # Per-request RoPE phases gathered from the fixed table (no host-int slicing). main_freqs = freqs_cis[start_pos].unsqueeze(1) # [G, 1, rd//2] @@ -640,8 +366,9 @@ def dspark_attention_forward_batched( main_kv_input = F.linear(main_x, wkv) # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm_rope_batched(F.linear(x, wq_a), q_norm_w, eps, 0, blk_freqs) - q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [G, block, h, head_dim] + q = _q_projection( + q_latent, q_norm_w, eps, wq_b, use_fused_q_input_projection, blk_freqs + ).unflatten(-1, (n_heads, head_dim)) q = _rmsnorm_rope_batched( q, None, @@ -653,85 +380,98 @@ def dspark_attention_forward_batched( apply_rmsnorm=q_b_norm_enabled, ) - block_kv_input = F.linear(x, wkv) - - # ``persist=True`` writes through to the worker-owned rolling window; - # single-shot callers keep the original functional clone behavior. - write_target = kv_cache if persist else kv_cache.clone() - main_rope_freqs = torch.view_as_real(main_freqs).reshape(-1, rd // 2, 2).contiguous() - inverse_rope_freqs = torch.view_as_real(blk_freqs).contiguous() - block_rope_freqs = inverse_rope_freqs.reshape(-1, rd // 2, 2) - - # The DSV4 fused kernel requires explicit physical-window validity; None - # keeps the reference path. - use_fused_dsv4_dspark_attention = ( - valid_len is not None - and IS_CUTLASS_DSL_AVAILABLE - and is_fused_dsv4_dspark_attention_supported( - q, write_target, valid_len, attn_sink, inverse_rope_freqs - ) - and is_fused_dspark_attention_preparation_supported( + if IS_CUTLASS_DSL_AVAILABLE and _get_dspark_arch_str() is not None and block in (5, 6): + main_rope_freqs = torch.view_as_real(main_freqs).reshape(-1, rd // 2, 2).contiguous() + inverse_rope_freqs = torch.view_as_real(blk_freqs).contiguous() + cute_dsl_dspark_rmsnorm_rope_page_write( main_kv_input, - block_kv_input, kv_norm_w, main_rope_freqs, - block_rope_freqs, - write_target, - slots, - start_pos, - ) - ) - if use_fused_dsv4_dspark_attention: - slots_i32, cache_seqs = cute_dsl_dspark_rmsnorm_rope_cache_write( - main_kv_input, - kv_norm_w, - main_rope_freqs, - write_target, - slots, + kv_cache, + block_tables, start_pos, + capacities, eps, ) draft_block = cute_dsl_dspark_rmsnorm_rope_draft_block( - block_kv_input, - kv_norm_w, - block_rope_freqs, - eps, + block_kv_input, kv_norm_w, inverse_rope_freqs.reshape(-1, rd // 2, 2), eps ) - o = fused_dsv4_dspark_attention( + if use_fused_output_projection is None: + use_fused_output_projection = ( + o_a_proj is not None + and _can_use_fused_fp8_output_projection(wo_b, o_a_proj, n_groups, o_lora_rank) + ) + if use_fused_output_projection: + if o_a_proj is None or o_a_scale is None: + raise ValueError("The fused DSpark FP8 output path requires o_a weights and scales") + grouped_o_fp8, grouped_o_scale = torch.ops.trtllm.fused_dsv4_dspark_attention_fp8( + q, + draft_block, + kv_cache, + block_tables, + start_pos, + valid_len, + capacities, + attn_sink, + inverse_rope_freqs, + softmax_scale, + n_groups, + ) + return _fused_fp8_output_projection_from_quantized( + grouped_o_fp8, + grouped_o_scale, + batch=g, + block=block, + wo_b=wo_b, + n_groups=n_groups, + o_lora_rank=o_lora_rank, + o_a_proj=o_a_proj, + o_a_scale=o_a_scale, + ) + o = torch.ops.trtllm.fused_dsv4_dspark_attention( q, draft_block, - write_target, - slots_i32, - cache_seqs, + kv_cache, + block_tables, + start_pos, valid_len, + capacities, attn_sink, inverse_rope_freqs, softmax_scale, ) else: main_kv = _rmsnorm_rope_batched(main_kv_input, kv_norm_w, eps, rd, main_freqs) - kv = _rmsnorm_rope_batched(block_kv_input, kv_norm_w, eps, rd, blk_freqs) - main_kv_flat = main_kv.squeeze(1).to(write_target.dtype) - - slot_pos = start_pos % window_size # [G] - write_target[slots, slot_pos] = main_kv_flat - cache_rows = write_target[slots] # [G, window, head_dim] - kv_full = torch.cat([cache_rows, kv], dim=1) # [G, window + block, head_dim] - topk = get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len) - o = dspark_sparse_attn( - q, kv_full, attn_sink, topk, softmax_scale - ) # [G, block, h, head_dim] - o = _rmsnorm_rope_batched( - o, - kv_norm_w, - eps, - rd, - blk_freqs, - num_heads=n_heads, - apply_weight=False, - apply_rmsnorm=False, - inverse_rope=True, - ) + draft_block = _rmsnorm_rope_batched(block_kv_input, kv_norm_w, eps, rd, blk_freqs) + write_dspark_context( + main_kv, + start_pos[:, None], + torch.ones_like(start_pos[:, None], dtype=torch.bool), + kv_cache, + block_tables, + capacities, + ) + + # Gather only the retained context window; masked rows may have no pages. + positions = start_pos[:, None] - window_size + torch.arange(window_size, device=x.device) + valid = (positions >= (start_pos - valid_len).clamp_min(0)[:, None]) & ( + positions < capacities[:, None] + ) + page_size = kv_cache.shape[1] + columns = (positions // page_size).clamp(0, block_tables.shape[1] - 1) + pages = block_tables.gather(1, columns).long() + valid = valid & (pages >= 0) + context = kv_cache[pages.clamp_min(0), positions % page_size] + context = torch.where(valid[..., None], context, 0) + kv = torch.cat((context, draft_block), dim=1) + + scores = torch.einsum("bshd,btd->bsht", q.float(), kv.float()) * softmax_scale + mask = torch.cat((valid, torch.ones((g, block), device=x.device, dtype=torch.bool)), dim=1) + scores = scores.masked_fill(~mask[:, None, None, :], float("-inf")) + sink = attn_sink.float()[None, None, :, None].expand(g, block, -1, 1) + probs = torch.cat((scores, sink), dim=-1).softmax(-1)[..., :-1] + o = torch.einsum("bsht,btd->bshd", probs, kv.float()).to(q.dtype) + o = _rope_last_dims_batched(o, rd, blk_freqs, inverse=True) # Grouped low-rank O projection. o = o.reshape(g, block, n_groups, -1) @@ -1516,110 +1256,48 @@ def _draft_quant_config_dict(model_config, base: int, num_stages: int): changed = True return new_qcd if changed else None - @torch.inference_mode() - def write_context_windows( - self, - main_hidden: torch.Tensor, - positions: torch.Tensor, - stage_windows: torch.Tensor, - ) -> None: - """Write captured-context ``main_kv`` into the rolling per-stage KV windows. - - Replicates exactly the per-position context write that - :func:`dspark_attention_forward` performs each generation step - (``main_kv = RoPE_pos(rmsnorm(wkv @ main_x))`` written at - ``window[pos % window_size]``), but for an arbitrary set of - ``(captured-hidden, position)`` pairs. Used to (a) seed a request's - window from its prompt at prefill and (b) back-fill the intermediate - accepted tokens of a multi-accept step — both of which the per-step - generation path would otherwise leave as holes, starving the draft - attention of context (acceptance-rate only; verified decoding keeps - output correctness regardless). - - Args: - main_hidden: ``[M, num_capture * hidden]`` captured target hiddens. - positions: ``[M]`` absolute window positions (used for BOTH the RoPE - phase and the slot ``pos % window_size``). By the generation-path - convention this is ``committed_position + 1``. Must hold at most - ``window_size`` entries with distinct slots (the caller passes a - contiguous, deduplicated range) so the scatter is well defined. - stage_windows: ``[num_stages, window_size, head_dim]`` window for one - request's slot; updated in place. - """ - if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: - return - M = int(main_hidden.shape[0]) - if M == 0: - return - win = int(self._attn_params["window_size"]) - rd = int(self._attn_params["rope_head_dim"]) - eps = float(self._attn_params["eps"]) - positions = positions.to(main_hidden.device).long() - # main_x is stage-invariant (stage 0's projection), matching forward_embed. - stage0 = self.mtp_layers[0] - main_x = stage0.main_norm(stage0.main_proj(main_hidden)) # [M, hidden] - freqs = self._dspark_freqs_table(main_x.device)[positions] - slots = positions % win # [M] - mx = main_x.unsqueeze(0) # [1, M, hidden] for the per-position RoPE layout - for s, stage in enumerate(self.mtp_layers): - a = stage._dspark_attn - kv = _rmsnorm(F.linear(mx, a["wkv"]), a["kv_norm_w"], eps) # [1, M, head_dim] - kv = _rope_last_dims(kv, rd, freqs) # [1, M, head_dim] - stage_windows[s, slots] = kv[0].to(stage_windows.dtype) + def _project_context_wkv(self, main_x: torch.Tensor) -> Optional[torch.Tensor]: + """Quantize once and project every stage's WKV with one stacked GEMM.""" + if self._context_wkv_fp8 is None: + return None + rows = main_x.reshape(-1, main_x.shape[-1]) + input_fp8, input_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(rows) + weight, weight_scale = self._context_wkv_fp8 + output = torch.ops.trtllm.cute_dsl_dspark_context_wkv_gemm_rubin( + input_fp8, weight, input_scale, weight_scale + ) + head_dim = self._attn_params["head_dim"] + return output.view(*main_x.shape[:-1], len(self.mtp_layers), head_dim) - def write_context_windows_batched( + def write_context_pages( self, main_hidden: torch.Tensor, positions: torch.Tensor, - slots: torch.Tensor, mask: torch.Tensor, - kv_windows: torch.Tensor, + pools: tuple[torch.Tensor, ...], + block_tables: torch.Tensor, + capacities: torch.Tensor, ) -> None: - """CUDA-graph-safe batched + masked variant of :meth:`write_context_windows`. - - Back-fills the intermediate accepted tokens of a multi-accept step into the - rolling per-stage KV windows for ALL gen requests at once, with a fixed - ``[G, M]`` shape (``M = max interim per request``) and a validity mask - (invalid entries are no-ops via a read-modify-write), so nothing depends on - the per-request accept count. Same per-position math as the scalar version - (``RoPE_pos(rmsnorm(wkv @ main_x))`` written at ``window[pos % window]``), - but indexed/scattered through ``slots`` into the shared persistent buffer. + """Project captured context and store only valid tokens in managed pages. - Args: - main_hidden: ``[G, M, num_capture * hidden]`` captured target hiddens - (rows beyond a request's interim count are masked). - positions: ``[G, M]`` absolute window positions (RoPE phase + slot). - slots: ``[G]`` row index of each request into ``kv_windows``. - mask: ``[G, M]`` bool — which ``(g, m)`` entries are real interim writes. - kv_windows: ``[N, num_stages, window_size, head_dim]`` persistent buffer; - updated in place. + Inputs have fixed [batch, tokens] leading dimensions. Positions are RoPE + frames (logical token position + 1), including for chunked prefill. """ - if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: - return - G, M = positions.shape - if G == 0 or M == 0: + if main_hidden.numel() == 0 or getattr(self.mtp_layers[0], "_dspark_attn", None) is None: return - win = int(self._attn_params["window_size"]) rd = int(self._attn_params["rope_head_dim"]) eps = float(self._attn_params["eps"]) - positions = positions.long() - slots = slots.long() safe_positions = torch.where(mask, positions, torch.zeros_like(positions)) - freqs = self._dspark_freqs_table(main_hidden.device)[safe_positions] # [G, M, rd//2] - cols = positions % win # [G, M] - rows = slots[:, None].expand(-1, M) # [G, M] - mask3 = mask.unsqueeze(-1) # [G, M, 1] + freqs = self._dspark_freqs_table(main_hidden.device)[safe_positions] stage0 = self.mtp_layers[0] - main_x = stage0.main_norm(stage0.main_proj(main_hidden)) # [G, M, hidden] - for s, stage in enumerate(self.mtp_layers): + main_x = stage0.main_norm(stage0.main_proj(main_hidden)) + head_dim = int(self._attn_params["head_dim"]) + projected_wkv = self._project_context_wkv(main_x) + for s, (stage, pool) in enumerate(zip(self.mtp_layers, pools)): a = stage._dspark_attn kv = _rmsnorm(F.linear(main_x, a["wkv"]), a["kv_norm_w"], eps) # [G, M, head_dim] kv = _rope_last_dims_batched(kv, rd, freqs) # [G, M, head_dim] - win_s = kv_windows[:, s] # [N, win, head_dim] view onto the base buffer - # Read-modify-write so masked-out (g, m) entries keep their current - # value — a graph-safe masked scatter (no dynamic-shape compaction). - cur = win_s[rows, cols] # [G, M, head_dim] - win_s[rows, cols] = torch.where(mask3, kv.to(win_s.dtype), cur) + write_dspark_context(kv, positions, mask, pool[:, 0, 0], block_tables, capacities) def forward_embed( self, main_hidden: torch.Tensor, bonus_token_ids: torch.Tensor @@ -1650,12 +1328,13 @@ def _forward_stage( stage: "DSv4DSparkBlock", h: torch.Tensor, main_x: torch.Tensor, - start_pos, + start_pos: torch.Tensor, freqs_cis: torch.Tensor, moe_input_ids: torch.Tensor, - stage_window: Optional[torch.Tensor] = None, - slots: Optional[torch.Tensor] = None, - valid_len: Optional[torch.Tensor] = None, + kv_cache: torch.Tensor, + block_tables: torch.Tensor, + capacities: torch.Tensor, + valid_len: torch.Tensor, all_rank_num_tokens: Optional[List[int]] = None, pre_mix: Optional[torch.Tensor] = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: @@ -1668,15 +1347,8 @@ def _forward_stage( ``hc_pre → attn_norm → DSparkAttention → hc_post`` then ``hc_pre → ffn_norm → MoE → hc_post``. - ``stage_window`` is this stage's persistent rolling captured-context KV - window ``[T, window_size, head_dim]`` owned by the worker across decode - steps; the attention writes the current ``main_kv`` into it in place. When - ``None`` (golden / single-shot) a zero window is allocated per call. - - When ``slots`` (a ``[G]`` int tensor) is given, the CUDA-graph-safe batched - attention (:func:`dspark_attention_forward_batched`) is used instead: it - takes ``start_pos`` as a ``[G]`` tensor and writes/reads ``stage_window`` - (then shaped ``[N, window_size, head_dim]``) through the ``slots`` index. + ``kv_cache`` is this stage's managed [pages, page_size, head_dim] pool. + Page tables and per-request lengths address captured context directly. """ T, block, _, hidden = h.shape @@ -1689,58 +1361,19 @@ def _forward_stage( residual, pre_mix ) layer_input = stage.input_layernorm(layer_input) # [T, block, hidden] - # Rolling-window cache: persist through the worker-owned ``stage_window`` - # for cross-step decode, else a fresh zero window for a single block. - persist = stage_window is not None - kv_cache = ( - stage_window - if persist - else torch.zeros( - T, - self._attn_params["window_size"], - self._attn_params["head_dim"], - dtype=torch.bfloat16, - device=h.device, - ) + attn = dspark_attention_forward( + layer_input, + main_x, + start_pos, + kv_cache, + block_tables, + capacities, + valid_len, + freqs_cis=freqs_cis, + **stage._dspark_attn, + **self._attn_params, ) - if slots is not None: - # Batched, CUDA-graph-safe path (start_pos is a [G] tensor; window is - # written/read through ``slots``). - attn = dspark_attention_forward_batched( - layer_input, - main_x, - start_pos, - kv_cache, - slots, - valid_len=valid_len, - freqs_cis=freqs_cis, - persist=True, - **stage._dspark_attn, - **self._attn_params, - ) - else: - attn = dspark_attention_forward( - layer_input, - main_x, - start_pos, - kv_cache, - freqs_cis=freqs_cis, - persist=persist, - **stage._dspark_attn, - **self._attn_params, - ) - if pre_mix is not None: - residual = stage.hc_attn.post_mapping( - x=attn, - residual=residual, - post_layer_mix=post_mix, - comb_res_mix=comb_mix, - ) - ffn_pre, post_mix, comb_mix, layer_input = stage.hc_ffn.pre_mapping_lagged( - residual, attn_pre - ) - layer_input = stage.post_attention_layernorm(layer_input) - elif stage.enable_fused_hc: + if stage.enable_fused_hc: residual, post_mix, comb_mix, layer_input = stage.hc_ffn.fused_hc( x_prev=attn, residual_prev=residual, @@ -1795,109 +1428,28 @@ def _forward_stage( return (h, ffn_pre) if pre_mix is not None else h def forward( - self, - main_hidden: torch.Tensor, - bonus_token_ids: torch.Tensor, - start_pos: int, - *, - kv_windows: Optional[torch.Tensor] = None, - temperature: float = 0.0, - return_confidence: bool = False, - return_logits: bool = False, - all_rank_num_tokens: Optional[List[int]] = None, - ) -> tuple: - """Full block-draft forward: chain the ``num_stages`` DSpark stages. - - Mirrors the reference ``Transformer.forward_spec`` (generation path, - ``start_pos > 0``): ``forward_embed`` builds the block input + captured - context, each stage runs the captured-context backbone, and - ``forward_head`` emits the block draft tokens + per-position confidence. - - Args: - main_hidden: ``[T, num_capture * hidden]`` captured target context. - bonus_token_ids: ``[T]`` last accepted token per request. - start_pos: absolute decode position (must be > 0). - kv_windows: optional persistent per-stage rolling captured-context KV - windows ``[T, num_stages, window_size, head_dim]`` owned by the - worker; updated in place each call. ``None`` allocates fresh zero - windows (single-shot golden / test path). - Returns: - ``(draft_tokens [T, block], confidence [T, block] or None)``. - """ - assert start_pos > 0, "DSpark draft runs at generation (start_pos > 0)" - if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: - raise RuntimeError( - "DSpark attention weights not cached; call " - "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." - ) - x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) - main_x = main_x.unsqueeze(1) # [T, 1, hidden] for the MQA K/V projection - freqs_cis = self._dspark_freqs_table(x.device) - moe_input_ids = draft_ids.reshape(-1) - - h = x - pre_mix = self._initial_pre_mix(x) - for s, stage in enumerate(self.mtp_layers): - stage_window = kv_windows[:, s] if kv_windows is not None else None - h = self._forward_stage( - stage, - h, - main_x, - start_pos, - freqs_cis, - moe_input_ids, - stage_window, - all_rank_num_tokens=all_rank_num_tokens, - pre_mix=pre_mix, - ) - if self.uses_lagged_pre: - h, pre_mix = h - - return self.forward_head( - h, - bonus_token_ids, - temperature=temperature, - return_confidence=return_confidence, - return_logits=return_logits, - pre_mix=pre_mix, - ) - - def forward_batched( self, main_hidden: torch.Tensor, bonus_token_ids: torch.Tensor, start_pos: torch.Tensor, *, - kv_windows: torch.Tensor, - slots: torch.Tensor, - valid_len: Optional[torch.Tensor] = None, + kv_pages: tuple[torch.Tensor, ...], + block_tables: torch.Tensor, + capacities: torch.Tensor, + valid_len: torch.Tensor, temperature: float = 0.0, return_confidence: bool = False, return_logits: bool = False, all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: - """CUDA-graph-safe batched block-draft forward (all gen requests at once). - - Same computation as :meth:`forward`, but every host-int / data-dependent - operation is tensorized so the whole path can be captured into the target's - CUDA graph (DSpark is a one-engine drafter — its worker runs inside the - graph). ``start_pos`` is a ``[G]`` tensor (one absolute decode position per - gen request); the rolling captured-context windows are written/read through - ``slots`` into the worker-owned ``kv_windows`` buffer; RoPE phases are - gathered from a fixed table. ``forward_head`` always proposes the full - block; ``return_confidence`` only adds a fixed-shape scoring matmul, so - both settings stay capturable. - - Args: - main_hidden: ``[G, num_capture * hidden]`` captured target context. - bonus_token_ids: ``[G]`` last accepted token per gen request. - start_pos: ``[G]`` int tensor of absolute decode positions (> 0). - kv_windows: ``[N, num_stages, window_size, head_dim]`` persistent rolling - windows; written in place through ``slots``. - slots: ``[G]`` int tensor mapping each request to its ``kv_windows`` row. - valid_len: ``[G]`` count of actually written rolling-window entries. - Returns: - ``(draft_tokens [G, block], confidence [G, block] or None)``. + """Draft all generation requests against manager-owned context pages. + + ``main_hidden`` is [batch, num_capture * hidden], ``bonus_token_ids`` + and ``start_pos`` are [batch]. Each stage updates its pool in + ``kv_pages`` through the shared [batch, max_pages] block table. + ``valid_len`` counts initialized context tokens; ``capacities`` is + zero for synthetic rows. Returns draft tokens and proposal lengths, + optionally including the per-position logits. """ if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: raise RuntimeError( @@ -1912,7 +1464,6 @@ def forward_batched( h = x pre_mix = self._initial_pre_mix(x) for s, stage in enumerate(self.mtp_layers): - stage_window = kv_windows[:, s] # [N, window_size, head_dim] h = self._forward_stage( stage, h, @@ -1920,8 +1471,9 @@ def forward_batched( start_pos, freqs_cis, moe_input_ids, - stage_window, - slots, + kv_pages[s][:, 0, 0], + block_tables, + capacities, valid_len, all_rank_num_tokens=all_rank_num_tokens, pre_mix=pre_mix, @@ -2085,22 +1637,18 @@ def embed_tokens(self): def forward(self, main_hidden, bonus_token_ids, start_pos, **kwargs): return self.dspark_model.forward(main_hidden, bonus_token_ids, start_pos, **kwargs) - def forward_batched(self, main_hidden, bonus_token_ids, start_pos, **kwargs): - """CUDA-graph-safe batched draft forward (delegates to the draft model).""" - return self.dspark_model.forward_batched(main_hidden, bonus_token_ids, start_pos, **kwargs) - def run_moe_lockstep_noop(self, all_rank_num_tokens, device): """Empty-batch MoE barrier lockstep (delegates to the draft model).""" return self.dspark_model.run_moe_lockstep_noop(all_rank_num_tokens, device) - def write_context_windows(self, main_hidden, positions, stage_windows): - """Seed / back-fill the rolling KV windows (delegates to the draft model).""" - return self.dspark_model.write_context_windows(main_hidden, positions, stage_windows) + def configure_generation_autotuner(self, max_batch): + """Delegate generation autotuner capacity to the DSpark model.""" + self.dspark_model.configure_generation_autotuner(max_batch) - def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_windows): - """Batched + masked window back-fill (delegates to the draft model).""" - return self.dspark_model.write_context_windows_batched( - main_hidden, positions, slots, mask, kv_windows + def write_context_pages(self, main_hidden, positions, mask, pools, block_tables, capacities): + """Project captured context directly into the managed draft pools.""" + return self.dspark_model.write_context_pages( + main_hidden, positions, mask, pools, block_tables, capacities ) def _prepare_confidence_head_for_load(self, weights: Dict) -> None: @@ -3660,14 +3208,10 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", # Captured-context attention primitives. - "get_dspark_topk_idxs", - "get_dspark_topk_idxs_batched", - "dspark_sparse_attn", "precompute_dspark_freqs_cis", "apply_dspark_rotary", "apply_dspark_rotary_batched", "dspark_attention_forward", - "dspark_attention_forward_batched", # Block draft I/O. "build_draft_input_ids", "dspark_propose", diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 08609658b69d..994b5dd2c427 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -24,7 +24,10 @@ import tensorrt_llm.bindings.executor as trtllm from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, is_sm_100f, prefer_pinned, - str_dtype_to_binding, torch_dtype_to_str) + str_dtype_to_binding, torch_dtype_to_binding, + torch_dtype_to_str) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.executor import DecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams # isort: off @@ -992,6 +995,7 @@ def __init__( self._disable_overlap_scheduler = llm_args.disable_overlap_scheduler self._draft_config = draft_config self._skip_est = skip_est + self._validate_standalone_draft_cache() # Admission cap (tokens of summed context attended-KV) that the fp8 context-MLA workspace reservation # covers, computed in configure_kv_cache_capacity and carried to the KV manager so the scheduler # reads it directly instead of re-deriving it from pool layout. None until reserved (or w == 0). @@ -1244,11 +1248,15 @@ def _get_kv_size_per_token(self, model_config = self._model_engine.model.model_config use_separate_draft_kv_cache = ( self._should_create_separate_draft_kv_cache()) + draft_kwargs = {} + if self._uses_unified_standalone_draft_cache(): + draft_kwargs = self._get_unified_draft_cache_kwargs() total = self._per_manager_cache_cost( self._kv_cache_manager_cls, model_config, kv_cache_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache) + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + **draft_kwargs) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) draft_cost = self._get_draft_cache_cost( @@ -1264,7 +1272,6 @@ def _get_draft_cache_cost( kv_cache_config: KvCacheConfig, *, use_separate_draft_kv_cache: bool, - include_request_cost: bool = False, ) -> Optional[CacheCost]: """Return the draft manager's standalone cache cost, if it has one. @@ -1280,14 +1287,20 @@ def _get_draft_cache_cost( helix_cp_scale = self._mapping.cp_size def scaled(cost: CacheCost) -> CacheCost: - # Rescale only the slope: every other cost component (e.g. the - # fixed intercept) is per-request rank-local bytes and stays - # unscaled. dataclasses.replace keeps any component the rescale - # does not touch, where rebuilding the CacheCost from named - # fields would silently drop it. - return dataclasses.replace(cost, slope=cost.slope * helix_cp_scale) - - if use_separate_draft_kv_cache: + return CacheCost(slope=cost.slope * helix_cp_scale, + intercept=cost.intercept) + + if self._draft_model_engine is not None: + draft_model_config = self._draft_model_engine.model.model_config + draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( + self._draft_model_engine, kv_cache_config) + draft_cost = self._per_manager_cache_cost( + draft_kv_cache_manager_cls, + draft_model_config, + kv_cache_config, + mapping=draft_mapping) + return scaled(draft_cost) + elif use_separate_draft_kv_cache: # One-model draft with separate KV cache layout. # Pass num_layers explicitly since the HF config may report a # different layer count than what is actually used at runtime @@ -1315,27 +1328,30 @@ def scaled(cost: CacheCost) -> CacheCost: draft_kv_cache_config) if self._speculative_config.spec_dec_mode.is_external_drafter(): # External drafter: layers start from 0, normal PP distribution - return scaled( - self._per_manager_cache_cost( - draft_kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - mapping=draft_mapping, - tokens_per_block=draft_tokens_per_block, - is_draft=True, - include_request_cost=include_request_cost)) + # Resolve draft manager class from draft config — may differ + # from target (e.g. hybrid target + plain transformer draft). + draft_kv_cache_manager_cls = get_kv_cache_manager_cls( + effective_draft_config, + draft_kv_cache_config, + is_disagg=self._is_disagg) + draft_cost = self._per_manager_cache_cost( + draft_kv_cache_manager_cls, + effective_draft_config, + draft_kv_cache_config, + mapping=draft_mapping, + tokens_per_block=draft_tokens_per_block) + return scaled(draft_cost) elif self._mapping.is_last_pp_rank(): # EAGLE3/MTP: draft layers only on last PP rank - return scaled( - self._per_manager_cache_cost( - draft_kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - mapping=draft_mapping, - tokens_per_block=draft_tokens_per_block, - num_layers=self._get_num_draft_layers(), - is_draft=True, - include_request_cost=include_request_cost)) + draft_cost = self._per_manager_cache_cost( + self._kv_cache_manager_cls, + effective_draft_config, + draft_kv_cache_config, + mapping=draft_mapping, + tokens_per_block=draft_tokens_per_block, + num_layers=self._get_num_draft_layers(), + is_draft=True) + return scaled(draft_cost) return None def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -1597,6 +1613,8 @@ def _get_token_num_for_estimation(self) -> int: if spec_cfg is not None: num_extra_tokens_per_seq += spec_cfg.tokens_per_gen_step - 1 num_extra_tokens_per_seq += get_num_extra_kv_tokens(spec_cfg) + if self._uses_unified_standalone_draft_cache(): + num_extra_tokens_per_seq += self._get_draft_scratch_tokens() if self._dummy_reqs is None: self._dummy_reqs = self._create_dummy_context_requests( @@ -1641,6 +1659,8 @@ def _get_token_num_for_estimation(self) -> int: self._model_engine.max_seq_len, self._kv_cache_config, ) + if self._uses_unified_standalone_draft_cache(): + num_pool_groups += 1 num_cache_blocks *= num_pool_groups # Dummy context requests use the configured maximum beam width. Scale @@ -2083,12 +2103,14 @@ def _create_kv_cache_manager( kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( model_engine, kv_cache_config) - # When using separate draft KV cache in one-model speculative decoding, - # use layer_mask to include only target layers. The draft layers should - # only be in the separate draft KV cache manager. - # We still pass spec_config so that num_extra_kv_tokens is calculated. + # The target and draft layers have independent storage properties. + # spec_config also determines the target's extra-token capacity. spec_dec_layer_mask = None - if self._should_create_separate_draft_kv_cache(): + draft_cache_kwargs = (self._get_unified_draft_cache_kwargs() + if self._uses_unified_standalone_draft_cache() + else {}) + if (self._should_create_separate_draft_kv_cache() + or draft_cache_kwargs): num_target_layers = model_engine.model.model_config.pretrained_config.num_hidden_layers spec_dec_layer_mask = [True] * num_target_layers @@ -2117,6 +2139,7 @@ def _create_kv_cache_manager( self._llm_args.kv_cache_config.kv_events_config, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, + **draft_cache_kwargs, ) if not self._skip_est: @@ -2142,6 +2165,106 @@ def _create_kv_cache_manager( return kv_cache_manager + def _is_standalone_dspark(self) -> bool: + spec_config = self._speculative_config + return (spec_config is not None + and spec_config.spec_dec_mode.is_dspark() + and not spec_config.draft_is_embedded_in_target + and not spec_config._use_shared_kv_cache) + + def _is_embedded_dspark(self) -> bool: + spec_config = self._speculative_config + return (spec_config is not None + and spec_config.spec_dec_mode.is_dspark() + and spec_config.draft_is_embedded_in_target) + + def _uses_unified_standalone_draft_cache(self) -> bool: + if (not (self._is_standalone_dspark() or self._is_embedded_dspark()) + or not self._is_kv_cache_manager_v2): + return False + # Disaggregation validates unified support; aggregate may use legacy state. + return (self._is_disagg + or self._unified_draft_cache_unsupported_reason() is None) + + def _validate_standalone_draft_cache(self) -> None: + """Validate manager-owned DSpark history for disaggregated serving.""" + if (not (self._is_standalone_dspark() or self._is_embedded_dspark()) + or not self._is_disagg): + return + if not self._is_kv_cache_manager_v2: + raise ValueError("DSpark disaggregation requires " + "kv_cache_config.use_kv_cache_manager_v2=True " + "on both workers.") + reason = self._unified_draft_cache_unsupported_reason() + if reason is not None: + raise ValueError(reason) + + def _unified_draft_cache_unsupported_reason(self) -> Optional[str]: + """Shared admission requirements for unified aggregate and disagg KV.""" + if (self._is_standalone_dspark() + and is_mla(self._draft_config.pretrained_config)): + return "Unified DSpark KV cache does not yet support MLA drafters." + return None + + def _get_draft_mapping(self) -> Mapping: + """Use the drafter's dense TP mapping when target attention uses Helix.""" + return (self._mapping.repurpose_helix_cp_to_tp() + if self._mapping.has_cp_helix() else self._mapping) + + def _get_draft_scratch_tokens(self) -> int: + return 0 if self._is_embedded_dspark( + ) else self._speculative_config.max_draft_len + 1 + + def _get_unified_draft_cache_kwargs(self) -> dict: + """Pass draft storage properties separately from its history requirements.""" + draft = self._model_engine.model.draft_model + if self._is_embedded_dspark(): + return dict( + draft_num_layers=draft.num_stages, + draft_num_kv_heads=1, + draft_head_dim=int(draft._attn_params["head_dim"]), + draft_dtype=torch_dtype_to_binding(torch.bfloat16), + draft_kv_factor=1, + draft_scratch_tokens=0, + draft_attention_backend="DSv4", + draft_window_size=int(draft._attn_params["window_size"]), + ) + config = self._draft_config.pretrained_config + num_heads = config.num_attention_heads + num_kv_heads = getattr(config, "num_key_value_heads", num_heads) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + head_dim = config.hidden_size // num_heads + draft_mapping = self._get_draft_mapping() + attention_tp_size = (1 if draft_mapping.enable_attention_dp else + draft_mapping.tp_size) + attention_backend = self._speculative_config.attention_backend + if attention_backend == "AUTO": + attention_backend = draft.dflash_attention_backend + draft_model_config = getattr(draft, "config", None) + dtype = draft.fc.weight.dtype if hasattr(draft, + "fc") else torch.bfloat16 + if dtype not in (torch.float16, torch.bfloat16): + raise ValueError( + "Standalone draft KV supports FP16 and BF16 storage") + if attention_backend == "TRTLLM": + # The TRTLLM drafter attention reads FP8 pages with FP8 queries; + # the drafter's history is stored at half the BF16 footprint. + dtype = torch.float8_e4m3fn + return dict( + draft_num_layers=config.num_hidden_layers, + draft_num_kv_heads=(num_kv_heads + attention_tp_size - 1) // + attention_tp_size, + draft_head_dim=head_dim, + draft_dtype=torch_dtype_to_binding(dtype), + draft_total_num_kv_heads=num_kv_heads, + draft_scratch_tokens=self._get_draft_scratch_tokens(), + draft_attention_backend=attention_backend, + draft_max_position_embeddings=(getattr( + draft_model_config, "max_position_embeddings", None) + if draft_model_config else None), + ) + def _should_create_separate_draft_kv_cache(self) -> bool: """ Check if we need a separate draft KV cache manager for one-model mode. @@ -2154,15 +2277,22 @@ def _should_create_separate_draft_kv_cache(self) -> bool: if self._speculative_config is None: # No drafter at all, so there is nothing to give a manager to. return False - # Narrower than is_external_drafter(): PARD and DRAFT_TARGET_ONE_MODEL - # never reach the arena this carve-out exists for. - spec_dec_mode = self._speculative_config.spec_dec_mode - is_standalone_drafter = (spec_dec_mode.is_dflash() - or spec_dec_mode.is_dspark()) - if self._mapping.enable_attention_dp and not is_standalone_drafter: - # This bail suits MTP, whose draft layers are target-shaped and - # appendable to the target pool. A standalone drafter has nothing to - # append, so it would be stranded on the private arena instead. + if self._uses_unified_standalone_draft_cache(): + return False + is_external_drafter = ( + self._speculative_config.spec_dec_mode.is_external_drafter()) + if self._mapping.enable_attention_dp and not is_external_drafter: + # Attention-DP uses the target KV cache manager for MTP draft layers + # so draft cache state follows the DP-sharded target requests. That + # works only because MTP draft layers are target-shaped and can be + # appended to the target pool (get_pp_layers). An external drafter + # has its own architecture and layer count, so there is nothing to + # append and the bail would strand it on the private context arena + # -- dense in max_seq_len, and with KV heads unsharded under + # attention DP that is the allocation that OOMs at long ISL. Its + # own manager follows the DP sharding regardless: same mapping (so + # `tp_size = 1 if enable_attention_dp` applies to it too) and block + # tables copied per rank from attn_metadata.request_ids. logger.info( "Attention DP is enabled, separate draft KV cache is not supported." ) @@ -2405,6 +2535,12 @@ def _create_one_model_draft_kv_cache_manager( # the sparse_attention_config. Get it from effective_draft_config which # falls back to the target model's config for MTP mode. sparse_attn_config = effective_draft_config.sparse_attention_config + # The standalone drafter is a plain dense model built against the + # repurposed mapping under helix (CP ranks become TP ranks; each rank + # keeps its full drafter KV). Its paged KV manager must therefore use + # the repurposed, CP-free mapping — helix round-robin ledger semantics + # apply only to the TARGET KV (and KVCacheManagerV2 rejects + # is_draft x helix outright). draft_mapping = self._get_draft_mapping() draft_tokens_per_block = self._get_draft_tokens_per_block() if draft_tokens_per_block != self._tokens_per_block: @@ -3308,21 +3444,24 @@ def _create_kv_cache_manager( head_dim: Optional[Union[int, List[int]]] = None, kv_cache_type=None, is_disagg: bool = False, - disable_overlap_scheduler: bool = False, - cold_page_codec_provider: Optional[object] = None, - kv_events_config: Optional[KVEventsConfig] = None, - joint_kv_cache_reuse: bool = False, - max_cuda_graph_batch_size: Optional[int] = None) -> KVCacheManager: + draft_num_layers: int = 0, + draft_num_kv_heads: int = 0, + draft_head_dim: int = 0, + draft_dtype: DataType = DataType.BF16, + draft_kv_factor: int = 2, + draft_total_num_kv_heads: Optional[int] = None, + draft_scratch_tokens: int = 0, + draft_window_size: Optional[int] = None, + draft_attention_backend: Optional[str] = None, + draft_max_position_embeddings: Optional[int] = None, + joint_kv_cache_reuse: bool = False) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config """ - if cold_page_codec_provider is not None and not issubclass( - kv_cache_manager_cls, KVCacheManagerV2): - raise ValueError( - "Cold-page quantization requires the resolved KV cache manager " - f"to be KVCacheManagerV2; selected {kv_cache_manager_cls.__name__}") - + if draft_num_layers and not issubclass(kv_cache_manager_cls, + KVCacheManagerV2): + raise ValueError("Unified draft cache layers require KVCacheManagerV2.") if (estimating_kv_cache and issubclass(kv_cache_manager_cls, KVCacheManagerV2) and kv_cache_config.pool_ratio is None @@ -3463,27 +3602,18 @@ def _create_kv_cache_manager( "cold_page_codec_provider"] = cold_page_codec_provider manager_extra_kwargs["kv_events_config"] = kv_events_config manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse - manager_extra_kwargs[ - "disable_overlap_scheduler"] = disable_overlap_scheduler - # Vocab size also enables multimodal event decoding and its per-block - # digest scan. Leave it unset for text-only engines. Without an engine - # (e.g. separate one-model draft caches), preserve config resolution: - # a text sub-config alone cannot identify a multimodal deployment. - needs_multimodal_keys = model_engine is None or is_multimodal( - model_engine.model, model_engine.input_processor) - manager_extra_kwargs["vocab_size"] = (resolve_vocab_size(config) if - needs_multimodal_keys else None) - if (needs_multimodal_keys and manager_extra_kwargs["vocab_size"] is None - and kv_cache_config.enable_block_reuse): - logger.warning( - "Could not resolve vocab_size from the model config; " - "multimodal requests will fail when block reuse is on. " - "Disable block reuse to serve them.") - elif kv_events_config is not None and kv_events_config.enable_kv_cache_events: - logger.warning( - "kv_cache_config.kv_events_config is set but streaming KV event " - "publishing requires KV cache manager V2; events will not be " - f"published for {kv_cache_manager_cls.__name__}.") + manager_extra_kwargs.update( + draft_num_layers=draft_num_layers, + draft_num_kv_heads=draft_num_kv_heads, + draft_head_dim=draft_head_dim, + draft_dtype=draft_dtype, + draft_kv_factor=draft_kv_factor, + draft_total_num_kv_heads=draft_total_num_kv_heads, + draft_scratch_tokens=draft_scratch_tokens, + draft_window_size=draft_window_size, + draft_attention_backend=draft_attention_backend, + draft_max_position_embeddings=draft_max_position_embeddings, + ) if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py new file mode 100644 index 000000000000..234128fca0dc --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -0,0 +1,5326 @@ +# 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 hashlib +import math +import os +import sys +from collections import OrderedDict, defaultdict, deque +from dataclasses import dataclass, field, replace +from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union + +import numpy as np +import torch +from strenum import StrEnum + +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind +from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp +from tensorrt_llm._utils import ( + TensorWrapper, + convert_to_torch_tensor, + get_size_in_bytes, + prefer_pinned, +) +from tensorrt_llm.bindings.internal.batch_manager import KvCacheIterationStats, KvCacheStats +from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + IndexMapper, + copy_batch_block_offsets_to_device, +) +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.runtime.kv_cache_hash import get_effective_kv_cache_event_hash_algo +from tensorrt_llm.runtime.kv_cache_manager_v2 import _BACKEND as KV_CACHE_MANAGER_V2_BACKEND +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, + BAD_PAGE_INDEX, + CACHE_LEVEL1, + DEFAULT_BEAM_INDEX, + GPU_LEVEL, + AttentionLayerConfig, + AttnLifeCycle, + BatchDesc, + BufferConfig, + CacheLevel, + CacheTierConfig, + CuError, + DataRole, + DiskCacheTierConfig, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheDesc, + KVCacheEventManager, + KVCacheIterationStatsDelta, + LayerId, + PageIndexMode, + PlannedDropHandle, + PoolGroupPeakBlockStats, + ReuseScope, + SwaScratchReuseConfig, + TokenIdExt, + _cpp_introspection, + _introspection, + _KVCache, + exact_div, + gen_multimodal_cache_key_tokens, + typed_range, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as KVCacheManagerPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfMemoryError as KVCacheOutOfMemoryError +from tensorrt_llm.sampling_params import SamplingParams + +from ..._utils import binding_to_torch_dtype, mpi_rank, nvtx_range, str_dtype_to_torch +from ...logger import logger +from ...mapping import CpType, Mapping +from ..utils import maybe_compile +from .config_utils import uses_vswa_kv_cache_layout +from .connectors.kv_cache_connector import KvCacheConnectorManager +from .kv_cache_stats import ( + KVCacheV2IterationStatsReport, + KVCacheV2LifeCycleIterationStats, + KVCacheV2PoolGroupIterationStats, + KVCacheV2SsmLifeCycleIterationStats, + KVCacheV2SsmSnapshotIterationStats, +) +from .llm_request import LlmRequest, LlmRequestState, SamplingConfig, get_draft_token_length +from .resource_manager import ( + BaseResourceManager, + CacheTypeCpp, + DataType, + KVCacheManager, + ModelConfigCpp, + ModelConfigPython, + _populate_dummy_mrope_config, + get_pp_layers, + request_context, +) +from .scheduler import ScheduledRequests +from .standalone_draft_cache import StandaloneDraftHistory + +if KV_CACHE_MANAGER_V2_BACKEND == "cpp": + from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheDomain + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata + +BASE_GENERATION_TOKEN_COUNT = 1 + +KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS +KV_CACHE_ITERATION_STATS_REUSE_FIELDS = ( + "iter_reused_blocks", + "iter_full_reused_blocks", + "iter_partial_reused_blocks", + "iter_missed_blocks", +) +KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS = tuple( + field_name + for field_name in KV_CACHE_ITERATION_STATS_DELTA_FIELDS + if field_name not in KV_CACHE_ITERATION_STATS_REUSE_FIELDS +) + + +def _create_gpu_cache_tier_config(quota: int, enable_ugpu: bool) -> GpuCacheTierConfig: + if KV_CACHE_MANAGER_V2_BACKEND == "cpp": + if enable_ugpu: + logger.warning_once( + "uGPU KV-cache localization is not supported by the C++ KV cache manager V2; " + "disabling it.", + key="disable_ugpu_for_cpp_kv_cache_manager_v2", + ) + return GpuCacheTierConfig(quota=quota) + return GpuCacheTierConfig(quota=quota, enable_ugpu=enable_ugpu) + + +def _get_num_ugpus(impl) -> int: + """Return the effective uGPU count supported by the selected backend.""" + if KV_CACHE_MANAGER_V2_BACKEND == "cpp": + return 1 + return impl.num_ugpus + + +def _get_kv_cache_ugpu_id(kv_cache) -> int | None: + """Return the cache-local uGPU id supported by the selected backend.""" + if KV_CACHE_MANAGER_V2_BACKEND == "cpp": + return None + return kv_cache.ugpu_id + + +def _create_backend_kv_cache( + impl, + lora_task_id: int | None, + salt: int | None, + input_tokens: Sequence[TokenIdExt] | None, + request_id: int, + expected_prompt_length: int | None, + ugpu_id: int | None, +): + if KV_CACHE_MANAGER_V2_BACKEND == "cpp": + assert ugpu_id is None, "The C++ KV cache manager V2 does not support uGPU localization" + return impl.create_kv_cache( + ReuseScope(lora_id=lora_task_id, salt=salt), + input_tokens, + id=request_id, + expected_prompt_length=expected_prompt_length, + ) + return impl.create_kv_cache( + ReuseScope(lora_id=lora_task_id, salt=salt, ugpu_id=ugpu_id), + input_tokens, + id=request_id, + expected_prompt_length=expected_prompt_length, + ugpu_id=ugpu_id, + ) + + +class Role: + KEY = DataRole("key") + VALUE = DataRole("value") + KEY_BLOCK_SCALE = DataRole("key_block_scale") + VALUE_BLOCK_SCALE = DataRole("value_block_scale") + # Sparse-attention per-layer index-K cache (MiniMax-M3 and similar + # sparse-block-selection backends). Registered as a native V2 + # BufferConfig on sparse layers via the extra_buffers_per_layer hook on + # _build_base_config, so allocation, free, slot reuse, and prefix reuse + # share the lifecycle of the main K/V buffers for the same layer. + INDEX_KEY = DataRole("index_key") + MLA_V_SCALE = DataRole("mla_v_scale") + MLA_V_PACKED = DataRole("mla_v_packed") + MLA_HP_TAIL = DataRole("mla_hp_tail") + ALL = DataRole("all") + + +class BlockReusePolicy(StrEnum): + ALL_REUSABLE = "all_reusable" + PER_REQUEST = "per_request" + PER_CONVERSATION = "per_conversation" + + +def _request_conversation_id(request: LlmRequest) -> Optional[str]: + if request.is_dummy_request: + return None + conversation_params = request.py_conversation_params + if conversation_params is None: + return None + conversation_id = conversation_params.conversation_id.strip() + return conversation_id or None + + +@dataclass(slots=True) +class _ConversationState: + current_request_id: Optional[int] = None + planned_drop_handles: deque[PlannedDropHandle] = field(default_factory=deque) + + +class ConversationManager: + """Track the current request and drop plan for each conversation.""" + + def __init__(self, max_num_turns: int) -> None: + self._max_num_turns = max_num_turns + self._conversation_states: Dict[str, _ConversationState] = {} + + def save_drop_plan(self, request: LlmRequest, kv_cache: _KVCache) -> None: + """Save a completed context's drop plan and apply the preceding plan on success.""" + request_id = request.py_request_id + conversation_id = _request_conversation_id(request) + if conversation_id is None: + return + + state = self._conversation_states[conversation_id] + if state.current_request_id != request_id: + return + + drop_handle = kv_cache.plan_committed_block_drop() + if drop_handle is None: + logger.warning( + f"Committed blocks for request {request_id} in conversation " + f"{conversation_id} have been dropped." + ) + else: + state.planned_drop_handles.append(drop_handle) + if len(state.planned_drop_handles) > self._max_num_turns: + state.planned_drop_handles.popleft().drop() + + self.finish_request(request) + + def prepare_request(self, request: LlmRequest) -> None: + """Register a context request unless its conversation has another active one.""" + conversation_id = _request_conversation_id(request) + if conversation_id is None: + return + request_id = request.py_request_id + state = self._conversation_states.setdefault(conversation_id, _ConversationState()) + current_request_id = state.current_request_id + if current_request_id is not None and current_request_id != request_id: + logger.warning( + f"Conversation {conversation_id} already has current request " + f"{current_request_id}. Request {request_id} will ignore " + "conversation params." + ) + return + + state.current_request_id = request_id + + def finish_request(self, request: LlmRequest) -> None: + """Clear a request as active while preserving any saved drop plan.""" + conversation_id = _request_conversation_id(request) + if conversation_id is None: + return + state = self._conversation_states.get(conversation_id) + if state is None or state.current_request_id != request.py_request_id: + return + + state.current_request_id = None + if not state.planned_drop_handles: + self._conversation_states.pop(conversation_id) + + def clear(self) -> None: + """Clear state after reusable KV-cache blocks have been cleared.""" + self._conversation_states.clear() + + +def _estimate_full_attn_size_per_token( + layer_sizes: Sequence[int], attention_windows: Sequence[Optional[int]] +) -> int: + return sum( + layer_size + for layer_size, window_size in zip(layer_sizes, attention_windows) + if window_size is None or window_size <= 0 + ) + + +def _compute_auto_host_tier_quota( + quota: int, + local_ranks: int, + mem_available: float, + memlock_limit: float, +) -> int: + """Compute the auto-provisioned host cache tier quota for a single rank. + + The host tier backs the MAX_UTILIZATION scheduler's suspend/resume path, + so it must be positive. It defaults to the device quota, capped by the + per-node available-memory budget shared with co-located ranks and by the + pinnable-memory (RLIMIT_MEMLOCK) limit. + + Args: + quota: Device (GPU) KV cache quota in bytes; must be positive. + local_ranks: Number of ranks co-located on this physical node. + mem_available: Available host memory in bytes, or ``float("inf")`` + if unknown. + memlock_limit: RLIMIT_MEMLOCK soft limit in bytes, or + ``float("inf")`` if unlimited or unknown. + + Returns: + Host cache tier quota in bytes (always positive). + """ + candidates = [quota] + if mem_available != float("inf"): + candidates.append(int(mem_available / local_ranks * 0.5)) + if memlock_limit != float("inf"): + candidates.append(int(memlock_limit * 0.8)) + host_quota = min(candidates) + if host_quota <= 0: + logger.warning( + f"KV cache manager v2 auto host tier sizing computed a " + f"non-positive quota ({host_quota}); falling back to the " + f"device quota {quota / (1 << 30):.2f}GiB" + ) + host_quota = quota + return host_quota + + +def _sync_host_tier_quota(host_quota: int, mapping: Mapping) -> int: + """Reduce the auto-provisioned host cache tier quota to the fleet minimum. + + ``_compute_auto_host_tier_quota`` reads rank-local host state + (``SC_AVPHYS_PAGES``, ``RLIMIT_MEMLOCK``), so co-scheduled ranks can arrive + at divergent host quotas (observed up to 10x). Divergent host-tier + retention makes per-rank MAX_UTILIZATION schedulers disagree about which + suspended requests can resume, which wedges collectives on + non-attention-DP TP. Reducing to the fleet minimum makes the + most-constrained rank set the value for everyone, mirroring the device + quota sync. A single-rank job needs no collective and is returned as-is. + + Args: + host_quota: This rank's locally-computed host tier quota in bytes. + mapping: Parallelism mapping; ``world_size`` selects whether a + collective is needed. + + Returns: + The globally-agreed host tier quota in bytes (the fleet ``MIN``). + """ + if mapping.world_size > 1: + host_quota = Distributed.get(mapping).allreduce(host_quota, op=ReduceOp.MIN) + return host_quota + + +def _estimate_swa_cache_size( + layer_sizes: Sequence[int], + attention_windows: Sequence[Optional[int]], + tokens_per_block: int, + *, + context: bool, + scratch: bool, + generation_capacity_headroom: Optional[int] = None, +) -> tuple[int, int]: + tokens_per_block = int(tokens_per_block) + size_per_token = 0 + size_per_request = 0 + scratch_keys = set() + for layer_size, window_size in zip(layer_sizes, attention_windows): + if window_size is not None and window_size > 0: + if generation_capacity_headroom is None: + window_tokens = math.ceil(window_size / tokens_per_block) * tokens_per_block + else: + window_blocks = ( + math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + + 1 + ) + window_tokens = window_blocks * tokens_per_block + if not context: + size_per_request += window_tokens * layer_size + elif not scratch: + size_per_token += layer_size + else: + scratch_key = (int(window_size), layer_size) + if scratch_key in scratch_keys: + size_per_request += window_tokens * layer_size + else: + scratch_keys.add(scratch_key) + size_per_token += layer_size + return size_per_token, size_per_request + + +def _get_generation_kv_capacity(spec_config, *, is_draft: bool) -> tuple[int, int]: + """Return draft-token reserve and total capacity headroom over history.""" + from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens + + if spec_config is None: + return 0, BASE_GENERATION_TOKEN_COUNT + reserve = spec_config.max_total_draft_tokens + if ( + is_draft + and getattr(spec_config, "use_dynamic_tree", False) + and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0 + ): + draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len + reserve = max(reserve, draft_loop_tokens) + dynamic_reserve = reserve - spec_config.max_total_draft_tokens + headroom = ( + get_num_extra_kv_tokens(spec_config) + + spec_config.tokens_per_gen_step + + dynamic_reserve + + reserve # K3's overlap scheduler can be one verification round ahead. + ) + return reserve, headroom + + +def _estimate_cache_size_components( + layer_sizes: Sequence[int], + attention_windows: Sequence[int | None], + tokens_per_block: int, + *, + scratch: bool, + generation_capacity_headroom: int, + standalone_draft_reserve: int = 0, +) -> tuple[int, int, int]: + """Return context/generation bytes per token and generation bytes per request. + + Static profiling and runtime quota conversion must charge the same SWA + retention pages and context scratch space. The standalone draft reserve + is a shared allocation envelope, including windowed attention groups; + it does not represent committed history. Resume-watermark normalization + is separate from these usable-capacity costs. Callers estimate target and + draft layers separately so only target layers can share context scratch. + """ + full_attn_size = _estimate_full_attn_size_per_token(layer_sizes, attention_windows) + context_swa_size, _ = _estimate_swa_cache_size( + layer_sizes, + attention_windows, + tokens_per_block, + context=True, + scratch=scratch, + ) + generation_swa_size, generation_swa_per_request = _estimate_swa_cache_size( + layer_sizes, + attention_windows, + tokens_per_block, + context=False, + scratch=False, + generation_capacity_headroom=generation_capacity_headroom, + ) + return ( + full_attn_size + context_swa_size, + full_attn_size + generation_swa_size, + generation_swa_per_request + standalone_draft_reserve * sum(layer_sizes), + ) + + +def _get_draft_layer_sizes( + *, + draft_num_layers: int = 0, + draft_num_kv_heads: int = 0, + draft_head_dim: int = 0, + draft_dtype: DataType = DataType.BF16, + draft_kv_factor: int = 2, + **kwargs, +) -> tuple[int, ...]: + if draft_num_layers < 0: + raise ValueError("Draft layer count must be nonnegative") + if draft_num_layers: + if min(draft_num_kv_heads, draft_head_dim) <= 0: + raise ValueError("Standalone draft cache dimensions must be positive") + if draft_dtype not in (DataType.HALF, DataType.BF16, DataType.FP8): + raise ValueError("Standalone draft KV supports FP16, BF16 and FP8 storage") + if draft_kv_factor not in (1, 2): + raise ValueError("KV storage requires one or two planes") + size = get_size_in_bytes(draft_num_kv_heads * draft_head_dim * draft_kv_factor, draft_dtype) + return (size,) * draft_num_layers + + +def _estimate_draft_cache_size_components( + layer_sizes: Sequence[int], + tokens_per_block: int, + *, + generation_capacity_headroom: int, + scratch_tokens: int = 0, + window_size: int | None = None, + helix_cp_size: int = 1, +) -> tuple[int, int, int]: + """Draft storage cost, with history and scratch requirements supplied separately.""" + retention_window = window_size + 1 if window_size is not None else None + context, generation, per_request = _estimate_cache_size_components( + layer_sizes, + [retention_window] * len(layer_sizes), + tokens_per_block, + scratch=False, + generation_capacity_headroom=generation_capacity_headroom, + standalone_draft_reserve=scratch_tokens, + ) + extra_bytes_per_token = (helix_cp_size - 1) * sum(layer_sizes) + return context + extra_bytes_per_token, generation + extra_bytes_per_token, per_request + + +def _extend_swa_windows_for_reuse( + attention_windows: Sequence[Optional[int]], + reuse_match_backoff: int, + max_seq_len: Optional[int], +) -> List[Optional[int]]: + """Extend physical SWA retention to cover a backed-off reuse match.""" + if reuse_match_backoff <= 0: + return list(attention_windows) + + retention_windows: List[Optional[int]] = [] + for window_size in attention_windows: + if window_size is None or window_size <= 0: + retention_windows.append(window_size) + continue + + retention_window = int(window_size) + reuse_match_backoff + retention_windows.append( + None + if max_seq_len is not None and retention_window >= int(max_seq_len) + else retention_window + ) + return retention_windows + + +def _get_static_cache_size_layer_components( + model_config: ModelConfigPython, + mapping: Mapping, + num_layers: Optional[int] = None, + **kwargs, +) -> tuple[List[int], List[Optional[int]]]: + config = model_config.pretrained_config + max_seq_len = kwargs.get("max_seq_len") + kv_cache_config = kwargs.get("kv_cache_config") + + num_key_value_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) + if isinstance(num_key_value_heads, Iterable): + num_key_value_heads = sum(num_key_value_heads) / len(num_key_value_heads) + + mla = hasattr(config, "kv_lora_rank") and config.kv_lora_rank is not None + if mla: + head_dim = config.kv_lora_rank + config.qk_rope_head_dim + kv_factor = 1 + else: + tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size + head_dim = getattr(config, "head_dim", None) + if not isinstance(head_dim, int): + head_dim = config.hidden_size // config.num_attention_heads + head_dim = head_dim * num_key_value_heads // tp_size + kv_factor = 2 + + cache_size_per_token = kv_factor * head_dim + quant_config = model_config.quant_config + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): + layer_size = cache_size_per_token + elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache(): + layer_size = math.ceil(cache_size_per_token / 2) + math.ceil(cache_size_per_token / 16) + else: + assert quant_config is None or (not quant_config.quant_mode.has_kv_cache_quant()), ( + "Quantized kv cache is not expected" + ) + layer_size = cache_size_per_token * 2 + + num_attention_layers = KVCacheManager._resolve_num_attention_layers( + model_config, mapping, num_layers + ) + layer_sizes = [layer_size] * num_attention_layers + window_pattern = kv_cache_config.max_attention_window if kv_cache_config is not None else None + + def get_window_size(layer_idx: int) -> Optional[int]: + if window_pattern is None or not isinstance(window_pattern, (list, tuple)): + return None + window_size = window_pattern[layer_idx % len(window_pattern)] + if window_size is None or window_size <= 0: + return None + window_size = int(window_size) + if max_seq_len is not None and window_size == int(max_seq_len): + return None + return window_size + + attention_windows = [get_window_size(layer_idx) for layer_idx in range(num_attention_layers)] + return layer_sizes, attention_windows + + +class _MmRunMetadata(NamedTuple): + """CPU tensors for exact multimodal run lookup. + + `num_runs` is the total number of flat multimodal runs across all + multimodal items, i.e. `item_run_cu_offsets[-1]`. All members are 1-D CPU + int64 tensors with shape `[num_runs]`. + """ + + # Shape [num_runs]; full-prompt start index for each flat run. + run_positions: torch.Tensor + # Shape [num_runs]; one-past-last full-prompt index for each flat run. + run_ends: torch.Tensor + # Shape [num_runs]; logical multimodal item index for each flat run. + run_item_indices: torch.Tensor + # Shape [num_runs]; item-local token offset where each flat run begins. + run_item_offsets: torch.Tensor + + +def _hash_to_digest(hash_ints: Sequence[int]) -> bytes: + # Convert 8 x int32 hash chunks to the 32-byte digest used by C++ block + # keys. The byte order matches getNthByte(), which extracts MSB first. + try: + hash_len = len(hash_ints) + except TypeError as exc: + raise ValueError("Expected 8 int32 hash values, got non-sized input") from exc + if hash_len != 8: + raise ValueError(f"Expected 8 int32 hash values, got {hash_len}") + if not all(isinstance(value, int) for value in hash_ints): + raise ValueError("Expected multimodal hash values to be integers") + return b"".join(v.to_bytes(4, "big", signed=True) for v in hash_ints) + + +def _ensure_int64_cpu_tensor(values: Sequence[int] | torch.Tensor) -> torch.Tensor: + # Block-reuse augmentation is Python-side index math. The metadata is + # produced by host mm preprocessing and carried in + # MultimodalInput metadata. A non-CPU tensor here means the upstream + # contract drifted and would introduce an unexpected sync in the + # block-reuse path, so fail loudly. + if isinstance(values, torch.Tensor): + if values.device.type != "cpu": + raise ValueError( + f"multimodal block-reuse metadata must be CPU-resident, got {values.device}" + ) + return values.to(dtype=torch.int64) + return torch.as_tensor(values, dtype=torch.int64) + + +_MM_RUN_METADATA_UNSET = object() + + +def _resolve_multimodal_run_metadata(req: LlmRequest) -> Optional[_MmRunMetadata]: + # Memoize on the request. Block reuse resolves this on the first context + # chunk, again on every commit during chunked prefill, and on prefetch, yet + # the result is a pure function of the request's immutable multimodal run + # layout. Caching on the request bounds the cache to the request lifetime, + # so no separate teardown is needed. + cached = getattr(req, "py_mm_run_metadata", _MM_RUN_METADATA_UNSET) + if cached is not _MM_RUN_METADATA_UNSET: + return cached + + # Worked example for one logical multimodal item split by text: + # + # prompt index: 0 1 2 3 4 5 + # prompt token: T0 MM0:0 MM0:1 T1 MM0:2 MM0:3 + # flat run: run0 run0 run1 run1 + # + # Inputs encode the prompt layout as: + # + # item_run_cu_offsets = [0, 2] # item 0 owns flat runs [0, 2) + # run_positions = [1, 4] # run0 starts at prompt 1, run1 at 4 + # run_lengths = [2, 2] # token counts for run0 and run1 + # + # The derived arrays below let the augmentation loop map any overlapping + # flat run back to the item digest and to the item-local token offset: + # + # item_run_counts = [2] # item 0 has two runs + # cumulative_run_lengths = [0,2,4] # tokens before run0, run1, and end + # item_starts = [0] # item 0 starts at flat run 0 + # run_item_indices = [0,0] # run0 and run1 both belong to item 0 + # run_item_offsets = [0,2] # run1 begins at MM0 token offset 2 + # + # So a chunk slice covering prompt [4, 6) uses MM0's digest with token + # offset 2; a slice covering prompt [5, 6) uses token offset 3. + + # Shape [num_items + 1]; item i owns flat runs [offsets[i], offsets[i + 1]). + item_run_cu_offsets = req.multimodal_item_run_cu_offsets + # Shape [num_runs]; full-prompt start index for each flat run. + run_positions = req.multimodal_run_positions + # Shape [num_runs]; token count for each flat run. + run_lengths = req.multimodal_run_lengths + + if all(field is None for field in (item_run_cu_offsets, run_positions, run_lengths)): + req.py_mm_run_metadata = None + return None + + if item_run_cu_offsets is None or run_positions is None or run_lengths is None: + raise ValueError( + "multimodal run metadata must be validated before block reuse and provided together" + ) + + item_run_cu_offsets = _ensure_int64_cpu_tensor(item_run_cu_offsets) + run_positions = _ensure_int64_cpu_tensor(run_positions) + run_lengths = _ensure_int64_cpu_tensor(run_lengths) + + # Shape [num_items]; number of flat runs owned by each logical item. + item_run_counts = item_run_cu_offsets[1:] - item_run_cu_offsets[:-1] + + # Shape [num_runs + 1]; item-token count before each flat run, plus end. + cumulative_run_lengths = torch.cat( + (torch.zeros(1, dtype=torch.int64), torch.cumsum(run_lengths, dim=0)) + ) + + # Shape [num_items]; first flat-run index for each logical item. + item_starts = item_run_cu_offsets[:-1] + + # Shape [num_runs]; logical multimodal item index for each flat run. + run_item_indices = torch.repeat_interleave( + torch.arange(item_run_counts.numel(), dtype=torch.int64), item_run_counts + ) + + # Shape [num_runs]; item-local token offset where each flat run begins. + run_item_offsets = ( + cumulative_run_lengths[:-1] - cumulative_run_lengths[item_starts][run_item_indices] + ) + # Shape [num_runs]; one-past-last full-prompt index for each flat run. + run_ends = run_positions + run_lengths + + metadata = _MmRunMetadata( + run_positions=run_positions, + run_ends=run_ends, + run_item_indices=run_item_indices, + run_item_offsets=run_item_offsets, + ) + req.py_mm_run_metadata = metadata + return metadata + + +def _augment_tokens_with_mm_run_metadata( + vocab_size: int, + result: list[TokenIdExt], + multimodal_hashes: Sequence[Sequence[int]], + metadata: _MmRunMetadata, + chunk_start: int, + chunk_end: int, +) -> list[TokenIdExt]: + # Only rewrite multimodal runs that overlap the materialized prompt slice. + overlap_mask = (metadata.run_ends > chunk_start) & (metadata.run_positions < chunk_end) + # Shape [num_overlapping_runs]; flat-run indices that touch this chunk. + overlap_run_indices = torch.nonzero(overlap_mask).flatten() + if overlap_run_indices.numel() == 0: + return result + + # `result` is indexed relative to this chunk, but multimodal cache-key + # tokens are indexed relative to the logical multimodal item. The rewrite + # below therefore computes, for each selected run segment: + # - chunk_result_offset = prompt position - chunk_start + # - item_token_offset = item-local run start + intra-run prompt offset + # Shape [num_overlapping_runs]; item index for each selected flat run. + overlap_run_item_indices = metadata.run_item_indices[overlap_run_indices] + # Materialize the selected run metadata once before the Python loop. These + # are CPU tensors by construction, and the loop below is request hot-path + # bookkeeping. + overlap_run_positions = metadata.run_positions[overlap_run_indices] + prompt_overlap_starts = torch.clamp(overlap_run_positions, min=chunk_start) + prompt_overlap_ends = torch.clamp(metadata.run_ends[overlap_run_indices], max=chunk_end) + item_token_offsets = ( + metadata.run_item_offsets[overlap_run_indices] + + prompt_overlap_starts + - overlap_run_positions + ) + chunk_result_offsets = prompt_overlap_starts - chunk_start + lengths = prompt_overlap_ends - prompt_overlap_starts + + current_item_idx: Optional[int] = None + digest = b"" + for item_idx, chunk_result_offset, item_token_offset, length in zip( + overlap_run_item_indices.tolist(), + chunk_result_offsets.tolist(), + item_token_offsets.tolist(), + lengths.tolist(), + strict=True, + ): + if item_idx != current_item_idx: + current_item_idx = item_idx + digest = _hash_to_digest(multimodal_hashes[item_idx]) + # Feed the coarse item property (content digest) and granular run + # properties (item-local offset and span length) into the key + # generator, so cache keys reflect the actual multimodal tokens being + # rewritten. + result[chunk_result_offset : chunk_result_offset + length] = ( + gen_multimodal_cache_key_tokens( + vocab_size, digest, length, token_offset=item_token_offset + ) + ) + + return result + + +def _augment_tokens_with_contiguous_mm_metadata( + vocab_size: int, + result: list[TokenIdExt], + multimodal_hashes: Sequence[Sequence[int]], + multimodal_positions: Sequence[int] | torch.Tensor, + multimodal_lengths: Sequence[int] | torch.Tensor, + chunk_start: int, + chunk_end: int, +) -> list[TokenIdExt]: + # Legacy metadata assumes every item is one contiguous prompt span. Exact + # run metadata is preferred for video layouts that interleave text tokens. + positions = _ensure_int64_cpu_tensor(multimodal_positions) + lengths = _ensure_int64_cpu_tensor(multimodal_lengths) + + item_ends = positions + lengths + overlap_mask = (item_ends > chunk_start) & (positions < chunk_end) + overlap_item_indices = torch.nonzero(overlap_mask).flatten() + for item_idx_tensor in overlap_item_indices: + item_idx = int(item_idx_tensor) + pos = positions[item_idx].item() + length = lengths[item_idx].item() + overlap_start = max(pos, chunk_start) + overlap_end = min(pos + length, chunk_end) + source_offset = overlap_start - pos + result_offset = overlap_start - chunk_start + overlap_length = overlap_end - overlap_start + result[result_offset : result_offset + overlap_length] = gen_multimodal_cache_key_tokens( + vocab_size, + _hash_to_digest(multimodal_hashes[item_idx]), + overlap_length, + token_offset=source_offset, + ) + + return result + + +def _locate_accepted_draft_tokens(requests: List[LlmRequest]): + num_accepted_draft_tokens = [] + accepted_draft_tokens_indices = [] + rewind_draft_token_separate_adjustments = [] + # for context requests, the py_num_accepted_draft_tokens = 0, and py_num_accepted_draft_tokens_indices = [] + for seq in requests: + num_accepted_draft_tokens.append(seq.py_num_accepted_draft_tokens) + rewind_draft_token_separate_adjustments.append( + seq.py_rewind_draft_token_separate_adjustment + ) + accepted_draft_tokens_indices.extend(seq.py_num_accepted_draft_tokens_indices) + batch_size = len(requests) + num_accepted_draft_tokens_offset = torch.zeros(batch_size + 1, dtype=torch.int32, device="cuda") + num_accepted_draft_tokens_offset[1:] = torch.cumsum( + torch.tensor(num_accepted_draft_tokens, dtype=torch.int32), dim=0 + ) + accepted_draft_tokens_indices = torch.tensor( + accepted_draft_tokens_indices, dtype=torch.int32, device="cuda" + ) + rewind_draft_token_separate_adjustments = torch.tensor( + rewind_draft_token_separate_adjustments, dtype=torch.int32, device="cuda" + ) + return ( + num_accepted_draft_tokens_offset, + accepted_draft_tokens_indices, + rewind_draft_token_separate_adjustments, + ) + + +def _update_kv_cache_draft_token_location( + cache_manager, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata", + kv_cache_dtype_byte_size: float, +): + run_kv_cache_relocation = False + for request in scheduled_batch.generation_requests: + if request.state != LlmRequestState.GENERATION_COMPLETE: + if ( + request.py_num_accepted_draft_tokens > 0 + and len(request.py_num_accepted_draft_tokens_indices) > 0 + ): + run_kv_cache_relocation = True + if not run_kv_cache_relocation: + return + if getattr(cache_manager, "draft_layer_ids", ()): + raise ValueError("Unified standalone draft KV does not support tree-token relocation") + requests = scheduled_batch.all_requests() + ( + accepted_draft_token_offsets, + packed_accepted_draft_tokens_indices, + rewind_draft_token_separate_adjustments, + ) = _locate_accepted_draft_tokens(requests) + past_key_value_lengths = attn_metadata.kv_lens_cuda[: len(requests)] + if ( + attn_metadata.kv_cache_block_offsets is not None + and attn_metadata.host_kv_cache_pool_pointers is not None + and attn_metadata.host_kv_cache_pool_mapping is not None + ): + use_paged_kv_cache = True + else: + use_paged_kv_cache = False + assert use_paged_kv_cache, "Only paged kv cache is supported" + assert len(cache_manager.max_attention_window_vec) == 1, ( + "Currently, only one max attention window size is supported." + ) + + if use_paged_kv_cache: + assert len(set(cache_manager.num_kv_heads_per_layer)) == 1, ( + "update_kv_cache_draft_token_location requires uniform num_kv_heads across all layers, " + f"but got {cache_manager.num_kv_heads_per_layer}" + ) + torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( + accepted_draft_token_offsets, + packed_accepted_draft_tokens_indices, + past_key_value_lengths, + True, + cache_manager.num_layers, + # Use TP-sharded num_kv_heads (per-rank) instead of the unsharded + # total so the C++ kernel computes correct strides and grid dims. + cache_manager.num_kv_heads_per_layer[0], + int(cache_manager.head_dim * kv_cache_dtype_byte_size), + cache_manager.max_total_draft_tokens, + cache_manager.max_attention_window_vec[0], + rewind_draft_token_separate_adjustments, + None, + cache_manager.kv_cache_pool_pointers, + attn_metadata.kv_cache_block_offsets, + cache_manager.max_blocks_per_seq, + cache_manager.tokens_per_block, + None, + ) + + +@maybe_compile(options={"max-autotune": True}) +def _copy_swa_block_offsets_with_scratch_compiled( + block_offsets: torch.Tensor, + copy_idx: torch.Tensor, + pool_ids: torch.Tensor, + scales: torch.Tensor, + layer_offsets: torch.Tensor, + scratch_pages: torch.Tensor, + block_positions: torch.Tensor, + scratch_begs: torch.Tensor, + scratch_ends: torch.Tensor, + scratch_slots: torch.Tensor, + num_contexts: torch.Tensor, + output: torch.Tensor, +) -> None: + base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] + converted = torch.where( + base == BAD_PAGE_INDEX, + BAD_PAGE_INDEX, + base * scales[:, :, None, None] + layer_offsets[:, :, None, None], + ) + + context_positions = torch.arange( + scratch_begs.shape[1], + dtype=torch.int32, + device=scratch_begs.device, + ) + active_context = context_positions < num_contexts + scratch_mask_by_pool = ( + (block_positions >= scratch_begs[:, :, None]) + & (block_positions < scratch_ends[:, :, None]) + & active_context[None, :, None] + ) + range_index = torch.where( + scratch_mask_by_pool, + block_positions - scratch_begs[:, :, None], + 0, + ) + total_offset = range_index[pool_ids] * scratch_pages[:, :, None, None] + slot_idx = (total_offset // scales[:, :, None, None]).clamp(max=scratch_slots.shape[-1] - 1) + slot_id = scratch_slots[pool_ids].gather(-1, slot_idx.long()) + offset = total_offset % scales[:, :, None, None] + scratch_index = ( + slot_id * scales[:, :, None, None] + + (offset + layer_offsets[:, :, None, None]) % scales[:, :, None, None] + ) + scratch_mask = scratch_mask_by_pool[pool_ids] + converted = torch.where(scratch_mask, scratch_index, converted) + + output.copy_(converted.permute(0, 2, 1, 3)) + + +def _rewind_context_cursor(req: LlmRequest, reuse: int) -> None: + """Pull a context cursor back to a shallower prefix. ``setPrepopulatedPromptLen`` + only moves forward, so leave the forward case to it -- it re-floors chunk size too. + """ + if reuse < req.context_current_position: + req.context_current_position = reuse + + +class KVCacheManagerV2(BaseResourceManager): + _draft_layer_sizes: tuple[int, ...] = () + cp_as_tp_layer_ids: frozenset[int] = frozenset() + draft_window_size: int | None = None + draft_layer_ids: tuple[int, ...] = () + _standalone_draft_reserve: int = 0 + _draft_registered_as_model_layers: bool = False + # Read by KvCacheCreator: a subclass that overrides the context + # commit/history protocol opts out of generic reuse-match backoff. + _supports_reuse_match_backoff = True + + def __init__( + self, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: Union[int, List[int]], + tokens_per_block: int, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config=None, + layer_mask: Optional[List[bool]] = None, + vocab_size: int = None, + max_num_tokens: int = 8192, + model_config: Optional[ModelConfigCpp] = None, + max_beam_width: int = 1, + is_draft: bool = False, + kv_connector_manager: Optional[KvCacheConnectorManager] = None, + execution_stream: Optional[torch.cuda.Stream] = None, + is_disagg: bool = False, + enable_stats: bool = False, + num_reserved_index_slots: int = 1, + is_estimating_kv_cache: bool = False, + enable_ugpu: bool = False, + joint_kv_cache_reuse: bool = False, + draft_num_layers: int = 0, + draft_num_kv_heads: int = 0, + draft_head_dim: int = 0, + draft_dtype: DataType = DataType.BF16, + draft_kv_factor: int = 2, + draft_total_num_kv_heads: int | None = None, + draft_scratch_tokens: int = 0, + draft_window_size: int | None = None, + draft_attention_backend: str | None = None, + draft_max_position_embeddings: int | None = None, + **kwargs, + ) -> None: + self._draft_layer_sizes = _get_draft_layer_sizes( + draft_num_layers=draft_num_layers, + draft_num_kv_heads=draft_num_kv_heads, + draft_head_dim=draft_head_dim, + draft_dtype=draft_dtype, + draft_kv_factor=draft_kv_factor, + ) + if self._draft_layer_sizes and KV_CACHE_MANAGER_V2_BACKEND != "cpp": + raise ValueError("Unified standalone draft KV requires the C++ V2 cache backend") + if draft_total_num_kv_heads is not None and draft_total_num_kv_heads < draft_num_kv_heads: + raise ValueError("Total draft KV heads must cover the rank-local heads") + self._draft_total_num_kv_heads = draft_total_num_kv_heads + if draft_scratch_tokens < 0: + raise ValueError("Standalone draft scratch capacity must be nonnegative") + if draft_window_size is not None and draft_window_size <= 0: + raise ValueError("Draft history window must be positive") + self.draft_window_size = draft_window_size + self.draft_attention_backend = draft_attention_backend + self.draft_max_position_embeddings = draft_max_position_embeddings + self.draft_layer_ids: tuple[int, ...] = () + self.draft_history: dict[int, StandaloneDraftHistory] = {} + self._draft_dummy_request_ids: set[int] = set() + self._standalone_draft_reserve = draft_scratch_tokens + self.mapping = mapping + self.dtype = dtype + self.is_disagg = is_disagg + + assert kv_connector_manager is None, ( + "kv_connector_manager is not supported for KVCacheManagerV2" + ) + assert max_beam_width == 1, "max_beam_width must be 1 for KVCacheManagerV2" + assert not (mapping.cp_config.get("cp_type") == CpType.STAR), ( + "Star attention is not supported for KVCacheManagerV2" + ) + + self.kv_cache_type = kv_cache_type + self.pp_layers, self.num_layers = get_pp_layers( + num_layers, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + ) + self.is_draft = is_draft + + # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) + # can distinguish the throwaway estimation-phase managers from the + # final ones: the estimation cache is sized with no headroom for + # retained dummy requests. + self.is_estimating_kv_cache = is_estimating_kv_cache + + # Set True by a compression manager; generation-step resize then leaves history untouched. + self.kv_compression_manages_history: bool = False + self.enable_swa_scratch_reuse = ( + kv_cache_config.enable_swa_scratch_reuse and not self.is_draft + ) + block_reuse_config = kv_cache_config.block_reuse_config + self.block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) + self.num_local_layers = len(self.pp_layers) + self.layer_offsets = {idx: offset for offset, idx in enumerate(self.pp_layers)} + self.max_beam_width = max_beam_width + + tp_size = mapping.tp_size + if mapping.enable_attention_dp: + tp_size = 1 + + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.tokens_per_block = tokens_per_block + # Helix super-block ledger: bookkeeping runs in GLOBAL tokens, one + # ledger block = cp_size physical pages (one per rank), while + # self.tokens_per_block stays the PHYSICAL kernel page size. + self._has_cp_helix = mapping.has_cp_helix() + self._helix_cp_rank = mapping.cp_rank if self._has_cp_helix else 0 + self._helix_cp_size = mapping.cp_size if self._has_cp_helix else 1 + self._ledger_tokens_per_block = tokens_per_block * self._helix_cp_size + if self._has_cp_helix: + if kv_cache_config.enable_block_reuse: + raise ValueError( + "KVCacheManagerV2 does not support block reuse with " + "helix context parallelism: a ledger block's bytes span " + "all CP ranks, so publishing it requires a group-wide " + "protocol. Set kv_cache_config.enable_block_reuse=False." + ) + if is_draft: + raise ValueError( + "KVCacheManagerV2 does not support a draft cache " + "manager with helix context parallelism." + ) + if mapping.enable_attention_dp: + raise ValueError( + "KVCacheManagerV2 does not support attention-DP with " + "helix context parallelism: disagg transfer-completion " + "consensus is skipped under attention-DP, so the " + "scheduler's request view would not be rank-invariant " + "across the CP group." + ) + self.max_seq_len = max_seq_len + self.max_batch_size = max_batch_size + self.max_num_tokens = max_num_tokens + self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 + from ..speculative import draft_prompt_lookahead, get_num_extra_kv_tokens + + self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + # Mirror V1: expose max_draft_len so the native disagg AuxBuffer + # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. + self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 + self.max_total_draft_tokens = ( + spec_config.max_total_draft_tokens if spec_config is not None else 0 + ) + # KvCacheCreator owns the pairing decision; never re-derive it here, or + # the two predicates drift apart and so do the two pools. + self.enable_joint_kv_cache_reuse = joint_kv_cache_reuse + # A draft manager only publishes to its radix tree when it is paired. + self._can_publish_block_reuse = joint_kv_cache_reuse or not self.is_draft + # Unsupported adapters keep their main-like reuse endpoint. + self.reuse_match_backoff = draft_prompt_lookahead(spec_config) or 0 + if not self._supports_reuse_match_backoff: + self.reuse_match_backoff = 0 + # Mirror V1's KV reserve sizing (see V1 __init__ for rationale). + self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( + _get_generation_kv_capacity(spec_config, is_draft=self.is_draft) + ) + self._generation_kv_capacity_headroom += self._standalone_draft_reserve + + self.event_buffer_max_size = kv_cache_config.event_buffer_max_size + self.enable_stats = enable_stats + kv_cache_event_hash_algo = get_effective_kv_cache_event_hash_algo( + kv_cache_config.kv_cache_event_hash_algo, + use_kv_cache_manager_v2=True, + ) + + self._stream = ( + execution_stream if execution_stream is not None else torch.cuda.current_stream() + ) + logger.info(f"[KVCacheManager] execution_stream: {self._stream}") + + # Determine max_attention_window_vec + if kv_cache_config.max_attention_window is not None: + self.max_attention_window_vec = ( + kv_cache_config.max_attention_window.copy() + ) # Make a copy to avoid modifying original + # Clamp all window sizes to max_seq_len before calculating the + # number of KV cache blocks. This prevents the KV cache pool from + # being skewed by the largest window values. + self.max_attention_window_vec = [ + min(max_seq_len, w) for w in self.max_attention_window_vec + ] + + self.max_attention_window_vec = [ + None if w == max_seq_len else w for w in self.max_attention_window_vec + ] + + else: + self.max_attention_window_vec = [None] + + # A D-token match backoff needs D more physical SWA history at the + # committed endpoint so commit_min_snapshot still covers the rewind. + if kv_cache_config.enable_block_reuse and self.reuse_match_backoff: + self.max_attention_window_vec = _extend_swa_windows_for_reuse( + self.max_attention_window_vec, + self.reuse_match_backoff, + max_seq_len, + ) + + event_window_size = max( + self.max_seq_len if window_size is None else int(window_size) + for window_size in self.max_attention_window_vec + ) + self.event_manager: Optional[KVCacheEventManager] = None + if self.event_buffer_max_size > 0: + if mapping.enable_attention_dp: + self.event_manager = KVCacheEventManager( + self.event_buffer_max_size, + window_size=event_window_size, + attention_dp_rank=mapping.rank, + attention_dp_gather=Distributed.get(mapping).allgather, + hash_algo=kv_cache_event_hash_algo, + ) + elif mpi_rank() == 0: + self.event_manager = KVCacheEventManager( + self.event_buffer_max_size, + window_size=event_window_size, + hash_algo=kv_cache_event_hash_algo, + ) + + if isinstance(num_kv_heads, int): + self.num_kv_heads_per_layer = [ + (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_local_layers) + ] + self.total_num_kv_heads_per_layer = [ + (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_layers) + ] + else: + assert len(num_kv_heads) == self.num_layers + + def append_to_kv_heads_per_layer( + num_kv_heads_per_layer: List[int], kv_head: Optional[int] + ): + if kv_head is not None: + num_kv_heads_per_layer.append((kv_head + tp_size - 1) // tp_size) + else: + num_kv_heads_per_layer.append(0) + + self.num_kv_heads_per_layer = [] + if self.num_local_layers > 0: + for i in self.pp_layers: + kv_head = num_kv_heads[i] + append_to_kv_heads_per_layer(self.num_kv_heads_per_layer, kv_head) + + self.total_num_kv_heads_per_layer = [] + for i in range(self.num_layers): + kv_head = num_kv_heads[i] + append_to_kv_heads_per_layer(self.total_num_kv_heads_per_layer, kv_head) + + # Build per-layer head_dim (similar to num_kv_heads_per_layer) + if isinstance(head_dim, int): + self.head_dim_per_layer = [head_dim for _ in range(self.num_local_layers)] + else: + assert len(head_dim) == self.num_layers, ( + f"head_dim list length ({len(head_dim)}) must match num_layers ({self.num_layers})" + ) + self.head_dim_per_layer = [] + if self.num_local_layers > 0: + for i in self.pp_layers: + self.head_dim_per_layer.append(head_dim[i]) + if len(set(self.head_dim_per_layer)) > 1: + logger.info( + f"Per-layer head_dim: {len(self.head_dim_per_layer)} layers, " + f"unique values={set(self.head_dim_per_layer)}" + ) + + self.dtype_per_layer = [self.dtype] * self.num_local_layers + self.kv_factor_per_layer = [self.kv_factor] * self.num_local_layers + + self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) + + max_util_for_resume = kv_cache_config.max_util_for_resume + quota = sys.maxsize + if ( + kv_cache_config.max_gpu_total_bytes is not None + and kv_cache_config.max_gpu_total_bytes > 0 + ): + quota = int(kv_cache_config.max_gpu_total_bytes) + logger.info(f"max_gpu_total_bytes is provided. New quota is {quota / (1 << 30)}GiB") + if kv_cache_config.max_tokens is not None: + quota_from_max_tokens = int( + math.ceil( + self._get_quota_from_max_tokens(kv_cache_config.max_tokens) + / max_util_for_resume + ) + ) + quota = min(quota, quota_from_max_tokens) + logger.info( + f"max_tokens {kv_cache_config.max_tokens} is provided. " + f"Allowed quota from max_tokens is {quota_from_max_tokens / (1 << 30)}GiB. " + f"New quota is {quota / (1 << 30)}GiB" + ) + + assert quota < sys.maxsize, ( + "Quota not set. Check kv_cache_config.max_tokens or kv_cache_config.max_gpu_total_bytes" + ) + + # Sync resumable token capacity across ranks so the scheduler produces + # identical batches. Normalize to tokens because cache costs vary + # across PP ranks, including fixed per-rank costs. + if mapping.world_size > 1: + dist = Distributed.get(mapping) + resumable_quota = int(quota * max_util_for_resume) + max_tokens = self._get_max_tokens_from_quota(resumable_quota) + max_tokens = dist.allreduce(max_tokens, op=ReduceOp.MIN) + # inf max_tokens means all layers are SWA and every rank quota can + # fit all SWA fixed cache. + if not math.isinf(max_tokens): + # allreduce(MIN) must never increase the local quota. The + # token↔quota round-trip is not identity when SWA layers + # dominate (full_attn_size_per_token==0), so clamp to guard + # against a bogus inflation (nvbugs/6418103). + synced_quota = int( + math.ceil(self._get_quota_from_max_tokens(max_tokens) / max_util_for_resume) + ) + quota = min(quota, synced_quota) + + logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") + + cache_tiers: List[CacheTierConfig] = [ + _create_gpu_cache_tier_config(int(quota), enable_ugpu) + ] + if kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: + host_quota = kv_cache_config.host_cache_size + else: + # The V2 MAX_UTILIZATION scheduler relies on suspend/resume to + # evict and later restore KV cache pages. Without a secondary + # tier, suspended held pages cannot migrate out of GPU, so + # suspension cannot free capacity and scheduling can deadlock. + # + # Automatically provision a host tier matching the GPU quota so + # suspend/resume works out of the box. Cap at available host + # memory and pinnable memory limit to avoid allocation failures. + import resource + + # Rank-aware host tier sizing: divide the per-node memory budget + # by the number of ranks sharing this physical node. Each rank + # independently provisions a host tier; without this, N + # co-located ranks each reserving a device-quota-sized block can + # OOM the host (observed on GB300 NVL72 with 4 ranks/node and + # ~170GiB device quota each on a 975GiB node). + local_ranks = max(1, Distributed.get(mapping).local_world_size) + + try: + mem_available = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_AVPHYS_PAGES") + except (ValueError, OSError): + mem_available = float("inf") + try: + _soft, _hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) + memlock_limit = _soft if _soft != resource.RLIM_INFINITY else float("inf") + except (ValueError, OSError): + memlock_limit = float("inf") + host_quota = _compute_auto_host_tier_quota( + quota, local_ranks, mem_available, memlock_limit + ) + logger.info( + f"KV cache manager v2 auto host tier sizing: " + f"{local_ranks} co-located rank(s) on this node, " + f"available host memory {mem_available / (1 << 30):.2f}GiB" + ) + # Reduce the rank-local auto quota to the fleet minimum; see + # _sync_host_tier_quota for why divergence wedges collectives. + host_quota = _sync_host_tier_quota(host_quota, mapping) + if host_quota > 0: + cache_tiers.append(HostCacheTierConfig(quota=int(host_quota))) + logger.info( + f"KV cache manager v2 host cache quota set to {host_quota / (1 << 30):.2f}GiB" + ) + disk_cache_size = kv_cache_config.disk_cache_size + if disk_cache_size is not None and disk_cache_size > 0: + disk_cache_path = kv_cache_config.disk_cache_path + assert disk_cache_path is not None + cache_tiers.append( + DiskCacheTierConfig(quota=int(disk_cache_size), path=disk_cache_path) + ) + logger.info( + f"KV cache manager v2 disk cache quota set to {disk_cache_size / (1 << 30):.2f}GiB at {disk_cache_path}" + ) + + self.vocab_size = vocab_size + + config = self._build_base_config( + kv_cache_config, + tokens_per_block=tokens_per_block, + cache_tiers=cache_tiers, + ) + config = self._build_cache_config(config) + config = self._append_standalone_draft_layers( + config, + num_kv_heads=draft_num_kv_heads, + head_dim=draft_head_dim, + dtype=draft_dtype, + kv_factor=draft_kv_factor, + ) + config = self._configure_unified_draft_pool_ratio(config) + has_host_cache_tier = any( + isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers + ) + + self.kv_cache_manager_py_config = config + + try: + self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) + except (CuError, KVCacheOutOfMemoryError): + if has_host_cache_tier: + logger.warning( + "Failed to initialize KV cache manager with host cache " + "tier (cuMemHostRegister may have failed). " + "Retrying without host cache tier." + ) + cache_tiers_without_host = [ + tier for tier in config.cache_tiers if not isinstance(tier, HostCacheTierConfig) + ] + config = replace(config, cache_tiers=cache_tiers_without_host) + self.kv_cache_manager_py_config = config + self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) + else: + raise + self.can_evict = len(config.cache_tiers) > 1 + if self.event_manager is not None: + self.event_manager.set_layer_group_window_sizes( + self._get_event_window_sizes_by_layer_group() + ) + self.event_manager.add_created_event( + self._get_event_num_blocks_per_cache_level(config.cache_tiers, tokens_per_block), + self._get_event_layer_group_ids(), + ) + + # Both backends build layer_grouping on demand, and the layer order + # within a group is not part of its contract. Cache a stable physical- + # layout representative for each role from the public pool descriptors. + + # Cache the fork-join attention flag at construction time so the + # manager's per-uGPU metadata invariants stay stable for its lifetime. + self._fork_join_attn = self.num_ugpus > 1 and enable_ugpu + self.num_pools = len(self.impl.layer_grouping) + self._pool_layer_ids_by_role: Dict[Tuple[int, DataRole], LayerId] = {} + for pool_group in self.impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + pool_id = int(variant.layer_group_id) + for coalesced in variant.coalesced_buffers: + for buffer_id in coalesced.buffer_ids: + self._pool_layer_ids_by_role.setdefault( + (pool_id, buffer_id.role), buffer_id.layer_id + ) + # num_pools is the logical layer-group count. With SWA scratch reuse, + # scratch slot IDs are only valid with per-layer page indices, so the + # attention op sees one virtual pool per local layer while the + # underlying manager can still group layers. + if self.enable_swa_scratch_reuse: + self.num_attention_op_pools = self.num_local_layers + else: + self.num_attention_op_pools = self.num_pools + + num_layers = len(config.layers) + self.layer_to_pool_mapping_dict: dict[int, int] = { + layer_id: self.impl.get_layer_group_id(layer_id) + for layer_id in typed_range(LayerId(num_layers)) + } + + self.kv_cache_map: dict[int, _KVCache] = {} + + # Tracks the draft length allocated by try_allocate_generation per + # request. Used by extend_capacity_for_tokens to compute the exact + # padding delta instead of blindly extending, which would cause + # unbounded capacity growth. + self._allocated_draft_lens: dict[int, int] = {} + + # Overlap-scheduler slack (the extra draft_len in + # _required_gen_capacity's 1 + 2 * draft_len growth) granted to each + # request and not yet reclaimed. The sampler's py_rewind_len only + # covers the rejected part of the recurring 1 + draft_len growth, so + # update_resources must additionally trim this slack — otherwise it + # compounds, leaking draft_len tokens of capacity per generation + # iteration until the request's block count overruns + # max_blocks_per_seq ("User-provided base page indices is too short"). + self._pending_overlap_slack: dict[int, int] = {} + + # Defensive cap for get_num_available_tokens: when host cache is + # enabled, clamp_max_seq_len_for_mem may return a value that spans + # both GPU and host tiers. Storing the explicit max_tokens (if set) + # lets us cap the result to GPU-only capacity so callers like CUDA + # graph warmup don't over-allocate beyond the GPU pool. + # None when max_tokens is not explicitly configured — other config + # paths (max_gpu_total_bytes, free_gpu_memory_fraction) are already + # bounded by the GPU quota passed to GpuCacheTierConfig. + self._gpu_max_tokens = kv_cache_config.max_tokens + + max_num_tokens = self.get_num_available_tokens(token_num_upper_bound=max_seq_len) + + if max_seq_len > max_num_tokens: + logger.warning( + f"max_seq_len {max_seq_len} is greater than max_num_tokens {max_num_tokens} " + "that can be allocated in kv cache manager, setting " + f"max_seq_len to {max_num_tokens}" + ) + # max_num_tokens is a float from clamp_max_seq_len_for_mem; cast + # so downstream int-only consumers (torch.randint size, range) + # stay int. + self.max_seq_len = int(max_num_tokens) + + # Pad max_blocks_per_seq to next multiple of 4 (copy_block_offsets kernel). + # Account for max single-sequence capacity = seq_len + extra KV tokens + + # _kv_reserve_draft_tokens (see __init__) + 1 base decode token. + max_seq_capacity = ( + self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1 + ) + self.max_blocks_per_seq = ( + max_seq_capacity + self._ledger_tokens_per_block - 1 + ) // self._ledger_tokens_per_block + if self.max_blocks_per_seq % 4 != 0: + self.max_blocks_per_seq = ((self.max_blocks_per_seq + 3) // 4) * 4 + + self.enable_block_reuse = kv_cache_config.enable_block_reuse + self.enable_partial_reuse = kv_cache_config.enable_partial_reuse + self.disk_prefetch_num_reqs = kv_cache_config.disk_prefetch_num_reqs + enable_conversation_manager = ( + self.enable_block_reuse + and self.block_reuse_policy == BlockReusePolicy.PER_CONVERSATION + and self._can_publish_block_reuse + ) + self.conversation_manager = ( + ConversationManager(block_reuse_config.max_num_turns) + if enable_conversation_manager + else None + ) + + # With pipeline parallelism, multiple microbatches can be in-flight + # simultaneously, so we need slots for all concurrent sequences. + # Reserve stable slots for persistent request IDs such as CUDA-graph + # padding requests. The default preserves the main-branch allocation. + # In disaggregated mode, use a coefficient of 2: at any moment up to + # `max_num_sequences` requests can be actively generating while another + # up to `max_num_sequences` requests are still in KV transfer + # (TRANS_IN_PROGRESS) and continue to hold their index slots. The 2x + # capacity lets the next batch of active requests acquire slots without + # waiting for the previous batch's transfers to finish. + max_num_sequences = max_batch_size * mapping.pp_size + assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" + index_mapper_capacity = ( + max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots + ) + logger.info( + f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " + f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " + f"num_reserved_index_slots={num_reserved_index_slots}, " + f"max_beam_width={max_beam_width})" + ) + self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) + self._early_freed_index_requests: set[int] = set() + self._prepare_page_table_tensor(index_mapper_capacity) + + self._log_kv_cache_pool_lifecycle_mapping() + + def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: + """Return the roles represented by the two page-table index lanes. + + When present, role B must be addressable from role A using a constant + page-index offset. + """ + layer_id = int(self.impl.layer_grouping[pool_id][0]) + return Role.KEY, Role.VALUE if self.kv_factor_per_layer[layer_id] == 2 else None + + def _get_block_scale_role(self, role_a: DataRole, layer_id: int) -> Optional[DataRole]: + """Select block scales for a role and local layer in the pool mapping.""" + if self.dtype_per_layer[layer_id] != DataType.NVFP4 or role_a != Role.KEY: + return None + return Role.KEY_BLOCK_SCALE + + def _build_pool_mapping_tensors(self): + """Build the (kv_cache_pool_pointers, kv_cache_pool_mapping) tensors. + + An overridable hook for subclasses whose pools coalesce extra + per-layer buffers alongside K/V. + """ + has_block_scales = DataType.NVFP4 in self.dtype_per_layer + kv_cache_pool_pointers_list = [] + kv_cache_pool_mapping_list = [] + block_scale_pool_pointers_list = [] + if self.enable_swa_scratch_reuse: + for layer_id in typed_range(LayerId(self.num_local_layers)): + pool_id = self.impl.get_layer_group_id(layer_id) + role_a, _ = self._get_pool_roles(pool_id) + kv_cache_pool_pointers_list.append( + [ + self.impl.get_mem_pool_base_address( + layer_id, role_a, PageIndexMode.PER_LAYER + ), + 0, + ] + ) + if has_block_scales: + block_scale_role = self._get_block_scale_role(role_a, layer_id) + block_scale_pool_pointers_list.append( + [ + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, PageIndexMode.PER_LAYER + ) + if block_scale_role is not None + else 0, + 0, + ] + ) + kv_cache_pool_mapping_list.append([int(layer_id), 0]) + else: + for pool_id in range(self.num_pools): + role_a, _ = self._get_pool_roles(pool_id) + layer_id = self._pool_layer_ids_by_role[(pool_id, role_a)] + key_base_addr = self.impl.get_mem_pool_base_address( + layer_id, role_a, PageIndexMode.SHARED + ) + kv_cache_pool_pointers_list.append([key_base_addr, 0]) + if has_block_scales: + # The KEY/scale pointers are a (rep-layer, 0-offset) origin + # against which each layer's kv_cache_pool_mapping offset is + # resolved. The block-scale origin must reproduce the SAME + # per-layer offset() as KEY, so mirror the KEY base for the + # same representative layer and shift it back by that layer's + # offset. For the base manager offset(rep) == 0, so this is + # just the rep layer's scale base; for address-ranked + # subclasses (MiniMax-M3) offset(rep) may be non-zero, and the + # shift lands the origin on the pool's slot-0 scale address. + # This keeps block_scale_offset == offset without depending on + # the non-contractual layer_grouping order. + block_scale_role = self._get_block_scale_role(role_a, layer_id) + if block_scale_role is not None: + rep_offset = self._kv_pool_mapping_offset(layer_id, pool_id, key_base_addr) + scale_stride = ( + self.get_layer_bytes_per_token(layer_id, block_scale_role) + * self.kv_factor_per_layer[layer_id] + * self.tokens_per_block + ) + scale_base_addr = ( + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, PageIndexMode.SHARED + ) + - rep_offset * scale_stride + ) + else: + scale_base_addr = 0 + block_scale_pool_pointers_list.append([scale_base_addr, 0]) + + for layer_id in typed_range(LayerId(self.num_local_layers)): + layer_group_id = self.impl.get_layer_group_id(layer_id) + role_a, role_b = self._get_pool_roles(layer_group_id) + index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] + if role_a == Role.KEY: + offset = self._kv_pool_mapping_offset(layer_id, layer_group_id, index_base_addr) + else: + addr_offset = ( + self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED) + - index_base_addr + ) + offset_divisor = self.impl.get_page_stride(layer_id, role_a) + if role_b is not None: + offset_divisor *= self.kv_factor_per_layer[layer_id] + offset = exact_div( + addr_offset, + offset_divisor, + ) + + if self.dtype_per_layer[layer_id] != DataType.NVFP4 or role_a != Role.KEY: + block_scale_offset = None + else: + block_scale_role = self._get_block_scale_role(role_a, layer_id) + if block_scale_role is None: + block_scale_offset = None + else: + block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] + block_scale_addr_offset = ( + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, PageIndexMode.SHARED + ) + - block_scale_base_addr + ) + block_scale_offset = exact_div( + block_scale_addr_offset, + self.get_layer_bytes_per_token(layer_id, block_scale_role) + * self.kv_factor_per_layer[layer_id] + * self.tokens_per_block, + ) + + if block_scale_offset is not None: + assert block_scale_offset == offset, ( + "Block scale offset and offset should be the same" + ) + + kv_cache_pool_mapping_list.append([layer_group_id, offset]) + + if has_block_scales: + for pool_id, block_scale_pool_pointers in enumerate(block_scale_pool_pointers_list): + pool_pointers = kv_cache_pool_pointers_list[pool_id] + kv_cache_pool_pointers_list[pool_id] = [ + [pool_pointers[0], block_scale_pool_pointers[0]], + [pool_pointers[1], block_scale_pool_pointers[1]], + ] + + kv_cache_pool_pointers = torch.tensor( + kv_cache_pool_pointers_list, + dtype=torch.int64, + device="cpu", + pin_memory=prefer_pinned(), + ) + kv_cache_pool_mapping = torch.tensor( + kv_cache_pool_mapping_list, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + return kv_cache_pool_pointers, kv_cache_pool_mapping + + def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + self.kv_cache_pool_pointers, self.kv_cache_pool_mapping = self._build_pool_mapping_tensors() + self.index_scales = torch.empty( + self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu" + ) + self.kv_offset = torch.empty( + self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu" + ) + for pool_id in range(self.num_pools): + role_a, role_b = self._get_pool_roles(pool_id) + layer_id = self._pool_layer_ids_by_role[(pool_id, role_a)] + self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, role_a) + if role_b is not None: + self.kv_offset[pool_id] = exact_div( + self.impl.get_mem_pool_base_address(layer_id, role_b, PageIndexMode.SHARED) + - self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), + self.impl.get_page_stride(layer_id, role_a), + ) + else: + self.kv_offset[pool_id] = 0 + # Plain-int mirror of index_scales so the per-step block-table build + # does not index a tensor per request (see get_batch_cache_indices*). + self._index_scale_ints: List[int] = self.index_scales.tolist() + + # Keep unused block offsets as safe block index 0. + self.host_kv_cache_block_offsets = torch.zeros( + self.num_pools, + index_mapper_capacity * self.max_beam_width, + 2, # key and value + self.max_blocks_per_seq, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + if self.enable_swa_scratch_reuse: + self._prepare_swa_scratch_copy_tensors(index_mapper_capacity) + + def _kv_pool_mapping_offset( + self, layer_id: LayerId, layer_group_id: int, key_base_addr: int + ) -> int: + """Per-layer offset recorded in ``kv_cache_pool_mapping``. + + The default derives the layer's position within its pool from the K + base address, assuming every layer contributes exactly K(+V) to the + pool slot so the layer stride is uniform. Managers whose pool slots + may interleave extra per-layer buffers between layers (non-uniform + layer strides — e.g. MiniMax M3 when the index-K buffer coalesces + into the K/V pool) must override this with a positional formula. + """ + addr_offset = ( + self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) + - key_base_addr + ) + return exact_div( + addr_offset, + self.get_layer_bytes_per_token(layer_id, Role.KEY) + * self.get_layer_kv_factor(self.pp_layers[layer_id]) + * self.tokens_per_block, + ) + + def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: + layer_sizes = [] + attention_windows = [] + pattern_len = len(self.max_attention_window_vec) + for local_layer_idx in range(self.num_local_layers): + layer_sizes.append( + self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=Role.ALL) + ) + if self._is_standalone_draft_layer(local_layer_idx): + attention_windows.append(self.draft_retention_window_size) + else: + attention_windows.append( + self.max_attention_window_vec[self.pp_layers[local_layer_idx] % pattern_len] + ) + # Quota sizing precedes specialized target-layer construction. Include + # draft storage then; after append it is already in the per-layer arrays. + if self._draft_layer_sizes and not self._draft_registered_as_model_layers: + layer_sizes.extend(self._draft_layer_sizes) + attention_windows.extend( + [self.draft_retention_window_size] * len(self._draft_layer_sizes) + ) + return layer_sizes, attention_windows + + def _get_generation_request_capacity(self) -> int: + """Return the resident generation requests used for SWA sizing.""" + return self.max_batch_size + + def _get_runtime_cache_size_components(self) -> tuple[int, int, int]: + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + num_target_layers = len(layer_sizes) - (len(self._draft_layer_sizes)) + context_size, generation_size, fixed_size = _estimate_cache_size_components( + layer_sizes[:num_target_layers], + attention_windows[:num_target_layers], + self.tokens_per_block, + scratch=self.enable_swa_scratch_reuse, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + standalone_draft_reserve=self._standalone_draft_reserve, + ) + if self._draft_layer_sizes: + draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( + self._draft_layer_sizes, + self.tokens_per_block, + scratch_tokens=self._standalone_draft_reserve, + window_size=self.draft_window_size, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + helix_cp_size=self._helix_cp_size, + ) + context_size += draft_context + generation_size += draft_generation + fixed_size += draft_fixed + return context_size, generation_size, fixed_size + + def _get_max_tokens_from_quota(self, quota: int) -> float: + """Rank-local byte quota -> token capacity (GLOBAL tokens under helix).""" + tokens = self._get_max_tokens_from_quota_impl(quota) + if self._has_cp_helix and not math.isinf(tokens): + # Floor to whole physical pages before scaling: a ledger block + # allocates one full page on every CP rank, so a partial + # trailing page in the rank-local budget is never usable. + tokens = int(tokens) // self.tokens_per_block * self.tokens_per_block + tokens *= self._helix_cp_size + return tokens + + def _get_max_tokens_from_quota_impl(self, quota: int) -> float: + if not self._draft_layer_sizes: + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + full_attn_size_per_token = _estimate_full_attn_size_per_token( + layer_sizes, attention_windows + ) + context_swa_size_per_token, _ = _estimate_swa_cache_size( + layer_sizes, + attention_windows, + self.tokens_per_block, + context=True, + scratch=self.enable_swa_scratch_reuse, + ) + ( + generation_swa_size_per_token, + generation_swa_size_per_request, + ) = _estimate_swa_cache_size( + layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False + ) + size_per_batch = ( + self._get_generation_request_capacity() * generation_swa_size_per_request + ) + if quota < size_per_batch: + return 0 + context_size_per_token = full_attn_size_per_token + context_swa_size_per_token + context_limit_quota = self.max_num_tokens * context_size_per_token + size_per_batch + if quota <= context_limit_quota: + if context_size_per_token <= 0: + return float("inf") + return (quota - size_per_batch) / context_size_per_token + + generation_size_per_token = full_attn_size_per_token + generation_swa_size_per_token + if generation_size_per_token <= 0: + return float("inf") + return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token + + ( + context_size_per_token, + generation_size_per_token, + generation_size_per_request, + ) = self._get_runtime_cache_size_components() + size_per_batch = self._get_generation_request_capacity() * generation_size_per_request + if quota < size_per_batch: + return 0 + context_limit_quota = self.max_num_tokens * context_size_per_token + size_per_batch + if quota <= context_limit_quota: + if context_size_per_token <= 0: + return float("inf") + return (quota - size_per_batch) / context_size_per_token + + if generation_size_per_token <= 0: + return float("inf") + return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token + + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + """Token capacity (GLOBAL tokens under helix) -> rank-local byte quota.""" + if self._has_cp_helix: + # Round up to whole ledger blocks first: allocation is page- + # granular on every rank, so a request of N global tokens costs + # ceil(N / ledger_tpb) full physical pages per rank. + blocks = -(-int(max_tokens) // self._ledger_tokens_per_block) + max_tokens = blocks * self.tokens_per_block + return self._get_quota_from_max_tokens_impl(max_tokens) + + def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: + if not self._draft_layer_sizes: + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + full_attn_size_per_token = _estimate_full_attn_size_per_token( + layer_sizes, attention_windows + ) + ( + context_swa_size_per_token, + _, + ) = _estimate_swa_cache_size( + layer_sizes, + attention_windows, + self.tokens_per_block, + context=True, + scratch=self.enable_swa_scratch_reuse, + ) + ( + generation_swa_size_per_token, + generation_swa_size_per_request, + ) = _estimate_swa_cache_size( + layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False + ) + context_tokens = min(max_tokens, self.max_num_tokens) + generation_tokens = max_tokens - context_tokens + generation_quota = ( + max_tokens * full_attn_size_per_token + + generation_tokens * generation_swa_size_per_token + + self._get_generation_request_capacity() * generation_swa_size_per_request + ) + context_extra_quota = context_tokens * context_swa_size_per_token + return int(generation_quota + context_extra_quota) + + ( + context_size_per_token, + generation_size_per_token, + generation_size_per_request, + ) = self._get_runtime_cache_size_components() + context_tokens = min(max_tokens, self.max_num_tokens) + generation_tokens = max_tokens - context_tokens + return int( + context_tokens * context_size_per_token + + generation_tokens * generation_size_per_token + + self._get_generation_request_capacity() * generation_size_per_request + ) + + def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: + """Draft bytes within the unified manager's usable token capacity.""" + assert self._draft_layer_sizes + if self._has_cp_helix: + # Each target ledger block retains all of its draft subpages. + max_tokens = ( + math.ceil(max_tokens / self._ledger_tokens_per_block) + * self._ledger_tokens_per_block + ) + context_size, generation_size, fixed_size = _estimate_draft_cache_size_components( + self._draft_layer_sizes, + self.tokens_per_block, + scratch_tokens=self._standalone_draft_reserve, + window_size=self.draft_window_size, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + ) + context_tokens = ( + max_tokens if self.max_num_tokens is None else min(max_tokens, self.max_num_tokens) + ) + return math.ceil( + context_tokens * context_size + + (max_tokens - context_tokens) * generation_size + + self._get_generation_request_capacity() * fixed_size + ) + + def _get_event_num_blocks_per_cache_level( + self, + cache_tiers: List[CacheTierConfig], + tokens_per_block: int, + ) -> List[int]: + bytes_per_block = self.get_cache_bytes_per_token() * tokens_per_block + if bytes_per_block <= 0: + return [] + return [int(tier.quota // bytes_per_block) for tier in cache_tiers] + + def _get_event_layer_group_ids(self) -> List[int]: + return [int(layer_group_id) for layer_group_id in range(len(self.impl.layer_grouping))] + + def _get_event_window_sizes_by_layer_group(self) -> Dict[int, int]: + # Assumes every layer in a group shares the same sliding_window_size, + # which is how `impl.layer_grouping` partitions layers today. Only the + # first layer's window is read; if the grouping policy ever permits + # mixed windows in one group, this needs to fan out per-layer. + + def get_event_window_size(layer_id: int) -> int: + layer_config = self.kv_cache_manager_py_config.layers[layer_id] + window_size = getattr(layer_config, "sliding_window_size", None) + return self.max_seq_len if window_size is None else int(window_size) + + return { + int(layer_group_id): get_event_window_size(int(layer_ids[0])) + for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping) + } + + def _format_kv_cache_pool_lifecycle_entry(self, layer_id: LayerId, role: DataRole) -> str: + for pool_group in self.impl.pool_group_descs: + for variant in pool_group.slot_desc.variants: + for coalesced in variant.coalesced_buffers: + for buffer_id in coalesced.buffer_ids: + if int(buffer_id.layer_id) == int(layer_id) and str(buffer_id.role) == str( + role + ): + return ( + f"role={role!s}, " + f"pool_group_id={int(pool_group.pool_group_index)}, " + f"layer_group_id={int(variant.layer_group_id)}" + ) + return f"role={role!s}, pool_group_id=?, layer_group_id=?" + + def _format_kv_cache_pool_group_sizes(self) -> list[str]: + """Describe how many slots each pool group actually got. + + The aggregate KV utilization reported per iteration cannot answer + "which pool group is the binding constraint", because it sums every + pool group: a run can sit at 13% aggregate while one small pool group + is at its admission ceiling and refusing to activate requests. The + per-group slot count is what makes that case legible, and comparing it + against ``max_batch_size`` is what makes it actionable -- a pool group + sized at roughly one slot per declared request has no room for the + per-request slot count to vary, which windowed life cycles do as their + window slides across block boundaries. + + Read through ``_introspection`` so the numbers are the same under + either the Python or the C++ backend. + """ + try: + stats = _introspection.storage_statistics(self.impl, GPU_LEVEL) + except (AttributeError, RuntimeError) as exc: + # Introspection is a debugging aid, never a correctness + # dependency: a manager that cannot report its own layout must + # still start. + logger.debug(f"KV cache pool group sizes unavailable: {exc}") + return [] + + def flat_sizes(value) -> list[int]: + """Per-pool byte sizes, whatever shape the backend reported them in. + + `slot_size` is a scalar under one backend and a list under another, + and wrapping a list in another list made `sum()` raise + ``TypeError: unsupported operand type(s) for +: 'int' and 'list'`` + -- which is how a debugging aid took the startup log line with it. + """ + if isinstance(value, int): + return [value] + sizes: list[int] = [] + for item in value: + sizes.extend(flat_sizes(item)) + return sizes + + entries = [] + for pool_group_id, stat in enumerate(stats): + reported = getattr(stat, "slot_sizes", None) + if reported is None: + reported = stat.slot_size + slot_sizes = flat_sizes(reported) + entries.append( + f"pool_group_id={pool_group_id}, " + f"num_slots={stat.total}, " + f"slot_size={slot_sizes}, " + f"bytes={stat.total * sum(slot_sizes)}" + ) + return entries + + def _log_kv_cache_pool_lifecycle_mapping(self) -> None: + entries = OrderedDict() + for layer in self.kv_cache_manager_py_config.layers: + for buffer in layer.buffers: + entries.setdefault( + self._format_kv_cache_pool_lifecycle_entry(layer.layer_id, buffer.role), None + ) + + if not entries: + return + + logger.info(f"{type(self).__name__} role-to-pool/lifecycle mapping:") + for entry in entries: + logger.info(entry) + + for entry in self._format_kv_cache_pool_group_sizes(): + logger.info(entry) + + def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: + pool_ids = torch.empty( + self.num_attention_op_pools, + 2, + dtype=torch.long, + device="cpu", + ) + scales = torch.empty( + self.num_attention_op_pools, + 2, + dtype=torch.int32, + device="cpu", + ) + layer_offsets = torch.empty_like(scales) + scratch_pages = torch.empty_like(scales) + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + pool_id = self.layer_to_pool_mapping_dict[layer_id] + role_a, role_b = self._get_pool_roles(pool_id) + roles = [ + role_a, + role_a if role_b is None else role_b, + ] + for role_idx, role in enumerate(roles): + converter = self.impl.get_page_index_converter(layer_id, role) + if converter.expansion != 1 and self._is_standalone_draft_layer(layer_id): + pool_ids[local_layer_idx, role_idx] = pool_id + scales[local_layer_idx, role_idx] = 1 + layer_offsets[local_layer_idx, role_idx] = 0 + scratch_pages[local_layer_idx, role_idx] = 0 + continue + if converter.expansion != 1: + raise NotImplementedError( + "SWA scratch block-table conversion does not support " + f"expanded page indices yet: layer={layer_id}, role={role}, " + f"expansion={converter.expansion}" + ) + pool_ids[local_layer_idx, role_idx] = pool_id + scales[local_layer_idx, role_idx] = int(converter.scale) + layer_offsets[local_layer_idx, role_idx] = int(converter.layer_offset) + scratch_pages[local_layer_idx, role_idx] = int(converter.scratch_pages_per_block) + + staging_capacity = index_mapper_capacity * self.max_beam_width + device = torch.device("cuda", torch.cuda.current_device()) + self._device_kv_cache_block_offsets_input = torch.empty_like( + self.host_kv_cache_block_offsets, + device=device, + ) + self._device_attention_op_block_offsets_staging = torch.empty( + self.num_attention_op_pools, + staging_capacity, + 2, + self.max_blocks_per_seq, + dtype=torch.int32, + device=device, + ) + self._device_copy_idx_staging = torch.zeros( + staging_capacity, + dtype=torch.long, + device=device, + ) + self._device_num_contexts = torch.empty((), dtype=torch.int32, device=device) + self._device_attention_op_pool_ids = pool_ids.to(device=device) + self._device_attention_op_scales = scales.to(device=device) + self._device_attention_op_layer_offsets = layer_offsets.to(device=device) + self._device_attention_op_scratch_pages = scratch_pages.to(device=device) + self._device_block_positions = torch.arange( + self.max_blocks_per_seq, + dtype=torch.int32, + device=device, + ) + + min_scale = int(scales.min().item()) if scales.numel() > 0 else 1 + max_scratch_pages = int(scratch_pages.max().item()) if scratch_pages.numel() > 0 else 1 + self._max_scratch_slots = max( + 1, + (self.max_blocks_per_seq * max_scratch_pages + min_scale - 1) // min_scale, + ) + scratch_slots_shape = ( + self.num_pools, + staging_capacity, + self._max_scratch_slots, + ) + self._host_scratch_begs_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + self._host_scratch_ends_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + self._host_scratch_slots_staging = torch.zeros( + scratch_slots_shape, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + self._device_scratch_begs_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + device=device, + ) + self._device_scratch_ends_staging = torch.zeros( + self.num_pools, + staging_capacity, + dtype=torch.int32, + device=device, + ) + self._device_scratch_slots_staging = torch.zeros( + scratch_slots_shape, + dtype=torch.int32, + device=device, + ) + + def _copy_idx_to_device(self, copy_idx: torch.Tensor) -> torch.Tensor: + self._device_copy_idx_staging[: copy_idx.size(0)].copy_(copy_idx, non_blocking=True) + return self._device_copy_idx_staging + + def _copy_scratch_metadata_to_device( + self, + request_ids: List[int], + num_contexts: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + host_begs = self._host_scratch_begs_staging[:, :num_contexts] + host_ends = self._host_scratch_ends_staging[:, :num_contexts] + host_slots = self._host_scratch_slots_staging[:, :num_contexts, :] + host_begs.zero_() + host_ends.zero_() + host_slots.zero_() + + for pool_id in range(self.num_pools): + for context_idx, request_id in enumerate(request_ids[:num_contexts]): + desc = self.kv_cache_map[request_id].get_scratch_desc(pool_id) + if desc is None: + continue + slot_ids = desc.slot_ids + if len(slot_ids) > self._max_scratch_slots: + raise RuntimeError( + f"Scratch slot count {len(slot_ids)} exceeds staging capacity " + f"{self._max_scratch_slots}" + ) + host_begs[pool_id, context_idx] = int(desc.range.beg) + host_ends[pool_id, context_idx] = int(desc.range.end) + for slot_idx, slot_id in enumerate(slot_ids): + host_slots[pool_id, context_idx, slot_idx] = int(slot_id) + + self._device_scratch_begs_staging.copy_( + self._host_scratch_begs_staging, + non_blocking=True, + ) + self._device_scratch_ends_staging.copy_( + self._host_scratch_ends_staging, + non_blocking=True, + ) + self._device_scratch_slots_staging.copy_( + self._host_scratch_slots_staging, + non_blocking=True, + ) + return ( + self._device_scratch_begs_staging, + self._device_scratch_ends_staging, + self._device_scratch_slots_staging, + ) + + def _copy_batch_block_offsets_per_layer( + self, + dst_tensor: torch.Tensor, + request_ids: List[int], + copy_idx: torch.Tensor, + num_contexts: int, + num_seqs: int, + ) -> None: + device_copy_idx = self._copy_idx_to_device(copy_idx) + self._device_kv_cache_block_offsets_input.copy_( + self.host_kv_cache_block_offsets, + non_blocking=True, + ) + scratch_begs, scratch_ends, scratch_slots = self._copy_scratch_metadata_to_device( + request_ids, + num_contexts, + ) + self._device_num_contexts.fill_(num_contexts) + _copy_swa_block_offsets_with_scratch_compiled( + self._device_kv_cache_block_offsets_input, + device_copy_idx, + self._device_attention_op_pool_ids, + self._device_attention_op_scales, + self._device_attention_op_layer_offsets, + self._device_attention_op_scratch_pages, + self._device_block_positions, + scratch_begs, + scratch_ends, + scratch_slots, + self._device_num_contexts, + self._device_attention_op_block_offsets_staging, + ) + dst_tensor[: self.num_attention_op_pools, :num_seqs].copy_( + self._device_attention_op_block_offsets_staging[ + : self.num_attention_op_pools, + :num_seqs, + ], + non_blocking=True, + ) + + def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: + """Return the primary cache roles allocated for one local layer.""" + roles = [Role.KEY, Role.VALUE][: self.kv_factor_per_layer[local_layer_idx]] + if self.dtype_per_layer[local_layer_idx] == DataType.NVFP4: + assert self.head_dim_per_layer[local_layer_idx] % 2 == 0, ( + "NVFP4 cache head_dim must be divisible by two" + ) + roles += [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE][ + : self.kv_factor_per_layer[local_layer_idx] + ] + return roles + + def _build_base_config( + self, + kv_cache_config: KvCacheConfig, + *, + tokens_per_block: int, + cache_tiers: List[CacheTierConfig], + ) -> KVCacheManagerConfigPy: + """Build the general cache configuration used by most models. + + Models that need a custom configuration should subclass + :class:`KVCacheManagerV2` and override :meth:`_build_cache_config` to + update the necessary fields. + """ + scratch_reuse_config = None + if self.enable_swa_scratch_reuse: + # Unwritten speculative and draft reserves must stay outside the + # scratch range when a context chunk advances its history. + scratch_reuse_config = SwaScratchReuseConfig( + max_rewind_len=self.num_extra_kv_tokens + self._standalone_draft_reserve + ) + + typical_step = None + constraints = [] + if kv_cache_config.pool_ratio is None: + typical_seq_len = self._get_typical_seq_len(kv_cache_config) + if typical_seq_len is not None and typical_seq_len > self.max_seq_len: + raise ValueError( + f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or " + f"equal to max_seq_len ({self.max_seq_len})" + ) + + if typical_seq_len is not None: + # Model one context request and enough generation requests to fill + # the resident request capacity without over-provisioning windowed + # cache pools. + generation_request_capacity = self._get_generation_request_capacity() + context_capacity = ( + self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len + ) + self.num_extra_kv_tokens + generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) + typical_step = BatchDesc( + [KVCacheDesc(capacity=context_capacity, history_length=0)] + + [ + KVCacheDesc( + capacity=typical_seq_len, + history_length=generation_history_length, + ) + ] + * (generation_request_capacity - 1) + ) + + # CUDA graph generation warmup uses one request at max_seq_len and + # enough minimal decode requests to fill the resident capacity. + min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens + constraints.append( + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_seq_len, + history_length=self.max_seq_len - 1, + ) + ] + + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)] + * (generation_request_capacity - 1) + ) + ) + # General and chunked-prefill warmup uses one fresh context request + # at the per-iteration token budget. + if self.max_num_tokens is not None: + constraints.append( + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_num_tokens + self.num_extra_kv_tokens, + history_length=0, + ) + ] + ) + ) + + if kv_cache_config.dtype == "nvfp4": + for layer_idx, hd in enumerate(self.head_dim_per_layer): + buffer_roles = self._get_buffer_roles_for_layer(layer_idx) + if Role.KEY_BLOCK_SCALE not in buffer_roles: + continue + assert hd % 2 == 0, ( + f"head_dim must be divisible by 2 for nvfp4 kv cache, but layer {layer_idx} has head_dim={hd}" + ) + + # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional + # per-layer BufferConfig entries — for example a sparse index-K + # buffer — without overriding the K/V/NVFP4 scale wiring above. + # The dict maps local layer id -> list of extra BufferConfig. Each + # extra buffer's role must be unique within the layer (asserted by + # AttentionLayerConfig.__post_init__) and its size must be in bytes + # per block (= bytes_per_token * tokens_per_block). + extra_buffers_per_layer = ( + self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} + ) + + layer_configs: List[AttentionLayerConfig] = [] + for layer_id in typed_range(LayerId(self.num_local_layers)): + buffer_roles = self._get_buffer_roles_for_layer(layer_id) + buffers = [ + BufferConfig( + role=role, + size=self.get_layer_bytes_per_token(local_layer_idx=layer_id, data_role=role) + * tokens_per_block, + ) + for role in buffer_roles + ] + for extra in extra_buffers_per_layer.get(int(layer_id), ()): + assert extra.role not in buffer_roles, ( + f"extra buffer role {extra.role!r} for layer " + f"{int(layer_id)} duplicates a standard K/V/scale role" + ) + buffers.append(extra) + layer_configs.append( + AttentionLayerConfig( + layer_id=layer_id, + buffers=buffers, + sliding_window_size=self.max_attention_window_vec[ + self.pp_layers[layer_id] % len(self.max_attention_window_vec) + ], + num_sink_tokens=None, + ) + ) + + return KVCacheManagerConfigPy( + # Used by the backend only for token<->block arithmetic and + # radix hashing; BufferConfig.size above stays the physical page. + tokens_per_block=self._ledger_tokens_per_block, + cache_tiers=cache_tiers, + layers=layer_configs, + typical_step=typical_step, + constraints=constraints, + max_util_for_resume=kv_cache_config.max_util_for_resume, + enable_partial_reuse=kv_cache_config.enable_partial_reuse, + # Keep the lookahead evidence and its backoff in the same tree match. + # A paired scheduler caps the claim by retaining D evidence past the + # common usable depth, so this still trims exactly once. + reuse_match_backoff=self.reuse_match_backoff, + enable_stats=self.enable_stats, + swa_scratch_reuse=scratch_reuse_config, + commit_min_snapshot=( + kv_cache_config.enable_block_reuse + and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE + ), + initial_pool_ratio=kv_cache_config.pool_ratio, + ) + + def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + """Customize the general cache config for a specialized cache manager.""" + return config + + def _append_standalone_draft_layers( + self, + config: KVCacheManagerConfigPy, + *, + num_kv_heads: int, + head_dim: int, + dtype: DataType, + kv_factor: int, + register_model_layers: bool = True, + ) -> KVCacheManagerConfigPy: + if not self._draft_layer_sizes: + return config + first_global_id = max(self.pp_layers, default=-1) + 1 + first_local_id = len(config.layers) + # FP4 MLA adds native HP-tail layers outside the model-layer arrays. + register_model_layers = register_model_layers and first_local_id == len(self.pp_layers) + self._draft_registered_as_model_layers = register_model_layers + # Specialized managers can replace model layers with multiple native layers. + padding = first_local_id - len(self.dtype_per_layer) + assert padding >= 0 + num_draft_layers = len(self._draft_layer_sizes) + self.dtype_per_layer.extend([self.dtype] * padding + [dtype] * num_draft_layers) + self.kv_factor_per_layer.extend([self.kv_factor] * padding + [kv_factor] * num_draft_layers) + self.num_kv_heads_per_layer.extend([0] * padding + [num_kv_heads] * num_draft_layers) + self.head_dim_per_layer.extend([0] * padding + [head_dim] * num_draft_layers) + self.cp_as_tp_layer_ids = frozenset( + range(first_local_id, first_local_id + num_draft_layers) + ) + self.draft_layer_ids = tuple( + range(first_global_id, first_global_id + len(self._draft_layer_sizes)) + ) + layers = list(config.layers) + for offset, global_id in enumerate(self.draft_layer_ids): + local_id = first_local_id + offset + layers.append( + AttentionLayerConfig( + layer_id=LayerId(local_id), + buffers=[ + BufferConfig( + role=role, + size=KVCacheManagerV2.get_layer_bytes_per_token(self, local_id, role) + * self.tokens_per_block, + tokens_per_block_override=( + self.tokens_per_block if self._helix_cp_size > 1 else None + ), + ) + for role in self._get_buffer_roles_for_layer(local_id) + ], + sliding_window_size=self.draft_retention_window_size, + cache_domain=CacheDomain.STANDALONE_DRAFT, + ) + ) + self.layer_offsets[global_id] = local_id + if register_model_layers: + self.pp_layers.append(global_id) + self.total_num_kv_heads_per_layer.append( + self._draft_total_num_kv_heads or num_kv_heads + ) + if register_model_layers: + self.num_local_layers = len(self.pp_layers) + self.num_layers += len(self._draft_layer_sizes) + + def reserve_scratch(batch: BatchDesc) -> BatchDesc: + return BatchDesc( + [ + KVCacheDesc( + capacity=desc.capacity + + (self._standalone_draft_reserve if desc.capacity else 0), + history_length=desc.history_length, + ) + for desc in batch.kv_caches + ], + system_prompt_length=batch.system_prompt_length, + ) + + return replace( + config, + layers=layers, + constraints=[reserve_scratch(batch) for batch in config.constraints], + typical_step=reserve_scratch(config.typical_step) + if config.typical_step is not None + else None, + ) + + def _get_typical_seq_len(self, kv_cache_config: KvCacheConfig) -> int | None: + """Return the configured typical sequence length, if any.""" + return kv_cache_config.avg_seq_len + + def _configure_unified_draft_pool_ratio( + self, config: KVCacheManagerConfigPy + ) -> KVCacheManagerConfigPy: + """Preserve target ratios within their share of the unified byte quota.""" + if not self.draft_layer_ids: + return config + target_ratios = config.initial_pool_ratio + if target_ratios is None: + return config + + # Native lifecycle equality includes window, rounded sink blocks, sparsity, + # and cache domain. All SSM layers share one lifecycle, represented by None. + target_life_cycles: list[AttnLifeCycle | None] = [] + for layer in config.layers: + if isinstance(layer, AttentionLayerConfig): + if layer.cache_domain == CacheDomain.STANDALONE_DRAFT: + continue + life_cycle = AttnLifeCycle.make( + layer.sliding_window_size, + layer.num_sink_tokens, + config.tokens_per_block, + cache_domain=layer.cache_domain, + ) + else: + life_cycle = None + if not any( + type(existing) is type(life_cycle) and existing == life_cycle + for existing in target_life_cycles + ): + target_life_cycles.append(life_cycle) + + if len(target_ratios) != len(target_life_cycles): + raise ValueError( + "kv_cache_config.pool_ratio must have one entry per target layer group " + f"({len(target_life_cycles)}), got {len(target_ratios)}" + ) + # Use the same model-specific capacity accounting as quota selection; + # specialized targets can include compressed KV or recurrent snapshots. + usable_quota = int(config.cache_tiers[GPU_LEVEL].quota * config.max_util_for_resume) + max_tokens = self._get_max_tokens_from_quota(usable_quota) + if math.isinf(max_tokens): + max_tokens = self.max_num_tokens or self.max_seq_len + if usable_quota <= 0 or self._get_quota_from_max_tokens(max_tokens) > usable_quota: + raise ValueError("Unified KV cache quota cannot hold the target and draft reserves") + draft_ratio = self._get_draft_quota_from_max_tokens(max_tokens) / usable_quota + if not 0 < draft_ratio < 1: + raise ValueError("Unified KV cache quota must leave room for target and draft pools") + + # Appended draft layers share one lifecycle, after all target groups. + return replace( + config, + initial_pool_ratio=[ratio * (1 - draft_ratio) for ratio in target_ratios] + + [draft_ratio], + ) + + def _extra_buffers_per_layer( + self, *, tokens_per_block: int + ) -> Optional[dict[int, List[BufferConfig]]]: + """Return per-local-layer extra BufferConfig entries to register + alongside the standard K/V/NVFP4 scale buffers. + + Default implementation returns ``None``. Subclasses override this + to register additional buffers — for example, MiniMax-M3 registers + a sparse index-K buffer for each sparse local layer. Each + ``BufferConfig.size`` is interpreted as bytes per block (i.e., + ``bytes_per_token * tokens_per_block``), matching the standard + buffers built in :meth:`_build_base_config`. The block storage + groups buffers by lifecycle and size with an opaque role key, so + new roles do not require C++ changes. + """ + return None + + def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: + """Map native cache roles to disaggregation mapper kinds. + + ``Role.ALL`` is the required fallback for roles without an explicit + entry. The default is the head-major (HND) ``INDEXED`` layout written + by the TRTLLM attention kernels — correct for V1 and standard V2 + managers. ``Role.INDEX_KEY`` defaults to ``REPLICATED``: every + index-key side cache shipped so far (DSA indexer-K on V1, MiniMax M3 + on V2) computes its projection replicated across TP ranks, so the + cache bytes are identical per rank; the entry is inert unless a + subclass actually registers ``INDEX_KEY`` buffers via + ``_extra_buffers_per_layer``. Model-specific managers may declare + logical layouts without requiring the shared extractor to inspect + private attributes or role names. MiniMax M3, for example, maps + ordinary K/V to ``NHD`` and keeps index-key ``REPLICATED``. + + This declaration does not influence storage pooling: V2 storage + coalesces buffers purely by ``(life_cycle, buffer size)``, so roles + with different transfer semantics may share a pool slot when their + per-block sizes coincide (e.g. MiniMax M3 at TP degrees where + K == V == INDEX_KEY bytes per block). The disagg page-table builder + splits each physical pool into one logical view per mapper kind, so + transfer correctness never depends on the coalescing outcome. + + Pool memory is layout-agnostic; this declaration describes what the + manager's paired attention backend actually writes. A static mapping + is valid only for a fixed manager/backend pair. A manager whose backend + selects the layout at runtime must derive the mapping from that + backend's configuration. + """ + return {Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED} + + def get_disagg_transfer_roles(self) -> Optional[frozenset[DataRole]]: + """Return the native roles included in disaggregated transfer. + + ``None`` preserves the default policy and transfers every registered + role. Specialized managers may omit deterministic derived roles that + are rebuilt on the receiver after canonical cache data arrives. + """ + return None + + @property + def blocks_in_primary_pool(self) -> int: + """ + Get the number of blocks in the primary pool. + """ + return self.impl.get_page_index_upper_bound(0, Role.KEY) + + def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: + return self._get_layer_buffer_view(self.layer_offsets[layer_idx], kv_layout) + + def _get_layer_buffer_view(self, local_layer_idx: int, kv_layout: str) -> torch.Tensor: + return self._get_kv_buffer_view( + local_layer_idx, + dtype=self.dtype_per_layer[local_layer_idx], + kv_factor=self.kv_factor_per_layer[local_layer_idx], + num_kv_heads=self.num_kv_heads_per_layer[local_layer_idx], + head_dim=self.head_dim_per_layer[local_layer_idx], + kv_layout=kv_layout, + page_expansion=self._helix_cp_size if local_layer_idx in self.cp_as_tp_layer_ids else 1, + ) + + def _get_kv_buffer_view( + self, + layer_id: int, + *, + dtype: DataType | torch.dtype, + kv_factor: int, + num_kv_heads: int, + head_dim: int, + kv_layout: str, + page_expansion: int = 1, + ) -> torch.Tensor: + """View shared pool pages for a manager-local layer and its KV geometry.""" + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported KV layout: {kv_layout}") + addr_key = self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) + if kv_factor == 2: + addr_value = self.impl.get_mem_pool_base_address( + layer_id, Role.VALUE, PageIndexMode.SHARED + ) + page_size_key = self.impl.get_page_stride(layer_id, Role.KEY) + page_size_value = self.impl.get_page_stride(layer_id, Role.VALUE) + if ( + addr_key + page_size_key * page_expansion != addr_value + or page_size_key != page_size_value + ): + raise ValueError("K/V buffers must have adjacent equal-sized pages") + + element_per_container = 1 + if dtype == DataType.NVFP4: + element_per_container = 2 + dtype = torch.int8 + + dimensions = ( + [self.tokens_per_block, num_kv_heads] + if kv_layout == "NHD" + else [num_kv_heads, self.tokens_per_block] + ) + head_dim //= element_per_container + page_upper = self.impl.get_page_index_upper_bound(layer_id, Role.KEY) + num_pages = page_upper // kv_factor + strides = None + if page_expansion > 1: + # Each ledger slot contains all of its K subpages, then all V + # subpages. Expanded KEY indices skip the V spans between slots. + num_pages = page_upper - (kv_factor - 1) * page_expansion + page_elements = self.tokens_per_block * num_kv_heads * head_dim + strides = [ + page_elements, + page_expansion * page_elements, + dimensions[1] * head_dim, + head_dim, + 1, + ] + return convert_to_torch_tensor( + TensorWrapper( + addr_key, + dtype, + [ + num_pages, + kv_factor, + *dimensions, + head_dim, + ], + strides=strides, + ) + ) + + def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: + """Return the storage dtype for a layer, or the manager default.""" + return ( + self.dtype if layer_idx is None else self.dtype_per_layer[self.layer_offsets[layer_idx]] + ) + + def _is_standalone_draft_layer(self, local_layer_idx: int) -> bool: + return local_layer_idx in self.cp_as_tp_layer_ids + + def get_layer_kv_factor(self, layer_idx: int) -> int: + return self.kv_factor_per_layer[self.layer_offsets[layer_idx]] + + def get_draft_buffers(self, local_layer_idx: int, kv_layout: str = "HND") -> torch.Tensor: + """View one draft layer using the same storage contract as every other layer.""" + if not 0 <= local_layer_idx < len(self.draft_layer_ids): + raise ValueError("No standalone draft cache layer at this index") + layer_idx = self.draft_layer_ids[local_layer_idx] + return self._get_layer_buffer_view(self.layer_offsets[layer_idx], kv_layout) + + @property + def draft_max_blocks_per_seq(self) -> int: + """Maximum physical draft pages, including every Helix ledger subpage.""" + return self.max_blocks_per_seq * self._helix_cp_size + + def get_draft_num_blocks(self, request_id: int) -> int: + """Allocated draft pages covering the request's global token positions.""" + return self.kv_cache_map[request_id].num_blocks * self._helix_cp_size + + @property + def draft_retention_window_size(self) -> int | None: + return self.draft_window_size + 1 if self.draft_window_size is not None else None + + def get_draft_transfer_identity(self) -> dict: + local_layer_idx = self.layer_offsets[self.draft_layer_ids[0]] + return { + "num_layers": len(self.draft_layer_ids), + "num_kv_heads": self._draft_total_num_kv_heads + or self.num_kv_heads_per_layer[local_layer_idx], + "head_dim": self.head_dim_per_layer[local_layer_idx], + "dtype": str(binding_to_torch_dtype(self.dtype_per_layer[local_layer_idx])), + "attention_backend": ( + "VANILLA" + if self.draft_attention_backend == "TRTLLM" + else self.draft_attention_backend + ), + "kv_factor": self.kv_factor_per_layer[local_layer_idx], + "window_size": self.draft_window_size, + } + + def _draft_history_uses_kv_history(self) -> bool: + """Embedded DSpark writes the complete live suffix of target history.""" + return getattr(self, "draft_attention_backend", None) == "DSv4" + + def get_draft_block_indices( + self, + request_ids: List[int], + histories: Optional[Sequence[StandaloneDraftHistory]] = None, + ) -> List[List[int]]: + """Validated rank-local mappings for allocated draft pages, without padding.""" + for request_id in request_ids: + cache = self.kv_cache_map.get(request_id) + if cache is None or not cache.is_active: + raise ValueError(f"Standalone draft request {request_id} has no active cache") + expanded = self._helix_cp_size > 1 + layer_idx = self.draft_layer_ids[0] + batch_indices = self.get_batch_cache_indices(request_ids, layer_idx, raw_indices=expanded) + if expanded: + converter = self.impl.get_page_index_converter(self.layer_offsets[layer_idx], Role.KEY) + batch_indices = [converter(indices, PageIndexMode.SHARED) for indices in batch_indices] + for row, indices in enumerate(batch_indices): + cache = self.kv_cache_map[request_ids[row]] + if ( + len(indices) != self.get_draft_num_blocks(request_ids[row]) + or len(indices) > self.draft_max_blocks_per_seq + ): + raise ValueError("Standalone draft cache has an incomplete or oversized page table") + first_required_block = 0 + if self.draft_window_size is not None: + history = None + if not self._draft_history_uses_kv_history(): + history = ( + histories[row] + if histories is not None + else self.get_draft_history(request_ids[row]) + ) + history_start = ( + history.position - history.valid_length + if history is not None + else max(0, cache.history_length - self.draft_window_size) + ) + first_required_block = history_start // self.tokens_per_block + if any(index == BAD_PAGE_INDEX for index in indices[first_required_block:]): + raise ValueError( + "Draft cache contains missing pages in its required history or scratch" + ) + return batch_indices + + def get_draft_block_table( + self, + request_ids: List[int], + histories: Optional[Sequence[StandaloneDraftHistory]] = None, + ) -> torch.Tensor: + """Current rank-local page mappings; unused tail entries point to page zero.""" + batch_indices = self.get_draft_block_indices(request_ids, histories) + table = torch.zeros((len(request_ids), self.draft_max_blocks_per_seq), dtype=torch.int32) + for row, indices in enumerate(batch_indices): + table[row, : len(indices)] = torch.tensor(indices, dtype=torch.int32) + return table + + def get_draft_history(self, request_id: int) -> Optional[StandaloneDraftHistory]: + return self.draft_history.get(request_id) + + def set_draft_history(self, request_id: int, valid_length: int, position: int) -> None: + cache = self.kv_cache_map.get(request_id) + if cache is None or not cache.is_active: + raise ValueError( + f"Standalone draft request {request_id} has no active cache allocation" + ) + history = StandaloneDraftHistory(valid_length, position) + if history.position > cache.capacity: + raise ValueError("Standalone draft history exceeds allocated capacity") + if self.draft_window_size is not None and history.valid_length > self.draft_window_size: + raise ValueError("Draft history exceeds its retention window") + self.draft_history[request_id] = history + + def export_draft_history(self, request_id: int) -> Optional[dict]: + if not self.draft_layer_ids: + return None + if self._draft_history_uses_kv_history(): + cache = self.kv_cache_map[request_id] + if not cache.is_active: + raise ValueError("Embedded DSpark transfer requires an active cache") + return { + "valid_length": min(cache.history_length, self.draft_window_size), + "position": cache.history_length, + "layout": self.get_draft_transfer_identity(), + } + history = self.get_draft_history(request_id) + if history is None: + raise ValueError( + f"Standalone draft request {request_id} has no valid history to transfer" + ) + return { + "valid_length": history.valid_length, + "position": history.position, + "layout": self.get_draft_transfer_identity(), + } + + def restore_draft_history(self, request_id: int, metadata: dict) -> None: + if not self.draft_layer_ids or metadata.get("layout") != self.get_draft_transfer_identity(): + raise ValueError("Standalone draft transfer layout does not match the receiving worker") + valid_length = metadata.get("valid_length") + position = metadata.get("position") + if self._draft_history_uses_kv_history(): + cache = self.kv_cache_map[request_id] + # Receive admission establishes the global prompt watermark; the + # transfer completion path validates that the received KV covers it. + if ( + type(position) is not int + or type(valid_length) is not int + or position != cache.history_length + or valid_length != min(position, self.draft_window_size) + ): + raise ValueError("Embedded DSpark transfer does not match the admitted KV history") + self.get_draft_block_table([request_id]) + return + history = StandaloneDraftHistory(valid_length, position) + # Validate receiver-local allocation before publishing history. + self.get_draft_block_table([request_id], [history]) + self.set_draft_history(request_id, valid_length, position) + + def get_index_k_buffer( + self, + layer_idx: int, + *, + num_heads: int = 1, + head_dim: int, + dtype: Union[torch.dtype, "DataType", str] = torch.bfloat16, + kv_layout: str = "NHD", + ) -> Optional[torch.Tensor]: + """Return a torch view over the V2-managed paged ``Role.INDEX_KEY`` + buffer for ``layer_idx``, or ``None`` when the layer has no + INDEX_KEY buffer registered (e.g. dense layers in a sparse model, + or non-local layers on the current PP rank). + + For ``kv_layout="NHD"``, the view has shape ``[num_pages, + tokens_per_block, num_heads, head_dim]``. For ``kv_layout="HND"``, + it has shape ``[num_pages, num_heads, tokens_per_block, head_dim]``. + Sparse modeling code decomposes the per-token slot id used by the + main paged K/V cache into ``(page, within_page)`` and indexes the + token axis selected by ``kv_layout``. + + Because :class:`BufferConfig` only carries an opaque byte ``size`` + per block, the dtype and head shape are caller-side contracts. + The caller must pass the same ``num_heads``, ``head_dim``, and + ``dtype`` it used to compute the registered + ``size = num_heads * head_dim * dtype_bytes * tokens_per_block`` + in ``_extra_buffers_per_layer``. The accessor validates + ``num_heads * head_dim * dtype_bytes * tokens_per_block == + page_stride`` so a wiring mismatch fails loudly instead of + returning a view with silently wrong stride. + + Follows the same ``TensorWrapper`` / ``convert_to_torch_tensor`` + pattern as :meth:`get_buffers`. The returned tensor is a + zero-copy view over V2-managed pool memory and stays valid for + the lifetime of the cache manager — writes through the view + propagate to the pool, and successive calls return views over + the same backing storage. + """ + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported kv_layout: {kv_layout}") + if layer_idx not in self.layer_offsets: + return None + layer_offset = self.layer_offsets[layer_idx] + try: + addr = self.impl.get_mem_pool_base_address(layer_offset, Role.INDEX_KEY) + except (KeyError, IndexError): + # INDEX_KEY not registered for this layer (default V2 manager + # registers only K/V/scale; sparse subclasses register + # INDEX_KEY only on sparse layers via + # ``_extra_buffers_per_layer``). + # + # The python backend raises ``KeyError`` from the missing dict + # lookup; the C++ backend's ``getBufferAttr`` throws + # ``std::out_of_range`` for an unknown buffer id, which nanobind + # maps to ``IndexError``. Catch both so the "role not registered + # -> None" contract holds on either backend. + return None + # INDEX_KEY is registered; the remaining lookups must succeed. Keep them + # outside the try so genuine failures surface instead of returning None. + page_stride = self.impl.get_page_stride(layer_offset, Role.INDEX_KEY) + page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.INDEX_KEY) + converter = self.impl.get_page_index_converter(layer_offset, Role.INDEX_KEY) + + if isinstance(dtype, DataType): + torch_dtype = binding_to_torch_dtype(dtype) + elif isinstance(dtype, str): + torch_dtype = str_dtype_to_torch(dtype) + else: + torch_dtype = dtype + + elem_bytes = torch.tensor([], dtype=torch_dtype).element_size() + expected_stride = num_heads * head_dim * elem_bytes * self.tokens_per_block + assert page_stride == expected_stride, ( + f"INDEX_KEY page stride mismatch for layer {layer_idx}: " + f"V2 reports page_stride={page_stride}, but the caller " + f"supplied num_heads={num_heads} * head_dim={head_dim} * " + f"elem_bytes={elem_bytes} * tokens_per_block=" + f"{self.tokens_per_block} = {expected_stride}. Re-check " + f"the BufferConfig.size used to register Role.INDEX_KEY " + f"in _extra_buffers_per_layer." + ) + + # Multi-layer coalescing: when INDEX_KEY shares the V2 storage + # pool with K/V (production MiniMax-M3 at TP=8 makes K, V, and + # INDEX_KEY all 256 bytes/token so V2 coalesces them by + # ``(life_cycle_id, single_buffer_size)``), ``page_upper`` is + # not the per-layer page count — it is the max page index from + # this buffer's base in the coalesced pool. Decompose it into + # ``num_slots = (page_upper + layer_offset_pages) // scale`` so + # the per-layer view has the right shape. ``scale`` here is the + # buffers-per-slot count; for a non-coalesced INDEX_KEY pool it + # equals 1 and ``num_slots == page_upper``, matching the legacy + # ``shape[0] == page_upper`` behavior. + scale = int(converter.scale) + layer_offset_pages = int(converter.layer_offset) + num_slots_total = page_upper + layer_offset_pages + assert num_slots_total % scale == 0, ( + f"V2 storage inconsistency for INDEX_KEY of layer " + f"{layer_idx}: page_upper + layer_offset_pages = " + f"{num_slots_total} is not divisible by scale = {scale}." + ) + num_slots = num_slots_total // scale + if kv_layout == "NHD": + page_shape = [self.tokens_per_block, num_heads, head_dim] + else: + page_shape = [num_heads, self.tokens_per_block, head_dim] + + if scale == 1: + # Non-coalesced INDEX_KEY pool: the per-buffer stride is + # the entire page, so either layout is a contiguous view. + shape = [page_upper, *page_shape] + return convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, shape)) + + # Coalesced pool: include the buffers-per-slot dimension, then + # slice ``[:, 0]`` to extract this layer's INDEX_KEY data while + # preserving dim-0 stride = ``scale * page_stride`` bytes. + full_slot_shape = [num_slots, scale, *page_shape] + full_view = convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, full_slot_shape)) + return full_view[:, 0] + + def get_num_available_tokens( + self, *, token_num_upper_bound: int, batch_size: int = 1, max_num_draft_tokens: int = 0 + ) -> int: + """Clamp ``token_num_upper_bound`` to the allocatable token capacity. + + Unit note: under helix the backend runs on the ledger + ``tokens_per_block``, so the returned capacity (like + ``token_num_upper_bound`` and ``max_seq_len``) is in GLOBAL ledger + tokens - the coordinate request lengths are expressed in. Callers + that additionally bound the result by per-forward budgets (e.g. + ``max_num_tokens``) stay consistent because a helix context forward + replicates all tokens on every rank, so both bounds constrain the + same request-length variable. + """ + extra_tokens = ( + self.num_extra_kv_tokens + max_num_draft_tokens + self._standalone_draft_reserve + ) + # Token num upper bound is the maximum number of tokens that can be allocated in the kv cache manager. + # We need to add extra tokens to the token num upper bound to account for the extra tokens. + clamped = ( + self.impl.clamp_max_seq_len_for_mem(batch_size, token_num_upper_bound + extra_tokens) + - extra_tokens + ) + # clamp_max_seq_len_for_mem considers all tiers (GPU + host). When + # max_tokens is explicitly set, cap by GPU-only capacity so callers + # (e.g. CUDA graph warmup) don't exceed the GPU pool. + if self._gpu_max_tokens is not None: + clamped = min(clamped, self._gpu_max_tokens - extra_tokens) + return clamped + + def _get_attention_pool_num_blocks(self, local_layer_idx: int) -> int: + """Count ledger blocks without counting expanded physical subpages.""" + layer_id = LayerId(local_layer_idx) + expansion = self.impl.get_page_index_converter(layer_id, Role.KEY).expansion + return self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // ( + self.get_layer_kv_factor(self.pp_layers[local_layer_idx]) * expansion + ) + + def get_num_free_blocks(self) -> int: + # NOTE This method is used to get the number of blocks in the primary pool not the FREE blocks. + # However, since we only use this function when the kv cache manager is empty, so it is safe to do so. + assert len(self.kv_cache_map) == 0, ( + "get_num_free_blocks is only used when the kv cache manager is empty" + ) + max_num_blocks = max( + [ + self._get_attention_pool_num_blocks(local_layer_idx) + for local_layer_idx in range(self.num_local_layers) + ] + ) + return max_num_blocks + + def _get_buffer_views_for_invalid_value_check( + self, layer_idx: int, kv_layout: str = "NHD" + ) -> list[torch.Tensor]: + """Return contiguous tensor views covering the mapped KV pages to scan.""" + if self.num_ugpus <= 1: + return [self.get_buffers(layer_idx, kv_layout=kv_layout).flatten(0, 1)] + + layer_offset = self.layer_offsets[layer_idx] + storage = self.impl._storage + lc_id = storage._layer_to_life_cycle_ids[LayerId(layer_offset)] + attr = storage.get_buffer_attr(LayerId(layer_offset), Role.KEY) + pg_idx = storage.get_pool_group_index(lc_id) + gpu_storage = storage._levels[GPU_LEVEL].storage + pool_group = gpu_storage._pool_groups[pg_idx] + slot_size = pool_group.slot_size[attr.pool_index] + element_per_container = 1 + dtype = self.dtype + if dtype == DataType.NVFP4: + element_per_container = 2 + dtype = torch.int8 + + views: list[torch.Tensor] = [] + for ugpu_id in range(self.num_ugpus): + local_page_upper_bound = ( + exact_div(slot_size, attr.size) * pool_group.num_slots(ugpu_id) + - exact_div(attr.offset, attr.size) + ) * attr.expansion + if kv_layout == "NHD": + shape = [ + local_page_upper_bound // self.kv_factor, + self.kv_factor, + self.tokens_per_block, + self.num_kv_heads_per_layer[layer_offset], + self.head_dim // element_per_container, + ] + else: + shape = [ + local_page_upper_bound // self.kv_factor, + self.kv_factor, + self.num_kv_heads_per_layer[layer_offset], + self.tokens_per_block, + self.head_dim // element_per_container, + ] + addr_key = ( + int(gpu_storage.slot_address(pg_idx, attr.pool_index, 0, ugpu_id)) + attr.offset + ) + views.append( + convert_to_torch_tensor(TensorWrapper(addr_key, dtype, shape)).flatten(0, 1) + ) + return views + + @property + def num_ugpus(self) -> int: + """Number of uGPUs. Returns 1 for non-localized configurations.""" + return _get_num_ugpus(self.impl) + + def get_per_ugpu_free_slots(self) -> list[int]: + """Free slot count per uGPU, summed across all pool groups.""" + return self.impl.get_per_ugpu_free_slots() + + def pick_ugpu(self, request_id: int) -> int | None: + """Default uGPU placement for newly-created requests.""" + if self.num_ugpus <= 1: + return None + return request_id % self.num_ugpus + + def _request_ugpu_id(self, req: LlmRequest) -> int | None: + """Return the request-owned uGPU id after validating the invariant.""" + ugpu_id = getattr(req, "py_ugpu_id", None) + if self.num_ugpus <= 1: + assert ugpu_id is None or ugpu_id == 0, ( + "Non-localized requests must not use a nonzero ugpu_id" + ) + return None + + assert ugpu_id is not None, ( + f"Request {req.py_request_id} must carry py_ugpu_id before localized KV-cache placement" + ) + assert 0 <= ugpu_id < self.num_ugpus, ( + f"ugpu_id must be in [0, {self.num_ugpus}), got {ugpu_id}" + ) + return ugpu_id + + def get_ugpu(self, request_id: int) -> int: + """Return the uGPU index assigned to an allocated request.""" + if self.num_ugpus <= 1: + return 0 + kv_cache = self.kv_cache_map.get(request_id) + assert kv_cache is not None, ( + f"get_ugpu: request {request_id} has no allocated KV cache " + f"(num_ugpus={self.num_ugpus}, " + f"kv_cache_map size={len(self.kv_cache_map)})" + ) + assert _get_kv_cache_ugpu_id(kv_cache) is not None, ( + f"get_ugpu: request {request_id} has no concrete ugpu_id " + f"(num_ugpus={self.num_ugpus}, " + f"kv_cache_map size={len(self.kv_cache_map)})" + ) + return _get_kv_cache_ugpu_id(kv_cache) + + @property + def fork_join_attn(self) -> bool: + """Whether to drive single-forward fork-join attention.""" + return self._fork_join_attn + + def commit_scheduled_kv_cache_stats(self, scheduled_batch: ScheduledRequests) -> None: + if self.is_draft or not self.enable_stats: + return + dirty_req_ids = self.impl.get_dirty_stats_kv_cache_ids() + for req in scheduled_batch.all_requests(): + if req.py_request_id in dirty_req_ids: + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + continue + request_stats = kv_cache.commit_pending_stats() + if not req.is_dummy and not request_stats.empty: + req.update_kv_cache_perf_metrics( + request_stats.alloc_total_blocks, + request_stats.alloc_new_blocks, + request_stats.reused_blocks, + request_stats.missed_blocks, + ) + + # ---- Scheduling API (called by KVCacheV2Scheduler) ---- + + def is_request_active(self, request_id: int) -> bool: + """Return True if *request_id* has a live, non-suspended KV cache.""" + kv_cache = self.kv_cache_map.get(request_id) + return kv_cache is not None and kv_cache.is_active + + def _effective_draft_len(self, req: LlmRequest) -> int: + """Draft token length to use for next-step KV capacity calculation. + + During the context-to-generation transition, ``py_draft_tokens`` is + still empty when the scheduler asks for capacity. Prefer transferred + context draft tokens; if there are none, reserve the generation + worker's configured draft capacity before its first forward pass. + """ + draft_len = get_draft_token_length(req) + if ( + draft_len == 0 + and req.is_disagg_generation_transmission_complete + and req.context_phase_params is not None + ): + ctx_draft_tokens = req.context_phase_params.draft_tokens + if ctx_draft_tokens is not None: + draft_len = len(ctx_draft_tokens) + if draft_len == 0 and not self.is_draft and not req.py_disable_speculative_decoding: + draft_len = self.max_total_draft_tokens + return draft_len + + def _helix_local_len(self, global_len: int) -> int: + """Tokens of the first ``global_len`` owned by this CP rank + (continuation round-robin: page b lives on rank b %% cp).""" + phys = self.tokens_per_block + full, rem = divmod(global_len, self._ledger_tokens_per_block) + return full * phys + min(max(rem - self._helix_cp_rank * phys, 0), phys) + + def _set_helix_rank_fields(self, req: LlmRequest) -> None: + """Derive the per-rank helix fields from the global position. + + The decode-step index is manager-owned (committed in + ``try_allocate_generation`` on successful resize) rather than derived + from ``py_decoding_iter``: the sampler advances that counter after + scheduling under the overlap loop, so a schedule-time read is one + step behind and would repeat the first decode position, overwriting + the first generated token's KV. Multi-token verify groups advance + py_helix_decode_group_index per committed group, so this formula + stays exact under DSpark speculation. + """ + step = req.py_helix_decode_group_index + 1 + pos = req.total_input_len_cp + step - 1 + owner = (pos // self.tokens_per_block) % self._helix_cp_size + active = owner == self._helix_cp_rank + req.py_helix_is_inactive_rank = not active + # Convention shared with model_engine: the active rank's seqlen + # includes the in-flight token (past_seen = seqlen - 1 there). + req.seqlen_this_rank_cp = self._helix_local_len(pos) + (1 if active else 0) + + def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int: + """Compute generation KV cache capacity for a request. + + Grows *current_capacity* by 1 + draft tokens, plus another draft_len + of slack for one-model speculative decoding: under the overlap + scheduler the device-side KV position runs one verification round + ahead of host bookkeeping (up to draft_len accepted-but-uncommitted + tokens), and the MTP draft layers then append draft KV beyond that + position. Without the slack, the draft KV append can address a block + ordinal the host has not allocated yet whenever the overrun crosses a + tokens_per_block boundary — an illegal memory access in the MLA rope + generation kernel (observed with DSV4 DEP+MTP3 at 128k, where the + first window-slide boundary crossing after prefill faults). + + The slack must stay a *constant* offset over the request's lifetime: + py_rewind_len only rewinds the rejected part of the recurring + 1 + draft_len growth, so every call site records the extra draft_len + in _pending_overlap_slack and update_resources trims it back. + """ + return current_capacity + 1 + 2 * self._generation_draft_slots(req) + + def _generation_draft_slots(self, req: LlmRequest) -> int: + """Physical draft width reserved for one iteration. Dynamic-tree draft pools + reserve wider than target pools; each reclaims its own in ``update_resources``. + """ + slots = self._effective_draft_len(req) + if self.is_draft: + slots = max(slots, self._kv_reserve_draft_tokens) + return slots + + def try_allocate_generation(self, req: LlmRequest) -> bool: + """Try to allocate one additional KV cache slot for a generation request. + + Resumes from suspended state if needed, then resizes capacity by 1 (+ + draft tokens). Returns True on success, False if allocation failed. + """ + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return False + + if not kv_cache.is_active: + if not kv_cache.resume(self._stream.cuda_stream): + return False + self._restore_page_index_bufs(req.py_request_id, kv_cache) + + request_id = req.py_request_id + draft_slots = self._generation_draft_slots(req) + self._allocated_draft_lens.pop(request_id, None) + is_helix_req = self._has_cp_helix and not req.is_dummy_request + if is_helix_req: + self._set_helix_rank_fields(req) + if not kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)): + return False + if is_helix_req: + # Commit only on success so a same-pass retry recomputes the + # same step instead of skipping one. + req.py_helix_decode_group_index += 1 + self._allocated_draft_lens[request_id] = draft_slots + if draft_slots > 0: + self._pending_overlap_slack[request_id] = ( + self._pending_overlap_slack.get(request_id, 0) + draft_slots + ) + return True + + def revert_allocate_generation(self, req: LlmRequest) -> None: + """Undo the capacity growth from try_allocate_generation. + + When attention DP causes can_queue=False after scheduling, the + forward pass is skipped but the scheduler already grew each + generation request's KV cache capacity by 1 (+draft tokens). + This method shrinks capacity back to undo that spurious growth + so it does not accumulate across iterations and overflow the + host page-index buffer. + + Keyed on ``_allocated_draft_lens``: no entry means no growth to give back. + Recording the width keeps the paired target/draft reverts in step. + """ + request_id = req.py_request_id + draft_slots = self._allocated_draft_lens.pop(request_id, None) + if draft_slots is None: + return + kv_cache = self.kv_cache_map.get(request_id) + if kv_cache is None or not kv_cache.is_active: + return + if self._has_cp_helix and not req.is_dummy_request and req.py_helix_decode_group_index > 0: + # The forward pass for this step is skipped; give the step back. + req.py_helix_decode_group_index -= 1 + # Mirror the 1 + 2 * draft_slots growth in _required_gen_capacity. + reverted_cap = kv_cache.capacity - 1 - 2 * draft_slots + if reverted_cap < 0: + return + # The reverted growth included draft_slots of overlap slack; deduct it + # so update_resources does not trim slack that no longer exists. + if draft_slots > 0: + remaining = self._pending_overlap_slack.get(request_id, 0) - draft_slots + if remaining > 0: + self._pending_overlap_slack[request_id] = remaining + else: + self._pending_overlap_slack.pop(request_id, None) + if not kv_cache.resize(reverted_cap): + raise RuntimeError( + f"Failed to revert KV cache capacity for request " + f"{req.py_request_id} from {kv_cache.capacity} to " + f"{reverted_cap}" + ) + + def revert_allocate_context(self, req: LlmRequest) -> bool: + """Undo this iteration's context resize. False means the cache was dropped, + not shrunk (history outran pre-resize capacity); the caller drops any draft pool. + """ + pre_cap = getattr(req, "py_ctx_pre_resize_cap", None) + if pre_cap is None: + return True + req.py_ctx_pre_resize_cap = None + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None or not kv_cache.is_active: + return True + if pre_cap >= kv_cache.capacity: + return True + if kv_cache.history_length > pre_cap: + self.free_resources(req) + # The setter asserts prepopulated < promptLen; a Helix CP rank that + # holds none of the prompt has promptLen 0 and nothing to reset. + if req.prompt_len > 0 and req.prepopulated_prompt_len > 0: + req.set_prepopulated_prompt_len(0, self.tokens_per_block) + # setPrepopulatedPromptLen only moves forward, so rewind the same + # fields LlmRequest::pause() resets for a recompute pause. + req.context_current_position = 0 + req.context_chunk_size = req.prompt_len + req.estimated_reusable_tokens = 0 + return False + history_length = min(kv_cache.history_length, pre_cap) + if not kv_cache.resize(pre_cap, history_length): + raise RuntimeError( + f"Failed to revert KV cache capacity for context " + f"request {req.py_request_id} from {kv_cache.capacity} " + f"to {pre_cap}" + ) + if pre_cap > 0: + kv_cache.suspend() + return True + + def _restore_page_index_bufs(self, request_id: int, kv_cache) -> None: + """Re-connect host page-index buffers after resume(). + + suspend() clears the base_page_index_buf pointers (sets them to + None) so the KV cache stops writing page indices to the host + buffer. After resume(), the KV cache has re-locked pages but + copy_batch_block_offsets still reads from the host buffer, so we + must re-connect the buffers to avoid stale/zero page indices that + would cause illegal memory accesses during the forward pass. + """ + index = self.index_mapper.get_index(request_id) + for i in range(self.max_beam_width): + for pool_idx in range(self.num_pools): + buffer: torch.Tensor = self.host_kv_cache_block_offsets[ + pool_idx, index * self.max_beam_width + i, 0 + ] + kv_cache.set_base_page_index_buf(i, pool_idx, memoryview(buffer.numpy())) + + def _resume_and_restore(self, req_id: int, kv_cache) -> bool: + """Resume a suspended KV cache and restore its page index buffers. + + Returns True if the cache is (or becomes) active, False on failure. + """ + if kv_cache.is_active: + return True + if not kv_cache.resume(self._stream.cuda_stream): + return False + self._restore_page_index_bufs(req_id, kv_cache) + return True + + def _draft_context_exceeds_limit(self, req: LlmRequest) -> bool: + """Whether draft prefill skips this prompt, preventing unified block reuse.""" + if ( + not self.draft_layer_ids + or self.draft_window_size is not None + or self.draft_max_position_embeddings is None + ): + return False + prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len + return prompt_len > min(self.max_seq_len, self.draft_max_position_embeddings) + + def _context_reuse_tokens( + self, req: LlmRequest, reuse_limit: int | None = None + ) -> Sequence[TokenIdExt]: + """Radix-tree evidence for one context prefix. Raw prompt, so block hashes + stay comparable with KV cache events, the cache-aware router and disagg peers. + """ + if self._draft_context_exceeds_limit(req): + return [] + all_tokens = self._reuse_token_source(req) + # The target must recompute the last prompt token, but a one-model draft + # also reads D prompt tokens ahead. Include those tokens as match evidence; + # the core drops D before claiming in both reuse protocols. + reuse_end = min( + len(all_tokens), + max(len(all_tokens) - 1, 0) + self.reuse_match_backoff, + ) + if reuse_limit is not None: + # ``reuse_limit`` caps the claimed depth, not the evidence. Retain + # the D following tokens that prove the draft KV at that depth. + evidence_limit = 0 if reuse_limit <= 0 else reuse_limit + self.reuse_match_backoff + reuse_end = min(reuse_end, evidence_limit) + return self._augment_tokens_for_block_reuse(all_tokens, req, end=reuse_end) + + def _record_branch_snapshot_point( + self, req: LlmRequest, kv_cache: _KVCache, num_lookup_tokens: Optional[int] + ) -> None: + """Hook for recurrent-state managers; attention-only caches have no snapshots.""" + + def probe_context_reuse(self, req: LlmRequest) -> int | None: + """Reusable prefix depth for this pool, without claiming pages. + + The lookup includes this pool's prompt lookahead as evidence, and the core + applies ``reuse_match_backoff`` in the same match. A paired caller takes + the common usable depth. ``None``: nothing to probe -- reuse is off, or the + request already holds a settled cache claim. + """ + if not self.enable_block_reuse or req.py_request_id in self.kv_cache_map: + return None + return self.impl.probe_reuse( + ReuseScope(lora_id=req.lora_task_id, salt=self._derive_reuse_salt(req.cache_salt)), + self._context_reuse_tokens(req), + ) + + def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) -> int | None: + """Create/resume a context cache without touching request cursors. ``reuse_limit`` + caps the claimed depth while retaining its lookahead evidence; the returned + match lets the scheduler check both pools agree. + """ + if req.is_first_context_chunk: + if self.conversation_manager is not None: + self.conversation_manager.prepare_request(req) + kv_cache = self.kv_cache_map.get(req.py_request_id) + # Stays None when the cache was resumed rather than freshly + # matched, which is also the case where no branch point applies. + num_lookup_tokens = None + if kv_cache is None: + if self.enable_block_reuse: + # Empty (not None) for a zero limit: lookup is off, but + # newly computed blocks may still be committed. + tokens = self._context_reuse_tokens(req, reuse_limit) + else: + tokens = None + ugpu_id = self._request_ugpu_id(req) + # Taken from the augmented sequence rather than derived as + # prompt_len - 1, because augmentation can change the length + # for multimodal requests. + num_lookup_tokens = len(tokens) if tokens is not None else None + kv_cache = self._create_kv_cache( + req.py_request_id, + req.lora_task_id, + tokens, + cache_salt=req.cache_salt, + is_dummy=req.is_dummy, + expected_prompt_length=( + req.total_input_len_cp if self._has_cp_helix else req.prompt_len + ) + - 1, + ugpu_id=ugpu_id, + ) + if kv_cache is None: + return None + assert _get_kv_cache_ugpu_id(kv_cache) == ugpu_id + kv_cache.cuda_stream = self._stream.cuda_stream + + if not self.enable_block_reuse or self._draft_context_exceeds_limit(req): + kv_cache.stop_committing() + else: + self._record_branch_snapshot_point(req, kv_cache, num_lookup_tokens) + + if req.is_disagg_generation_init_state: + # Disagg generation receives prompt KV from the context worker; + # scratch blocks are only valid for local prefill chunks. + kv_cache.enable_swa_scratch_reuse = False + if not self._resume_and_restore(req.py_request_id, kv_cache): + return None + if ( + self.draft_layer_ids + and not self._draft_history_uses_kv_history() + and req.py_request_id not in self.draft_history + ): + # The reuse match covers every cache domain. Initialize once, + # after resume succeeds, so overlap/chunk retries cannot rewind + # draft history that has already advanced on the device. + reused = kv_cache.num_committed_tokens + valid_length = reused + if self.draft_window_size is not None: + valid_length = min(valid_length, self.draft_window_size) + self.set_draft_history(req.py_request_id, valid_length, reused) + return kv_cache.num_committed_tokens + + # Subsequent chunk: cache must exist from first chunk. It may be + # suspended (e.g. evicted between chunks), so resume it here. + kv_cache = self.kv_cache_map.get(req.py_request_id) + assert kv_cache is not None, ( + f"KV cache missing for non-first context chunk, request {req.py_request_id}" + ) + if not self._resume_and_restore(req.py_request_id, kv_cache): + return None + return kv_cache.num_committed_tokens + + def prepare_context(self, req: LlmRequest) -> bool: + """Create _KVCache, handle block reuse, and resume. Does NOT resize.""" + assert not req.is_disagg_generation_init_state, ( + f"req {req.py_request_id}: use prepare_disagg_gen_init" + ) + reused = self.prepare_context_cache(req) + if reused is None: + return False + # First chunk only: num_committed_tokens holds at the initial prefix + # until context end, so reapplying later would rewind the cursor. + if req.is_first_context_chunk and self.enable_block_reuse: + _rewind_context_cursor(req, reused) + req.set_prepopulated_prompt_len(reused, self.tokens_per_block) + return True + + def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: + """Resize KV cache to cover context_current_position + num_tokens. + + Returns True on success, False if resize failed (first chunk is + suspended on failure). + """ + assert not req.is_disagg_generation_init_state, ( + f"req {req.py_request_id}: use prepare_disagg_gen_init" + ) + if self._has_cp_helix and not req.is_dummy_request: + raise ValueError( + "resize_context is not helix-aware: its rank-local chunk " + "target would under-size the global ledger. Helix requests " + "are disagg-generation-only and must never take the context " + "path." + ) + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return False + + target = ( + req.context_current_position + + num_tokens + + self.num_extra_kv_tokens + + self._standalone_draft_reserve + ) + capacity = max(kv_cache.capacity, target) + pre_cap = kv_cache.capacity + + success = kv_cache.resize(capacity) + if not success: + if req.is_first_context_chunk: + kv_cache.suspend() + return False + req.py_ctx_pre_resize_cap = pre_cap if capacity > pre_cap else None + return True + + def try_allocate_draft_context(self, req: LlmRequest, num_tokens: int) -> bool: + """Resize a draft context cache during admission. Joint reuse shares one + cursor; the draft manager owns only its radix tree and capacity. + """ + assert self.is_draft + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None or not self._resume_and_restore(req.py_request_id, kv_cache): + return False + + target = req.context_current_position + num_tokens + self.num_extra_kv_tokens + return kv_cache.resize(max(kv_cache.capacity, target)) + + def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: + """Prepare KV cache for a disagg generation init request. + + Allocates capacity for the full prompt (+ draft) and sets + ``kv_cache.history_length`` to ``prompt_len``. Returns True on + success, False if preparation or resize failed (cache is suspended + on resize failure). + """ + reused = self.prepare_context_cache(req) + if reused is None: + return False + if self.enable_block_reuse: + _rewind_context_cursor(req, reused) + req.set_prepopulated_prompt_len(reused, self.tokens_per_block) + + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return False + + # prompt_len is the full incoming prompt length, robust to block + # reuse (which may leave a non-zero context_current_position). + # Helix requests carry the rank-local strided slice in prompt_len; + # the global ledger sizes off the full prompt instead. + prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len + target = ( + prompt_len + + get_draft_token_length(req) + + self.num_extra_kv_tokens + + self._standalone_draft_reserve + ) + capacity = max(kv_cache.capacity, target) + pre_cap = kv_cache.capacity + + success = kv_cache.resize(capacity, prompt_len) + if not success: + if req.is_first_context_chunk: + kv_cache.suspend() + return False + req.py_ctx_pre_resize_cap = pre_cap if capacity > pre_cap else None + return True + + def get_history_length(self, req: LlmRequest) -> int | None: + """Return the cache's current history_length, or None if no cache. + + Exposes the per-request SWA history watermark so callers + (e.g., the disagg transceiver) can verify scheduler/cache contracts + without reaching into ``kv_cache_map`` directly. + """ + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return None + return kv_cache.history_length + + def extend_capacity_for_tokens(self, request: LlmRequest) -> None: + """Extend KV cache capacity for the CUDA-graph padding delta. + + ``try_allocate_generation`` allocated capacity for the schedule-reduced + draft length. After padding restores ``py_draft_tokens`` to the static + max, we must extend by exactly the difference so that the subsequent + rewind (which operates on the padded length) does not underflow. + + The delta is computed from ``_allocated_draft_lens`` (recorded by + ``try_allocate_generation``) vs the current draft length (post-padding). + """ + allocated = self._allocated_draft_lens.get(request.py_request_id) + if allocated is None: + return + current_draft_len = get_draft_token_length(request) + # Growth is 1 + 2 * draft_len (see _required_gen_capacity), so the + # padding delta scales by 2 as well. + delta = 2 * (current_draft_len - allocated) + if delta <= 0: + return + # Half of the delta tops up the overlap slack; record it so + # update_resources reclaims the full slack for this iteration. + self._pending_overlap_slack[request.py_request_id] = self._pending_overlap_slack.get( + request.py_request_id, 0 + ) + (current_draft_len - allocated) + kv_cache = self.kv_cache_map[request.py_request_id] + new_capacity = kv_cache.capacity + delta + success = kv_cache.resize(new_capacity) + if not success: + raise ValueError( + f"Failed to extend capacity of KV cache for request " + f"{request.py_request_id} by {delta} tokens " + f"(target capacity {new_capacity})" + ) + self._allocated_draft_lens[request.py_request_id] = current_draft_len + + def suspend_request(self, req: LlmRequest) -> None: + """Suspend a request's KV cache, allowing pages to migrate to a secondary tier.""" + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is not None and kv_cache.is_active: + kv_cache.suspend() + + def resume_request(self, req: LlmRequest) -> bool: + """Resume a previously-suspended KV cache for *req*. + + Returns True if the cache is (or becomes) active on GPU, False if + resume was refused (e.g. GPU pressure above max_util_for_resume) + or no cache exists for the request. + """ + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return False + return self._resume_and_restore(req.py_request_id, kv_cache) + + # ---- prepare_resources ---- + + @nvtx_range("prepare_resources_kv_cache_manager_v2") + def prepare_resources(self, scheduled_batch: ScheduledRequests): + if self.is_draft: + # Mirror the main manager. Under one-model spec decoding the + # scheduler may already have created the context cache. + self._prepare_draft_resources(scheduled_batch) + return + + def _mirror_draft_kv_cache(self, req: LlmRequest): + """The draft manager's entry for ``req``, created on first sight. + + Mirrors what the target manager does in ``_prepare_context_impl``: + create when missing rather than fail, passing ``is_dummy`` through. The + draft mirror never looks up block reuse (``tokens=None``) and never + commits, so this cannot share the target's implementation. + + Reached from both the context and the generation loop, because not + every generation request passed through context on this worker: + attention DP injects its idle placeholder straight as a generation + request when a rank has no real work (ATTENTION_DP_DUMMY_REQUEST_ID, + see py_executor._adp_dummy_is_gen), and a disaggregated generation + worker receives prompt KV rather than prefilling it. + + Returns None when the IndexMapper is saturated. The caller must treat + that as fatal: the request stays in the scheduled batch, so skipping it + only defers the failure to copy_batch_block_offsets(), which asserts in + C++ on the unmapped request ID. Slots are released early on a + context-only worker (release_index_slot), so saturation here means the + draft mirror is out of sync with the target, not merely busy. + """ + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is not None: + return kv_cache + ugpu_id = self._request_ugpu_id(req) + kv_cache = self._create_kv_cache( + req.py_request_id, + req.lora_task_id, + None, + cache_salt=req.cache_salt, + is_dummy=req.is_dummy, + ugpu_id=ugpu_id, + ) + if kv_cache is None: + return None + assert _get_kv_cache_ugpu_id(kv_cache) == ugpu_id + kv_cache.stop_committing() + return kv_cache + + def _draft_pool_diagnostic(self) -> str: + """Draft-pool occupancy, for the resize-failure messages below. + + The draft manager mirrors the target's tokens but is sized from its own + byte budget, and the capacity scheduler admits on the TARGET pool alone + (`scheduler_v2` touches `draft_kv_cache_manager` only to suspend/free). + A draft pool smaller in tokens than the target therefore cannot + backpressure -- it can only raise, and the raise kills every rank. When + that happens the first question is always "how big was the draft pool + and how full was it", so answer it in the message rather than leaving + it to post-hoc arithmetic over the budget-split log line. + + Best-effort: never let a diagnostic mask the failure it describes. + """ + try: + live = sum(c.capacity for c in self.kv_cache_map.values()) + return ( + f" [draft pool: {len(self.kv_cache_map)} live caches holding " + f"{live} tokens, gpu_max_tokens={self._gpu_max_tokens}]" + ) + except Exception: # noqa: BLE001 - diagnostic only + return "" + + def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): + """Create/resize KV caches in the draft V2 manager for scheduled requests. + + The main V2 scheduler manages capacity in the primary KV cache + manager. The draft manager mirrors scheduled requests so that its + IndexMapper contains the correct request IDs. With block reuse, the + scheduler pre-creates context caches to coordinate the reusable prefix. + """ + # Joint V2 reuse deliberately shares the target request progress. A + # separate draft engine still uses the legacy draft request view. + with request_context(not self.enable_joint_kv_cache_reuse, scheduled_batch): + for req in scheduled_batch.context_requests: + kv_cache = self._mirror_draft_kv_cache(req) + if kv_cache is None: + raise RuntimeError( + f"Draft KV cache mirror exhausted its IndexMapper on " + f"context request {req.py_request_id}" + ) + if not self._resume_and_restore(req.py_request_id, kv_cache): + raise RuntimeError( + f"Failed to resume draft KV cache for request {req.py_request_id}" + ) + draft_len = get_draft_token_length(req) if req.is_last_context_chunk else 0 + capacity = ( + req.context_current_position + + req.context_chunk_size + + draft_len + + self.num_extra_kv_tokens + ) + if not kv_cache.resize(capacity): + raise RuntimeError( + f"Draft KV cache context resize failed for request " + f"{req.py_request_id}: could not resize to {capacity} tokens" + f"{self._draft_pool_diagnostic()}" + ) + + for req in scheduled_batch.generation_requests: + kv_cache = self._mirror_draft_kv_cache(req) + if kv_cache is None: + raise RuntimeError( + f"Draft KV cache mirror exhausted its IndexMapper on " + f"generation request {req.py_request_id}" + ) + if not self._resume_and_restore(req.py_request_id, kv_cache): + raise RuntimeError( + f"Failed to resume draft KV cache for request {req.py_request_id}" + ) + # A paired scheduler already grew both pools atomically; the + # recorded width says so, so do not charge this step twice. + if req.py_request_id in self._allocated_draft_lens: + continue + new_cap = self._required_gen_capacity(req, kv_cache.capacity) + if not kv_cache.resize(new_cap): + raise RuntimeError( + f"Draft KV cache generation resize failed for request " + f"{req.py_request_id}: could not resize to {new_cap} tokens" + f"{self._draft_pool_diagnostic()}" + ) + slack = self._effective_draft_len(req) + if slack > 0: + self._pending_overlap_slack[req.py_request_id] = ( + self._pending_overlap_slack.get(req.py_request_id, 0) + slack + ) + + def _reuse_token_source(self, req: LlmRequest) -> Sequence[int]: + """Beam-0 tokens for block reuse, in the form the active backend consumes. + + The C++ backend ingests a zero-copy int32 view (get_tokens_view) — no per-token + PyLong allocation. The pure-Python backend cannot consume a numpy array, so it gets + a plain list (get_tokens), preserving pre-optimization behavior. ``_cpp_introspection`` + is None exactly when the Python backend is active. + """ + if _cpp_introspection is not None: + return req.get_tokens_view(DEFAULT_BEAM_INDEX) + return req.get_tokens(DEFAULT_BEAM_INDEX) + + def _augment_tokens_for_block_reuse( + self, tokens: Sequence[int], req: LlmRequest, start: int = 0, end: int | None = None + ) -> Sequence[TokenIdExt]: + """Augment token sequence with multimodal content digests for block reuse. + + Multimodal placeholder tokens (e.g. image_token_id) share the same ID + regardless of the underlying content. This method replaces each + multimodal token region with TokenIdExt values produced by + gen_multimodal_cache_key_tokens(), embedding the content digest + (Blake3 hash) into the token sequence so that the radix tree can + distinguish blocks belonging to different images/videos. + + When *start*/*end* are given, they define the chunk bounds; only + `tokens[start:end]` is materialized and returned. This avoids + re-augmenting the full prompt on every chunk during chunked prefill. + + For text-only requests this is a no-op. + """ + if end is None: + end = len(tokens) + chunk_start = start + chunk_end = end + is_sliced = chunk_start != 0 or chunk_end != len(tokens) + + if ( + req.multimodal_hashes is None + or req.multimodal_positions is None + or req.multimodal_lengths is None + ): + return tokens[chunk_start:chunk_end] if is_sliced else tokens + + # Multimodal path: materialize a Python-int list (digest bytes get spliced in below), + # which flows through the per-element binding fallback. tokens may be a zero-copy numpy + # int32 view (get_tokens_view) — use tolist() so elements are Python ints, not np.int32. + chunk = tokens[chunk_start:chunk_end] + result: list[TokenIdExt] = chunk.tolist() if hasattr(chunk, "tolist") else list(chunk) + run_metadata = _resolve_multimodal_run_metadata(req) + if run_metadata is not None: + return _augment_tokens_with_mm_run_metadata( + self.vocab_size, result, req.multimodal_hashes, run_metadata, chunk_start, chunk_end + ) + + return _augment_tokens_with_contiguous_mm_metadata( + self.vocab_size, + result, + req.multimodal_hashes, + req.multimodal_positions, + req.multimodal_lengths, + chunk_start, + chunk_end, + ) + + def _stats_window_size(self, window_size: Optional[int]) -> int: + return self.max_seq_len if window_size is None else int(window_size) + + def _stats_life_cycle_window_size(self, life_cycle) -> Optional[int]: + if not isinstance(life_cycle, AttnLifeCycle): + return None + return self._stats_window_size(life_cycle.window_size) + + def _get_storage_statistics(self, cache_level: CacheLevel): + return _introspection.storage_statistics(self.impl, cache_level) + + def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str]]: + # life cycle (== layer group) -> pool group is static structure exposed by + # the public pool_group_descs API; no introspection needed. + pool_groups_by_life_cycle = { + int(variant.layer_group_id): int(pool_group.pool_group_index) + for pool_group in self.impl.pool_group_descs + for variant in pool_group.slot_desc.variants + } + + metadata: dict[int, tuple[int, Optional[int], str]] = {} + for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): + if not layer_ids: + continue + layer = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] + is_attention = isinstance(layer, AttentionLayerConfig) + metadata[life_cycle_id] = ( + int(pool_groups_by_life_cycle[life_cycle_id]), + self._stats_window_size(layer.sliding_window_size) if is_attention else None, + "attention" if is_attention else "ssm", + ) + return metadata + + def _storage_pool_groups_by_window(self) -> dict[int, set[int]]: + pool_groups_by_window: dict[int, set[int]] = defaultdict(set) + for pool_group_id, window_size, _ in self._stats_life_cycle_metadata().values(): + if window_size is not None: + pool_groups_by_window[window_size].add(pool_group_id) + return pool_groups_by_window + + def _get_and_reset_iteration_peak_block_stats(self, cache_level: CacheLevel): + get_peak_stats = getattr(self.impl, "get_and_reset_iteration_peak_block_stats", None) + if get_peak_stats is not None: + return get_peak_stats(cache_level) + return [ + PoolGroupPeakBlockStats( + available=stats.available, + unavailable=stats.total - stats.available, + evictable=stats.evictable, + ) + for stats in self._get_storage_statistics(cache_level) + ] + + @staticmethod + def _windows_by_pool_group( + pool_groups_by_window: dict[int, set[int]], + ) -> dict[int, tuple[int, ...]]: + windows_by_pool_group: dict[int, set[int]] = defaultdict(set) + for window_size, pool_group_ids in pool_groups_by_window.items(): + for pool_group_id in pool_group_ids: + windows_by_pool_group[pool_group_id].add(window_size) + return { + pool_group_id: tuple(sorted(window_sizes)) + for pool_group_id, window_sizes in windows_by_pool_group.items() + } + + @staticmethod + def _filter_iteration_stats_delta(delta, field_names) -> KVCacheIterationStatsDelta: + filtered = KVCacheIterationStatsDelta() + for field_name in field_names: + setattr(filtered, field_name, getattr(delta, field_name)) + return filtered + + @staticmethod + def _add_iteration_stats_delta( + bucket: dict[int, KVCacheIterationStatsDelta], key: int, delta: KVCacheIterationStatsDelta + ) -> None: + if delta.empty: + return + if key not in bucket: + bucket[key] = delta.copy() + return + bucket[key].add(delta) + + @staticmethod + def _iteration_cache_hit_rate(stats) -> float: + total = stats.iter_reused_blocks + stats.iter_missed_blocks + if stats.iter_reused_blocks == 0 or total == 0: + return 0.0 + return stats.iter_reused_blocks / total + + @staticmethod + def _apply_iteration_stats_delta( + stats, delta, field_names=KV_CACHE_ITERATION_STATS_DELTA_FIELDS + ) -> None: + if delta is None: + return + for field_name in field_names: + setattr(stats, field_name, getattr(delta, field_name)) + stats.iter_cache_hit_rate = KVCacheManagerV2._iteration_cache_hit_rate(stats) + + def _build_iteration_stats( + self, + pool_group_ids: Iterable[int], + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + delta, + field_names=KV_CACHE_ITERATION_STATS_DELTA_FIELDS, + ): + pool_group_ids = tuple(pool_group_ids) + stats = KvCacheIterationStats() + stats.primary_max_num_blocks = sum( + primary_stats[pool_group_id].total for pool_group_id in pool_group_ids + ) + stats.primary_free_num_blocks = sum( + primary_stats[pool_group_id].available for pool_group_id in pool_group_ids + ) + stats.primary_used_num_blocks = stats.primary_max_num_blocks - stats.primary_free_num_blocks + stats.primary_evictable_num_blocks = sum( + primary_stats[pool_group_id].evictable for pool_group_id in pool_group_ids + ) + stats.primary_peak_free_num_blocks = sum( + primary_peak_stats[pool_group_id].available for pool_group_id in pool_group_ids + ) + stats.primary_peak_used_num_blocks = sum( + primary_peak_stats[pool_group_id].unavailable for pool_group_id in pool_group_ids + ) + stats.primary_peak_evictable_num_blocks = sum( + primary_peak_stats[pool_group_id].evictable for pool_group_id in pool_group_ids + ) + stats.secondary_max_num_blocks = sum( + level_stats[pool_group_id].total + for level_stats in secondary_stats_by_level + for pool_group_id in pool_group_ids + ) + stats.secondary_free_num_blocks = sum( + level_stats[pool_group_id].available + for level_stats in secondary_stats_by_level + for pool_group_id in pool_group_ids + ) + stats.secondary_used_num_blocks = ( + stats.secondary_max_num_blocks - stats.secondary_free_num_blocks + ) + stats.secondary_evictable_num_blocks = sum( + level_stats[pool_group_id].evictable + for level_stats in secondary_stats_by_level + for pool_group_id in pool_group_ids + ) + stats.secondary_peak_free_num_blocks = sum( + peak_stats[pool_group_id].available + for peak_stats in secondary_peak_stats_by_level + for pool_group_id in pool_group_ids + ) + stats.secondary_peak_used_num_blocks = sum( + peak_stats[pool_group_id].unavailable + for peak_stats in secondary_peak_stats_by_level + for pool_group_id in pool_group_ids + ) + stats.secondary_peak_evictable_num_blocks = sum( + peak_stats[pool_group_id].evictable + for peak_stats in secondary_peak_stats_by_level + for pool_group_id in pool_group_ids + ) + self._apply_iteration_stats_delta(stats, delta, field_names) + return stats + + def _collect_iteration_stats_deltas( + self, raw_iteration_stats, life_cycle_metadata + ) -> tuple[dict, dict, dict, dict]: + reuse_deltas_by_window: dict[int, KVCacheIterationStatsDelta] = {} + reuse_deltas_by_life_cycle: dict[int, KVCacheIterationStatsDelta] = {} + pool_group_deltas_by_window: dict[int, KVCacheIterationStatsDelta] = {} + pool_group_deltas: dict[int, KVCacheIterationStatsDelta] = {} + + for life_cycle_id, delta in raw_iteration_stats.items(): + pool_group_id, window_size, _ = life_cycle_metadata[int(life_cycle_id)] + + pool_group_delta = self._filter_iteration_stats_delta( + delta, KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS + ) + self._add_iteration_stats_delta(pool_group_deltas, pool_group_id, pool_group_delta) + if window_size is not None: + self._add_iteration_stats_delta( + pool_group_deltas_by_window, window_size, pool_group_delta + ) + + reuse_delta = self._filter_iteration_stats_delta( + delta, KV_CACHE_ITERATION_STATS_REUSE_FIELDS + ) + if reuse_delta.empty: + continue + reuse_deltas_by_life_cycle[int(life_cycle_id)] = reuse_delta.copy() + if window_size is not None: + self._add_iteration_stats_delta(reuse_deltas_by_window, window_size, reuse_delta) + + return ( + reuse_deltas_by_window, + reuse_deltas_by_life_cycle, + pool_group_deltas_by_window, + pool_group_deltas, + ) + + def _build_window_iteration_stats( + self, + window_size: int, + pool_groups_by_window: dict[int, set[int]], + windows_by_pool_group: dict[int, tuple[int, ...]], + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + pool_group_delta, + reuse_delta, + ): + pool_group_ids = tuple( + pool_group_id + for pool_group_id in pool_groups_by_window.get(window_size, set()) + if windows_by_pool_group.get(pool_group_id) == (window_size,) + ) + stats = self._build_iteration_stats( + pool_group_ids, + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + pool_group_delta, + KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS, + ) + self._apply_iteration_stats_delta(stats, reuse_delta, KV_CACHE_ITERATION_STATS_REUSE_FIELDS) + return stats + + def _build_pool_group_iteration_stats( + self, + pool_group_id: int, + windows_by_pool_group: dict[int, tuple[int, ...]], + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + pool_group_delta, + ) -> KVCacheV2PoolGroupIterationStats: + primary_pool_group_stats = primary_stats[pool_group_id] + slot_size = ( + primary_pool_group_stats.slot_sizes + if hasattr(primary_pool_group_stats, "slot_sizes") + else primary_pool_group_stats.slot_size + ) + return KVCacheV2PoolGroupIterationStats( + pool_group_id=pool_group_id, + slot_size=tuple(slot_size), + window_sizes=windows_by_pool_group.get(pool_group_id, ()), + stats=self._build_iteration_stats( + (pool_group_id,), + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + pool_group_delta, + KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS, + ), + ) + + def _build_attention_life_cycle_iteration_stats( + self, + life_cycle_id: int, + life_cycle_metadata, + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + reuse_delta, + ) -> KVCacheV2LifeCycleIterationStats: + pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] + assert kind == "attention" + return KVCacheV2LifeCycleIterationStats( + life_cycle_id=life_cycle_id, + pool_group_id=pool_group_id, + window_size=window_size, + kind=kind, + stats=self._build_iteration_stats( + (), + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + reuse_delta, + KV_CACHE_ITERATION_STATS_REUSE_FIELDS, + ), + ) + + @staticmethod + def _build_ssm_life_cycle_iteration_stats( + life_cycle_id: int, + life_cycle_metadata, + snapshot_delta, + ) -> KVCacheV2SsmLifeCycleIterationStats: + pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] + assert kind == "ssm" + assert window_size is None + return KVCacheV2SsmLifeCycleIterationStats( + life_cycle_id=life_cycle_id, + pool_group_id=pool_group_id, + snapshot_stats=KVCacheV2SsmSnapshotIterationStats( + iter_snapshot_lookups=snapshot_delta.iter_snapshot_lookups, + iter_snapshot_hits=snapshot_delta.iter_snapshot_hits, + iter_snapshot_misses=snapshot_delta.iter_snapshot_misses, + iter_reused_tokens=snapshot_delta.iter_reused_tokens, + iter_unreused_tokens=snapshot_delta.iter_unreused_tokens, + iter_aligned_snapshot_hits=snapshot_delta.iter_aligned_snapshot_hits, + iter_unaligned_snapshot_hits=snapshot_delta.iter_unaligned_snapshot_hits, + ), + ) + + def get_kv_cache_stats(self): + kv_cache_stats = KvCacheStats() + pool_group_stats = self._get_storage_statistics(GPU_LEVEL) + max_num_blocks = sum(stat.total for stat in pool_group_stats) + free_num_blocks = sum(stat.available for stat in pool_group_stats) + committed_stats = self.impl.get_committed_stats() + + kv_cache_stats.max_num_blocks = max_num_blocks + kv_cache_stats.free_num_blocks = free_num_blocks + kv_cache_stats.used_num_blocks = max_num_blocks - free_num_blocks + kv_cache_stats.tokens_per_block = self.tokens_per_block + kv_cache_stats.alloc_total_blocks = committed_stats.alloc_total_blocks + kv_cache_stats.alloc_new_blocks = committed_stats.alloc_new_blocks + kv_cache_stats.reused_blocks = committed_stats.reused_blocks + kv_cache_stats.missed_blocks = committed_stats.missed_blocks + total = kv_cache_stats.reused_blocks + kv_cache_stats.missed_blocks + kv_cache_stats.cache_hit_rate = ( + 0.0 + if kv_cache_stats.reused_blocks == 0 or total == 0 + else kv_cache_stats.reused_blocks / total + ) + kv_cache_stats.num_free_blocks_per_window_size = { + window_size: sum( + pool_group_stats[pool_group_id].available for pool_group_id in pool_group_ids + ) + for window_size, pool_group_ids in self._storage_pool_groups_by_window().items() + } + kv_cache_stats.allocated_bytes = self.impl.get_quota(GPU_LEVEL) + + return kv_cache_stats + + def flush_iteration_events(self): + if self.event_manager is not None: + self.event_manager.flush_iteration_events() + + def get_latest_events(self, timeout_ms: Optional[float] = None): + if self.event_manager is None: + return [] + return self.event_manager.get_latest_events(timeout_ms) + + def get_iteration_stats(self): + if not self.enable_stats: + return None + + life_cycle_metadata = self._stats_life_cycle_metadata() + pool_groups_by_window = self._storage_pool_groups_by_window() + windows_by_pool_group = self._windows_by_pool_group(pool_groups_by_window) + raw_iteration_stats = self.impl.get_and_reset_iteration_stats() + raw_ssm_snapshot_iteration_stats = self.impl.get_and_reset_ssm_snapshot_iteration_stats() + primary_peak_stats = self._get_and_reset_iteration_peak_block_stats(GPU_LEVEL) + num_cache_levels = len(self.impl.cache_tier_list) + secondary_peak_stats_by_level = [ + self._get_and_reset_iteration_peak_block_stats(CacheLevel(level)) + for level in range(1, num_cache_levels) + ] + ( + reuse_deltas_by_window, + reuse_deltas_by_life_cycle, + pool_group_deltas_by_window, + pool_group_deltas, + ) = self._collect_iteration_stats_deltas(raw_iteration_stats, life_cycle_metadata) + + windows = set(pool_groups_by_window) + windows.update(reuse_deltas_by_window) + windows.update(pool_group_deltas_by_window) + primary_stats = self._get_storage_statistics(GPU_LEVEL) + secondary_stats_by_level = [ + self._get_storage_statistics(CacheLevel(level)) for level in range(1, num_cache_levels) + ] + + stats_by_window = { + window_size: self._build_window_iteration_stats( + window_size, + pool_groups_by_window, + windows_by_pool_group, + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + pool_group_deltas_by_window.get(window_size), + reuse_deltas_by_window.get(window_size), + ) + for window_size in sorted(windows) + } + + all_pool_group_ids = set(range(len(primary_stats))) + pool_group_ids = sorted( + all_pool_group_ids | set(windows_by_pool_group) | set(pool_group_deltas) + ) + stats_by_pool_group = { + pool_group_id: self._build_pool_group_iteration_stats( + pool_group_id, + windows_by_pool_group, + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + pool_group_deltas.get(pool_group_id), + ) + for pool_group_id in pool_group_ids + } + + stats_by_life_cycle = { + life_cycle_id: self._build_attention_life_cycle_iteration_stats( + life_cycle_id, + life_cycle_metadata, + primary_stats, + secondary_stats_by_level, + primary_peak_stats, + secondary_peak_stats_by_level, + reuse_delta, + ) + for life_cycle_id, reuse_delta in sorted(reuse_deltas_by_life_cycle.items()) + } + for life_cycle_id, snapshot_delta in sorted(raw_ssm_snapshot_iteration_stats.items()): + assert int(life_cycle_id) not in stats_by_life_cycle + stats_by_life_cycle[int(life_cycle_id)] = self._build_ssm_life_cycle_iteration_stats( + int(life_cycle_id), + life_cycle_metadata, + snapshot_delta, + ) + stats_by_life_cycle = dict(sorted(stats_by_life_cycle.items())) + + return KVCacheV2IterationStatsReport( + stats_by_window, stats_by_pool_group, stats_by_life_cycle + ) + + def get_block_ids_per_seq(self, request_ids: List[int]) -> torch.Tensor: + block_ids_per_seq = self.get_batch_cache_indices(request_ids) + block_ids_per_seq_tensors = [ + torch.tensor( + [i // self.num_local_layers if i != BAD_PAGE_INDEX else 0 for i in sublist], + dtype=torch.int, + ) + for sublist in block_ids_per_seq + ] + padded_tensor = torch.nn.utils.rnn.pad_sequence( + block_ids_per_seq_tensors, batch_first=True, padding_value=0 + ) + return padded_tensor + + def add_dummy_requests( + self, + request_ids: List[int], + # Note that token_nums should be past_kv_len + input_len (without + # spec decoding). The draft tokens will be added in this function, + # so we don't need to take care of it in the caller. When preparing + # token_nums, we should not take the draft tokens into account, so + # don't use the kv_cache_manager.max_seq_len, which includes both + # extra tokens and draft tokens. + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + prepare_resource: bool = True, + max_num_draft_tokens: int = 0, + kv_reserve_draft_tokens: Optional[int] = None, + use_mrope: bool = False, + max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, + num_extra_decoding_steps: int = 0, + draft_kv_cache_manager: Optional["BaseResourceManager"] = None, + ugpu_ids: Optional[List[int | None]] = None, + ): + if ugpu_ids is not None: + assert len(ugpu_ids) == len(request_ids), ( + "ugpu_ids must have the same length as request_ids" + ) + + _kv_draft = ( + kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens + ) + + beam_width = max_beam_width + requests = [] + + def release_resources( + current_request: LlmRequest, free_draft_resources: bool = False + ) -> None: + for req in requests: + self.free_resources(req) + self.free_resources(current_request) + if draft_kv_cache_manager is not None: + for req in requests: + draft_kv_cache_manager.free_resources(req) + if free_draft_resources: + draft_kv_cache_manager.free_resources(current_request) + + for i, req_id in enumerate(request_ids): + # exact choice of n can be ignored for dummy requests + sampling_params = SamplingParams( + n=beam_width, best_of=beam_width, use_beam_search=beam_width > 1 + ) + # Here 1+max_num_draft_tokens is used to extend the prompt length to + # a non-zero number to skip illegal memory access issue in MLA kernel + # during warmup. + token_num = token_nums[i] if token_nums is not None else 1 + max_num_draft_tokens + if self._has_cp_helix: + # Keep the frozen dummy fields self-consistent (the active + # rank's past_seen = seqlen - 1 must stay >= 1). + token_num = max(token_num, 2) + # token_num - 1 is the past history length in generation. + history_hint = max(0, token_num - 1) if is_gen else None + encoder_output_len = encoder_output_lens[i] if encoder_output_lens is not None else None + encoder_input_tokens = ( + [1] * encoder_output_len if encoder_output_len is not None else None + ) + # Using 1 instead of 0 prevents NaN during warmup in e.g. Deepseek + input_tokens = [1 for _ in range(token_num)] + request_ugpu_id = ( + self.pick_ugpu(req_id) if ugpu_ids is None or ugpu_ids[i] is None else ugpu_ids[i] + ) + req = LlmRequest( + request_id=req_id, + max_new_tokens=1, + input_tokens=input_tokens, + sampling_config=SamplingConfig(sampling_params._get_sampling_config()), + is_streaming=False, + encoder_input_tokens=encoder_input_tokens, + ugpu_id=request_ugpu_id, + encoder_output_len=encoder_output_len, + ) + req.is_dummy_request = True + req.paged_kv_block_ids = [] + if prepare_resource: + # Dummy/warmup request. ``stop_committing()`` below blocks all + # writes to the radix tree, so the choice of branch does not + # affect committed state. ``cache_salt`` is left defaulted + # to None to avoid coupling synthetic data to any salted branch. + ugpu_id = self._request_ugpu_id(req) + kv_cache = self._create_kv_cache( + req.py_request_id, + req.lora_task_id, + input_tokens, + is_dummy=req.is_dummy, + ugpu_id=ugpu_id, + ) + # Saturated IndexMapper (e.g. disagg gen trans in progress) + # returns None; retry next iter. + if kv_cache is None: + logger.warning_once( + f"Unable to create KV cache for dummy request " + f"{req.py_request_id}; IndexMapper capacity is " + f"temporarily exhausted", + key="dummy_kv_cache_create_failed", + ) + release_resources(req) + return None + assert _get_kv_cache_ugpu_id(kv_cache) == ugpu_id + assert kv_cache.num_committed_tokens == 0 + success = kv_cache.resume(self._stream.cuda_stream) + if not success: + logger.warning_once( + f"Unable to resume KV cache for dummy request {req.py_request_id}", + key="dummy_kv_cache_resume_failed", + ) + release_resources(req) + return None + kv_cache.stop_committing() + dummy_capacity = ( + token_num + + self.num_extra_kv_tokens + + num_extra_decoding_steps + + self._standalone_draft_reserve + ) + if is_gen: + kv_cache.enable_swa_scratch_reuse = False + # Need to hint the committed history to activate stale-block + # optimization and match the solver's pool budget. + success = kv_cache.resize(dummy_capacity, history_length=history_hint) + if not success: + logger.warning_once( + f"Unable to resize KV cache for dummy request " + f"{req.py_request_id} to capacity {dummy_capacity} " + f"(history_length={history_hint})", + key="dummy_kv_cache_resize_failed", + ) + release_resources(req) + return None + draft_kv_cache = None + if draft_kv_cache_manager is not None: + ugpu_id = self._request_ugpu_id(req) + draft_kv_cache = draft_kv_cache_manager._create_kv_cache( + req.py_request_id, + req.lora_task_id, + input_tokens, + is_dummy=req.is_dummy, + ugpu_id=ugpu_id, + ) + # Dummy path: see comment above, no salt. + if draft_kv_cache is None: + logger.warning_once( + f"Unable to create draft KV cache for dummy " + f"request {req.py_request_id}", + key="dummy_draft_kv_cache_create_failed", + ) + release_resources(req) + return None + assert _get_kv_cache_ugpu_id(draft_kv_cache) == ugpu_id + success = draft_kv_cache.resume(draft_kv_cache_manager._stream.cuda_stream) + if not success: + logger.warning_once( + f"Unable to resume draft KV cache for dummy " + f"request {req.py_request_id}", + key="dummy_draft_kv_cache_resume_failed", + ) + release_resources(req, free_draft_resources=True) + return None + draft_kv_cache.stop_committing() + success = draft_kv_cache.resize(dummy_capacity) + if not success: + logger.warning_once( + f"Unable to resize draft KV cache for dummy " + f"request {req.py_request_id} to capacity " + f"{dummy_capacity}", + key="dummy_draft_kv_cache_resize_failed", + ) + release_resources(req, free_draft_resources=True) + return None + + if is_gen: + req.state = LlmRequestState.GENERATION_IN_PROGRESS + req.prompt_len = token_num - 1 + req.py_prompt_len = req.prompt_len + req.py_draft_tokens = [1] * max_num_draft_tokens + if self._has_cp_helix: + # Frozen fields (V1 parity): dummies never pass the + # per-step derivation — last CP rank active, shared + # synthetic global length. + if self._helix_cp_rank == self._helix_cp_size - 1: + req.py_helix_is_inactive_rank = False + req.prompt_len = token_num - 1 + else: + req.py_helix_is_inactive_rank = True + req.prompt_len = token_num + req.py_prompt_len = req.prompt_len + req.seqlen_this_rank_cp = req.prompt_len + req.total_input_len_cp = token_num * self._helix_cp_size - 1 + req.py_decoding_iter = 1 + if prepare_resource: + new_capacity = kv_cache.capacity + _kv_draft + 1 + success = kv_cache.resize(new_capacity, history_length=history_hint) + if not success: + logger.warning_once( + f"Unable to extend generation dummy request " + f"{req.py_request_id} to capacity {new_capacity} " + f"(history_length={history_hint})", + key="generation_dummy_kv_cache_resize_failed", + ) + release_resources(req, free_draft_resources=draft_kv_cache is not None) + return None + if draft_kv_cache is not None: + success = draft_kv_cache.resize(new_capacity) + if not success: + logger.warning_once( + f"Unable to extend generation draft dummy " + f"request {req.py_request_id} to capacity " + f"{new_capacity}", + key="generation_dummy_draft_kv_cache_resize_failed", + ) + release_resources(req, free_draft_resources=True) + return None + + if use_mrope: + _populate_dummy_mrope_config(req, token_num, is_gen) + requests.append(req) + + return requests + + def try_commit_blocks(self, request: LlmRequest) -> None: + should_block_reuse = ( + self.enable_block_reuse + and self._can_publish_block_reuse + and not request.is_dummy_request + and not self._draft_context_exceeds_limit(request) + ) + if not should_block_reuse: + return + + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return + + commit_end = request.context_current_position + if commit_end > kv_cache.num_committed_tokens: + tokens = self._augment_tokens_for_block_reuse( + self._reuse_token_source(request), + request, + start=kv_cache.num_committed_tokens, + end=commit_end, + ) + # TODO: On a disaggregated prefill server, pass is_end=True for + # the last context chunk to improve performance. + kv_cache.commit(tokens) + if request.context_remaining_length == 0: + kv_cache.stop_committing() + + def release_index_slot(self, request_id: int) -> None: + """Release IndexMapper slot early while keeping KV cache blocks allocated. + + After prefill completes on a context-only worker, the IndexMapper slot + (used for host_kv_cache_block_offsets during model forward) is no longer + needed. Releasing it early allows new requests to be scheduled while + the KV cache blocks are still being transferred via NIXL/UCX. + """ + kv_cache = self.kv_cache_map.get(request_id) + if self.is_draft and (kv_cache is None or request_id in self._early_freed_index_requests): + # The draft mirror only holds a slot for requests it actually + # mirrored, and the target may release the same request twice. + return + if kv_cache is not None: + for i in range(self.max_beam_width): + for pool_idx in range(self.num_pools): + kv_cache.set_base_page_index_buf(i, pool_idx, None) + self.index_mapper.remove_sequence(request_id) + self._early_freed_index_requests.add(request_id) + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + if self.conversation_manager is not None: + self.conversation_manager.finish_request(request) + self._allocated_draft_lens.pop(request.py_request_id, None) + if self.draft_layer_ids: + self.draft_history.pop(request.py_request_id, None) + self._draft_dummy_request_ids.discard(request.py_request_id) + self._pending_overlap_slack.pop(request.py_request_id, None) + kv_cache = self.kv_cache_map.pop(request.py_request_id, None) + if kv_cache is None: + self.impl.clear_stats_excluded(request.py_request_id) + return + kv_cache.discard_pending_stats() + kv_cache.close() + self.impl.clear_stats_excluded(request.py_request_id) + if request.py_request_id in self._early_freed_index_requests: + self._early_freed_index_requests.discard(request.py_request_id) + else: + self.index_mapper.remove_sequence(request.py_request_id) + + def get_layer_page_index_scale(self, layer_idx: int) -> int: + """Page-index scale of this layer's KV buffer. Layers in one pool can + have different scales (e.g. different head_dim), so per-layer callers + must not use the pool-level scale.""" + return int(self.impl.get_page_index_scale(self.layer_offsets[layer_idx], Role.KEY)) + + def get_batch_cache_indices( + self, + request_ids: List[int], + layer_idx: Optional[int] = None, + num_blocks_per_seq: Optional[Sequence[int]] = None, + *, + raw_indices: bool = False, + ) -> List[List[int]]: + """Return cache indices for a layer, or pool 0 when no layer is given. + + Set raw_indices for slot-major views that need base slot IDs without + page-index scaling or KV aggregation. + """ + if layer_idx is None: + pool_id = 0 + index_scale = 1 if raw_indices else None + else: + pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] + index_scale = 1 if raw_indices else self.get_layer_page_index_scale(layer_idx) + return self._get_batch_cache_indices_by_pool_id( + request_ids, + pool_id=pool_id, + is_kv_aggregate=not raw_indices, + num_blocks_per_seq=num_blocks_per_seq, + index_scale=index_scale, + kv_factor=self.get_layer_kv_factor(layer_idx) if layer_idx is not None else None, + ) + + def _get_batch_cache_indices_by_pool_id( + self, + request_ids: List[int], + *, + pool_id: int = 0, + is_kv_aggregate: bool = True, + num_blocks_per_seq: Optional[Sequence[int]] = None, + index_scale: Optional[int] = None, + kv_factor: Optional[int] = None, + ) -> List[List[int]]: + if is_kv_aggregate: + # Div by kv_factor to index kv cache with size + # [num_blocks, kv_factor, tokens_per_block, num_kv_heads, head_dim] + div_factor = ( + self.kv_factor_per_layer[int(self.impl.layer_grouping[pool_id][0])] + if kv_factor is None + else kv_factor + ) + else: + div_factor = 1 + + if index_scale is None: + index_scale = int(self.index_scales[pool_id]) + res = [] + + for req_idx, req_id in enumerate(request_ids): + kv_cache = self.kv_cache_map[req_id] + # Zero-copy page-index buffers are padded to max_blocks_per_seq. + # Only convert blocks owned by this request; attention callers + # discard the padded tail immediately. + num_blocks = kv_cache.num_blocks + if num_blocks_per_seq is not None: + num_blocks = min(num_blocks, num_blocks_per_seq[req_idx]) + base_page_indices = kv_cache.get_base_page_indices(pool_id)[:num_blocks] + res.append( + [ + base_page_index * index_scale // div_factor + if base_page_index != BAD_PAGE_INDEX + else BAD_PAGE_INDEX + for base_page_index in base_page_indices + ] + ) + + return res + + def get_batch_cache_indices_flat( + self, + request_ids: List[int], + num_blocks: List[int], + layer_idx: Optional[int] = None, + ) -> torch.Tensor: + """Concatenated per-request block tables, trimmed to real widths. + + Equivalent to concatenating + ``get_batch_cache_indices(request_ids, layer_idx)[i][:num_blocks[i]]`` + over all requests, but never materializes the padded-to-capacity + per-request lists: the page indices are host data maintained in place + by the KV cache, so only ``num_blocks[i]`` entries per request are + gathered and a single vectorized transform runs over the result. + + Returns a CPU int32 tensor (pinned when supported) ready for an + async H2D copy. + """ + if layer_idx is None: + pool_id = 0 + scale = self._index_scale_ints[pool_id] + else: + pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] + # Layers sharing a pool can still require different page-index + # scales (e.g. Gemma4 sliding/global head_dim). Use the per-layer + # scale so this flat block table matches get_batch_cache_indices() + # and never feeds out-of-range page ids to FlashInfer. + scale = self.get_layer_page_index_scale(layer_idx) + div_factor = ( + self.get_layer_kv_factor(layer_idx) + if layer_idx is not None + else self.kv_factor_per_layer[int(self.impl.layer_grouping[pool_id][0])] + ) + + out_tensor = torch.empty(sum(num_blocks), dtype=torch.int32, pin_memory=prefer_pinned()) + out = out_tensor.numpy() + offset = 0 + for req_id, n in zip(request_ids, num_blocks): + out[offset : offset + n] = np.frombuffer( + self.kv_cache_map[req_id].get_base_page_indices(pool_id), + dtype=np.int32, + count=n, + ) + offset += n + + # One batched transform over the real widths; BAD_PAGE_INDEX entries + # (e.g. evicted out-of-window SWA blocks) stay untouched, matching + # get_batch_cache_indices. + valid = out != BAD_PAGE_INDEX + np.copyto(out, out * scale // div_factor, where=valid) + return out_tensor + + def get_cache_bytes_per_token(self) -> int: + return sum( + self.get_layer_bytes_per_token(local_id, Role.ALL) + * (self._helix_cp_size if local_id in self.cp_as_tp_layer_ids else 1) + for local_id in range(len(self.dtype_per_layer)) + ) + + def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): + if data_role not in (Role.ALL, *self._get_buffer_roles_for_layer(local_layer_idx)): + return 0 + if self.dtype_per_layer[local_layer_idx] not in ( + DataType.FP8, + DataType.HALF, + DataType.BF16, + DataType.FLOAT, + DataType.NVFP4, + ): + raise ValueError(f"Cannot support {self.dtype_per_layer[local_layer_idx]} KV cache.") + + if data_role == Role.ALL: + kv_factor = self.kv_factor_per_layer[local_layer_idx] + elif data_role in [Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: + if data_role in [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: + assert self.dtype_per_layer[local_layer_idx] == DataType.NVFP4, ( + "NVFP4 is the only supported dtype for block quant data roles" + ) + if data_role == Role.VALUE: + assert self.kv_factor_per_layer[local_layer_idx] == 2, ( + "VALUE data role is not supported for SELFKONLY cache type" + ) + kv_factor = 1 + else: + raise ValueError(f"Invalid data role: {data_role}") + + cache_size_per_token = ( + kv_factor + * self.num_kv_heads_per_layer[local_layer_idx] + * self.head_dim_per_layer[local_layer_idx] + ) + + cache_size_bytes_per_token = get_size_in_bytes( + cache_size_per_token, self.dtype_per_layer[local_layer_idx] + ) + + if data_role in [Role.KEY, Role.VALUE]: + return cache_size_bytes_per_token + + quant_size_per_token = 0 + + if self.dtype_per_layer[local_layer_idx] == DataType.NVFP4: + quant_size_per_token = self.calculate_scaling_factor_size_bytes( + cache_size_per_token, + quant_vector_size=16, + scaling_factor_dtype=DataType.FP8, + ) + + if data_role in [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: + return quant_size_per_token + + # Role.ALL combines both + return cache_size_bytes_per_token + quant_size_per_token + + @staticmethod + def calculate_scaling_factor_size_bytes( + cache_size: int, quant_vector_size: int, scaling_factor_dtype: DataType + ) -> int: + assert cache_size % quant_vector_size == 0, ( + "NVFP4 cache size must be divisible by quant vector size" + ) + return get_size_in_bytes(cache_size // quant_vector_size, scaling_factor_dtype) + + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: + pool_handled = set() + for layer_id, layer_offset in self.layer_offsets.items(): + pool_id = self.layer_to_pool_mapping_dict[layer_offset] + if pool_id in pool_handled: + continue + for buffer in self._get_buffer_views_for_invalid_value_check(layer_id): + yield buffer + pool_handled.add(pool_id) + + @torch.inference_mode() + 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() + ) + + for buffer in self._iter_cache_buffers_for_invalid_check(): + if torch.is_floating_point(buffer) or torch.is_complex(buffer): + # Process page-major physical views in bounded chunks. Packed + # integer cache data cannot contain NaN or Inf. + 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_() + 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) + + def shutdown(self): + for kv_cache in self.kv_cache_map.values(): + kv_cache.close() + self.kv_cache_map.clear() + if self.draft_layer_ids: + self.draft_history.clear() + self._draft_dummy_request_ids.clear() + self.impl.shutdown() + if self.conversation_manager is not None: + self.conversation_manager.clear() + + def get_max_resource_count(self) -> int: + # TODO: implement this + return 1 + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + # TODO: implement this + # context_token_count = request.orig_prompt_len + # num_context_blocks = context_token_count // self.tokens_per_block + # remaining_tokens = context_token_count + request.max_new_tokens - num_context_blocks * self.tokens_per_block + # need_blocks = num_context_blocks + math.ceil( + # remaining_tokens / self.tokens_per_block) + # return need_blocks + return 0 + + # TODO: refactor get_cache_size_per_token and get_cache_bytes_per_token to use the same logic + @classmethod + def get_cache_size_per_token( + cls, + model_config: ModelConfigPython, + mapping: Mapping, + num_layers: Optional[int] = None, + draft_scratch_tokens: int = 0, + draft_window_size: int | None = None, + **kwargs, + ): + layer_sizes, attention_windows = _get_static_cache_size_layer_components( + model_config, mapping, num_layers=num_layers, **kwargs + ) + # Upstream takes these as named parameters; here they ride **kwargs. + kv_cache_config = kwargs.get("kv_cache_config") + reuse_backoff_enabled = ( + kv_cache_config is not None + and kv_cache_config.enable_block_reuse + and cls._supports_reuse_match_backoff + ) + if reuse_backoff_enabled: + from ..speculative import draft_prompt_lookahead + + # Mirror the runtime W+D retention window before page rounding. + backoff = draft_prompt_lookahead(kwargs.get("spec_config")) or 0 + attention_windows = _extend_swa_windows_for_reuse( + attention_windows, + backoff, + kwargs.get("max_seq_len"), + ) + draft_layer_sizes = _get_draft_layer_sizes(**kwargs) + if draft_layer_sizes: + _, headroom = _get_generation_kv_capacity( + kwargs.get("spec_config"), is_draft=kwargs.get("is_draft", False) + ) + headroom += draft_scratch_tokens + context, generation, fixed = _estimate_cache_size_components( + layer_sizes, + attention_windows, + kwargs["tokens_per_block"], + scratch=bool(kv_cache_config and kv_cache_config.enable_swa_scratch_reuse), + generation_capacity_headroom=headroom, + standalone_draft_reserve=draft_scratch_tokens, + ) + draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( + draft_layer_sizes, + kwargs["tokens_per_block"], + scratch_tokens=draft_scratch_tokens, + window_size=draft_window_size, + generation_capacity_headroom=headroom, + helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, + ) + return ( + generation + draft_generation, + (fixed + draft_fixed) * int(kwargs.get("max_batch_size") or 0) + + (context + draft_context - generation - draft_generation) + * int(kwargs.get("max_num_tokens") or 0), + ) + full_attn_size_per_token = _estimate_full_attn_size_per_token( + layer_sizes, attention_windows + ) + swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( + layer_sizes, + attention_windows, + kwargs["tokens_per_block"], + context=False, + scratch=False, + ) + max_batch_size = int(kwargs.get("max_batch_size") or 0) + return ( + full_attn_size_per_token + swa_size_per_token, + swa_size_per_request * max_batch_size, + ) + + def update_context_resources(self, scheduled_batch: ScheduledRequests): + """Update KV cache for context requests in the current batch. + + This is separated from update_resources (which handles generation + requests only) because the overlap executor needs context KV cache + updates to happen before next batch scheduling. Otherwise, the scheduler + would under-estimate available KV cache for sliding-window attention + layer. In non-overlap scheduler, you should call it together with + update_resources(). + """ + for req in scheduled_batch.context_requests: + if req.py_request_id not in self.kv_cache_map: + continue + kv_cache = self.kv_cache_map[req.py_request_id] + # In the overlap scheduler, iteration N+1's eviction may + # suspend a ctx request's KV cache while iteration N's + # update still needs to process it. The request will be resumed by + # the scheduler next iteration. + if not kv_cache.is_active: + continue + should_block_reuse = ( + self.enable_block_reuse + and self._can_publish_block_reuse + and not req.is_dummy_request + and not self._draft_context_exceeds_limit(req) + ) + is_all_reusable = self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE + should_resize = not should_block_reuse or not is_all_reusable + should_commit = is_all_reusable or req.context_remaining_length == 0 + + if should_resize: + self._resize_context_history(req, kv_cache, req.context_current_position) + if should_commit: + self.try_commit_blocks(req) + if req.context_remaining_length == 0: + if self.conversation_manager is not None: + self.conversation_manager.save_drop_plan(req, kv_cache) + # Scratch blocks are only for prefill chunks. Disable them at + # the context/generation boundary so generation uses normal KV + # pages before the first generation allocation. + kv_cache.enable_swa_scratch_reuse = False + + def _resize_context_history(self, req: LlmRequest, kv_cache, history_length: int) -> None: + """Advance a context cache's history marker, or raise.""" + if not kv_cache.resize(None, history_length): + raise ValueError( + "Failed to resize history length of KV cache for request " + f"{req.py_request_id} to {history_length} tokens at context update" + ) + + def update_resources( + self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None, + ): + if not self.is_draft: + _update_kv_cache_draft_token_location( + self, scheduled_batch, attn_metadata, kv_cache_dtype_byte_size + ) + # Context request KV cache updates are handled by + # update_context_resources, called separately from the executor loop. + for req in scheduled_batch.generation_requests: + if req.py_request_id not in self.kv_cache_map: + continue + kv_cache = self.kv_cache_map[req.py_request_id] + # In the overlap scheduler, the scheduler for iteration N+1 + # may suspend a gen request's KV cache (via self-eviction or + # victim eviction) while iteration N's update_resources still + # needs to process it. Skip suspended caches — the request + # will be resumed by the scheduler on the next iteration. + if not kv_cache.is_active: + continue + accepted_draft_len = req.py_num_accepted_draft_tokens + runtime_draft_len = req.py_rewind_len + accepted_draft_len + # No entry: this manager did not grow through the scheduler (e.g. + # disagg), so fall back to _generation_draft_slots' width. + allocated_draft_slots = self._allocated_draft_lens.get( + req.py_request_id, + max(runtime_draft_len, self._kv_reserve_draft_tokens) + if self.is_draft + else runtime_draft_len, + ) + # Reclaim every reserved slot that was not occupied by an accepted + # draft token. + rewind_len = max(allocated_draft_slots - accepted_draft_len, 0) + # _required_gen_capacity grows by 1 + 2 * slots; rewind_len gives + # back the unused half of the first slots term, this the second. + overlap_slack = self._pending_overlap_slack.pop(req.py_request_id, 0) + new_capacity = ( + None + if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) + else max( + kv_cache.capacity - rewind_len - overlap_slack, + req.max_beam_num_tokens - 1, + ) + ) + embedded_draft = KVCacheManagerV2._draft_history_uses_kv_history(self) + if self.draft_layer_ids and not embedded_draft and new_capacity is not None: + draft_history = self.get_draft_history(req.py_request_id) + if draft_history is not None: + # Rejected target verification slots do not revoke draft + # feature history. Retain it and its next-forward scratch. + new_capacity = max( + new_capacity, draft_history.position + self._standalone_draft_reserve + ) + history_length = ( + None + # Reuse (history's consumer) is disabled under helix, and + # max_beam_num_tokens mixes rank-local and global counts. + if self.kv_compression_manages_history or self._has_cp_helix + else req.max_beam_num_tokens - 1 + ) + if embedded_draft and self._has_cp_helix: + # HELIX stores rank-local prompt tokens but generated tokens + # are replicated. The cache ledger uses global positions. + history_length = ( + req.total_input_len_cp + req.max_beam_num_tokens - req.py_prompt_len - 1 + ) + if embedded_draft and new_capacity is not None and history_length is not None: + new_capacity = max(new_capacity, history_length) + success = kv_cache.resize(new_capacity, history_length) + if not success: + raise ValueError( + f"Failed to resize KV cache for request {req.py_request_id} " + f"to capacity {new_capacity} and history length " + f"{history_length} tokens at generation update" + ) + self._allocated_draft_lens.pop(req.py_request_id, None) + + def copy_batch_block_offsets( + self, + dst_tensor: torch.Tensor, + request_ids: List[int], + beam_width: int, + num_contexts: int, + num_seqs: int, + max_blocks: Optional[int] = None, + ): + # max_blocks is accepted for signature parity with KVCacheManager; the + # device-side copy op here already scales with allocated blocks only. + assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" + + copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + assert copy_idx.shape[0] == num_seqs + + if self.enable_swa_scratch_reuse: + self._copy_batch_block_offsets_per_layer( + dst_tensor, request_ids, copy_idx, num_contexts, num_seqs + ) + return + + copy_batch_block_offsets_to_device( + self.host_kv_cache_block_offsets, + dst_tensor, + copy_idx, + self.index_scales, + self.kv_offset, + torch.cuda.current_stream().cuda_stream, + ) + + @staticmethod + def _derive_reuse_salt(cache_salt: str | None) -> int | None: + """Derive ``ReuseScope.salt`` (int|None) from the ``cache_salt`` string. + + Deterministic so the same string yields the same reuse namespace across + processes (matches C++ blockKey hashing on cacheSalt). Shared by cache + creation, prefetch, and the cache-aware router probe so they all hit the + same radix-tree namespace. + """ + if cache_salt is None: + return None + digest = hashlib.sha256(cache_salt.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "little") + + def _create_kv_cache( + self, + request_id: int, + lora_task_id: int | None, + input_tokens: Sequence[TokenIdExt] | None, + *, + cache_salt: str | None = None, + is_dummy: bool = False, + expected_prompt_length: int | None = None, + ugpu_id: int | None = None, + ): + assert request_id not in self.kv_cache_map, ( + f"KV cache for request {request_id} already exists" + ) + if self.num_ugpus > 1: + assert ugpu_id is not None, ( + "Localized KV cache managers require a concrete ugpu_id for KV-cache creation" + ) + assert 0 <= ugpu_id < self.num_ugpus, ( + f"ugpu_id must be in [0, {self.num_ugpus}), got {ugpu_id}" + ) + else: + assert ugpu_id is None, "Non-localized KV cache managers require ugpu_id=None" + if self.index_mapper.num_free_slots() == 0: + logger.warning( + "No free IndexMapper slots for request %s " + "(%d/%d slots in use, likely held by DISAGG_GENERATION_TRANS_IN_PROGRESS requests). " + "Skipping KV cache creation; request will retry next iteration.", + request_id, + self.index_mapper.size(), + self.index_mapper.size(), + ) + return None + salt_int = self._derive_reuse_salt(cache_salt) + kv_cache = _create_backend_kv_cache( + self.impl, + lora_task_id, + salt_int, + input_tokens, + request_id, + expected_prompt_length, + ugpu_id, + ) + self.kv_cache_map[request_id] = kv_cache + if is_dummy: + self.impl.mark_stats_excluded(request_id) + kv_cache.discard_pending_stats() + if self.draft_layer_ids: + self._draft_dummy_request_ids.add(request_id) + index = self.index_mapper.add_new_sequence(request_id) + for i in range(self.max_beam_width): + for pool_idx in range(self.num_pools): + buffer: torch.Tensor = self.host_kv_cache_block_offsets[ + pool_idx, index * self.max_beam_width + i, 0 + ] + kv_cache.set_base_page_index_buf(i, pool_idx, memoryview(buffer.numpy())) + return kv_cache + + def probe_prefix_match_length(self, input_tokens, lora_task_id=None, cache_salt=None): + """Probe the v2 KV cache radix tree for prefix match length. + + Returns the number of prefix tokens already cached on this rank, + without acquiring page ownership. Mirrors the v1 ``KVCacheManager`` + adapter so ``KVCacheAwareADPRouter`` works on both KV cache backends. + + ``cache_salt`` (and ``lora_task_id``) must match the values used by + ``_create_kv_cache``; the salt is derived from the ``cache_salt`` + string the same way, so the probe queries the same reuse namespace. + Otherwise the router would see an incorrect match length. + """ + if not self.enable_block_reuse: + return 0 + if len(input_tokens) == 0: + return 0 + salt_int = self._derive_reuse_salt(cache_salt) + return self.impl.probe_reuse( + ReuseScope(lora_id=lora_task_id, salt=salt_int), + input_tokens, + ) + + def prefetch_for_context_tokens(self, requests: list) -> bool: + """Prefetch radix-tree blocks from disk→host for upcoming context requests. + + Returns True if all prefetches succeeded, False if any failed. + """ + if not self.enable_block_reuse: + return False + # Prefetch via a transient KV cache that holds the reuse-matched blocks, + # prefetches disk->host, then closes. Holding blocks costs no GPU space + # (never resumed) and close() needs no stream sync. The transient cache + # is NOT registered in kv_cache_map / IndexMapper. + success = True + for req in requests: + tokens = self._context_reuse_tokens(req) + # Use the same salt derivation as _create_kv_cache so the transient + # cache hits the same radix-tree blocks. + salt_int = self._derive_reuse_salt(req.cache_salt) + kv_cache = self.impl.create_kv_cache( + ReuseScope(lora_id=req.lora_task_id, salt=salt_int), tokens + ) + # Prefetch to the first tier below GPU (host if present, otherwise + # disk). prefetch() is a best-effort hint either way. + if not kv_cache.prefetch(CACHE_LEVEL1): + logger.warning("prefetch failed for request %s", req.py_request_id) + success = False + kv_cache.close() + return success + + def reset_reuse_state(self): + self.impl.clear_reusable_blocks() + if self.conversation_manager is not None: + self.conversation_manager.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py new file mode 100644 index 000000000000..26139027838b --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -0,0 +1,4723 @@ +# 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 math +import os +from abc import ABC, abstractmethod +from dataclasses import dataclass, replace +from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, + Optional, Protocol, Sequence, Tuple, Union, cast) + +import torch +import triton +import triton.language as tl + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + BlockReusePolicy, KVCacheManagerV2, Role, + _estimate_draft_cache_size_components, _get_draft_layer_sizes) +from tensorrt_llm._torch.pyexecutor.kv_cache_stats import \ + KVCacheV2IterationStatsReport +from tensorrt_llm._torch.pyexecutor.llm_request import ( + ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest) +from tensorrt_llm._torch.pyexecutor.resource_manager import ( + BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, + PoolConfiguration, get_pp_layers) +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, + nvtx_range, prefer_pinned, + torch_dtype_to_binding) +from tensorrt_llm.bindings.internal.batch_manager import ( + LinearAttentionMetadata, LinearCacheType) +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, + BatchDesc, BufferConfig, + DataRole, KVCacheDesc) +from tensorrt_llm.runtime.kv_cache_manager_v2 import \ + KVCacheManagerConfig as KVCacheManagerConfigPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, + SsmLayerConfig, + TokenIdExt, _KVCache) + +# Shared with the KV budget estimator so allocator and budgeting can never +# diverge on the sharding rule (config_utils is import-cycle-free). +from .config_utils import mamba_effective_tp_size as _mamba_effective_tp_size + +GB = 1 << 30 + + +class _PrefixReuseDiagnostics(Protocol): + + def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: + ... + + def _get_num_reusable_tokens_before_pruning(self) -> int: + ... + + +# Replay kernels pad the token/window dimension to at least 16 for tensor-core +# tiles, so history sizes below 16 are no faster when not writing and do +# expensive checkpoint writes more often. This minimum tensor-core tile size is +# present in all tensor-core generations, with newer GPUs adding larger tile +# operations, not smaller ones. In the other direction, history sizes above 16 +# currently pad to 32, which is substantially slower with the current kernel +# design. Keep the default floor at 16 while still allowing larger T, which +# degenerates to checkpointing every step instead of rejecting the request. If +# larger T becomes common, we could consider padding larger histories to powers +# of 2 or multiples of 16, but this is untested. For smaller T, we could explore +# combining larger histories with new kernel designs that stay efficient when +# the window is only partly full. +MIN_REPLAY_HISTORY_SIZE = 16 + +_KDA_BETA_CACHE_ALIGNMENT_BYTES = 16 +# KVCM copies 16-byte payloads; the replay kernel reads only the first int32. +_KDA_REPLAY_COUNT_BYTES = 16 + + +def _allocate_kda_beta_cache( + shape: Tuple[int, ...], + device: Optional[torch.device], +) -> torch.Tensor: + """Allocate a logically shaped bf16 beta cache with aligned physical rows.""" + logical_num_heads = shape[-1] + heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // + torch.bfloat16.itemsize) + padded_num_heads = ((logical_num_heads + heads_per_alignment - 1) // + heads_per_alignment) * heads_per_alignment + # CuTe requires every nested cache view to be 16-byte aligned. Keep the + # logical head count while padding the physical stride (six fp32 heads are + # 24 bytes, so an unpadded row makes alternating views only 8-byte aligned). + return torch.zeros( + *shape[:-1], + padded_num_heads, + dtype=torch.bfloat16, + device=device, + )[..., :logical_num_heads] + + +def _get_num_cuda_graph_padding_dummy_slots( + spec_config: Optional["DecodingBaseConfig"], + max_batch_size: int, +) -> int: + """Return the number of persistent CUDA-graph padding dummy IDs. + + This is computed before ``ModelEngine`` exists and covers draft lengths + reachable at every batch size, including the zero-length acceptance-rate + fallback. ``ModelEngine._compute_dynamic_draft_len_mapping`` is created + later and covers only configured CUDA-graph batch sizes, so it cannot size + this persistent ID set. + """ + if spec_config is None: + return 1 + + draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + supports_dynamic_draft_len = (spec_dec_mode is not None and hasattr( + spec_dec_mode, "support_dynamic_draft_len") + and spec_dec_mode.support_dynamic_draft_len()) + if draft_len_schedule and supports_dynamic_draft_len: + runtime_draft_lengths = set() + first_uncovered_batch_size = 1 + for batch_size_threshold, draft_len in draft_len_schedule.items(): + if first_uncovered_batch_size > max_batch_size: + break + if batch_size_threshold >= first_uncovered_batch_size: + runtime_draft_lengths.add(draft_len) + first_uncovered_batch_size = batch_size_threshold + 1 + if first_uncovered_batch_size <= max_batch_size: + runtime_draft_lengths.add(0) + else: + max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 + max_total_draft_tokens = (getattr(spec_config, "max_total_draft_tokens", + 0) or 0) + is_linear_tree = getattr( + spec_config, + "is_linear_tree", + max_draft_len == max_total_draft_tokens, + ) + static_draft_len = (max_draft_len + if is_linear_tree else max_total_draft_tokens) + runtime_draft_lengths = {static_draft_len or 0} + + if ((getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and + (getattr(spec_config, "acceptance_rate_threshold", 0) or 0) > 0): + runtime_draft_lengths.add(0) + return len(runtime_draft_lengths) + + +class MambaRole: + """V2 buffer roles owned only by the hybrid Mamba manager.""" + + SSM_STATE = DataRole("ssm_state") + CONV_STATE = DataRole("conv_state") + + +def get_tensor_size_bytes(tensor): + """Calculate tensor size in bytes.""" + if isinstance(tensor, torch.Tensor): + return tensor.element_size() * tensor.nelement() + elif isinstance(tensor, list): + return sum(get_tensor_size_bytes(t) for t in tensor) + return 0 + + +# Mamba SSM stochastic-rounding Philox seed plumbing. +# +# Both the replay kernel and flashinfer's selective_state_update consume a +# `rand_seed` int64 tensor. Historically the non-replay paths created this +# tensor via `torch.randint(..., (1,))` on every decode step, which (a) is +# non-deterministic across runs, (b) allocates a fresh CUDA tensor per call +# and is therefore unfriendly to CUDA-graph capture, and (c) has no notion +# of cache-slot identity (so slot reuse cannot rotate the stream). +# +# The functions below produce per-slot int64 seeds via SplitMix64 finalization +# of (counter, slot, rank). Adjacent inputs yield uncorrelated outputs so +# consecutive cache slots and consecutive request counters do not leave +# structural fingerprints in the Philox input stream. All outputs live in +# (0, 2**62) so they avoid Philox's degenerate seed=0 case while staying +# within int64. +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +def _splitmix64(x: int) -> int: + """SplitMix64 finalizer; pure function, no torch.""" + x = (x + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 30) + x = (x * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 27) + x = (x * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 31) + return x & _MAMBA_SSM_UINT64_MASK + + +def _compute_deterministic_mamba_seed(counter: int, slot: int, + rank_offset: int) -> int: + """Deterministic int64 seed in (0, 2**62) from (counter, slot, rank). + + Pure function (no RNG, no torch.randint). Identical inputs across + process invocations produce identical outputs, which is what the + acceptance criteria require for cross-run reproducibility. + """ + folded = (_MAMBA_SSM_SEED_BASE + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK) + folded &= _MAMBA_SSM_UINT64_MASK + value = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + if value == 0: + value = 1 + return value + + +def _allocate_mamba_seed_buffer(cache_size: int, rank_offset: int, + device: torch.device) -> torch.Tensor: + """Allocate a (cache_size,) int64 CUDA buffer of deterministic seeds. + + counter=0 at allocation time; per-slot reset on fresh request assignment + bumps a host-side counter and rewrites only the freshly-assigned slot. + """ + slot_seeds = [ + _compute_deterministic_mamba_seed(0, i, rank_offset) + for i in range(cache_size) + ] + return torch.tensor(slot_seeds, dtype=torch.int64, device=device) + + +def _host_slots_to_device(slots: List[int], + device: torch.device) -> torch.Tensor: + """Copy host slot ids to ``device`` without a host sync. + + A pageable copy would wait for the queued step. Each call stages through a + fresh pinned tensor, so no later host write can race the copy. + """ + return torch.tensor(slots, + dtype=torch.long, + pin_memory=device.type == "cuda" + and prefer_pinned()).to(device, non_blocking=True) + + +def _copy_from_host_snapshot(dst: torch.Tensor, src: torch.Tensor) -> None: + """Copy reusable host staging ``src`` into ``dst`` without a host sync. + + Disaggregated admission and the execution batch can both prepare in one + executor iteration, refilling ``src`` before the first copy has run. Each + copy therefore reads a fresh pinned snapshot; ``dst`` keeps its address. + """ + snapshot = torch.empty(src.shape, + dtype=src.dtype, + pin_memory=dst.device.type == "cuda" + and prefer_pinned()) + snapshot.copy_(src) + dst.copy_(snapshot, non_blocking=True) + + +def _mamba_rank_offset(mapping: Mapping) -> int: + """Distinct seed offset per (tp_rank, pp_rank) so independent ranks + don't draw identical streams when not coordinated.""" + return (mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + + mapping.rank * 1_009) + + +def use_py_mamba_cache_manager() -> bool: + """Check if PythonMambaCacheManager should be forced (agg mode override). + + Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. + + Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager + with PythonMambaCacheManager inside instead of the configured manager. + Disagg mode is unaffected — its compatibility routing selects Mixed or + Cpp based on the transceiver configuration. + """ + return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' + + +class ReplayStateUpdateMetadata(NamedTuple): + """Shared tensors and fixed sizes for replay state updates.""" + prev_num_accepted_tokens: torch.Tensor + cache_buf_idx: torch.Tensor + replay_step_width: int + replay_history_size: int + + +def _advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: Optional[torch.Tensor] = None, +) -> None: + """Advance replay bookkeeping after a speculative generation step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to( + replay_metadata.prev_num_accepted_tokens.dtype) + prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = (prev_num_accepted_tokens + + replay_metadata.replay_step_width + > replay_metadata.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, + accepted_tokens, + prev_num_accepted_tokens + accepted_tokens, + ) + cache_buf_idx = replay_metadata.cache_buf_idx[slots] + next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx) + if is_dummy_request is not None: + next_num_accepted_tokens = torch.where( + is_dummy_request, + prev_num_accepted_tokens, + next_num_accepted_tokens, + ) + next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, + next_cache_buf_idx) + replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens + replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx + + +class BaseMambaCacheManager(ABC): + """Abstract interface for accessing mamba/recurrent state caches.""" + + @abstractmethod + def get_state_indices(self, *args, **kwargs) -> torch.Tensor: + """Return slot indices of each request. + + Shape: [max_batch_size] + """ + ... + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + """Return replay metadata tensors and fixed replay sizes.""" + return None + + @abstractmethod + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + """Return conv states for specific layer. + + Shape: [slot_size, conv_dim, d_conv - 1] + """ + ... + + @abstractmethod + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + """Return SSM states for specific layer. + + Shape: [slot_size, num_heads, head_dim, d_state] + """ + ... + + @abstractmethod + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + ... + + @abstractmethod + def is_speculative(self) -> bool: + ... + + @abstractmethod + def mamba_layer_cache( + self, layer_idx: int + ) -> Union['PythonMambaCacheManager.State', + 'PythonMambaCacheManager.SpeculativeState', None]: + ... + + +class PythonMambaCacheManager(BaseResourceManager): + """Pure-Python mamba state manager with speculative decoding support. + + Manages only mamba states (conv + SSM) using PyTorch tensors on GPU. + Supports speculative decoding and disaggregated serving. + """ + + @dataclass(frozen=True, kw_only=True) + class State: + """Base state container for Mamba cache.""" + conv: torch.Tensor + # Represents last state, or "two back" if prev_num_accepted_tokens is > 0 + temporal: torch.Tensor + + # Subclasses override to list fields shared across layers (not indexed by layer). + _SHARED_FIELDS = frozenset() + + def at_layer_idx(self, layer: int): + kwargs = {} + for k, v in vars(self).items(): + if v is None: + kwargs[k] = None + elif k in self._SHARED_FIELDS: + kwargs[k] = v + else: + kwargs[k] = v[layer] + return type(self)(**kwargs) + + @dataclass(frozen=True, kw_only=True) + class SpeculativeState(State): + """Speculative state with intermediate states for draft tokens. + + Supports three SSM update paths (only one set of tensors is + allocated): + - Legacy: caches full intermediate SSM states (intermediate_ssm) + - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) + - KDA replay: per-slot draft-token caches consumed by the fused + ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted + drafts and commits states in place (kda_conv_*, kda_*_cache) + """ + _SHARED_FIELDS = frozenset({ + "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" + }) + + # Allocated for the legacy and Mamba2-replay paths; None for the + # KDA replay path (the kernel commits conv windows in place). + intermediate_conv_window: torch.Tensor | None = None + + # Legacy path: full intermediate SSM states at each step + intermediate_ssm: torch.Tensor | None = None + + # KDA replay path (fused multi-token verify, kimi_linear). + # Pool invariant under this path: `temporal` holds the state after + # the LAST GOLDEN token; the accepted drafts recorded in + # prev_num_accepted_tokens are pending in these caches and are + # replayed by the kernel at the start of the next verify round. + # Extended conv caches [layers, slots, dim, (W-1) + num_spec] fp32, + # dim-contiguous (stride(dim) == 1): columns [0, W-1) are the + # committed raw-input window, tail columns the pending drafts' raw + # inputs. + kda_conv_q: torch.Tensor | None = None + kda_conv_k: torch.Tensor | None = None + kda_conv_v: torch.Tensor | None = None + # Post-processed per-draft quantities for replay: k + # [layers, slots, num_spec, H*K] bf16, gate [.., num_spec, H*K] fp32 + # (feeds the cumulative-gate exponents), v [.., num_spec, H*V] bf16, + # beta [.., num_spec, H] bf16. + kda_k_cache: torch.Tensor | None = None + kda_g_cache: torch.Tensor | None = None + kda_v_cache: torch.Tensor | None = None + kda_beta_cache: torch.Tensor | None = None + + # Replay path: compact double-buffered cache + # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). + # 0 means temporal saved state is actually the last state, not two back. + prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers + cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers + # Per-cache-slot Philox seeds. Replay bumps them in-place per launch so + # CUDA graph replay uses fresh SR draws without allocating RNG tensors. + # (cache,) int64 - shared across layers + mamba_ssm_rand_seed: torch.Tensor | None = None + old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) + old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) + # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. + old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + + def __init__( + self, + d_state: int, + d_conv: int, + num_heads: int, + n_groups: int, + head_dim: int, + num_layers: int, + max_batch_size: int, + spec_state_size: int, + mapping: Mapping, + dtype: torch.dtype, + ssm_cache_dtype: torch.dtype, + layer_mask: Optional[List[bool]] = None, + speculative_num_draft_tokens: Optional[int] = None, + model_type: str = "nemotron_hybrid", + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + ) -> None: + + self.mamba_ssm_cache_dtype = ssm_cache_dtype + self.speculative_num_draft_tokens = speculative_num_draft_tokens + self.spec_state_size = spec_state_size + self._use_replay_state_update = use_replay_state_update + # KDA replay path (kimi_linear fused multi-token verify). Mutually + # exclusive with use_replay_state_update; requires speculative mode. + self._kda_replay_num_spec = kda_replay_num_spec + self._use_kda_replay_update = kda_replay_num_spec is not None + if self._use_kda_replay_update: + assert not use_replay_state_update, ( + "kda_replay_num_spec and use_replay_state_update are " + "mutually exclusive") + assert speculative_num_draft_tokens is not None, ( + "KDA replay caches require speculative decoding") + assert kda_replay_num_spec == speculative_num_draft_tokens, ( + f"KDA replay cache width ({kda_replay_num_spec}) must match " + f"the draft length ({speculative_num_draft_tokens}): the " + "fused verify kernel is compiled with a static NUM_SPEC") + self.replay_history_size: Optional[int] = None + self.replay_step_width: Optional[int] = None + # When True, allocate the per-slot Philox seed buffer even outside + # the replay path so the non-replay flashinfer SR kernel reads a + # persistent deterministic seed instead of a per-call torch.randint. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh cache-slot assignment. Combined + # with slot index and rank offset to produce reproducible per-slot + # seed values. Starts at 0 so the post-init "reset" stream is + # disjoint from the counter=0 stream used at allocation time. + self._seed_request_counter = 0 + + # get tp size + tp_size = _mamba_effective_tp_size(mapping) + + # derive mamba parameters for conv and ssm states + d_inner = head_dim * num_heads + conv_dim = d_inner + 2 * n_groups * d_state + nheads = num_heads + + # check that can be partitioned + assert nheads % tp_size == 0, "nheads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + + # partition conv_dim and nheads + d_inner_local = d_inner // tp_size + ng_ds_local = n_groups * d_state // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + d_inner = d_inner // tp_size + + # Per-section dims for conv_state. + # Qwen3-Next: [Q | K | V] = [ng*ds, ng*ds, d_inner] + # Nemotron_hybrid: [x | B | C] = [d_inner, ng*ds, ng*ds] + if model_type == "qwen3_next": + self.conv_section_dims = [ng_ds_local, ng_ds_local, d_inner_local] + elif model_type == "nemotron_hybrid": + self.conv_section_dims = [d_inner_local, ng_ds_local, ng_ds_local] + else: + raise ValueError(f"Unsupported model type: {model_type}") + + # conv and ssm states device + device = torch.device("cuda") + + pp_layers, num_layers = get_pp_layers( + num_layers, + mapping, + layer_mask=layer_mask, + ) + num_local_layers = len(pp_layers) + self.mamba_layer_offsets = { + idx: offset + for offset, idx in enumerate(pp_layers) + } + + conv_state_shape = (conv_dim, d_conv - 1) + ssm_state_shape = (nheads, head_dim, d_state) + + # create mamba conv and ssm states + conv_states = torch.zeros( + size=(num_local_layers, max_batch_size) + conv_state_shape, + dtype=dtype, + device=device, + ) + + ssm_states = torch.zeros( + size=(num_local_layers, max_batch_size) + ssm_state_shape, + dtype=self.mamba_ssm_cache_dtype, + device=device, + ) + + # Per-slot Philox seeds. Allocated whenever stochastic rounding can + # fire, even outside the replay path and even when MTP/spec is off, + # so all consumers (replay kernel, MTP non-replay flashinfer, + # non-MTP flashinfer) read a persistent deterministic seed from the + # cache manager instead of calling torch.randint per forward. + self._mamba_ssm_rand_seed: Optional[torch.Tensor] = None + if (self._use_replay_state_update + or self._mamba_ssm_stochastic_rounding): + self._mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + max_batch_size, self._seed_rank_offset, device) + + # create state container + if speculative_num_draft_tokens is not None: + T = speculative_num_draft_tokens + 1 + self.replay_step_width = T + + # Conv intermediate cache — legacy and Mamba2-replay paths only. + # The KDA replay kernel commits conv windows in place, so the + # per-step window scratch is not needed. + intermediate_conv_window_cache = None + if not self._use_kda_replay_update: + intermediate_conv_window_cache = torch.zeros( + size=(num_local_layers, self.spec_state_size, T) + + conv_state_shape, + dtype=dtype, + device=device, + ) + + # SSM speculative cache — path-specific tensors + spec_kwargs = {} + # Share the manager-level seed buffer through SpeculativeState + # so the MTP path can still read it via layer_cache. + if self._mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self._mamba_ssm_rand_seed + if self._use_kda_replay_update: + # KDA replay caches for the fused multi-token verify kernel + # (trtllm::kda_mtp_decode). Per cache slot (persistent + # across rounds — the next round replays from them), unlike + # the batch-row-indexed intermediate buffers. + M = self._kda_replay_num_spec + # conv_dim covers the [q | k | v] short-conv sections; the + # kernel consumes them as three dim-contiguous caches. + assert conv_dim % 3 == 0, ( + "KDA replay caches expect the [q | k | v] conv-state " + "sectioning (3 equal sections)") + section_dim = conv_dim // 3 + # d_conv is KDA's convolution width; the live pool stores + # its W - 1 committed input columns. + w_kernel = d_conv + extended_s = w_kernel - 1 + M + + def _dim_contiguous_conv_cache(): + # raw bf16 projection inputs: bf16 is exact and halves + # the per-draft replay traffic of the verify kernel + return torch.zeros(num_local_layers, + max_batch_size, + extended_s, + section_dim, + dtype=torch.bfloat16, + device=device).transpose(-1, -2) + + spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device) + spec_kwargs['kda_conv_q'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_conv_k'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_conv_v'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_k_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.bfloat16, + device=device) + spec_kwargs['kda_g_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.float32, + device=device) + spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.bfloat16, + device=device) + spec_kwargs['kda_beta_cache'] = _allocate_kda_beta_cache( + (num_local_layers, max_batch_size, M, nheads), device) + ssm_spec_cache = [ + spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], + spec_kwargs['kda_conv_v'], spec_kwargs['kda_k_cache'], + spec_kwargs['kda_g_cache'], spec_kwargs['kda_v_cache'], + spec_kwargs['kda_beta_cache'] + ] + spec_path_label = "kda-replay" + elif self._use_replay_state_update: + assert n_groups % tp_size == 0, \ + "replay state update requires n_groups divisible by tp_size" + n_groups_per_rank = n_groups // tp_size + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) + + # Compact replay cache. + spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device) + spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size, + dtype=torch.int32, + device=device) + spec_kwargs['old_x'] = torch.zeros(num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + nheads, + head_dim, + dtype=dtype, + device=device) + spec_kwargs['old_B'] = torch.zeros(num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + n_groups_per_rank, + d_state, + dtype=dtype, + device=device) + spec_kwargs['old_dt'] = torch.zeros(num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device) + spec_kwargs['old_dA_cumsum'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device) + ssm_spec_cache = [ + spec_kwargs['old_x'], spec_kwargs['old_B'], + spec_kwargs['old_dt'], spec_kwargs['old_dA_cumsum'] + ] + spec_path_label = "replay" + else: + # Legacy: full intermediate SSM states at each step + spec_kwargs['intermediate_ssm'] = torch.zeros( + size=(num_local_layers, self.spec_state_size, T) + + ssm_state_shape, + dtype=self.mamba_ssm_cache_dtype, + device=device) + ssm_spec_cache = [spec_kwargs['intermediate_ssm']] + spec_path_label = "legacy" + + self.mamba_cache = self.SpeculativeState( + conv=conv_states, + temporal=ssm_states, + intermediate_conv_window=intermediate_conv_window_cache, + **spec_kwargs, + ) + + logger.info( + f"Mamba Cache ({spec_path_label}) is allocated. " + f"max_mamba_cache_size: {max_batch_size}, " + f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " + f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB, " + f"ssm_spec_cache size: {get_tensor_size_bytes(ssm_spec_cache) / GB:.2f}GB, " + "intermediate_conv_window_cache size: " + f"{get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" + ) + else: + self.mamba_cache = self.State( + conv=conv_states, + temporal=ssm_states, + ) + + logger.info( + f"Mamba Cache is allocated. " + f"max_mamba_cache_size: {max_batch_size}, " + f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " + f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB" + ) + + # mamba cache available blocks + self.mamba_cache_free_blocks = [i for i in range(max_batch_size)] + + # mamba cache index, maps request_id -> state indices + self.mamba_cache_index: Dict[int, int] = {} + self._dummy_request_ids: set[int] = set() + # Batch-order mask aligned with state_indices; duplicate dummy request + # IDs mark every batch row even when they share one cache slot. + self._dummy_request_mask = torch.zeros(max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros(max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned()) + + # Permanent slot shared by every CUDA-graph padding sentinel id + # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per + # draft length). Pool sizing must include +1 headroom for this; + # see MixedMambaHybridCacheManager. + self._padding_slot: int = self.mamba_cache_free_blocks.pop() + + # Reserved slot for the attention-DP padding dummy. Keep it out of the + # free pool so dummy insertion never consumes real-request capacity. + self._attention_dp_dummy_slot: Optional[int] = ( + self.mamba_cache_free_blocks.pop() + if mapping.enable_attention_dp else None) + + # save intermediate state indices for requests + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device=device) + + # Physical tensor rows include reserved dummy slots. Resource capacity + # is the number of real request slots remaining after those + # reservations. + self._max_batch_size = max_batch_size + self._max_resource_count = len(self.mamba_cache_free_blocks) + + def get_max_resource_count(self) -> int: + """Return the maximum number of real requests that can be cached.""" + return self._max_resource_count + + def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: + """Return the prefix of *context_requests* that fits in the + available Mamba cache blocks. Requests that already have a + cached block do not consume a free block.""" + free = len(self.mamba_cache_free_blocks) + result = [] + for r in context_requests: + if r.py_request_id in self.mamba_cache_index: + result.append(r) + elif free > 0: + result.append(r) + free -= 1 + else: + break + return result + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + """For Mamba cache manager, we always need one slot per request.""" + return 1 + + @torch.inference_mode() + def _prepare_mamba_cache_blocks(self, request_ids: List[int]): + for r in request_ids: + if r in self.mamba_cache_index: + continue + if len(self.mamba_cache_free_blocks) == 0: + raise RuntimeError("run out of mamba cache blocks") + block = self.mamba_cache_free_blocks.pop() + self.mamba_cache_index[r] = block + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + # Fresh request: no drafts pending in the replay caches. + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + if self._mamba_ssm_rand_seed is not None: + # Deterministic per-slot rotation on fresh assignment. + # `block` is pulled from mamba_cache_free_blocks, which + # excludes _padding_slot by construction (see __init__), + # so padding sentinels never reach this branch. + self._seed_request_counter += 1 + self._mamba_ssm_rand_seed[block] = ( + _compute_deterministic_mamba_seed( + self._seed_request_counter, block, + self._seed_rank_offset)) + + @torch.inference_mode() + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + """Seed the fused-verify KDA replay conv caches from the conv pool. + + On a disaggregated generation server the ctx->gen transfer populates + only the base ``conv`` / ``temporal`` pools (see + ``disaggregation/resource/page.py``); the per-slot ``kda_conv_*`` + replay caches consumed by ``trtllm::kda_mtp_decode`` are normally + seeded by a *local* prefill/decode via + ``_sync_kda_replay_conv_window`` and would otherwise hold zeros (or a + previous occupant's window) for a transferred request — corrupting + the recurrent state from the first verify step onward. Call this + after the state transfer completes, before the first generation + forward. Mirrors ``_sync_kda_replay_conv_window``: the conv pool row + stores the committed ``W - 1`` raw inputs directly. Pending-draft + scratch is cleared (no drafts are pending for a freshly transferred + request). + """ + if not (self._use_kda_replay_update + and isinstance(self.mamba_cache, self.SpeculativeState)): + return + blocks = [ + self.mamba_cache_index[rid] for rid in request_ids + if rid in self.mamba_cache_index + ] + if not blocks: + return + conv = self.mamba_cache.conv # [L, slots, 3D, W - 1] + idx = torch.tensor(sorted(set(blocks)), + dtype=torch.long, + device=conv.device) + d = conv.shape[2] // 3 + committed = conv.shape[3] + cs = conv.index_select(1, idx) + for cache, section in ( + (self.mamba_cache.kda_conv_q, cs[:, :, :d]), + (self.mamba_cache.kda_conv_k, cs[:, :, d:2 * d]), + (self.mamba_cache.kda_conv_v, cs[:, :, 2 * d:]), + ): + # cache: [L, slots, D, committed + num_spec]; zero the draft + # tail columns and seed the committed window in one copy. + seeded = torch.zeros( + (cache.shape[0], idx.numel()) + cache.shape[2:], + dtype=cache.dtype, + device=cache.device) + seeded[:, :, :, :committed] = section.to(cache.dtype) + cache.index_copy_(1, idx, seeded) + for buf in (self.mamba_cache.kda_k_cache, self.mamba_cache.kda_g_cache, + self.mamba_cache.kda_v_cache, + self.mamba_cache.kda_beta_cache): + buf.index_fill_(1, idx, 0) + self.mamba_cache.prev_num_accepted_tokens[idx] = 0 + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + context_ids = [ + i.py_request_id for i in scheduled_batch.context_requests + ] + generation_ids = [ + i.py_request_id for i in scheduled_batch.generation_requests + ] + request_ids = context_ids + generation_ids + self._prepare_mamba_cache_blocks(request_ids) + self._refresh_dummy_request_mask([req.is_dummy for req in requests]) + + def _is_padding_sentinel(self, request_id: int) -> bool: + # cuda_graph_runner caches one dummy per runtime_draft_len value + # (see _get_padded_batch), so any id in the range of dummy request IDs + # may be live concurrently. + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDA_GRAPH_DUMMY_REQUEST_ID + max_dl = self.speculative_num_draft_tokens or 0 + return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= + CUDA_GRAPH_DUMMY_REQUEST_ID) + + @torch.inference_mode() + def add_dummy_requests(self, request_ids: List[int], **kwargs): + # Sentinels alias to the shared _padding_slot; non-sentinel + # dummies (warmup, attention-DP idle padding) get their own + # slot and are freed individually. + if not request_ids: + return + self._dummy_request_ids.update(request_ids) + for r in request_ids: + if r in self.mamba_cache_index: + block = self.mamba_cache_index[r] + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + continue + if self._is_padding_sentinel(r): + block = self._padding_slot + elif (r == ATTENTION_DP_DUMMY_REQUEST_ID + and self._attention_dp_dummy_slot is not None): + block = self._attention_dp_dummy_slot + else: + if len(self.mamba_cache_free_blocks) == 0: + raise RuntimeError("run out of mamba cache blocks") + block = self.mamba_cache_free_blocks.pop() + self.mamba_cache_index[r] = block + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + + def free_resources(self, request: LlmRequest): + request_id = request.py_request_id + if request_id not in self.mamba_cache_index: + return + self._dummy_request_ids.discard(request_id) + block = self.mamba_cache_index.pop(request_id) + # Reserved slots must not re-enter the real-request free pool. + if block != self._padding_slot and \ + block != self._attention_dp_dummy_slot: + self.mamba_cache_free_blocks.append(block) + + def get_state_indices(self, request_ids: List[int], + is_padding: List[bool]) -> List[int]: + assert len(request_ids) == len(is_padding) + indices = [self.mamba_cache_index[rid] for rid in request_ids] + is_dummy = [ + rid in self._dummy_request_ids or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.as_tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, + non_blocking=True) + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).conv + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).temporal + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if not isinstance(self.mamba_cache, self.SpeculativeState): + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).intermediate_ssm + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + """Get intermediate conv states for speculative decoding.""" + if not isinstance(self.mamba_cache, self.SpeculativeState): + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx( + layer_offset).intermediate_conv_window + + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_x[layer_offset] + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_B[layer_offset] + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dt[layer_offset] + + def get_replay_old_dA_cumsum(self, + layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dA_cumsum[layer_offset] + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.cache_buf_idx + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.prev_num_accepted_tokens + + def is_speculative(self) -> bool: + return isinstance(self.mamba_cache, self.SpeculativeState) + + def mamba_layer_cache(self, + layer_idx: int) -> Union[State, SpeculativeState]: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.mamba_ssm_cache_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Return the persistent (cache_size,) int64 Philox seed buffer or + None when stochastic rounding is not active for this manager. + + Used by mamba2_mixer non-MTP paths that don't hold a SpeculativeState. + Callers must bump in-place (`.add_(1)` or slice-and-add) and pass a + view that matches the consuming kernel's expected shape. + """ + return self._mamba_ssm_rand_seed + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + if (not self._use_replay_state_update + or not isinstance(self.mamba_cache, self.SpeculativeState) + or self.mamba_cache.prev_num_accepted_tokens is None + or self.mamba_cache.cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=( + self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=self.mamba_cache.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) + + def shutdown(self): + """Release tensor memory.""" + # Clear mamba cache states + empty = torch.tensor([]) + + def _drop(tensor): + return empty if tensor is not None else None + + if isinstance(self.mamba_cache, self.SpeculativeState): + self.mamba_cache = self.SpeculativeState( + conv=empty, + temporal=empty, + intermediate_conv_window=_drop( + self.mamba_cache.intermediate_conv_window), + intermediate_ssm=_drop(self.mamba_cache.intermediate_ssm), + prev_num_accepted_tokens=_drop( + self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=_drop(self.mamba_cache.cache_buf_idx), + mamba_ssm_rand_seed=_drop(self.mamba_cache.mamba_ssm_rand_seed), + old_x=_drop(self.mamba_cache.old_x), + old_B=_drop(self.mamba_cache.old_B), + old_dt=_drop(self.mamba_cache.old_dt), + old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), + kda_conv_q=_drop(self.mamba_cache.kda_conv_q), + kda_conv_k=_drop(self.mamba_cache.kda_conv_k), + kda_conv_v=_drop(self.mamba_cache.kda_conv_v), + kda_k_cache=_drop(self.mamba_cache.kda_k_cache), + kda_g_cache=_drop(self.mamba_cache.kda_g_cache), + kda_v_cache=_drop(self.mamba_cache.kda_v_cache), + kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), + ) + else: + self.mamba_cache = self.State(conv=empty, temporal=empty) + + torch.cuda.empty_cache() + + @torch.compile(options={"max-autotune": True}) + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + num_accepted_draft_tokens = num_accepted_tokens[ + num_contexts:num_contexts + num_gens] - 1 + # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. + accepted_positions = (accepted_leaf_positions if accepted_leaf_positions + is not None else num_accepted_draft_tokens) + state_indices_d = state_indices[num_contexts:num_contexts + num_gens] + src_state_indices = self.intermediate_state_indices[:num_gens] + + if self._use_kda_replay_update: + # KDA replay: the fused verify kernel already committed the SSM + # state and conv windows (in place, after the golden token) and + # cached this round's drafts. All that remains is recording how + # many of those drafts the sampler accepted, so the next round's + # kernel launch replays exactly that prefix. + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + prev = self.mamba_cache.prev_num_accepted_tokens + current = prev[state_indices_d] + accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) + prev[state_indices_d] = torch.where(is_dummy_request, current, + accepted) + return + + if self._use_replay_state_update: + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) + else: + # Legacy: copy accepted SSM state from intermediate cache. + ssm_states = self.mamba_cache.temporal + intermediate_ssm_cache = self.mamba_cache.intermediate_ssm + accepted_ssm_state = intermediate_ssm_cache[:, src_state_indices, + accepted_positions] + ssm_states[:, state_indices_d, :] = accepted_ssm_state + + # Conv: both paths save all intermediate conv windows, carry over the accepted one. + conv_states = self.mamba_cache.conv + intermediate_conv_window_cache = self.mamba_cache.intermediate_conv_window + accepted_conv_state = intermediate_conv_window_cache[:, + src_state_indices, + accepted_positions] + conv_states[:, state_indices_d, :] = accepted_conv_state + + +class MambaCacheManager(BaseResourceManager, BaseMambaCacheManager): + """Facade for standalone mamba state management (no KV cache). + + Delegates to PythonMambaCacheManager. + """ + + def __init__( + self, + d_state: int, + d_conv: int, + num_heads: int, + n_groups: int, + head_dim: int, + num_layers: int, + max_batch_size: int, + spec_state_size: int, + mapping: Mapping, + dtype: torch.dtype, + ssm_cache_dtype: torch.dtype, + layer_mask: Optional[List[bool]] = None, + stream: Optional[torch.cuda.Stream] = None, + speculative_num_draft_tokens: Optional[int] = None, + model_type: str = "nemotron_hybrid", + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + ) -> None: + max_num_sequences = max_batch_size * mapping.pp_size + + self._impl = PythonMambaCacheManager( + d_state=d_state, + d_conv=d_conv, + num_heads=num_heads, + n_groups=n_groups, + head_dim=head_dim, + num_layers=num_layers, + max_batch_size=max_num_sequences, + spec_state_size=spec_state_size, + mapping=mapping, + dtype=dtype, + ssm_cache_dtype=ssm_cache_dtype, + layer_mask=layer_mask, + speculative_num_draft_tokens=speculative_num_draft_tokens, + model_type=model_type, + use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, + ) + + def get_max_resource_count(self) -> int: + return self._impl.get_max_resource_count() + + def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: + return self._impl.filter_ctx_requests_by_capacity(context_requests) + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + return self._impl.get_needed_resource_to_completion(request) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + self._impl.prepare_resources(scheduled_batch) + + def free_resources(self, request: LlmRequest): + self._impl.free_resources(request) + + def add_dummy_requests(self, request_ids: List[int], **kwargs): + self._impl.add_dummy_requests(request_ids, **kwargs) + + def get_state_indices( + self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None + ) -> Union[torch.Tensor, List[int]]: + return self._impl.get_state_indices(request_ids, is_padding) + + @property + def mamba_cache_free_blocks(self) -> List[int]: + return self._impl.mamba_cache_free_blocks + + @property + def mamba_cache_index(self) -> Dict[int, int]: + return self._impl.mamba_cache_index + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self._impl.get_conv_states(layer_idx) + + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + self._impl.seed_kda_replay_caches_for_disagg_gen(request_ids) + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self._impl.get_ssm_states(layer_idx) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self._impl.get_mamba_ssm_cache_dtype() + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Delegate to the underlying Python manager.""" + return self._impl.get_mamba_ssm_rand_seed() + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + return self._impl.get_replay_state_update_metadata() + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_intermediate_ssm_states(layer_idx) + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_intermediate_conv_states(layer_idx) + + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_x(layer_idx) + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_B(layer_idx) + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_dt(layer_idx) + + def get_replay_old_dA_cumsum(self, + layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_dA_cumsum(layer_idx) + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + return self._impl.get_replay_cache_buf_idx() + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + return self._impl.get_replay_prev_num_accepted_tokens() + + def is_speculative(self) -> bool: + return self._impl.is_speculative() + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + return self._impl.mamba_layer_cache(layer_idx) + + def shutdown(self): + self._impl.shutdown() + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None): + # Non-speculative configs don't allocate intermediate state; the + # promotion is a clean no-op. + if not self._impl.is_speculative(): + return + self._impl.update_mamba_states(attn_metadata, num_accepted_tokens, + state_indices, accepted_leaf_positions) + + +class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): + """Shared interface and state plumbing for hybrid Mamba managers. + + Concrete storage, state-index, and resource lifecycles remain owned by the + Cpp and V2 implementations. + """ + + _supports_additional_snapshot_offsets = False + + def _setup_mtp_intermediate_states(self, spec_config, + max_batch_size: int) -> None: + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if spec_config is None or self.local_num_mamba_layers == 0: + return + + tokens_per_gen_step = spec_config.tokens_per_gen_step + if not self._use_replay_state_update: + self.intermediate_ssm_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, + tokens_per_gen_step + ] + self.ssm_state_shape, + dtype=self.ssm_state_dtype, + device="cuda", + ) + self.intermediate_conv_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step + ] + self.conv_state_shape, + dtype=self.conv_state_dtype, + device="cuda", + ) + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device="cuda") + + def _allocate_pool_replay_buffers( + self, + spec_config, + cache_size: int, + device: Optional[torch.device], + ) -> bool: + """Allocate replay tensors shared by the Cpp and V2 state pools.""" + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + if (self.local_num_mamba_layers == 0 + or (not self._use_replay_state_update + and not self._mamba_ssm_stochastic_rounding)): + return False + + assert device is not None + self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + cache_size, self._seed_rank_offset, device) + if spec_config is None or not self._use_replay_state_update: + return False + + history_size = self.replay_history_size + assert history_size is not None + nheads, head_dim, d_state = self.ssm_state_shape + common_shape = [self.local_num_mamba_layers, cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.cache_buf_idx = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.old_x = torch.zeros( + common_shape + [history_size, nheads, head_dim], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_shape + [history_size, self._n_groups_per_rank, d_state], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + return True + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + if self._dummy_request_mask is None: + return + + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.tensor(is_dummy, dtype=torch.bool)) + _copy_from_host_snapshot(self._dummy_request_mask, + self._dummy_request_mask_host) + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + if num_contexts == 0: + return + + context_slots = self.cuda_state_indices[:num_contexts].long() + if (self._use_replay_state_update + and self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None): + self.prev_num_accepted_tokens[context_slots] = 0 + self.cache_buf_idx[context_slots] = 0 + if self.old_x is not None: + self.old_x[:, context_slots] = 0 + if self.old_B is not None: + self.old_B[:, context_slots] = 0 + if self.old_dt is not None: + self.old_dt[:, context_slots] = 0 + if self.old_dA_cumsum is not None: + self.old_dA_cumsum[:, context_slots] = 0 + + if self.mamba_ssm_rand_seed is None: + return + self._seed_request_counter += 1 + counter = self._seed_request_counter + rank_offset = self._seed_rank_offset + host_slots = self._host_state_indices[:num_contexts].tolist() + new_seeds = [ + _compute_deterministic_mamba_seed(counter, slot, rank_offset) + for slot in host_slots + ] + self.mamba_ssm_rand_seed[context_slots] = torch.tensor( + new_seeds, + dtype=torch.int64, + device=self.mamba_ssm_rand_seed.device, + ) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + """Set reusable Mamba snapshot boundaries before scheduling.""" + if not self.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + state_config = self.kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + for request in requests: + snapshot_points = set() + if interval is not None and interval > 0: + snapshot_points.update( + range(interval, request.prompt_len + 1, interval)) + if self._supports_additional_snapshot_offsets: + for offset in state_config.additional_snapshot_offsets_from_start: + if offset <= request.prompt_len: + snapshot_points.add(offset) + for offset in state_config.additional_snapshot_offsets_from_end: + point = request.prompt_len - offset + if point > 0: + snapshot_points.add(point) + request.expect_snapshot_points = sorted(snapshot_points) + + def is_speculative(self) -> bool: + return self.spec_config is not None + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_ssm_states is None: + return None + return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_conv_states is None: + return None + return self.intermediate_conv_states[ + self.mamba_layer_offsets[layer_idx]] + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + conv = self.get_conv_states(layer_idx) + ssm = self.get_ssm_states(layer_idx) + if self.spec_config is not None: + layer_offset = self.mamba_layer_offsets[layer_idx] + spec_kwargs = {} + if self.mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed + if self._use_replay_state_update: + spec_kwargs['old_x'] = self.old_x[layer_offset] + spec_kwargs['old_B'] = self.old_B[layer_offset] + spec_kwargs['old_dt'] = self.old_dt[layer_offset] + spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] + spec_kwargs['cache_buf_idx'] = self.cache_buf_idx + spec_kwargs['prev_num_accepted_tokens'] = ( + self.prev_num_accepted_tokens) + else: + spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ + layer_offset] + return PythonMambaCacheManager.SpeculativeState( + conv=conv, + temporal=ssm, + intermediate_conv_window=self. + intermediate_conv_states[layer_offset], + **spec_kwargs, + ) + return PythonMambaCacheManager.State(conv=conv, temporal=ssm) + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', + None) + cache_buf_idx = getattr(self, 'cache_buf_idx', None) + if (not self._use_replay_state_update + or prev_num_accepted_tokens is None or cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.ssm_state_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + return getattr(self, 'mamba_ssm_rand_seed', None) + + +def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: + """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" + pool_size = max_batch_size + # One permanent slot is shared by every CUDA-graph padding sentinel. + pool_size += 1 + if mapping.enable_attention_dp: + # Attention-DP can insert a transient dummy request on an otherwise + # idle rank. Keep this headroom internal so scheduler-visible + # max_batch_size still limits real requests. + pool_size += 1 + return pool_size + + +class MixedMambaHybridCacheManager(KVCacheManager, MambaCacheManager, + MambaHybridCacheManager): + """Hybrid cache manager combining separate KVCacheManager and MambaCacheManager. + + Manages KV cache and mamba states in independent pools, with support of + speculative decoding and disaggregated serving. + Does not support block reuse / prefix caching for mamba states. + """ + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + # kv cache parameters + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + layer_mask: List[bool], + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + max_num_tokens: int = 8192, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + is_estimating_kv_cache: bool = False, + execution_stream: Optional[torch.cuda.Stream] = None, + model_type: str = "nemotron_hybrid", + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + # Per-pool configurations forwarded to the C++ KVCacheManager ctor. + # Lets a single manager host pools with mixed shapes (e.g. Gemma4 + # hybrid attention). See KVCacheManager.__init__. + pool_configurations: Optional[List[PoolConfiguration]] = None, + ) -> None: + + # mamba hybrid cache requires block reuse to be disabled in KV cache config + assert not kv_cache_config.enable_block_reuse, ( + "mamba hybrid cache requires block reuse to be disabled in KV cache config" + ) + + # Host-drafter spec modes (NGram) have no spec worker to call + # update_mamba_states; this manager promotes accepted states itself + # in update_resources. One-model modes (MTP/Eagle/DFlash) and the + # suffix-automaton worker promote from their spec workers and must + # NOT be promoted twice. + self._promote_states_in_update_resources = ( + spec_config is not None + and getattr(spec_config, "decoding_type", None) == "NGram") + + pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) + + MambaCacheManager.__init__( + self, + mamba_d_state, + mamba_d_conv, + mamba_num_heads, + mamba_n_groups, + mamba_head_dim, + mamba_num_layers, + pool_size, + max_batch_size, + mapping, + mamba_cache_dtype, + mamba_ssm_cache_dtype, + mamba_layer_mask, + execution_stream, + speculative_num_draft_tokens=(spec_config.tokens_per_gen_step - 1 + if spec_config is not None else None), + model_type=model_type, + use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, + ) + + # initialize kv cache manager + KVCacheManager.__init__( + self, + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_num_tokens=max_num_tokens, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + execution_stream=execution_stream, + is_draft=is_draft, + pool_configurations=pool_configurations, + ) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + MambaCacheManager.prepare_resources(self, scheduled_batch) + KVCacheManager.prepare_resources(self, scheduled_batch) + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + MambaCacheManager.free_resources(self, request) + KVCacheManager.free_resources(self, request, pin_on_release) + + def add_dummy_requests(self, request_ids: List[int], **kwargs): + MambaCacheManager.add_dummy_requests(self, request_ids) + return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) + + def shutdown(self): + MambaCacheManager.shutdown(self) + KVCacheManager.shutdown(self) + + def update_resources(self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None): + KVCacheManager.update_resources(self, scheduled_batch, attn_metadata, + kv_cache_dtype_byte_size) + self._maybe_promote_drafter_states(scheduled_batch, attn_metadata) + + def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): + """Commit accepted verification states for host-drafter spec modes. + + One-model spec workers (MTP/Eagle/DFlash) and the suffix-automaton + worker call update_mamba_states themselves right after on-device + acceptance. Host-drafter modes (NGram) have no spec worker: + acceptance lands on the requests as + ``py_num_accepted_draft_tokens`` during sampler update, and this + hook — running right after, alongside the KV rewind — promotes the + accepted step's intermediate state into the live pools (or, on the + KDA replay path, records the accepted-draft count for the next + round's replay). Without it, the pools would keep the + pre-verification state and the next forward would resume from a + stale prefix. + """ + if not self._promote_states_in_update_resources: + return + if not self.is_speculative() or attn_metadata is None: + return + gen_requests = scheduled_batch.generation_requests + if not gen_requests: + return + drafted = [ + r for r in gen_requests + if r.py_draft_tokens is not None and len(r.py_draft_tokens) > 0 + ] + if not drafted: + # Drafter skipped this step (e.g. speculation gated off): the + # forward ran the plain in-place decode path; nothing to promote. + return + assert len(drafted) == len(gen_requests), ( + "mixed drafted/undrafted generation batch is not supported for " + "hybrid state promotion (drafts are padded to the static max)") + device = self._impl.mamba_cache.temporal.device + num_contexts = len(scheduled_batch.context_requests) + num_accepted = torch.tensor( + [0] * num_contexts + + [r.py_num_accepted_draft_tokens + 1 for r in gen_requests], + dtype=torch.int32, + device=device) + # Batch-ordered slots (contexts then gens), matching the ordering + # the forward used for the intermediate scratch buffers. Requests + # that finished this step were already freed by response handling + # (which runs before update_resources) and are gone from the index; + # their rows must stay in place for alignment, so redirect them to + # the reserved padding slot (a harmless scratch write). + slot_index = self.mamba_cache_index + padding_slot = self._impl._padding_slot + state_indices = torch.tensor([ + slot_index.get(r.py_request_id, padding_slot) + for r in scheduled_batch.context_requests + gen_requests + ], + dtype=torch.int32, + device=device) + self.update_mamba_states(attn_metadata, num_accepted, state_indices) + + +@triton.jit +def _promote_mamba_state_kernel( + src_ptr, + dst_ptr, + src_idx_ptr, + accepted_position_source_ptr, + num_accepted_tokens_ptr, + blk_ptr, + replay_pnat_ptr, + replay_cache_buf_idx_ptr, + dummy_request_mask_ptr, + num_gens, + count, + src_s_layer, + src_s_row, + src_s_step, + dst_s_layer, + dst_s_block, + POSITION_SOURCE_IS_TOKEN_COUNT: tl.constexpr, + REPLAY_STEP_WIDTH: tl.constexpr, + REPLAY_HISTORY_SIZE: tl.constexpr, + UPDATE_REPLAY_BOOKKEEPING: tl.constexpr, + BLOCK: tl.constexpr, +): + # One program copies a BLOCK-sized tile of the contiguous inner state for a + # single (layer, gen) pair. grid = (num_layers * num_gens, ceil(count/BLOCK)). + pair = tl.program_id(0) + tile = tl.program_id(1) + layer = (pair // num_gens).to(tl.int64) + g = pair % num_gens + row = tl.load(src_idx_ptr + g).to(tl.int64) + accepted_position = tl.load(accepted_position_source_ptr + g) + if POSITION_SOURCE_IS_TOKEN_COUNT: + accepted_position -= 1 + acc = accepted_position.to(tl.int64) + blk = tl.load(blk_ptr + g).to(tl.int64) + + # The replay metadata update is batch-scoped, so layer 0's first copy tile + # owns it. This folds PNAT and active-buffer maintenance into the existing + # accepted-convolution-state promotion launch. + if UPDATE_REPLAY_BOOKKEEPING: + if (layer == 0) & (tile == 0): + num_accepted_tokens = tl.load(num_accepted_tokens_ptr + g) + is_dummy_request = tl.load(dummy_request_mask_ptr + g) + pnat = tl.load(replay_pnat_ptr + blk) + cache_buf_idx = tl.load(replay_cache_buf_idx_ptr + blk) + wrote_checkpoint = pnat + REPLAY_STEP_WIDTH > REPLAY_HISTORY_SIZE + next_pnat = tl.where(wrote_checkpoint, num_accepted_tokens, + pnat + num_accepted_tokens) + next_cache_buf_idx = tl.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx) + tl.store(replay_pnat_ptr + blk, next_pnat, mask=~is_dummy_request) + tl.store(replay_cache_buf_idx_ptr + blk, + next_cache_buf_idx, + mask=~is_dummy_request) + + # int64 throughout: per-layer strides are O(1e8), so layer*stride overflows + # int32 and would corrupt addresses / fault. + src_base = layer * src_s_layer + row * src_s_row + acc * src_s_step + dst_base = layer * dst_s_layer + blk * dst_s_block + offs = tile * BLOCK + tl.arange(0, BLOCK) + mask = offs < count + v = tl.load(src_ptr + src_base + offs, mask=mask) + tl.store(dst_ptr + dst_base + offs, v, mask=mask) + + +def _promote_mamba_state_triton( + dst: torch.Tensor, + intermediate: torch.Tensor, + src_state_indices: torch.Tensor, + accepted_position_source: torch.Tensor, + dst_state_indices: torch.Tensor, + num_accepted_tokens: Optional[torch.Tensor] = None, + position_source_is_token_count: bool = False, + replay_pnat: Optional[torch.Tensor] = None, + replay_cache_buf_idx: Optional[torch.Tensor] = None, + dummy_request_mask: Optional[torch.Tensor] = None, + replay_step_width: int = 0, + replay_history_size: int = 0, + BLOCK: int = 2048) -> None: + """Scatter each generation request's accepted draft-step recurrent state + from the per-request intermediate buffer into the unified C++ pool view. + + Args: + dst: ``[num_layers, num_blocks, *state_shape]`` view of the C++ pool. + The block dim (dim 1) is strided (the pool interleaves + ``ssm_bytes | conv_bytes`` per block), so ``dst`` is non-contiguous. + The kernel writes through ``dst``'s real strides, touching only this + state's bytes -- the case that defeats torch.compile / aot_autograd + (dtype-view mutation) and Inductor codegen (``XBLOCK`` on the uint8 + pool). + intermediate: ``[num_layers, max_batch_size, T, *state_shape]`` dense. + src_state_indices, accepted_position_source, dst_state_indices: + ``[num_gens]`` int. + num_accepted_tokens: ``[num_gens]`` accepted-token counts. Required + when replay bookkeeping is enabled. + position_source_is_token_count: Whether ``accepted_position_source`` + contains accepted-token counts instead of zero-based intermediate + state positions. When true, the kernel subtracts one while + computing the source position. + replay_pnat, replay_cache_buf_idx, dummy_request_mask: Optional replay + metadata. When supplied together, the kernel advances PNAT and the + active history buffer while promoting the convolution state. + + For each gen ``g``:: + + accepted_position = ( + accepted_position_source[g] - 1 + if position_source_is_token_count + else accepted_position_source[g] + ) + dst[:, dst_state_indices[g]] = intermediate[ + :, src_state_indices[g], accepted_position + ] + + Pure gather->scatter copy; bandwidth-bound (~85% of HBM peak), one launch. + """ + num_layers = dst.shape[0] + num_gens = src_state_indices.shape[0] + if num_gens == 0: + return + update_replay_bookkeeping = replay_pnat is not None + if update_replay_bookkeeping: + assert num_accepted_tokens is not None + assert replay_cache_buf_idx is not None + assert dummy_request_mask is not None + assert replay_step_width > 0 + assert replay_history_size >= replay_step_width + assert replay_pnat.dtype == torch.int32 + assert replay_cache_buf_idx.dtype == torch.int32 + assert dummy_request_mask.dtype == torch.bool + assert dummy_request_mask.numel() >= num_gens + else: + # Compile-time-disabled pointer arguments still require tensors. + num_accepted_tokens = accepted_position_source + assert replay_cache_buf_idx is None + assert dummy_request_mask is None + replay_pnat = dst_state_indices + replay_cache_buf_idx = dst_state_indices + dummy_request_mask = dst_state_indices + count = 1 + for s in dst.shape[2:]: + count *= s + # Grid dim order matters: dim 0 -> CUDA grid.x (limit 2^31-1), dim 1 -> + # grid.y (limit 65535). Put the (layer, gen) pairs in dim 0 since that is + # what grows with batch size (num_layers*num_gens stays far below 2^31 for + # any real config); the tile count in dim 1 is small (~count/BLOCK). Do NOT + # swap them: pairs would hit the 65535 y-limit at num_layers*num_gens>65535. + grid = (num_layers * num_gens, triton.cdiv(count, BLOCK)) + _promote_mamba_state_kernel[grid]( + intermediate, + dst, + src_state_indices, + accepted_position_source, + num_accepted_tokens, + dst_state_indices, + replay_pnat, + replay_cache_buf_idx, + dummy_request_mask, + num_gens, + count, + intermediate.stride(0), + intermediate.stride(1), + intermediate.stride(2), + dst.stride(0), + dst.stride(1), + POSITION_SOURCE_IS_TOKEN_COUNT=position_source_is_token_count, + REPLAY_STEP_WIDTH=replay_step_width, + REPLAY_HISTORY_SIZE=replay_history_size, + UPDATE_REPLAY_BOOKKEEPING=update_replay_bookkeeping, + BLOCK=BLOCK, + ) + + +def _mamba_snapshot_rule_counts( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], + tokens_per_block: int, +) -> Tuple[int, int]: + """Return reachable fixed rules and their partial-block upper bound.""" + if not kv_cache_config.enable_block_reuse: + return 0, 0 + + num_rules = 0 + num_unaligned_rules = 0 + state_config = kv_cache_config.mamba_state_config + for offset in set(state_config.additional_snapshot_offsets_from_start): + if max_seq_len is not None and offset > max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += int(offset % tokens_per_block != 0) + for offset in set(state_config.additional_snapshot_offsets_from_end): + # A from-end offset is reachable iff some valid prompt is longer than + # the offset. Its absolute alignment depends on that prompt. + if max_seq_len is not None and offset >= max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += 1 + return num_rules, num_unaligned_rules + + +def _mamba_regular_snapshot_interval( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], +) -> Optional[int]: + if not kv_cache_config.enable_block_reuse: + return None + interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval + if interval is None or interval <= 0: + return None + if max_seq_len is not None and interval > max_seq_len: + return None + return interval + + +def _get_local_mamba_cache_layout( + model_config, + mapping: Mapping, + *, + spec_config=None, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, +): + """Return normalized params and local Mamba/attention layer counts. + + Cache construction and affine sizing must follow the model's PP layout: + partition base layers first, then place appended speculative layers on the + last PP rank. The normalized params retain target masks and the appended + draft-layer count so estimation selects the same combined or per-manager + layout as runtime. + """ + from tensorrt_llm._torch.pyexecutor.config_utils import \ + extract_mamba_kv_cache_params + + params = extract_mamba_kv_cache_params( + model_config.pretrained_config, + spec_config=spec_config, + quant_config=model_config.quant_config, + ) + mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + combined_layer_mask = [ + is_mamba or is_attention for is_mamba, is_attention in zip( + mamba_layer_mask, full_attention_layer_mask) + ] + local_layer_indices, _ = get_pp_layers( + sum(combined_layer_mask), + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + local_mamba_layers = sum(mamba_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + local_attention_layers = sum(full_attention_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + return params, local_mamba_layers, local_attention_layers + + +def _estimate_mamba_hybrid_cache_cost( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + tokens_per_block: int, + max_seq_len: Optional[int], + num_reserved_dummy_slots: int, + include_explicit_snapshots: bool, + cap_partial_attention_snapshots: bool, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, + attention_slope_override: Optional[int] = None, + attention_page_bytes_override: int = 0, + draft_scratch_tokens: int = 0, + **kwargs, +) -> Tuple[int, int]: + spec_config = kwargs.get("spec_config") + params, local_mamba_layers, local_attention_layers = ( + _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + )) + attention_slope = attention_slope_override + if attention_slope is None: + attention_slope = (KVCacheManager.get_cache_size_per_token( + model_config, + mapping, + num_layers=local_attention_layers, + **kwargs, + ) if local_attention_layers > 0 else 0) + state_bytes_per_rank = (local_mamba_layers * + params.get_states_bytes_per_layer(mapping)) + max_resident_sequences = max_batch_size * mapping.pp_size + + if include_explicit_snapshots: + fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( + kv_cache_config, max_seq_len, tokens_per_block) + else: + fixed_rules = 0 + unaligned_fixed_rules = 0 + fixed_state_slots = (max_resident_sequences + num_reserved_dummy_slots + + max_resident_sequences * fixed_rules) + attention_block_bytes = attention_slope * tokens_per_block + + interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) + has_unaligned_periodic_snapshot = (interval is not None + and interval % tokens_per_block != 0) + if cap_partial_attention_snapshots: + # Snapshot alignment is unknown while estimating cache cost. Once a + # snapshot is possible, reserve one retained partial attention page + # per resident lineage. Dummy requests carry no attention capacity. + has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None + partial_attention_slots = (max_resident_sequences + if has_non_live_ssm_capacity else 0) + else: + partial_attention_slots = (max_resident_sequences * + unaligned_fixed_rules) + intercept = (fixed_state_slots * state_bytes_per_rank + + partial_attention_slots * + (attention_block_bytes + attention_page_bytes_override)) + + if interval is None: + regular_slope = 0 + else: + regular_slope = math.ceil(state_bytes_per_rank / interval) + if (has_unaligned_periodic_snapshot + and not cap_partial_attention_snapshots): + regular_slope += math.ceil(attention_block_bytes / interval) + draft_layer_sizes = _get_draft_layer_sizes(**kwargs) + if draft_layer_sizes: + # Hybrid target attention is full attention. Reserve capture/noise + # capacity only in its attention pools and the distinct draft pools; + # recurrent state retains the fixed/snapshot accounting above. + # Hybrid profiling bounds draft storage by full-history attention. + _, draft_slope, draft_fixed = _estimate_draft_cache_size_components( + draft_layer_sizes, + tokens_per_block, + generation_capacity_headroom=1, + scratch_tokens=draft_scratch_tokens, + helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, + ) + draft_fixed += draft_scratch_tokens * attention_slope + attention_slope += draft_slope + intercept += max_batch_size * draft_fixed + return attention_slope + regular_slope, intercept + + +class CppMambaHybridCacheManager(KVCacheManager, MambaHybridCacheManager): + """Hybrid cache manager storing mamba states inside the KVCacheManager pool. + + Both KV cache blocks and recurrent state blocks are managed by the unified + C++ KVCacheManager, enabling block reuse / prefix caching across attention + and mamba layers. This compatibility manager remains available through the + manager preference override and legacy disaggregated routing. + + """ + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + layer_mask: Optional[ + List[bool]] = None, # this is the full attention layer mask + is_estimating_kv_cache: bool = False, + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + model_type: str = "nemotron_hybrid", + **kwargs, + ) -> None: + # 3 kinds of layers: + # 1) Mamba layers (mamba_layer_mask is True) + # 2) Full attention layers (full_attention_layer_mask is True) + # 3) Not managed layers (both masks are False) + total_layers = len(mamba_layer_mask) + if layer_mask is None: + full_attention_layer_mask = [False] * total_layers + elif len(layer_mask) != total_layers: + raise ValueError( + f"layer_mask length ({len(layer_mask)}) must match " + f"mamba_layer_mask length ({total_layers})") + else: + full_attention_layer_mask = list(layer_mask) + layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] + for i in range(total_layers) + ] + # PP sharding is done across all layers. + # This is called again in the super().__init__, but we want it to run first + # to set up mtp states before the C++ backend is initialized. + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + ) + self.mamba_pp_layers = [ + layer_idx for layer_idx in self.pp_layers + if mamba_layer_mask[layer_idx] + ] + self.local_num_mamba_layers = len(self.mamba_pp_layers) + self.requests = [] + # Seed externally visible mamba fields before any early return so that + # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work + # on ranks with no local mamba layers. + self._use_replay_state_update = use_replay_state_update + self._use_gdn_cached_replay_all_layer_commit = ( + use_replay_state_update and model_type == "qwen3_next" + and self.local_num_mamba_layers > 0) + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None else None) + # Same allocation gate as PythonMambaCacheManager: the rand_seed + # buffer must exist whenever SR can fire, not only on the replay path. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh context-slot assignment; combined + # with the slot index and rank offset to produce reproducible per-slot + # seed values without any torch.randint. + self._seed_request_counter = 0 + self.ssm_state_dtype = mamba_ssm_cache_dtype + # Keep the shared Mamba interface valid on PP ranks that do not own a + # local Mamba layer. + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + + if self.local_num_mamba_layers == 0: + logger.info( + "No local mamba layers for this rank, skipping mamba cache initialization" + ) + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=full_attention_layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, + **kwargs, + ) + # PP ranks replay the same scheduling decisions, so a rank without + # local Mamba layers must still publish the configured boundaries. + self.kv_cache_config = kv_cache_config + self.linear_attention_metadata = LinearAttentionMetadata() + self.linear_attention_metadata.states_snapshot_interval = ( + kv_cache_config.mamba_state_config.periodic_snapshot_interval + if kv_cache_config.enable_block_reuse else 0) + return + + # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) + tp_size = _mamba_effective_tp_size(mapping) + d_inner = mamba_head_dim * mamba_num_heads + conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state + nheads = mamba_num_heads + assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if use_replay_state_update: + assert mamba_n_groups % tp_size == 0, \ + "replay state update requires mamba_n_groups divisible by tp_size" + self._n_groups_per_rank = mamba_n_groups // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + self.conv_state_dtype = mamba_cache_dtype + + # Store GLOBAL (pre-TP-division) mamba params for disagg RnnModelConfig. + # d_inner is computed at the top of __init__ before TP division. + d_inner_global = d_inner # = mamba_head_dim * mamba_num_heads (GLOBAL) + conv_dim_global = d_inner_global + 2 * mamba_n_groups * mamba_d_state + self._rnn_d_state = mamba_d_state + self._rnn_d_conv = mamba_d_conv + self._rnn_num_heads = mamba_num_heads # GLOBAL + self._rnn_n_groups = mamba_n_groups + self._rnn_head_dim = mamba_head_dim + self._rnn_hidden_size = d_inner_global # GLOBAL = head_dim * num_heads + self._rnn_conv_dim_size = conv_dim_global # GLOBAL conv dim + self._rnn_num_layers = mamba_num_layers + + # Conv section layout for section-aware split/concat in TP mismatch. + # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; + # we just need to tell C++ which ordering to use. + self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) + + total_bytes = self.ssm_bytes + self.conv_bytes + if total_bytes % self.ssm_state_dtype.itemsize != 0: + raise RuntimeError( + f"Total state bytes ({total_bytes}) not divisible by " + f"ssm_state_dtype size ({self.ssm_state_dtype.itemsize})") + if total_bytes % self.conv_state_dtype.itemsize != 0: + raise RuntimeError( + f"Total state bytes ({total_bytes}) not divisible by " + f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") + if self.ssm_bytes % self.conv_state_dtype.itemsize != 0: + raise RuntimeError( + f"SSM state bytes ({self.ssm_bytes}) not divisible by " + f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") + self.linear_attention_metadata = LinearAttentionMetadata() + self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value + self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes + self.linear_attention_metadata.states_snapshot_interval = ( + kv_cache_config.mamba_state_config.periodic_snapshot_interval + if kv_cache_config.enable_block_reuse else 0) + # RNN model params for disagg TP-mismatch split/concat. + conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} + self.linear_attention_metadata.rnn_num_heads = self._rnn_num_heads + self.linear_attention_metadata.rnn_head_dim = self._rnn_head_dim + self.linear_attention_metadata.rnn_d_state = self._rnn_d_state + self.linear_attention_metadata.rnn_d_conv = self._rnn_d_conv + self.linear_attention_metadata.rnn_n_groups = self._rnn_n_groups + self.linear_attention_metadata.rnn_conv_section_layout = conv_section_map.get( + self._rnn_conv_section_layout, 0) + self.linear_attention_metadata.rnn_ssm_bytes = self.ssm_bytes + self.linear_attention_metadata.rnn_ssm_dtype_size = self.ssm_state_dtype.itemsize + self.linear_attention_metadata.rnn_conv_dtype_size = self.conv_state_dtype.itemsize + kv_cache_config = kv_cache_config.model_copy(deep=True) + if kv_cache_config.enable_partial_reuse: + logger.warning( + "Partial reuse is not supported for mamba hybrid models, disabling partial reuse" + ) + kv_cache_config.enable_partial_reuse = False + + kv_cache_config.max_attention_window = [] + for i in range(len(layer_mask)): + if layer_mask[i]: + kv_cache_config.max_attention_window.append( + LinearCacheType.RECURRENT_STATES. + value if mamba_layer_mask[i] else max_seq_len) + + recurrent_states_window = LinearCacheType.RECURRENT_STATES.value + local_windows = { + recurrent_states_window + if mamba_layer_mask[layer_idx] else max_seq_len + for layer_idx in self.pp_layers + } + kwargs["pool_configurations"] = [ + PoolConfiguration( + window_size=window_size, + head_dim=head_dim, + dtype=torch_dtype_to_binding(self.ssm_state_dtype) + if window_size == recurrent_states_window else dtype, + ) for window_size in sorted(local_windows) + ] + + # Normalize num_kv_heads to a per-layer list and zero out mamba + # layer positions: those layers carry SSM/conv state instead of KV + # heads, so the parent KV cache should not allocate KV head storage + # for them. + if isinstance(num_kv_heads, int): + per_layer_kv_heads = [num_kv_heads] * total_layers + else: + if len(num_kv_heads) != total_layers: + raise ValueError( + f"num_kv_heads list length ({len(num_kv_heads)}) does not " + f"match total layers ({total_layers})") + per_layer_kv_heads = list(num_kv_heads) + for i, is_mamba in enumerate(mamba_layer_mask): + if is_mamba: + per_layer_kv_heads[i] = 0 + + self._setup_mtp_intermediate_states(spec_config, max_batch_size) + + # pass remaining arguments to super class + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=mamba_num_layers + num_layers, + num_kv_heads=per_layer_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, + linear_attention_metadata=self.linear_attention_metadata, + **kwargs, + ) + + assert self.local_num_mamba_layers > 0, "At least one mamba layer is required" + self.mamba_layer_offsets = {} + for idx, layer_id in enumerate(self.mamba_pp_layers): + self.mamba_layer_offsets[layer_id] = idx + + self.host_block_offsets = torch.zeros([ + self.impl.num_pools, self.max_batch_size, 2, self.max_blocks_per_seq + ], + dtype=torch.int32, + device="cpu") + self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ + self.layer_offsets[self.mamba_pp_layers[0]]][0] + + self.cuda_state_indices = torch.zeros([self.max_batch_size], + dtype=torch.int32, + device="cuda") + self._host_state_indices = torch.zeros([self.max_batch_size], + dtype=torch.int32, + pin_memory=prefer_pinned()) + self._row_indices = torch.arange(self.max_batch_size, + dtype=torch.long, + device="cpu") + self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} + self.kv_cache_config = kv_cache_config + self.is_estimating_kv_cache = is_estimating_kv_cache + + self._setup_states() + self._setup_replay_buffers(spec_config) + if use_replay_state_update and model_type == "qwen3_next": + logger.info_once( + "Configured GDN cached replay commit mode: small-batch fused, " + "large-batch all-layer", + key="gdn_cached_replay_commit_mode_fused", + ) + + @staticmethod + def get_cache_size_per_token( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs, + ): + """Affine memory model for the unified hybrid KV pool. + + Returns ``(slope_bytes_per_token, intercept_bytes)``: + + * ``slope`` = attention KV bytes per token plus conservative periodic + Mamba-state and partial-attention-snapshot costs. + * ``intercept`` = live and CUDA-graph dummy Mamba state. + + Memory budget -> max tokens then becomes + ``T = (budget - intercept) // slope`` instead of plain + ``T = budget // bytes_per_token``. + """ + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=1, + include_explicit_snapshots=False, + cap_partial_attention_snapshots=False, + **kwargs, + ) + + @property + def use_gdn_cached_replay_all_layer_commit(self) -> bool: + return self._use_gdn_cached_replay_all_layer_commit + + def _commit_gdn_cached_replay_history_layers( + self, + attn_metadata: "AttentionMetadata", + num_decodes: int, + ) -> None: + """Synchronously advance every local GDN checkpoint in one launch.""" + from tensorrt_llm._torch.modules.fla.cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers) + + if (not self._use_gdn_cached_replay_all_layer_commit + or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): + return + if (self.all_ssm_states is None or self.old_x is None + or self.old_B is None or self.old_dt is None + or self.replay_history_size is None): + raise RuntimeError( + "GDN cached replay all-layer commit requires replay state buffers." + ) + + mamba_metadata = attn_metadata.mamba_metadata + if mamba_metadata.replay_num_decodes != num_decodes: + raise RuntimeError( + "GDN replay metadata contains " + f"{mamba_metadata.replay_num_decodes} decode requests, " + f"but state update received {num_decodes}.") + commit_gdn_cached_replay_history_layers( + ssm_states=self.all_ssm_states, + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=mamba_metadata.replay_work_items[:num_decodes], + n_writes=mamba_metadata.replay_n_writes, + history_size=self.replay_history_size, + ) + + def shutdown(self): + # Release tensor views into the pool before the pool memory is freed, + # so their deleters don't see stale pointers. + self.all_ssm_states = None + self.all_conv_states = None + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + super().shutdown() + + def add_dummy_requests( + self, + request_ids: List[int], + # Note that token_nums should be past_kv_len + input_len (without + # spec decoding). The draft tokens will be added in this function, + # so we don't need to take care of it in the caller. When preparing + # token_nums, we should not take the draft tokens into account, so + # don't use the kv_cache_manager.max_seq_len, which includes both + # extra tokens and draft tokens. + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + prepare_resource: bool = True, + max_num_draft_tokens: int = 0, + kv_reserve_draft_tokens: Optional[int] = None, + use_mrope: bool = False, + max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, + # For capturable drafting loops. During normal inference, the draft model always + # has enough KV cache space to fit all of our draft tokens. During warmup, however, + # we need to make the KV cache manager aware that multiple autoregressive steps will + # occur. + num_extra_decoding_steps: int = 0, + draft_kv_cache_manager: Optional[KVCacheManager] = None, + ) -> List[LlmRequest]: + requests = super().add_dummy_requests( + request_ids=request_ids, + token_nums=token_nums, + is_gen=is_gen, + prepare_resource=prepare_resource, + max_num_draft_tokens=max_num_draft_tokens, + kv_reserve_draft_tokens=kv_reserve_draft_tokens, + use_mrope=use_mrope, + max_beam_width=max_beam_width, + encoder_output_lens=encoder_output_lens, + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if requests: + self.requests.extend(requests) + # Only process the newly added requests, not all of self.requests. + # self.requests may contain stale entries from the previous + # _prepare_resources call (e.g., disagg transfer-pending requests) + # that would exceed max_batch_size and cause out-of-bounds access. + self._setup_state_indices(requests) + return requests + + def update_resources(self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None): + super().update_resources(scheduled_batch, attn_metadata, + kv_cache_dtype_byte_size) + + @nvtx_range("hybrid_prepare_resources") + def _prepare_resources(self, scheduled_batch: ScheduledRequests): + self.requests = scheduled_batch.context_requests + \ + scheduled_batch.generation_requests + # Issue all async block onboards; defer the syncTransfers() (refresh_blocks) + # until just before forward so the CPU-side prep (attn metadata, input + # tensor assembly, draft model prep, etc.) can overlap with the in-flight + # cudaMemcpyAsync calls instead of being serialized behind them. + # copy_linear_attention_block returns True iff it actually issued a copy; + # we skip refresh_blocks entirely when nothing was scheduled. + self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( + self.requests) + self._setup_state_indices() + self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + self._prepare_resources(scheduled_batch) + + @nvtx_range("hybrid_flush_state_transfers") + def flush_state_transfers(self) -> None: + """Complete any deferred state-block onboards scheduled by + prepare_resources(). Must be called before forward() reads recurrent + state blocks. Cheap no-op when nothing was scheduled. + """ + if getattr(self, "_pending_state_transfers", False): + self.impl.refresh_blocks() + self._pending_state_transfers = False + + @nvtx_range("hybrid_update_mamba_states") + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None): + if self.local_num_mamba_layers == 0: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + + num_gens] + # Dynamic tree supplies explicit tree-node leaf positions. Linear MTP + # passes accepted-token counts and lets the promotion kernel subtract + # one, avoiding a separate GPU subtraction/cast. + if accepted_leaf_positions is None: + accepted_position_source = accepted_tokens + position_source_is_token_count = True + else: + accepted_position_source = accepted_leaf_positions.to(torch.int32) + position_source_is_token_count = False + # Match the API of MambaCacheManager.update_mamba_states: callers + # may pass per-request state slot indices explicitly (e.g. MTP via + # attn_metadata.mamba_metadata.state_indices). Fall back to this + # manager's own slot mapping when not provided. + if state_indices is None: + state_indices = self.get_state_indices() + state_indices_d = state_indices[num_contexts:num_contexts + + num_gens].to(torch.int32) + src_state_indices = self.intermediate_state_indices[:num_gens] + + # The accepted SSM/conv promotion is a bandwidth-bound gather->scatter + # into the dtype-reinterpreted, strided C++ pool view. torch.compile + # can't handle the dtype-view mutation (and Inductor chokes on the + # uint8 pool with "XBLOCK too large"), so a dedicated Triton kernel + # writes through the view's real strides (~85% of HBM peak, one launch + # per state). + if self._use_replay_state_update: + # Every GDN layer has finished reading the old checkpoint and + # writing its candidate history. Advance all local checkpoints + # in one launch before PNAT and the active history buffer change. + self._commit_gdn_cached_replay_history_layers( + attn_metadata, num_gens) + assert self._dummy_request_mask is not None + else: + # Legacy: copy the accepted SSM state from the intermediate buffer. + _promote_mamba_state_triton( + self.all_ssm_states, + self.intermediate_ssm_states, + src_state_indices, + accepted_position_source, + state_indices_d, + position_source_is_token_count=position_source_is_token_count) + + # Conv: both paths save all intermediate conv windows, carry over the + # accepted one. The replay path also advances PNAT and the active cache + # buffer in this same launch. + replay_dummy_request_mask = None + if self._use_replay_state_update: + replay_dummy_request_mask = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + _promote_mamba_state_triton( + self.all_conv_states, + self.intermediate_conv_states, + src_state_indices, + accepted_position_source, + state_indices_d, + num_accepted_tokens=accepted_tokens, + position_source_is_token_count=position_source_is_token_count, + replay_pnat=(self.prev_num_accepted_tokens + if self._use_replay_state_update else None), + replay_cache_buf_idx=(self.cache_buf_idx + if self._use_replay_state_update else None), + dummy_request_mask=replay_dummy_request_mask, + replay_step_width=(self.replay_step_width + if self._use_replay_state_update else 0), + replay_history_size=(self.replay_history_size + if self._use_replay_state_update else 0)) + + def get_num_available_tokens(self, + token_num_upper_bound: int, + max_num_draft_tokens: int = 0, + **kwargs) -> int: + # Base bound from attention KV cache pool (the parent's behaviour). + result = super().get_num_available_tokens(token_num_upper_bound, + max_num_draft_tokens, + **kwargs) + # Also bound by the recurrent-state pool capacity: each request needs + # roughly ceil(N / states_snapshot_interval) recurrent-state blocks + # (+1 for the corner case where N is a multiple of tokens_per_block). + # When block reuse is disabled, only one snapshot is needed per + # request, so no additional capping is required here. + interval = (self.linear_attention_metadata.states_snapshot_interval + if self.linear_attention_metadata is not None else 0) + # Attention-only PP ranks keep the interval so every rank publishes + # identical scheduling boundaries, but they have no recurrent-state + # pool whose capacity should constrain their attention KV cache. + if self.local_num_mamba_layers > 0 and interval and interval > 0: + stats = self.impl.get_kv_cache_stats() + rs_free = stats.num_free_blocks_per_window_size.get( + LinearCacheType.RECURRENT_STATES.value, 0) + # Reserve 1 block for the always-allocated last block (corner case + # / final live state) so we don't promise more tokens than the + # pool can actually back at allocation time. + usable_rs_blocks = max(0, rs_free - 1) + rs_token_cap = usable_rs_blocks * interval + result = min(result, rs_token_cap) + return max(result, 0) + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + if request in self.requests: + self.requests.remove(request) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) + + def _setup_state_indices(self, requests=None) -> None: + if self.local_num_mamba_layers == 0: + return + if requests is None: + requests = self.requests + block_indices = [] + for req in requests: + if req.is_context_finished: + next_step = self.get_num_tokens(req) - 1 + elif self.kv_cache_config.enable_block_reuse: + next_step = (req.context_current_position - 1 + + req.context_chunk_size) + else: + next_step = req.prompt_len - 1 + block_indices.append(next_step // self.tokens_per_block) + self.impl.copy_batch_block_offsets( + self.host_block_offsets, [req.py_request_id for req in requests], 1, + 0) + max_blocks = self.blocks_per_window[ + LinearCacheType.RECURRENT_STATES.value][0] + n = len(requests) + self._host_state_indices.zero_() + if n > 0: + # Vectorized gather: replace per-element Python loop with a single + # tensor index op. block_indices is a small Python list so + # torch.tensor() conversion here is O(n) at C level, much cheaper + # than n round-trips through Python indexing. + bi = torch.tensor(block_indices, dtype=torch.long) + rows = self._row_indices[:n] + # host_block_offsets: [num_pools, max_batch_size, 2, max_blocks_per_seq] + values = self.host_block_offsets[self.recurrent_states_pool_index, + rows, 0, bi] + invalid_mask = (values < 0) | (values >= max_blocks) + # The C++ recurrent-state manager uses null page-table entries for + # logical blocks that are not snapshot boundaries. Usually a + # context chunk ends exactly at an allocated snapshot (or at the + # final live-state block), but the scheduler may shorten a chunk + # further when KV capacity is tight. In that case the Mamba + # kernel must keep accumulating into the next allocated snapshot + # or final block, matching copyLinearAttentionBlock(), which also + # skips placeholders when it advances the live state. + for bad_i in invalid_mask.nonzero( + as_tuple=False).flatten().tolist(): + req = requests[bad_i] + if req.is_context_finished: + continue + last_prompt_block = (req.prompt_len - + 1) // self.tokens_per_block + row = self.host_block_offsets[self.recurrent_states_pool_index, + rows[bad_i], 0] + candidates = row[block_indices[bad_i]:last_prompt_block + 1] + valid_candidates = ((candidates >= 0) + & (candidates < max_blocks)).nonzero( + as_tuple=False) + if valid_candidates.numel() > 0: + values[bad_i] = candidates[valid_candidates[0, 0]] + + invalid_mask = (values < 0) | (values >= max_blocks) + if invalid_mask.any(): + bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) + req = requests[bad_i] + value = int(values[bad_i]) + raise RuntimeError( + f"Invalid recurrent state block index {value} " + f"(expected 0 <= index < {max_blocks}) for request {bad_i}, " + f"prompt_len={req.prompt_len}, " + f"is_context_finished={req.is_context_finished}, " + f"context_current_position={req.context_current_position}, " + f"prepopulated_token_num={req.prepopulated_prompt_len}, " + f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " + f"block_index for next step is {block_indices[bad_i]}, " + "\nblock_ids=" + f"{self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" + ) + self._host_state_indices[:n] = values + + self.cuda_state_indices.copy_(self._host_state_indices, + non_blocking=True) + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask( + is_dummy if requests is + self.requests else [req.is_dummy for req in self.requests]) + + # Build request_id → pool block offset mapping so that + # get_state_indices can return indices in arbitrary request order. + # Bulk tolist avoids a per-request tensor-index + .item() round-trip. + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy + + def get_state_indices(self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None) -> list: + if self.local_num_mamba_layers == 0: + # Mamba metadata is prepared on every PP rank even when this rank + # owns only attention layers. No local kernel consumes these + # indices, so avoid consulting state that is intentionally absent. + return [0] * len(request_ids) if request_ids is not None else [] + if request_ids is not None: + # Return indices in the order of the caller's request_ids, + # not the internal self.requests order. This is critical when + # the batch is reordered after prepare_resources (e.g. disagg + # serving sorts generation_requests by py_batch_idx). + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + return self.cuda_state_indices + + def _setup_states(self) -> None: + # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) + pool: torch.Tensor = self.impl.get_recurrent_states_pool().view( + torch.uint8).reshape(self.local_num_mamba_layers, -1, + self.ssm_bytes + self.conv_bytes) + num_blocks_in_pool = pool.shape[1] + self.all_ssm_states = pool[:, :, :self.ssm_bytes].view( + self.ssm_state_dtype).view( + [self.local_num_mamba_layers, num_blocks_in_pool] + + self.ssm_state_shape) + self.all_conv_states = pool[:, :, self.ssm_bytes:self.ssm_bytes + + self.conv_bytes].view( + self.conv_state_dtype).view([ + self.local_num_mamba_layers, + num_blocks_in_pool + ] + self.conv_state_shape) + self.all_ssm_states.zero_() + self.all_conv_states.zero_() + + def _setup_replay_buffers(self, spec_config) -> None: + cache_size = self.all_ssm_states.shape[1] + device = self.all_ssm_states.device + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): + return + + self._dummy_request_mask = torch.zeros(self.max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros( + self.max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) + + +class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): + """Hybrid Mamba cache manager backed by KVCacheManagerV2. + + Attention KV pages and Mamba recurrent-state pages are both owned by the + Python V2 cache manager. Mamba layers are represented as V2 SSM layers, + while this wrapper exposes the state tensors and slot indices expected by + the PyTorch Mamba kernels. + """ + + _supports_additional_snapshot_offsets = True + # Recurrent-state snapshots use a specialized commit/history protocol, so + # keep main-like reuse endpoints and the existing unpaired draft path. + _supports_reuse_match_backoff = False + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + layer_mask: Optional[List[bool]] = None, + is_estimating_kv_cache: bool = False, + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", + **kwargs, + ) -> None: + if conv_state_layout not in ("x_b_c", "q_k_v"): + raise ValueError( + f"Unsupported convolution state layout: {conv_state_layout!r}") + if "model_type" in kwargs: + # The V1 managers select the conv-state layout by model_type; this + # class takes it explicitly. Silently absorbing model_type here + # means a caller's layout request would be dropped on the floor. + raise TypeError( + "MambaHybridCacheManagerV2 does not accept 'model_type' " + f"(got {kwargs['model_type']!r}); pass " + "conv_state_layout='x_b_c' or 'q_k_v' instead") + total_layers = len(mamba_layer_mask) + if layer_mask is None: + full_attention_layer_mask = [False] * total_layers + elif len(layer_mask) != total_layers: + raise ValueError( + f"layer_mask length ({len(layer_mask)}) must match " + f"mamba_layer_mask length ({total_layers})") + else: + full_attention_layer_mask = list(layer_mask) + + combined_layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] + for i in range(total_layers) + ] + + self._mamba_layer_mask = list(mamba_layer_mask) + self._use_replay_state_update = use_replay_state_update + self._use_gdn_cached_replay_all_layer_commit = (use_replay_state_update + and conv_state_layout + == "q_k_v") + self._gdn_cached_replay_state_descriptors: Optional[torch.Tensor] = None + self._gdn_cached_replay_state_strides: Optional[Tuple[int, int, + int]] = None + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None else None) + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + self._seed_request_counter = 0 + self._recurrent_evicted_blocks_total = 0 + self._recurrent_onboarded_blocks_total = 0 + self._recurrent_dropped_blocks_total = 0 + self._recurrent_status_logged = False + # Branch points keyed by request id. prepare_expect_snapshot_points() + # overwrites the per-request list on every scheduler pass, so the point + # has to be re-merged from here rather than stored only on the request. + self._branch_snapshot_points: Dict[int, int] = {} + self._snapshot_pruned_tokens_total = 0 + self._page_pruned_tokens_total = 0 + self._branch_snapshots_taken_total = 0 + self._branch_snapshots_skipped_total: Dict[str, int] = {} + num_cuda_graph_padding_dummy_slots = ( + _get_num_cuda_graph_padding_dummy_slots(spec_config, + max_batch_size)) + self._num_reserved_dummy_slots = (num_cuda_graph_padding_dummy_slots + + int(mapping.enable_attention_dp)) + self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype + is not None else mamba_cache_dtype) + self.conv_state_dtype = mamba_cache_dtype + + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + self.mamba_pp_layers = [ + layer_idx for layer_idx in self.pp_layers + if mamba_layer_mask[layer_idx] + ] + self.local_num_mamba_layers = len(self.mamba_pp_layers) + self._use_gdn_cached_replay_all_layer_commit = ( + self._use_gdn_cached_replay_all_layer_commit + and self.local_num_mamba_layers > 0) + + if self.local_num_mamba_layers > 0: + tp_size = _mamba_effective_tp_size(mapping) + d_inner = mamba_head_dim * mamba_num_heads + grouped_state_dim = mamba_n_groups * mamba_d_state + conv_dim = d_inner + 2 * grouped_state_dim + nheads = mamba_num_heads + assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if kwargs.get("is_disagg", + False) and grouped_state_dim % tp_size != 0: + raise ValueError( + "Disaggregated Mamba transfer requires each convolution " + "state section to be divisible by tp_size") + if use_replay_state_update: + assert mamba_n_groups % tp_size == 0, \ + "replay state update requires mamba_n_groups divisible by tp_size" + self._n_groups_per_rank = mamba_n_groups // tp_size + d_inner_local = d_inner // tp_size + grouped_state_dim_local = grouped_state_dim // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + # TP-mismatch disaggregated transfers must split the flat + # convolution state at its true semantic boundaries. Mamba2 stores + # [x | B | C], while GDN stores [Q | K | V]. The large section is + # therefore first for Mamba2 and last for GDN. + if conv_state_layout == "x_b_c": + self.conv_section_dims = [ + d_inner_local, + grouped_state_dim_local, + grouped_state_dim_local, + ] + else: + self.conv_section_dims = [ + grouped_state_dim_local, + grouped_state_dim_local, + d_inner_local, + ] + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) + else: + logger.info( + "No local mamba layers for this rank, skipping mamba state views" + ) + self._n_groups_per_rank = 0 + self.conv_state_shape = [] + self.ssm_state_shape = [] + self.conv_section_dims = [] + self.ssm_bytes = 0 + self.conv_bytes = 0 + + if isinstance(num_kv_heads, int): + per_layer_kv_heads = [num_kv_heads] * total_layers + else: + if len(num_kv_heads) != total_layers: + raise ValueError( + f"num_kv_heads list length ({len(num_kv_heads)}) does not " + f"match total layers ({total_layers})") + per_layer_kv_heads = list(num_kv_heads) + for i, is_mamba in enumerate(mamba_layer_mask): + if is_mamba: + per_layer_kv_heads[i] = 0 + + self._setup_mtp_intermediate_states(spec_config, max_batch_size) + + kv_cache_config = kv_cache_config.model_copy(deep=True) + if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: + block_reuse_config = kv_cache_config.block_reuse_config + block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) + if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: + # SSM reuse is valid only at explicit snapshot boundaries. + kv_cache_config.block_reuse_config = block_reuse_config.model_copy( + update={"policy": BlockReusePolicy.PER_REQUEST.value}) + self._attention_cache_layout_policy = None + pretrained_config = kwargs.get("pretrained_config") + if (dtype == DataType.NVFP4 and kv_cache_type == CacheTypeCpp.SELFKONLY + and getattr(pretrained_config, "kv_lora_rank", None) is not None + and getattr(pretrained_config, "qk_rope_head_dim", + None) is not None): + from tensorrt_llm._torch.attention_backend.fp4_mla.cache_manager import \ + Fp4MlaV2CacheLayoutPolicy + + kv_cache_config = Fp4MlaV2CacheLayoutPolicy.configure_manager( + self, + kv_cache_config, + kv_cache_type, + dtype=dtype, + tokens_per_block=tokens_per_block, + head_dim=head_dim, + pretrained_config=pretrained_config, + spec_config=spec_config, + is_disagg=kwargs.get("is_disagg", False), + ) + self.kv_cache_config = kv_cache_config + + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=mamba_num_layers + num_layers, + num_kv_heads=per_layer_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=combined_layer_mask, + is_draft=is_draft, + is_estimating_kv_cache=is_estimating_kv_cache, + num_reserved_index_slots=self._num_reserved_dummy_slots, + **kwargs, + ) + + if self._attention_cache_layout_policy is not None: + self._attention_cache_layout_policy.finalize_manager(self) + + self.mamba_layer_offsets = { + layer_id: offset + for offset, layer_id in enumerate(self.mamba_pp_layers) + } + self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} + + state_index_capacity = (self.max_batch_size + + self._num_reserved_dummy_slots) + self.cuda_state_indices = torch.zeros([state_index_capacity], + dtype=torch.int32, + device="cuda") + self._host_state_indices = torch.zeros([state_index_capacity], + dtype=torch.int32, + pin_memory=prefer_pinned()) + + if self.local_num_mamba_layers > 0: + first_mamba_local_layer = self.layer_offsets[ + self.mamba_pp_layers[0]] + self.ssm_layer_group_id = self.impl.get_layer_group_id( + LayerId(first_mamba_local_layer)) + self._ssm_page_index_scale = self.impl.get_page_index_scale( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_pages = self.impl.get_page_index_upper_bound( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // + self._ssm_page_index_scale) + required_live_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + if num_ssm_slots < required_live_slots: + KVCacheManagerV2.shutdown(self) + raise ValueError( + "The V2 Mamba state pool has only " + f"{num_ssm_slots} slots but needs at least " + f"{required_live_slots} live/dummy slots. Increase the " + "KV cache budget or allocate a larger Mamba pool_ratio.") + self._setup_states() + self._setup_replay_buffers(spec_config) + if self._use_gdn_cached_replay_all_layer_commit: + state_layout = ("affine" + if self._gdn_cached_replay_state_strides + is not None else "indirect") + logger.info_once( + "Configured GDN cached replay commit mode for V2: " + "small-batch fused, large-batch all-layer; " + f"state layout: {state_layout}", + key="gdn_cached_replay_v2_commit_mode_fused", + ) + else: + self.ssm_layer_group_id = None + self._ssm_page_index_scale = 1 + self.all_ssm_states = [] + self.all_conv_states = [] + self._setup_replay_buffers(spec_config) + + @property + def use_gdn_cached_replay_all_layer_commit(self) -> bool: + return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) + + def _commit_gdn_cached_replay_history_layers( + self, + attn_metadata: "AttentionMetadata", + num_decodes: int, + ) -> None: + """Advance all V2 GDN checkpoints through their state views.""" + from tensorrt_llm._torch.modules.fla.cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers) + + if (not getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) + or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): + return + if (not self.all_ssm_states + or (self._gdn_cached_replay_state_descriptors is None + and self._gdn_cached_replay_state_strides is None) + or self.old_x is None or self.old_B is None + or self.old_dt is None or self.replay_history_size is None): + raise RuntimeError( + "GDN cached replay all-layer commit requires V2 replay state buffers." + ) + + mamba_metadata = attn_metadata.mamba_metadata + if mamba_metadata.replay_num_decodes != num_decodes: + raise RuntimeError( + "GDN replay metadata contains " + f"{mamba_metadata.replay_num_decodes} decode requests, " + f"but state update received {num_decodes}.") + state_strides = self._gdn_cached_replay_state_strides + commit_gdn_cached_replay_history_layers( + ssm_states=self.all_ssm_states[0], + ssm_state_descriptors=(self._gdn_cached_replay_state_descriptors), + ssm_state_num_layers=(state_strides[0] + if state_strides is not None else None), + ssm_state_layer_stride=(state_strides[1] + if state_strides is not None else None), + ssm_state_slot_stride=(state_strides[2] + if state_strides is not None else None), + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=mamba_metadata.replay_work_items[:num_decodes], + n_writes=mamba_metadata.replay_n_writes, + history_size=self.replay_history_size, + ) + + def _create_kv_cache( + self, + request_id: int, + lora_task_id: Optional[int], + input_tokens: Optional[Sequence[TokenIdExt]], + *, + cache_salt: Optional[str] = None, + is_dummy: bool = False, + expected_prompt_length: Optional[int] = None, + ugpu_id: Optional[int] = None, + ) -> Optional[_KVCache]: + kv_cache = super()._create_kv_cache( + request_id, + lora_task_id, + input_tokens, + cache_salt=cache_salt, + is_dummy=is_dummy, + expected_prompt_length=expected_prompt_length, + ugpu_id=ugpu_id, + ) + if (self.mapping.rank == 0 and kv_cache is not None + and input_tokens is not None and not is_dummy + and self.local_num_mamba_layers > 0): + # Prefix lookup excludes the final prompt token because it must be + # recomputed by prefill. Restore it for the request-length metric. + request_total_tokens = len(input_tokens) + 1 + prefix_reuse_diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) + logger.debug( + f"[MambaHybridCacheManagerV2] prefix reuse rank={self.mapping.rank} " + f"request_id={request_id} " + f"request_total_tokens={request_total_tokens} " + f"longest_attention_match_tokens=" + f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_hybrid_pruning()} " + f"content_divergence_tokens=" + f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_pruning()} " + f"latest_recurrent_snapshot_tokens=" + f"{kv_cache.num_committed_tokens}") + return kv_cache + + def _skip_branch_snapshot(self, reason: str) -> None: + self._branch_snapshots_skipped_total[reason] = ( + self._branch_snapshots_skipped_total.get(reason, 0) + 1) + + def _format_branch_snapshot_counters(self) -> str: + """Attribute reuse loss between missing snapshots and evicted pages. + + A flat cache hit rate means something different depending on whether + branch points were never found or were found and not reused, so the + skip reasons are reported alongside the totals. + """ + skipped = self._branch_snapshots_skipped_total + return (f"snapshot_pruned_tokens={self._snapshot_pruned_tokens_total} " + f"page_pruned_tokens={self._page_pruned_tokens_total} " + f"branch_snapshots_taken={self._branch_snapshots_taken_total} " + f"branch_snapshots_skipped={sum(skipped.values())} " + f"branch_snapshots_skipped_by_reason=" + f"{dict(sorted(skipped.items()))}") + + def _record_branch_snapshot_point(self, req: LlmRequest, kv_cache: _KVCache, + num_lookup_tokens: Optional[int]) -> None: + """Snapshot where this request leaves the reuse tree, for its siblings.""" + if not self.kv_cache_config.mamba_state_config.enable_branch_snapshot: + return + if req.is_dummy_request: + return + if num_lookup_tokens is None: + self._skip_branch_snapshot("no_fresh_match") + # A fresh lookup can record a point before resume() fails. On the + # next attempt, reapply that point against the cache's authoritative + # reuse depth: the request cursor is not rewound until preparation + # succeeds and may still describe a deeper, abandoned attempt. + if req.py_request_id in self._branch_snapshot_points: + self._apply_branch_snapshot_point( + req, context_current_position=kv_cache.num_committed_tokens) + return + + diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) + divergence = diagnostics._get_num_reusable_tokens_before_pruning() + reused = kv_cache.num_committed_tokens + if self.local_num_mamba_layers > 0: + hybrid_depth = ( + diagnostics._get_num_reusable_tokens_before_hybrid_pruning()) + # Loss attributable to a missing recurrent snapshot, versus loss + # attributable to attention pages having been evicted. These call + # for different fixes, so they are counted apart. + self._snapshot_pruned_tokens_total += max(0, hybrid_depth - reused) + self._page_pruned_tokens_total += max(0, divergence - hybrid_depth) + + # The whole lookup range matched, so there is no fork here. + if divergence >= num_lookup_tokens: + self._skip_branch_snapshot("no_divergence") + return + # Align down: tokens past a block boundary come from a partial-match + # child rather than a confirmed shared prefix, and an unaligned commit + # takes the partial-block snapshot path. + point = (divergence // self.tokens_per_block) * self.tokens_per_block + if point == 0: + self._skip_branch_snapshot("aligned_to_zero") + return + # Reuse already reached the fork, so a snapshot there exists. + if point <= reused: + self._skip_branch_snapshot("already_reused") + return + # A branch point at or beyond the prompt end cannot improve sibling + # reuse. + if point >= req.prompt_len: + self._skip_branch_snapshot("at_prompt_end") + return + + self._branch_snapshot_points[req.py_request_id] = point + self._apply_branch_snapshot_point(req, context_current_position=reused) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + super().prepare_expect_snapshot_points(requests) + for request in requests: + self._align_fp4_mla_snapshot_points(request) + if (not self.enable_block_reuse or not self.kv_cache_config. + mamba_state_config.enable_branch_snapshot): + return + for request in requests: + self._apply_branch_snapshot_point(request) + + def _align_fp4_mla_snapshot_points(self, request: LlmRequest) -> None: + if self.dtype != DataType.NVFP4: + return + # NVFP4 KV is quantized in 16-token tiles, so FP4 context resumes only + # on a full cache block (partial-block reuse is disabled). Capture the + # recurrent state at that boundary; rounding only the reused length + # would restore a later token's state. + request.expect_snapshot_points = sorted({ + point // self.tokens_per_block * self.tokens_per_block + for point in request.expect_snapshot_points + if point >= self.tokens_per_block + }) + + def _apply_branch_snapshot_point( + self, + request: LlmRequest, + context_current_position: Optional[int] = None, + ) -> None: + # prompt_len is unioned in unconditionally: try_commit_blocks derives + # commit_limit from max(snapshot_points), so a lone branch point below + # the prompt end would silently stop commits there and lose the + # prompt-end snapshot. It also makes the flag usable on its own. + points = set(request.expect_snapshot_points) + # Keep the flag usable without any configured snapshot placement. If a + # placement is configured, preserve it: adding prompt_len can replace a + # nearby recurrent checkpoint in the same radix-tree block because a + # block stores one recurrent page per life cycle. + if not points: + points.add(request.prompt_len) + point = self._branch_snapshot_points.get(request.py_request_id) + current_position = (request.context_current_position + if context_current_position is None else + context_current_position) + if point is not None and point > current_position: + points.add(point) + request.expect_snapshot_points = sorted(points) + self._align_fp4_mla_snapshot_points(request) + + def get_iteration_stats(self) -> Optional[KVCacheV2IterationStatsReport]: + """Log recurrent-cache movement; this is KDA state for Kimi K3.""" + report = super().get_iteration_stats() + if report is None: + return None + + pool_group_ids = sorted({ + pool_group_id + for pool_group_id, _, kind in + self._stats_life_cycle_metadata().values() if kind == "ssm" + }) + if not pool_group_ids: + return report + pool_group_reports = [ + report.by_pool_group[pool_group_id] + for pool_group_id in pool_group_ids + ] + + stats = [ + pool_group_report.stats for pool_group_report in pool_group_reports + ] + evicted_blocks = sum(stat.iter_offload_blocks for stat in stats) + evicted_bytes = sum(stat.iter_offload_bytes for stat in stats) + onboarded_blocks = sum(stat.iter_onboard_blocks for stat in stats) + onboarded_bytes = sum(stat.iter_onboard_bytes for stat in stats) + dropped_blocks = sum(stat.iter_host_dropped_blocks for stat in stats) + dropped_bytes = sum(stat.iter_host_dropped_bytes for stat in stats) + has_movement = bool(evicted_blocks or onboarded_blocks + or dropped_blocks) + if has_movement: + self._recurrent_evicted_blocks_total += evicted_blocks + self._recurrent_onboarded_blocks_total += onboarded_blocks + self._recurrent_dropped_blocks_total += dropped_blocks + if (self.mapping.rank == 0 + and (has_movement or not self._recurrent_status_logged)): + logger.debug( + f"[MambaHybridCacheManagerV2] recurrent cache status " + f"rank={self.mapping.rank} pool_group_ids={pool_group_ids} " + f"evicted_recurrent_blocks={evicted_blocks} " + f"evicted_recurrent_bytes={evicted_bytes} " + f"onboarded_recurrent_blocks={onboarded_blocks} " + f"onboarded_recurrent_bytes={onboarded_bytes} " + f"dropped_recurrent_blocks={dropped_blocks} " + f"dropped_recurrent_bytes={dropped_bytes} " + f"total_evicted_recurrent_blocks=" + f"{self._recurrent_evicted_blocks_total} " + f"total_onboarded_recurrent_blocks=" + f"{self._recurrent_onboarded_blocks_total} " + f"total_dropped_recurrent_blocks=" + f"{self._recurrent_dropped_blocks_total} " + f"gpu_used_recurrent_blocks=" + f"{sum(stat.primary_used_num_blocks for stat in stats)} " + f"gpu_free_recurrent_blocks=" + f"{sum(stat.primary_free_num_blocks for stat in stats)} " + f"gpu_evictable_recurrent_blocks=" + f"{sum(stat.primary_evictable_num_blocks for stat in stats)} " + f"host_used_recurrent_blocks=" + f"{sum(stat.secondary_used_num_blocks for stat in stats)} " + f"host_free_recurrent_blocks=" + f"{sum(stat.secondary_free_num_blocks for stat in stats)} " + f"{self._format_branch_snapshot_counters()}") + self._recurrent_status_logged = True + return report + + @staticmethod + def get_cache_size_per_token(model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs): + spec_config = kwargs.get("spec_config") + num_reserved_dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( + spec_config, max_batch_size) + int(mapping.enable_attention_dp)) + attention_slope_override = None + attention_intercept = 0 + attention_page_bytes_override = 0 + config = model_config.pretrained_config + quant_config = model_config.quant_config + if (getattr(config, "kv_lora_rank", None) is not None + and quant_config is not None + and quant_config.quant_mode.has_fp4_kv_cache()): + from tensorrt_llm._torch.attention_backend.fp4_mla.cache_manager import \ + Fp4MlaKVCacheManagerV2 + + _, _, local_attention_layers = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=kwargs.get("is_draft", False), + use_separate_draft_kv_cache=kwargs.get( + "use_separate_draft_kv_cache", False), + ) + attention_slope_override, attention_intercept = ( + Fp4MlaKVCacheManagerV2.get_cache_size_per_token( + model_config, + mapping, + num_layers=local_attention_layers, + tokens_per_block=tokens_per_block, + max_batch_size=max_batch_size, + spec_config=spec_config, + is_disagg=kwargs.get("is_disagg", False), + )) + max_resident_sequences = max_batch_size * mapping.pp_size + if max_resident_sequences > 0: + attention_page_bytes_override = (attention_intercept // + max_resident_sequences) + + slope, intercept = _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=num_reserved_dummy_slots, + include_explicit_snapshots=True, + cap_partial_attention_snapshots=True, + attention_slope_override=attention_slope_override, + attention_page_bytes_override=attention_page_bytes_override, + **kwargs, + ) + return slope, intercept + attention_intercept + + def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: + if self._is_standalone_draft_layer(local_layer_idx): + return False + layer_id = self.pp_layers[local_layer_idx] + return layer_id < len( + self._mamba_layer_mask) and self._mamba_layer_mask[layer_id] + + def _get_pool_roles(self, + pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_standalone_draft_layer(layer_id): + return KVCacheManagerV2._get_pool_roles(self, pool_id) + if layer_id >= self.num_local_layers: + if (getattr(self, "_attention_cache_layout_policy", + None) is not None + and LayerId(layer_id) in self._hp_manager_layer_ids): + return Role.MLA_HP_TAIL, None + raise ValueError(f"Unknown hybrid V2 virtual layer {layer_id}.") + if self._is_local_mamba_layer(layer_id): + return MambaRole.SSM_STATE, None + return super()._get_pool_roles(pool_id) + + def _max_resident_sequences(self) -> int: + return self.max_batch_size * self.mapping.pp_size + + def _mamba_state_bytes_per_slot(self) -> int: + return self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) + + def _num_ssm_snapshots_for_capacity( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + if capacity <= 0 or not kv_cache_config.enable_block_reuse: + return 0 + + fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, + self.max_seq_len, + self.tokens_per_block) + interval = _mamba_regular_snapshot_interval(kv_cache_config, + self.max_seq_len) + regular_snapshots = capacity // interval if interval is not None else 0 + return (self._max_resident_sequences() * fixed_rules + + regular_snapshots) + + def _num_ssm_states_per_typical_request( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + fixed_rules, _ = _mamba_snapshot_rule_counts( + kv_cache_config, + capacity, + self.tokens_per_block, + ) + # Additional snapshots are stable boundaries that must remain alive. + # Periodic snapshots are evictable cache entries and therefore do not + # increase the guaranteed state count represented by BatchDesc. + return 1 + fixed_rules + + def _typical_request_descs( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> List[KVCacheDesc]: + """Model one request with one descriptor per live SSM state.""" + num_states = self._num_ssm_states_per_typical_request( + capacity, kv_cache_config) + capacity_per_state, capacity_remainder = divmod(capacity, num_states) + capacities = [ + capacity_per_state + int(i < capacity_remainder) + for i in range(num_states) + ] + return [ + KVCacheDesc( + capacity=state_capacity, + history_length=max(0, state_capacity - 1), + ) for state_capacity in capacities + ] + + def _get_typical_request_capacity( + self, + kv_cache_config: KvCacheConfig, + ) -> int: + if kv_cache_config.avg_seq_len is not None: + return kv_cache_config.avg_seq_len + + fallback_capacity = max(1, self.max_seq_len // 2) + logger.warning( + "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " + "model using KV cache manager V2. Falling back to " + f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " + "'kv_cache_config.avg_seq_len' in the YAML configuration to the " + "workload's average total sequence length for an accurate KV/SSM " + "pool ratio.") + return fallback_capacity + + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + attention_quota = super()._get_quota_from_max_tokens(max_tokens) + num_request_lineages = self._max_resident_sequences() + snapshot_slots = self._num_ssm_snapshots_for_capacity( + max_tokens, self.kv_cache_config) + state_slots = (num_request_lineages + self._num_reserved_dummy_slots + + snapshot_slots) + state_quota = state_slots * self._mamba_state_bytes_per_slot() + # Once the plan contains any non-live SSM capacity, reserve one partial + # attention page per request lineage. This remains conservative when + # the plan contains fewer than one non-live slot per lineage. + extra_attention_quota = (num_request_lineages * + self._attention_cache_bytes_per_token() * + self.tokens_per_block + if snapshot_slots > 0 else 0) + return attention_quota + state_quota + extra_attention_quota + + def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: + draft_quota = super()._get_draft_quota_from_max_tokens(max_tokens) + snapshot_slots = self._num_ssm_snapshots_for_capacity( + max_tokens, self.kv_cache_config) + if snapshot_slots > 0: + # The snapshot plan's partial-page reserve also covers draft KV. + draft_quota += (self._max_resident_sequences() * + sum(self._draft_layer_sizes) * + self.tokens_per_block) + return draft_quota + + def _get_max_tokens_from_quota(self, quota: int) -> float: + if self._get_quota_from_max_tokens(0) > quota: + return 0 + + low = 0 + high = 1 + while self._get_quota_from_max_tokens(high) <= quota: + low = high + high *= 2 + if high >= 1 << 62: + return float("inf") + + while low + 1 < high: + mid = (low + high) // 2 + if self._get_quota_from_max_tokens(mid) <= quota: + low = mid + else: + high = mid + return low + + def _minimum_live_gpu_quota(self) -> int: + """Return the minimum quota for live states and one attention page.""" + attention_block_quota = (self._attention_cache_bytes_per_token() * + self.tokens_per_block) + num_state_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + return max( + self._get_quota_from_max_tokens(0), + state_quota + attention_block_quota, + ) + + def _build_cache_config( + self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + kv_cache_config = self.kv_cache_config + cache_tiers = config.cache_tiers + gpu_quota = cache_tiers[0].quota + minimum_live_quota = self._minimum_live_gpu_quota() + if minimum_live_quota > gpu_quota: + raise ValueError( + "The V2 Mamba GPU cache quota is too small for live recurrent " + f"states and attention pages: got {gpu_quota} bytes, need at " + f"least {minimum_live_quota} bytes.") + # _build_base_config already constructed every attention layer, + # including dtype-specific scale and subclass-provided side buffers. + # Preserve those configs and replace only the local Mamba layers. + layers = list(config.layers) + for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): + if self._mamba_layer_mask[global_layer_idx]: + layer_id = LayerId(local_layer_idx) + layers[local_layer_idx] = SsmLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig(role=MambaRole.SSM_STATE, + size=self.ssm_bytes), + BufferConfig(role=MambaRole.CONV_STATE, + size=self.conv_bytes), + ], + ) + + dummy_requests = [ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(self._num_reserved_dummy_slots) + ] + constraints = [ + replace( + batch, + kv_caches=[*batch.kv_caches, *dummy_requests], + ) for batch in config.constraints + ] + + typical_step = config.typical_step + if config.initial_pool_ratio is None: + typical_capacity = self._get_typical_request_capacity( + kv_cache_config) + request_descs = self._typical_request_descs(typical_capacity, + kv_cache_config) + typical_step = BatchDesc(request_descs * + self._max_resident_sequences() + + dummy_requests) + # The recurrent (SSM) state pool must hold one slot per resident + # sequence plus every reserved dummy slot. Unlike attention pages, a + # Mamba state is fixed-size per sequence, so this floor is independent + # of sequence length. The base config only emits constraints when + # ``avg_seq_len`` is set, and speculative decoding inflates the reserved + # dummy slots (CUDA-graph padding), so without an explicit floor the SSM + # pool can be undersized (see the live/dummy-slot check in _setup_states + # / __init__). Add a min-slots constraint of zero-capacity requests: + # these cost no attention pages but reserve one SSM slot each. + if any(isinstance(layer, SsmLayerConfig) for layer in layers): + ssm_floor_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + constraints = [ + *constraints, + BatchDesc([ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(ssm_floor_slots) + ]), + ] + result = replace( + config, + layers=layers, + typical_step=typical_step, + constraints=constraints, + # SSM lifecycles require minimum-snapshot commit semantics. The + # flag is harmless when reuse is disabled because no commits are + # attempted, while the runtime config still needs the invariant. + commit_min_snapshot=True, + ) + layout_policy = getattr(self, "_attention_cache_layout_policy", None) + if layout_policy is not None: + result = layout_policy._build_cache_config(self, result) + return result + + def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, + state_shape: List[int]) -> torch.Tensor: + addr = self.impl.get_mem_pool_base_address(LayerId(local_layer_idx), + role, PageIndexMode.SHARED) + num_pages = self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), role) + raw = convert_to_torch_tensor( + TensorWrapper(addr, dtype, [num_pages] + state_shape)) + page_index_scale = self.impl.get_page_index_scale( + LayerId(local_layer_idx), role) + num_slots = (num_pages + page_index_scale - 1) // page_index_scale + # V2 coalesces same-size per-layer buffers inside each slot. Kernels + # index Mamba states by logical slot id, so expose only this layer's + # sub-page from each coalesced slot instead of the raw page-index view. + return raw.as_strided( + [num_slots] + state_shape, + [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), + ) + + def _setup_states(self) -> None: + local_layer_ids = [ + self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers + ] + self.all_ssm_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, + self.ssm_state_dtype, self.ssm_state_shape) + for local_layer_idx in local_layer_ids + ] + self.all_conv_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, + self.conv_state_dtype, self.conv_state_shape) + for local_layer_idx in local_layer_ids + ] + if self._use_gdn_cached_replay_all_layer_commit: + expected_inner_strides = ( + self.ssm_state_shape[1] * self.ssm_state_shape[2], + self.ssm_state_shape[2], + 1, + ) + reference = self.all_ssm_states[0] + for state in self.all_ssm_states: + if (state.dtype != reference.dtype + or state.device != reference.device + or list(state.shape[1:]) != self.ssm_state_shape + or state.stride()[1:] != expected_inner_strides): + raise RuntimeError( + "GDN cached replay requires V2 SSM layers with matching " + "dtype, device, shape, and dense inner dimensions.") + pointers = [state.data_ptr() for state in self.all_ssm_states] + slot_stride = reference.stride(0) + layer_stride_bytes = (pointers[1] - + pointers[0] if len(pointers) > 1 else 0) + has_affine_layout = ( + layer_stride_bytes % reference.element_size() == 0 and all( + state.stride(0) == slot_stride + for state in self.all_ssm_states) + and all(pointer == pointers[0] + layer * layer_stride_bytes + for layer, pointer in enumerate(pointers))) + if has_affine_layout: + self._gdn_cached_replay_state_strides = ( + len(pointers), + layer_stride_bytes // reference.element_size(), + slot_stride, + ) + else: + logger.warning_once( + "V2 GDN state views are not affine; using indirect replay " + "checkpoint addressing", + key="gdn_cached_replay_v2_indirect_state_layout", + ) + self._gdn_cached_replay_state_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) + for state in self.all_ssm_states], + dtype=torch.int64, + device=reference.device, + ) + + def _setup_replay_buffers(self, spec_config) -> None: + cache_size = 0 + device = None + if self.local_num_mamba_layers > 0: + cache_size = self.all_ssm_states[0].shape[0] + assert all(t.shape[0] == cache_size for t in self.all_ssm_states) + device = self.all_ssm_states[0].device + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): + return + + mask_capacity = self._host_state_indices.shape[0] + self._dummy_request_mask = torch.zeros(mask_capacity, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros( + mask_capacity, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) + + def _attention_cache_bytes_per_token(self) -> int: + layout_policy = getattr(self, "_attention_cache_layout_policy", None) + if layout_policy is not None: + return layout_policy.get_cache_bytes_per_token(self) + # Mamba layers have zero KV heads, so the generic calculation naturally + # returns only bytes owned by local attention layers. + return super().get_cache_bytes_per_token() + + def get_cache_bytes_per_token(self) -> int: + cache_bytes = self._attention_cache_bytes_per_token() + + interval = ( + self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) + if (self.kv_cache_config.enable_block_reuse and interval is not None + and interval > 0): + cache_bytes += self._mamba_state_bytes_per_slot() // interval + if cache_bytes == 0 and self.local_num_mamba_layers > 0: + cache_bytes = self._mamba_state_bytes_per_slot() + return max(1, cache_bytes) + + def get_num_free_blocks(self) -> int: + assert len(self.kv_cache_map) == 0, ( + "get_num_free_blocks is only used when the kv cache manager is empty" + ) + if getattr(self, "_attention_cache_layout_policy", None) is not None: + return self.blocks_in_primary_pool + attention_pages = [] + ssm_pages = [] + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + if self._is_local_mamba_layer(local_layer_idx): + ssm_pages.append( + self.impl.get_page_index_upper_bound( + layer_id, MambaRole.SSM_STATE) // + self._ssm_page_index_scale) + else: + attention_pages.append( + self._get_attention_pool_num_blocks(local_layer_idx)) + if attention_pages: + return max(attention_pages) + return max(ssm_pages) if ssm_pages else 0 + + @property + def blocks_in_primary_pool(self) -> int: + layout_policy = getattr(self, "_attention_cache_layout_policy", None) + if layout_policy is not None: + return layout_policy.blocks_in_primary_pool.fget(self) + for local_layer_idx in range(self.num_local_layers): + if self._is_local_mamba_layer(local_layer_idx): + continue + return self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), Role.KEY) + return 0 + + def get_buffers(self, + layer_idx: int, + kv_layout: str = "NHD") -> Optional[torch.Tensor]: + local_layer_idx = self.layer_offsets[layer_idx] + if self._is_local_mamba_layer(local_layer_idx): + return None + return super().get_buffers(layer_idx, kv_layout) + + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: + layout_policy = getattr(self, "_attention_cache_layout_policy", None) + if layout_policy is not None: + yield from self._iter_fp4_mla_physical_pool_views() + else: + for global_layer_id, local_layer_id in self.layer_offsets.items(): + if self._is_local_mamba_layer(local_layer_id): + continue + # A layer group is a lifecycle, not a physical memory pool. + # Differently sized attention buffers can share one lifecycle, + # so scan every attention layer in this diagnostic path. + yield KVCacheManagerV2.get_buffers(self, global_layer_id) + + yield from self.all_ssm_states + yield from self.all_conv_states + + def add_dummy_requests( + self, + request_ids: List[int], + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + prepare_resource: bool = True, + max_num_draft_tokens: int = 0, + kv_reserve_draft_tokens: Optional[int] = None, + use_mrope: bool = False, + max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, + num_extra_decoding_steps: int = 0, + draft_kv_cache_manager: Optional[BaseResourceManager] = None, + ) -> List[LlmRequest]: + requests = super().add_dummy_requests( + request_ids=request_ids, + token_nums=token_nums, + is_gen=is_gen, + prepare_resource=prepare_resource, + max_num_draft_tokens=max_num_draft_tokens, + kv_reserve_draft_tokens=kv_reserve_draft_tokens, + use_mrope=use_mrope, + max_beam_width=max_beam_width, + encoder_output_lens=encoder_output_lens, + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if requests and prepare_resource: + self._setup_state_indices(requests) + return requests + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is not None and kv_cache.is_active: + self.try_commit_blocks(request, kv_cache) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + self._branch_snapshot_points.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + self._setup_state_indices(requests) + num_contexts = len(scheduled_batch.context_requests) + self._reset_context_mamba_slots(num_contexts) + + def _setup_state_indices(self, requests: List[LlmRequest]) -> None: + if self.local_num_mamba_layers == 0: + return + n = len(requests) + assert n <= self._host_state_indices.shape[0], ( + f"State-index batch size {n} exceeds max_batch_size " + f"{self._host_state_indices.shape[0]}") + self._host_state_indices.zero_() + if n > 0: + for i, req in enumerate(requests): + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + raise RuntimeError( + f"Missing V2 KV cache for request {req.py_request_id}") + base_index = kv_cache.get_ssm_block_base_index( + self.ssm_layer_group_id) + if base_index < 0: + raise RuntimeError( + f"Invalid SSM state block index {base_index} for " + f"request {req.py_request_id}") + self._host_state_indices[i] = base_index + + _copy_from_host_snapshot(self.cuda_state_indices, + self._host_state_indices) + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask(is_dummy) + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy + + def get_state_indices(self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None): + if self.local_num_mamba_layers == 0: + # Mamba metadata is still prepared on attention-only PP ranks, + # but no local kernel consumes these indices. Return harmless + # placeholders instead of consulting an intentionally empty map. + if request_ids is not None: + return [0] * len(request_ids) + return self.cuda_state_indices + if request_ids is not None: + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + return self.cuda_state_indices + + def get_max_resource_count(self) -> int: + return self.max_batch_size + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None): + if self.local_num_mamba_layers == 0: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + num_accepted_draft_tokens = ( + num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( + torch.int32) + # Dynamic tree selects a tree node rather than a linear draft depth. + accepted_positions = (accepted_leaf_positions.to(torch.int32) + if accepted_leaf_positions is not None else + num_accepted_draft_tokens) + if state_indices is None: + state_indices = self.get_state_indices() + state_indices_d = state_indices[num_contexts:num_contexts + + num_gens].to(torch.int32) + src_state_indices = self.intermediate_state_indices[:num_gens] + + if self._use_replay_state_update: + self._commit_gdn_cached_replay_history_layers( + attn_metadata, num_gens) + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) + else: + for layer_offset, dst in enumerate(self.all_ssm_states): + _promote_mamba_state_triton( + dst.unsqueeze(0), + self.intermediate_ssm_states[layer_offset:layer_offset + 1], + src_state_indices, + accepted_positions, + state_indices_d, + ) + + for layer_offset, dst in enumerate(self.all_conv_states): + _promote_mamba_state_triton( + dst.unsqueeze(0), + self.intermediate_conv_states[layer_offset:layer_offset + 1], + src_state_indices, + accepted_positions, + state_indices_d, + ) + + def _mark_context_position_as_history(self, request: LlmRequest, + kv_cache) -> None: + """Advance history without making later recurrent state reusable.""" + history_length = request.context_current_position + if history_length <= kv_cache.history_length: + return + capacity = max(kv_cache.capacity, history_length) + if not kv_cache.resize(capacity, history_length=history_length): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to {history_length} tokens") + + def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + if not should_block_reuse: + return + + if kv_cache is None: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return + + snapshot_points = request.expect_snapshot_points + commit_limit = (min(max(snapshot_points), request.prompt_len) + if snapshot_points else request.prompt_len) + commit_end = min(request.context_current_position, commit_limit) + if (request.context_current_position in request.expect_snapshot_points + and commit_end > kv_cache.num_committed_tokens): + tokens = self._augment_tokens_for_block_reuse( + request.get_tokens(DEFAULT_BEAM_INDEX), + request, + start=kv_cache.num_committed_tokens, + end=commit_end, + ) + kv_cache.commit(tokens) + if (self.local_num_mamba_layers > 0 + and self._branch_snapshot_points.get( + request.py_request_id) == commit_end): + self._branch_snapshots_taken_total += 1 + if request.context_current_position >= commit_limit: + self._mark_context_position_as_history(request, kv_cache) + if request.context_remaining_length == 0: + kv_cache.stop_committing() + + def update_context_resources(self, + scheduled_batch: ScheduledRequests) -> None: + for request in scheduled_batch.context_requests: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None or not kv_cache.is_active: + continue + + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + is_all_reusable = ( + self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE) + is_snapshot_boundary = (request.context_current_position + in request.expect_snapshot_points) + has_pending_snapshot = any( + point > request.context_current_position + for point in request.expect_snapshot_points) + # NVFP4: the BF16 HP-tail ring is positioned by the history + # length, so it needs the completed chunk's watermark even while a + # later recurrent-state snapshot is still pending. + should_resize = (self.dtype == DataType.NVFP4 + or not should_block_reuse or + (not is_all_reusable and not has_pending_snapshot)) + should_commit = (is_all_reusable or is_snapshot_boundary + or request.context_remaining_length == 0) + + if should_resize and not kv_cache.resize( + None, request.context_current_position): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to " + f"{request.context_current_position} tokens at context " + "update") + if should_commit: + self.try_commit_blocks(request, kv_cache) + if request.context_remaining_length == 0: + if self.conversation_manager is not None: + self.conversation_manager.save_drop_plan(request, kv_cache) + kv_cache.enable_swa_scratch_reuse = False + + def shutdown(self): + self._gdn_cached_replay_state_descriptors = None + self._gdn_cached_replay_state_strides = None + self.all_ssm_states = [] + self.all_conv_states = [] + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + self._branch_snapshot_points.clear() + super().shutdown() + + +def _is_disagg_context_server(is_disagg: bool) -> bool: + return is_disagg and os.getenv("TRTLLM_DISAGG_ROLE") == "context" + + +def _kda_replay_layout( + num_spec: int, num_heads: int, head_dim: int, + conv_window: int) -> dict[str, tuple[torch.dtype, list[int]]]: + """Physical per-slot shapes; beta rows retain the CuTe alignment.""" + dim = num_heads * head_dim + heads_per_alignment = _KDA_BETA_CACHE_ALIGNMENT_BYTES // torch.bfloat16.itemsize + aligned_heads = ((num_heads + heads_per_alignment - 1) // + heads_per_alignment) * heads_per_alignment + return { + "kda_conv_q": (torch.bfloat16, [conv_window + num_spec, dim]), + "kda_conv_k": (torch.bfloat16, [conv_window + num_spec, dim]), + "kda_conv_v": (torch.bfloat16, [conv_window + num_spec, dim]), + "kda_k_cache": (torch.bfloat16, [num_spec, dim]), + "kda_g_cache": (torch.float32, [num_spec, dim]), + "kda_v_cache": (torch.bfloat16, [num_spec, dim]), + "kda_beta_cache": (torch.bfloat16, [num_spec, aligned_heads]), + } + + +class KDAHybridCacheManagerV2(MambaHybridCacheManagerV2): + """V2 hybrid cache manager with Kimi KDA replay state. + + The live convolution pool stores the committed ``W - 1`` input window. + Fused verification extends that window with one column per draft token and + carries the accepted draft prefix between verification rounds. + """ + + def __init__( + self, + *args, + kda_replay_num_spec: Optional[int] = None, + **kwargs, + ) -> None: + spec_config = kwargs.get("spec_config") + use_replay_state_update = kwargs.get("use_replay_state_update", False) + conv_state_layout = kwargs.get("conv_state_layout", "x_b_c") + + self.kda_context_only = _is_disagg_context_server( + kwargs.get("is_disagg", False)) + if self.kda_context_only: + kda_replay_num_spec = None + self._kda_replay_layers: list[dict[str, torch.Tensor]] = [] + self._kda_replay_num_spec = kda_replay_num_spec + self._use_kda_replay_update = kda_replay_num_spec is not None + if self._use_kda_replay_update: + if use_replay_state_update: + raise ValueError( + "kda_replay_num_spec and use_replay_state_update are " + "mutually exclusive") + if spec_config is None: + raise ValueError( + "KDA replay caches require speculative decoding") + expected_num_spec = spec_config.tokens_per_gen_step - 1 + if kda_replay_num_spec != expected_num_spec: + raise ValueError( + "KDA replay cache width " + f"({kda_replay_num_spec}) must match the draft length " + f"({expected_num_spec})") + if conv_state_layout != "q_k_v": + raise ValueError( + "KDA replay requires conv_state_layout='q_k_v'") + self._record_kda_replay_in_update_resources = ( + self._use_kda_replay_update + and getattr(spec_config, "decoding_type", None) == "NGram") + + super().__init__(*args, **kwargs) + + @staticmethod + def get_cache_size_per_token(model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs): + from ..modules.kimi_kda._kda_kernels import is_kda_mtp_verify_available + + slope, intercept = MambaHybridCacheManagerV2.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + **kwargs) + spec_config = kwargs.get("spec_config") + if (_is_disagg_context_server(kwargs.get("is_disagg", False)) + or spec_config is None or not is_kda_mtp_verify_available()): + return slope, intercept + + params, local_layers, _ = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=kwargs.get("is_draft", False), + use_separate_draft_kv_cache=kwargs.get( + "use_separate_draft_kv_cache", False)) + if not local_layers: + return slope, intercept + + tp_size = _mamba_effective_tp_size(mapping) + layout = _kda_replay_layout(spec_config.tokens_per_gen_step - 1, + params.num_heads // tp_size, + params.head_dim, params.conv_kernel - 1) + replay_bytes_per_slot = ( + local_layers * sum(dtype.itemsize * math.prod(shape) + for dtype, shape in layout.values()) + + _KDA_REPLAY_COUNT_BYTES) + resident_sequences = max_batch_size * mapping.pp_size + dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( + spec_config, max_batch_size) + int(mapping.enable_attention_dp)) + fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, + max_seq_len, + tokens_per_block) + intercept += (resident_sequences * + (1 + fixed_rules) + dummy_slots) * replay_bytes_per_slot + + interval = _mamba_regular_snapshot_interval(kv_cache_config, + max_seq_len) + if interval is not None: + # Round the combined snapshot cost once, preserving the parent's + # attention cost and avoiding a second rounding of replay bytes. + state_bytes_per_slot = (local_layers * + params.get_states_bytes_per_layer(mapping)) + slope += (math.ceil( + (state_bytes_per_slot + replay_bytes_per_slot) / interval) - + math.ceil(state_bytes_per_slot / interval)) + return slope, intercept + + def _replay_layout(self) -> dict[str, tuple[torch.dtype, list[int]]]: + if not self._use_kda_replay_update or not self.local_num_mamba_layers: + return {} + if (len(self.conv_section_dims) != 3 + or len(set(self.conv_section_dims)) != 1): + raise ValueError("KDA replay caches require equal [Q | K | V] " + "convolution-state sections") + return _kda_replay_layout(self._kda_replay_num_spec, + self.ssm_state_shape[0], + self.ssm_state_shape[1], + self.conv_state_shape[1]) + + def _mamba_state_bytes_per_slot(self) -> int: + replay_bytes = sum(dtype.itemsize * math.prod(shape) + for dtype, shape in self._replay_layout().values()) + count_bytes = _KDA_REPLAY_COUNT_BYTES if replay_bytes else 0 + return (super()._mamba_state_bytes_per_slot() + + self.local_num_mamba_layers * replay_bytes + count_bytes) + + def _build_cache_config( + self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + config = super()._build_cache_config(config) + replay_buffers = [ + BufferConfig(role=DataRole(name), + size=dtype.itemsize * math.prod(shape)) + for name, (dtype, shape) in self._replay_layout().items() + ] + if not replay_buffers: + return config + first_layer_id = LayerId(self.layer_offsets[self.mamba_pp_layers[0]]) + layers = [] + for layer in config.layers: + if isinstance(layer, SsmLayerConfig): + buffers = [*layer.buffers, *replay_buffers] + if layer.layer_id == first_layer_id: + # All SSM layers share a slot lifecycle. Store the request's + # count once so KVCM moves it with the replay histories. + buffers.append( + BufferConfig(role=DataRole("prev_num_accepted_tokens"), + size=_KDA_REPLAY_COUNT_BYTES)) + layer = SsmLayerConfig(layer_id=layer.layer_id, buffers=buffers) + layers.append(layer) + return replace(config, layers=layers) + + def _setup_mtp_intermediate_states(self, spec_config, + max_batch_size: int) -> None: + if self.kda_context_only: + super()._setup_mtp_intermediate_states(None, max_batch_size) + return + if not self._use_kda_replay_update: + super()._setup_mtp_intermediate_states(spec_config, max_batch_size) + return + + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if spec_config is None or self.local_num_mamba_layers == 0: + return + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device="cuda") + + def _validate_replay_buffer_colocation(self, + replay_roles: Iterable[str]) -> None: + """Verify replay and recurrent buffers share a logical slot lifecycle.""" + # All layers reuse the same logical slot index. Check the core's + # allocated layout, not just equal capacities: buffers in one pool + # group and lifecycle are allocated and moved together. + required_buffers = {(self.layer_offsets[layer_id], str(role)) + for layer_id in self.mamba_pp_layers + for role in (MambaRole.SSM_STATE, + MambaRole.CONV_STATE, *replay_roles)} + required_buffers.add((self.layer_offsets[self.mamba_pp_layers[0]], + "prev_num_accepted_tokens")) + if not any( + required_buffers.issubset( + {(int(buffer_id.layer_id), str(buffer_id.role)) + for coalesced in variant.coalesced_buffers + for buffer_id in coalesced.buffer_ids}) + for pool_group in self.impl.pool_group_descs + for variant in pool_group.slot_desc.variants): + raise RuntimeError( + "KDA replay and recurrent-state buffers must share one KVCM " + "pool group and lifecycle to reuse the same logical slot index") + + def _allocate_pool_replay_buffers( + self, + spec_config, + cache_size: int, + device: Optional[torch.device], + ) -> bool: + allocated = super()._allocate_pool_replay_buffers( + spec_config, cache_size, device) + self._kda_replay_layers = [] + layout = self._replay_layout() + if not layout: + return allocated + self._validate_replay_buffer_colocation(layout) + for layer_id in self.mamba_pp_layers: + layer = {} + for name, (dtype, shape) in layout.items(): + view = self._get_state_buffer(self.layer_offsets[layer_id], + DataRole(name), dtype, shape) + if name in ("kda_conv_q", "kda_conv_k", "kda_conv_v"): + view = view.transpose(-1, -2) + elif name == "kda_beta_cache": + view = view[..., :self.ssm_state_shape[0]] + # Pool capacity can exceed physically resident memory. Only + # initialize rows after the request locks its pages. + layer[name] = view + self._kda_replay_layers.append(layer) + self.prev_num_accepted_tokens = self._get_state_buffer( + self.layer_offsets[self.mamba_pp_layers[0]], + DataRole("prev_num_accepted_tokens"), torch.int32, + [_KDA_REPLAY_COUNT_BYTES // torch.int32.itemsize])[..., 0] + return True + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + super()._reset_context_mamba_slots(num_contexts) + if num_contexts and self._use_kda_replay_update: + self._reset_kda_replay_slots( + self._host_state_indices[:num_contexts].tolist()) + + def _reset_kda_replay_slots(self, slots: Union[List[int], + torch.Tensor]) -> None: + if not self._kda_replay_layers: + return + device = self.prev_num_accepted_tokens.device + if isinstance(slots, torch.Tensor) and slots.device.type != "cpu": + slots = slots.to(device=device, dtype=torch.long, non_blocking=True) + else: + if isinstance(slots, torch.Tensor): + slots = slots.tolist() + if not slots: + return + slots = _host_slots_to_device(slots, device) + if slots.numel(): + self.prev_num_accepted_tokens.index_fill_(0, slots, 0) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + KVCacheManagerV2.prepare_resources(self, scheduled_batch) + if self.local_num_mamba_layers == 0: + return + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + num_contexts = len(scheduled_batch.context_requests) + self._setup_state_indices(requests, num_contexts=num_contexts) + self._reset_context_mamba_slots(num_contexts) + + def _setup_state_indices(self, + requests: List[LlmRequest], + num_contexts: Optional[int] = None) -> None: + if not self._use_kda_replay_update: + super()._setup_state_indices(requests) + return + + replay_requests = (requests + if num_contexts is None else requests[num_contexts:]) + old_state_values = [ + self._request_id_to_state_index.get(request.py_request_id, -1) + for request in replay_requests + ] + super()._setup_state_indices(requests) + if self.local_num_mamba_layers == 0: + return + + new_state_values = [ + self._request_id_to_state_index[request.py_request_id] + for request in replay_requests + ] + self._reset_kda_replay_slots([ + new for old, new in zip(old_state_values, new_state_values) + if old < 0 + ]) + + def _record_kda_replay_acceptance( + self, + state_indices: torch.Tensor, + num_accepted_draft_tokens: torch.Tensor, + is_dummy_request: Optional[torch.Tensor] = None, + ) -> None: + """Record the draft prefix that the next KDA verify must replay.""" + slots = state_indices.to(torch.long) + accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) + counts = self.prev_num_accepted_tokens + values = (torch.where(is_dummy_request, counts[slots], accepted) + if is_dummy_request is not None else accepted) + counts[slots] = values + + @torch.inference_mode() + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + """Seed replay caches from transferred ``W - 1`` conv states.""" + if (not self._use_kda_replay_update + or self.local_num_mamba_layers == 0): + return + slots = [ + self._request_id_to_state_index[request_id] + for request_id in request_ids + if request_id in self._request_id_to_state_index + ] + if not slots: + return + + device = self.prev_num_accepted_tokens.device + state_indices = _host_slots_to_device(sorted(set(slots)), device) + self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0) + committed_window = self.conv_state_shape[1] + section_offsets = [0] + for section_dim in self.conv_section_dims: + section_offsets.append(section_offsets[-1] + section_dim) + + for layer, conv_state in zip(self._kda_replay_layers, + self.all_conv_states): + for replay_buffer in layer.values(): + replay_buffer.index_fill_(0, state_indices, 0) + selected_conv = conv_state.index_select(0, state_indices) + for section_idx, name in enumerate( + ("kda_conv_q", "kda_conv_k", "kda_conv_v")): + replay_layer = layer[name] + section = selected_conv[:, section_offsets[section_idx]: + section_offsets[section_idx + 1], :] + replay_layer[:, :, :committed_window].index_copy_( + 0, state_indices, section.to(replay_layer.dtype)) + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + if self.kda_context_only: + return PythonMambaCacheManager.State( + conv=self.get_conv_states(layer_idx), + temporal=self.get_ssm_states(layer_idx)) + if not self._use_kda_replay_update: + return super().mamba_layer_cache(layer_idx) + + layer_offset = self.mamba_layer_offsets[layer_idx] + spec_kwargs = dict(self._kda_replay_layers[layer_offset]) + spec_kwargs["prev_num_accepted_tokens"] = self.prev_num_accepted_tokens + if self.mamba_ssm_rand_seed is not None: + spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed + return PythonMambaCacheManager.SpeculativeState( + conv=self.get_conv_states(layer_idx), + temporal=self.get_ssm_states(layer_idx), + intermediate_conv_window=None, + **spec_kwargs, + ) + + def update_resources(self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None) -> None: + KVCacheManagerV2.update_resources(self, scheduled_batch, attn_metadata, + kv_cache_dtype_byte_size) + if (not self._record_kda_replay_in_update_resources + or self.local_num_mamba_layers == 0): + return + + generation_requests = scheduled_batch.generation_requests + drafted_requests = [ + request for request in generation_requests + if request.py_draft_tokens is not None + and len(request.py_draft_tokens) > 0 + ] + if not drafted_requests: + return + if len(drafted_requests) != len(generation_requests): + raise RuntimeError( + "Mixed drafted/undrafted generation batch is not supported " + "for KDA replay bookkeeping") + + active_requests = [ + request for request in generation_requests + if request.py_request_id in self._request_id_to_state_index + ] + if not active_requests: + return + device = self.prev_num_accepted_tokens.device + state_indices = torch.tensor([ + self._request_id_to_state_index[request.py_request_id] + for request in active_requests + ], + dtype=torch.int32, + device=device) + accepted_drafts = torch.tensor([ + request.py_num_accepted_draft_tokens for request in active_requests + ], + dtype=torch.int32, + device=device) + is_dummy_request = torch.tensor([ + self._request_id_to_is_dummy.get(request.py_request_id, False) + for request in active_requests + ], + dtype=torch.bool, + device=device) + self._record_kda_replay_acceptance(state_indices, accepted_drafts, + is_dummy_request) + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None): + if self.kda_context_only: + return + if not self._use_kda_replay_update: + return super().update_mamba_states(attn_metadata, + num_accepted_tokens, + state_indices, + accepted_leaf_positions) + if self.local_num_mamba_layers == 0: + return + + num_contexts = attn_metadata.num_contexts + num_gens = attn_metadata.num_seqs - num_contexts + num_accepted_draft_tokens = ( + num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( + torch.int32) + if state_indices is None: + state_indices = self.get_state_indices() + state_indices_d = state_indices[num_contexts:num_contexts + + num_gens].to(torch.int32) + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[num_contexts:num_contexts + + num_gens] + self._record_kda_replay_acceptance( + state_indices_d, + num_accepted_draft_tokens, + is_dummy_request, + ) + + def shutdown(self): + self._kda_replay_layers.clear() + self.prev_num_accepted_tokens = None + super().shutdown() diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2372d00d16ea..b54cb6f5122c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -37,14 +37,25 @@ from ..distributed.communicator import init_pp_comm from ..distributed.ops import MNNVLAllReduce from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader -from ..models.modeling_multimodal_mixin import MultimodalModelMixin -from ..models.modeling_utils import DecoderModelForCausalLM, timing_metric -from ..moe.expert_statistic import ExpertStatistic -from ..moe.fused_moe.moe_comm_timeout_guard import set_moe_comm_warmup -from ..moe.fused_moe.moe_load_balancer import MoeLoadBalancer -from ..nccl_window_tensor_scope import install_eager_nccl_window_tensor_scopes -from ..route_capture import ROUTE_CAPTURE_ATTR, RouteCapture -from ..speculative import SpecMetadata, update_spec_config_from_loaded_model +from ..models.modeling_multimodal_encoder import MultimodalEncoderMixin +from ..models.modeling_multimodal_mixin import (MultimodalEncoderContractError, + MultimodalModelMixin, + _build_request_multimodal_input) +from ..models.modeling_multimodal_utils import filter_mm_token_from_input_ids +from ..models.modeling_speculative import SpecDecOneEngineForCausalLM +from ..models.modeling_utils import DecoderModelForCausalLM +from ..modules.fused_moe.moe_load_balancer import (MoeLoadBalancer, + MoeLoadBalancerIterContext) +from ..modules.mamba.mamba2_metadata import Mamba2Metadata +from ..peft.lora.cuda_graph_lora_manager import CudaGraphLoraManager +from ..speculative import (SpecMetadata, get_draft_kv_cache_manager, + get_num_extra_kv_tokens, get_spec_metadata, + prepare_attn_metadata_for_draft_replay, + restore_attn_metadata_after_draft_replay, + update_spec_config_from_loaded_model) +from ..speculative.drafting_loops import BaseDraftingLoopWrapper +from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata +from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..tensor_lru_cache import TensorLRUCache from ..utils import set_torch_compiling from .config_utils import is_hybrid_linear @@ -563,6 +574,100 @@ def _init_execution_services(self) -> None: self._runner = self._initialize_runner(runner_cls) + # Once encoder CUDA graphs are usable, enable mixed decoder graphs by + # default unless the user explicitly opts out. + encoder_decoder_cuda_graph_enabled = ( + self.encoder_cuda_graph_runner.enabled + and self.encoder_cuda_graph_runner.is_encoder_decoder + and bool(self.encoder_cuda_graph_runner.capture_keys)) + enable_encoder_decoder_mixed_cuda_graph = ( + encoder_decoder_cuda_graph_enabled + and self.cuda_graph_config is not None + and self.llm_args.enable_encoder_decoder_mixed_cuda_graph) + + # Create decoder CUDA graph config and runner. + cuda_graph_runner_config = CUDAGraphRunnerConfig( + use_cuda_graph=(not self._is_encode_only + and self.cuda_graph_config is not None), + cuda_graph_padding_enabled=self._cuda_graph_padding_enabled, + cuda_graph_batch_sizes=self._cuda_graph_batch_sizes, + max_cuda_graph_batch_size=self._max_cuda_graph_batch_size, + max_beam_width=self.max_beam_width, + spec_config=self.spec_config, + cuda_graph_mem_pool=self._cuda_graph_mem_pool, + dynamic_draft_len_mapping=self._dynamic_draft_len_mapping, + max_num_tokens=self.max_num_tokens, + use_mrope=self.use_mrope, + original_max_draft_len=self.original_max_draft_len, + original_max_total_draft_tokens=self. + original_max_total_draft_tokens, + is_draft_model=self.is_draft_model, + enable_attention_dp=self.enable_attention_dp, + is_encoder_decoder=self._is_encoder_decoder_model(), + batch_size=self.batch_size, + mapping=self.mapping, + dist=self.dist, + kv_cache_manager_key=self.kv_cache_manager_key, + sparse_attention_config=self.sparse_attention_config, + enable_encoder_decoder_mixed_cuda_graph=( + enable_encoder_decoder_mixed_cuda_graph), + ) + self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) + self.breakable_cuda_graph_runner = None + if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: + decoder_model = (self.model if isinstance( + self.model, DecoderModelForCausalLM) else getattr( + self.model, "llm", None)) + if not isinstance(decoder_model, DecoderModelForCausalLM): + raise ValueError( + "breakable prefill CUDA graph requires a decoder model body" + ) + self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner( + decoder_model.model) + + # Initialize CUDA Graph LoRA manager if LoRA is enabled + self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None + + # Setup the local cache indirection buffer only once and reuse it. + # This way it can also be used for CUDA graphs. + if self.use_beam_search: + self.cache_indirection_attention = torch.zeros( + (self.batch_size, self.max_beam_width, self.max_seq_len), + device="cuda", + dtype=torch.int32) + else: + self.cache_indirection_attention = None + + self.kv_cache_dtype_byte_size = self.get_kv_cache_dtype_byte_size() + + self._prepare_inputs_event: Optional[torch.cuda.Event] = None + + # Cache for enc-dec cross-attention stable generation steps. + # Populated on the first CUDA-graph generation step; cleared whenever + # the batch composition changes (new encoder request arrives). + self._cross_attn_stable_cached_tokens: Optional[List[int]] = None + self._cross_attn_stable_request_ids: Optional[List[int]] = None + + def register_forward_pass_callable(self, callable: Callable): + self.forward_pass_callable = callable + + def get_speculative_model(self) -> Optional[SpecDecOneEngineForCausalLM]: + """Return the one-engine speculative model, including inside a VLM.""" + if not self.is_spec_decode: + return None + model = self.model + if isinstance(model, MultimodalModelMixin): + model = model.language_model + if not isinstance(model, SpecDecOneEngineForCausalLM): + # Kimi's legacy multimodal wrapper exposes its decoder as llm + # without inheriting MultimodalModelMixin. + model = getattr(model, "llm", None) + return model if isinstance(model, SpecDecOneEngineForCausalLM) else None + + def get_kv_cache_dtype_byte_size(self) -> float: + """ + Returns the size (in bytes) occupied by kv cache type. + """ layer_quant_mode = self.model.model_config.quant_config.layer_quant_mode if layer_quant_mode.has_fp4_kv_cache(): self.kv_cache_dtype_byte_size = 1 / 2 @@ -1460,24 +1565,660 @@ def forward(self, resource_manager: Optional[ResourceManager] = None, new_tensors_device: Optional[SampleStateTensors] = None, cache_indirection_buffer: Optional[torch.Tensor] = None, - *, - enable_spec_decode: bool = False, - runtime_draft_len: int = 0): - if isinstance(batch, PackedInputs): - assert isinstance(self._runner, PackedModelRunner), ( - "a packed batch requires a packed-batch runner") - return self._runner.forward(batch) - assert resource_manager is not None, ( - "scheduled execution requires a resource manager") - inputs = make_scheduled_inputs(batch, - new_tensors_device, - cache_indirection_buffer, - enable_spec_decode=enable_spec_decode, - runtime_draft_len=runtime_draft_len) - # Only a warmup scope makes the pass dummy; padding requests in a - # serving batch do not. - return self._forward_scheduled( - inputs, + num_accepted_tokens_device: Optional[torch.Tensor] = None, + req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None): + if not self._disable_overlap_scheduler: + # Do not refill reusable host staging while the previous + # iteration's asynchronous H2D copies still consume it. This + # event precedes the model forward, so its synchronization retains + # GPU forward overlap while protecting every forward entry, + # including speculative/draft paths. + self.wait_for_input_copy() + + kv_cache_manager = resource_manager.get_resource_manager( + self.kv_cache_manager_key) + draft_kv_cache_manager = self._get_draft_kv_cache_manager( + resource_manager) + + attn_metadata = self._set_up_attn_metadata(kv_cache_manager, + draft_kv_cache_manager) + if isinstance(attn_metadata, TrtllmAttentionMetadata): + attn_metadata.trtllm_gen_jit_warmup = self._trtllm_gen_jit_warmup + if self.enable_spec_decode: + spec_resource_manager = resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + spec_tree_manager = None + if spec_resource_manager is not None and hasattr( + spec_resource_manager, 'spec_tree_manager'): + spec_tree_manager = spec_resource_manager.spec_tree_manager + spec_metadata = self._set_up_spec_metadata(spec_resource_manager, + no_cache=kv_cache_manager + is None) + # attn_metadata now depends on spec_metadata since it determines the shape/content of spec_dec parameter Tensors + is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode( + spec_resource_manager, self.is_draft_model, self.attn_backend, + self.model_is_wrapped) + # Propagate runtime_draft_len (already set on self by py_executor) + # to spec_metadata so downstream code (eagle3, interface, trtllm) can read it. + spec_metadata.runtime_draft_len = self.runtime_draft_len + spec_metadata.runtime_tokens_per_gen_step = ( + self.get_runtime_tokens_per_gen_step(self.runtime_draft_len)) + + # Parallel-draft modes advertise a per-gen-step width via + # tokens_per_gen_step (PARD: 2K, DFlash: K+1). Pass + # (tokens_per_gen_step - 1) so generation_lengths = tokens_per_gen_step + # and the XQA kernel computes the correct past_kv_len. + if spec_metadata.spec_dec_mode.is_parallel_draft(): + sd_max_draft_len = self.original_max_total_draft_tokens + sd_max_total = self.original_max_total_draft_tokens + else: + sd_max_draft_len = self.original_max_draft_len + sd_max_total = self._spec_dec_max_total_draft_tokens + + # Fill slot-ID buffer for update_spec_dec_param + if (spec_tree_manager is not None + and spec_tree_manager.use_dynamic_tree + and not self.is_draft_model): + spec_tree_manager.slot_storage.fill_all_slot_ids( + scheduled_requests.context_requests, + scheduled_requests.generation_requests, + ) + + attn_metadata.update_spec_dec_param( + batch_size=scheduled_requests.batch_size, + is_spec_decoding_enabled=is_spec_dec_mode, + is_spec_dec_tree=spec_metadata.is_spec_dec_tree, + is_spec_dec_dynamic_tree=spec_metadata.is_spec_dec_dynamic_tree, + max_draft_len=sd_max_draft_len, + max_total_draft_tokens=sd_max_total, + model_is_wrapped=self.model_is_wrapped, + spec_metadata=spec_metadata, + spec_tree_manager=spec_tree_manager, + num_contexts=scheduled_requests.num_context_requests) + else: + spec_resource_manager = None + spec_metadata = None + + moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', + None) + if kv_cache_manager is None: + inputs, gather_ids = self._prepare_tp_inputs_no_cache( + scheduled_requests, attn_metadata, spec_metadata, + resource_manager) + + with MoeLoadBalancerIterContext(moe_load_balancer): + # Special handling for multimodal encoder only mode + if self.llm_args.mm_encoder_only: + return self._forward_step_mm_encoder_only( + inputs, scheduled_requests) + else: + return self._forward_step( + inputs, + gather_ids=gather_ids, + gather_context_logits=gather_context_logits) + + graph_requests = scheduled_requests + promoted_context_request_ids: frozenset[int] = frozenset() + # Non-linear tree input preparation expands runtime_draft_len to the + # total tree width after graph selection. Only linear-tree zero-draft + # iterations can therefore safely reuse a zero-draft graph. + can_promote_spec_decode = (not self.enable_spec_decode + or (not self.is_draft_model + and self.runtime_draft_len == 0 + and self.spec_config is not None + and self.spec_config.is_linear_tree)) + # TODO: Generalize these conservative gates as actual-draft, beam, and + # context-parallel providers for decoder-only LLMs gain support for + # promoted final-context rows. Each relaxation must preserve whole-batch + # fallback on graph miss and prove parity with the provider's native + # q_len=1 path. Encoder-decoder and non-LLM engines remain out of scope. + if (scheduled_requests.num_context_requests > 0 + and self.cuda_graph_runner.enabled and can_promote_spec_decode + and not self.use_beam_search + and not self._is_encoder_decoder_model() + and not self._is_encode_only + and not self.llm_args.mm_encoder_only + and self.mapping.cp_size == 1): + graph_requests, promoted_context_request_ids = \ + _make_single_token_context_graph_batch( + scheduled_requests, + self._is_final_multimodal_context_decode_compatible) + + with self.cuda_graph_runner.pad_batch( + graph_requests, resource_manager, + self.runtime_draft_len) as padded_graph_requests: + # Callee already no-ops when use_mrope=False, but the Python call / + # frame setup itself is non-trivial under high concurrency. Gating + # at the caller avoids that overhead for non-mrope models. + if self.use_mrope: + self._pad_batch_seed_mrope_delta_cache(padded_graph_requests) + + # Refresh is_all_greedy_sample for the *current* batch BEFORE the + # CUDA graph key is built below. The key includes this flag to pick + # the argmax vs advanced-sampling graph variant; populate (inside + # _prepare_inputs) runs later and fills the matching GPU buffers. + # Without this pre-scan the key would use the previous iteration's + # stale value and could replay the advanced graph against + # unpopulated (greedy) buffers, hanging the run (e.g. MTP nextn>=2). + if spec_metadata is not None: + spec_metadata.update_is_all_greedy_sample( + padded_graph_requests.all_requests()) + self._sync_group_all_greedy_sample(spec_metadata) + + peft_cache_data_type = None + if getattr(self, "cuda_graph_lora_manager", None) is not None: + peft_cache_manager = resource_manager.get_resource_manager( + ResourceManagerType.PEFT_CACHE_MANAGER) + peft_cache_data_type = peft_cache_manager.data_type + + maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( + padded_graph_requests, + enable_spec_decode=self.enable_spec_decode, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_tokens_cuda=self.draft_tokens_cuda + if self.is_spec_decode else None, + new_tensors_device=new_tensors_device, + spec_resource_manager=spec_resource_manager, + promoted_context_request_ids=promoted_context_request_ids, + peft_cache_data_type=peft_cache_data_type, + ) + + can_run_graph = key is not None + if can_run_graph: + attn_metadata = maybe_attn_metadata + spec_metadata = maybe_spec_metadata + execution_requests = padded_graph_requests + execution_promoted_context_ids = promoted_context_request_ids + else: + attn_metadata = self.attn_metadata + if self.enable_spec_decode: + spec_metadata = self.spec_metadata + else: + spec_metadata = None + execution_requests = scheduled_requests + execution_promoted_context_ids = frozenset() + + # Fill slot-ID buffer for scatter inside draft loop + if (self.enable_spec_decode and spec_tree_manager is not None + and spec_tree_manager.use_dynamic_tree + and not self.is_draft_model): + spec_tree_manager.slot_storage.fill_all_slot_ids( + execution_requests.context_requests, + execution_requests.generation_requests, + ) + + inputs, gather_ids = self._prepare_inputs( + execution_requests, kv_cache_manager, attn_metadata, + spec_metadata, new_tensors_device, cache_indirection_buffer, + num_accepted_tokens_device, req_id_to_old_request, + resource_manager, can_run_graph, execution_promoted_context_ids) + speculative_model = self.get_speculative_model() + spec_worker = (speculative_model.spec_worker + if speculative_model is not None else None) + if spec_worker is not None: + spec_worker.prepare_managed_draft_cache( + speculative_model.draft_model, spec_metadata, attn_metadata, + resource_manager) + if execution_promoted_context_ids: + self.iter_states[ + 'num_ctx_requests'] = scheduled_requests.num_context_requests + self.iter_states['num_ctx_tokens'] = sum( + request.context_chunk_size + for request in scheduled_requests.context_requests) + self.iter_states[ + 'num_generation_tokens'] = scheduled_requests.num_generation_requests + self._prepare_inputs_event = torch.cuda.Event() + self._prepare_inputs_event.record() + + breakable_runner = self.breakable_cuda_graph_runner + + with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): + + def forward_step(): + with MoeLoadBalancerIterContext(moe_load_balancer): + return self._forward_step( + inputs, + gather_ids=gather_ids, + gather_context_logits=gather_context_logits) + + if not can_run_graph: + if (breakable_runner is not None + and breakable_runner.is_capturing): + return breakable_runner.capture_model_body(forward_step) + + num_tokens = inputs['input_ids'].shape[0] + can_run_breakable_graph = ( + breakable_runner is not None + and get_per_request_prefill_cuda_graph_flag() + and not gather_context_logits + and breakable_runner.has_graph(num_tokens)) + if can_run_breakable_graph and not breakable_runner.is_warming_up: + outputs = breakable_runner.execute( + num_tokens, forward_step) + else: + # real eager or BCG warmup or PCG + outputs = forward_step() + else: + needs_capture = self.cuda_graph_runner.needs_capture(key) + if needs_capture: + + def capture_forward_fn(inputs: Dict[str, Any]): + with MoeLoadBalancerIterContext(moe_load_balancer): + return self._forward_step( + inputs, + gather_ids=gather_ids, + gather_context_logits=gather_context_logits) + + def capture_postprocess_fn(inputs: Dict[str, Any]): + self._postprocess_inputs(inputs) + + capture_outputs = self.cuda_graph_runner.capture( + key, + capture_forward_fn, + inputs, + enable_spec_decode=self.enable_spec_decode, + postprocess_fn=capture_postprocess_fn) + + if self.cuda_graph_runner.is_warmup_only: + outputs = capture_outputs + elif needs_capture: + # Refresh attention metadata for the current batch's + # draft cache before replaying the captured graph. + saved_draft = prepare_attn_metadata_for_draft_replay( + attn_metadata, draft_kv_cache_manager) + try: + outputs = self.cuda_graph_runner.replay(key, inputs) + finally: + restore_attn_metadata_after_draft_replay( + attn_metadata, saved_draft) + else: + saved_draft = prepare_attn_metadata_for_draft_replay( + attn_metadata, draft_kv_cache_manager) + try: + with MoeLoadBalancerIterContext(moe_load_balancer): + outputs = self.cuda_graph_runner.replay( + key, inputs) + finally: + restore_attn_metadata_after_draft_replay( + attn_metadata, saved_draft) + + if (spec_worker is not None and not self.is_warmup + and not self.cuda_graph_runner.is_warmup_only): + draft_history_update = spec_worker.snapshot_managed_draft_history( + ) + if draft_history_update is not None: + # Graph output dictionaries persist across replays. Each + # overlapped iteration must retain its own host readback. + outputs = dict(outputs) + outputs['draft_history_update'] = draft_history_update + + if self.forward_pass_callable is not None: + self.forward_pass_callable() + + self._execute_logit_post_processors(scheduled_requests, outputs) + + return outputs + + def _get_spec_worker(self): + """Access the spec_worker from DecoderModelForCausalLM (one-model spec dec).""" + return getattr(self.model, 'spec_worker', None) + + def model_forward(self, **kwargs): + attrs = get_model_extra_attrs() + assert attrs is not None, "Model extra attrs is not set" + attrs["attention_metadata"] = weakref.ref(kwargs['attn_metadata']) + attrs.update(self.model.model_config.extra_attrs) + attrs["spec_metadata"] = kwargs.get('spec_metadata', None) + + if self._torch_compile_backend is not None: + # Register aux streams and events to model extra attrs. + # The streams and events are list which could be updated during compilation. + attrs["aux_streams"] = weakref.ref(self.backend_num_streams) + attrs["events"] = weakref.ref(self._torch_compile_backend.events) + attrs["global_stream"] = torch.cuda.current_stream() + + if is_trace_enabled("TLLM_TRACE_MODEL_FORWARD"): + return trace_func(self.model.forward)(**kwargs) + else: + return self.model.forward(**kwargs) + + @nvtx_range("_forward_step") + def _forward_step(self, + inputs: Dict[str, Any], + *, + gather_ids: Optional[torch.Tensor] = None, + gather_context_logits: bool = False) -> Dict[str, Any]: + inputs = self._preprocess_inputs(inputs) + if inputs.get('spec_metadata', None): + gather_ids = inputs['spec_metadata'].gather_ids + + # For simplicity, just return all the the logits if we have special gather_ids + # from speculative decoding. + outputs = self.model_forward( + **inputs, + return_context_logits=gather_ids is not None + or gather_context_logits, + ) + + if self.without_logits: + return outputs + + if isinstance(outputs, dict): + # If the model returns a dict, get the logits from it. All other keys are kept. + logits = outputs.get('logits', None) + # If the logits are not found, no further processing is needed. + if logits is None: + return outputs + else: + # If the model returns a single tensor, assume it is the logits and wrap it in a dict. + logits = outputs + outputs = {'logits': logits} + + # If we have special gather_ids, gather the logits + if gather_ids is not None: + outputs['logits'] = logits[gather_ids] + + return outputs + + @nvtx_range("_forward_step_mm_encoder_only") + def _forward_step_mm_encoder_only( + self, inputs: Dict[str, Any], + scheduled_requests: ScheduledRequests) -> Dict[str, Any]: + """Forward step for multimodal encoder only mode - returns mm_embeddings instead of logits.""" + # Get multimodal parameters from inputs + multimodal_params = inputs.get("multimodal_params", []) + if not multimodal_params or len(multimodal_params) == 0: + # Return empty embeddings if no multimodal data + return { + 'mm_embeddings': [], + 'mm_embedding_request_indices': [], + 'mm_embedding_lengths': [], + } + # Some ctx requests carry only mrope metadata (no actual vision + # content). Skip them so the encoder only runs on real image payloads. + mm_context_requests = [(request_idx, request) for request_idx, request + in enumerate(scheduled_requests.context_requests) + if request.py_multimodal_data is not None] + if len(mm_context_requests) != len(multimodal_params): + raise ValueError( + "mm_encoder_only expects one multimodal payload per context " + "request carrying py_multimodal_data") + mm_request_indices_with_payload = [] + mm_params_with_payload = [] + mm_embedding_lengths = [] + for (request_idx, + request), multimodal_param in zip(mm_context_requests, + multimodal_params): + if not _has_mm_payload_keys(request.py_multimodal_data): + # mrope-only warmup request (no actual vision content) -> skip. + continue + multimodal_embedding_lengths = get_multimodal_embedding_lengths( + request) + if multimodal_embedding_lengths is None: + # Vision payload keys present but no pre-computed embedding + # lengths — skip to avoid a downstream sum(None) TypeError. + continue + mm_request_indices_with_payload.append(request_idx) + mm_params_with_payload.append(multimodal_param) + mm_embedding_lengths.append(multimodal_embedding_lengths) + if not mm_params_with_payload: + return { + 'mm_embeddings': [], + 'mm_embedding_request_indices': [], + 'mm_embedding_lengths': [], + } + # For mm_encoder_only mode, we only run the vision encoder part + # The model should be a vision encoder (e.g., Qwen2VisionModelBase) + mm_embeddings = self.model.forward(mm_params_with_payload) + assert len( + mm_embeddings + ) == 1, "mm_embeddings should be a 1-element list, mix modality (video+image) is not supported" + + split_lengths = [sum(lengths) for lengths in mm_embedding_lengths] + mm_embeddings = list(torch.split(mm_embeddings[0], split_lengths, + dim=0)) + if len(mm_embeddings) != len(mm_embedding_lengths): + raise ValueError( + "mm_encoder_only produced an embedding batch that does not " + "match mm_embedding_lengths") + + # Extract mrope position data from multimodal_params if available + mrope_position_ids_list = [] + mrope_position_deltas_list = [] + for multimodal_param in mm_params_with_payload: + mrope_config = multimodal_param.multimodal_data.get( + 'mrope_config', {}) + mrope_position_ids = mrope_config.get('mrope_position_ids') + mrope_position_deltas = mrope_config.get('mrope_position_deltas') + if mrope_position_ids is not None: + mrope_position_ids_list.append(mrope_position_ids) + if mrope_position_deltas is not None: + mrope_position_deltas_list.append(mrope_position_deltas) + + # mrope lists must align 1:1 with multimodal_params (or be empty); + # the sampler indexes them by per-MM-result position into mm_embeddings. + assert (len(mrope_position_ids_list) == len(mrope_position_deltas_list) + and len(mrope_position_ids_list) + in (0, len(mm_params_with_payload))), ( + f"mrope alignment: got {len(mrope_position_ids_list)} ids, " + f"{len(mrope_position_deltas_list)} deltas, " + f"{len(mm_params_with_payload)} mm params") + + result = { + 'mm_embeddings': mm_embeddings, + 'logits': None, + 'mm_embedding_request_indices': mm_request_indices_with_payload, + 'mm_embedding_lengths': mm_embedding_lengths, + } + if mrope_position_ids_list: + result['mrope_position_ids'] = mrope_position_ids_list + if mrope_position_deltas_list: + result['mrope_position_deltas'] = mrope_position_deltas_list + + return result + + def _make_encoder_attn_metadata( + self, + sequence_lengths: List[int], + request_ids: List[int], + ): + """Build fresh, no-cache attention metadata for one packed encoder + batch. ``self.attn_metadata`` is not reused because that object is + bound to the decoder's KV-cache manager.""" + if len(sequence_lengths) != len(request_ids): + raise ValueError("Encoder sequence lengths and request IDs must " + "have the same length.") + sparse_metadata_params = ( + self.sparse_attention_config.to_sparse_metadata_params( + pretrained_config=self.model.model_config.pretrained_config) + if self.sparse_attention_config is not None else None) + encoder_attn_metadata = self.attn_backend.Metadata( + max_num_requests=self.encoder_batch_size, + max_num_tokens=self.encoder_max_num_tokens, + max_num_sequences=self.encoder_batch_size * self.max_beam_width, + kv_cache_manager=None, + mapping=self.mapping, + runtime_features=self.attn_runtime_features, + enable_flash_mla=self.model.model_config.enable_flash_mla, + enable_context_mla_with_cached_kv=False, + cache_indirection=None, + sparse_metadata_params=sparse_metadata_params, + num_heads_per_kv=1, + ) + assert isinstance( + encoder_attn_metadata, + (VanillaAttentionMetadata, TrtllmAttentionMetadata) + ), "Only vanilla and trtllm attention metadata are supported for the encoder pass" + + encoder_attn_metadata.seq_lens = torch.tensor( + sequence_lengths, + dtype=torch.int, + pin_memory=prefer_pinned(), + ) + encoder_attn_metadata.num_contexts = len(sequence_lengths) + encoder_attn_metadata.max_seq_len = self.max_seq_len + encoder_attn_metadata.request_ids = request_ids + encoder_attn_metadata.prepare_encoder_only() + return encoder_attn_metadata + + def _prepare_encoder_decoder_encoder_inputs( + self, + encoder_input_ids: List[int], + encoder_position_ids: List[int], + sequence_lengths: List[int], + request_ids: List[int], + resource_manager: Optional[ResourceManager] = None, + ) -> Dict[str, Any]: + num_tokens = len(encoder_input_ids) + if num_tokens != len(encoder_position_ids): + raise ValueError("Encoder input IDs and position IDs must have " + "the same length.") + assert num_tokens <= self.encoder_max_num_tokens, ( + f"encoder packed length ({num_tokens}) exceeds " + f"encoder_max_num_tokens ({self.encoder_max_num_tokens})") + + encoder_attn_metadata = self._make_encoder_attn_metadata( + sequence_lengths, request_ids) + encoder_input_ids_t = torch.tensor(encoder_input_ids, + dtype=torch.int, + pin_memory=prefer_pinned()) + encoder_position_ids_t = torch.tensor(encoder_position_ids, + dtype=torch.int, + pin_memory=prefer_pinned()) + encoder_graph_runner = self.encoder_cuda_graph_runner + encoder_batch_size = len(sequence_lengths) + use_graph_staging = ( + encoder_graph_runner.enabled and + (encoder_batch_size in encoder_graph_runner.supported_batch_sizes or + (encoder_graph_runner.padding_enabled and encoder_batch_size + <= encoder_graph_runner.max_supported_batch_size))) + + return { + 'encoder_input_ids': + (encoder_input_ids_t if use_graph_staging else + encoder_input_ids_t.to('cuda', non_blocking=True)), + 'encoder_position_ids': + ((encoder_position_ids_t + if use_graph_staging else encoder_position_ids_t.to( + 'cuda', non_blocking=True)).unsqueeze(0)), + 'encoder_attn_metadata': + encoder_attn_metadata, + 'encoder_seq_lens': + sequence_lengths, + 'encoder_input_ids_host': + encoder_input_ids_t, + 'encoder_position_ids_host': + encoder_position_ids_t, + 'resource_manager': + resource_manager, + } + + @nvtx_range("_prepare_tp_inputs_encoder_features") + def _prepare_tp_inputs_encoder_features( + self, + encoder_requests: List[LlmRequest], + resource_manager: Optional[ResourceManager] = None, + ): + """Pack encoder inputs for feature-driven audio encoders (Whisper). + + The encoder input is a per-request feature tensor (an opaque audio + tensor, e.g. Whisper's 30 s-padded waveform) rather than token ids, and + the packed sequence lengths are the post-encoder position counts + (``encoder_output_len``), not the raw feature length. + """ + features: List[torch.Tensor] = [] + sequence_lengths: List[int] = [] + request_ids: List[int] = [] + + for request in encoder_requests: + request_features = request.py_encoder_input_features + if request_features is None: + raise ValueError( + f"Encoder request {request.py_request_id} has no " + "encoder_input_features; feature- and token-driven " + "encoder requests cannot share one batch.") + features.append(request_features) + sequence_lengths.append(int(request.encoder_output_len)) + request_ids.append(request.py_request_id) + + num_tokens = sum(sequence_lengths) + assert num_tokens <= self.encoder_max_num_tokens, ( + f"encoder packed length ({num_tokens}) exceeds " + f"encoder_max_num_tokens ({self.encoder_max_num_tokens})") + + encoder_attn_metadata = self._make_encoder_attn_metadata( + sequence_lengths, request_ids) + + inputs = { + 'input_features': + torch.cat(features, dim=0).to('cuda', non_blocking=True), + 'encoder_attn_metadata': + encoder_attn_metadata, + 'encoder_seq_lens': + sequence_lengths, + 'resource_manager': + resource_manager, + } + return inputs + + @nvtx_range("_prepare_tp_inputs_encoder") + def _prepare_tp_inputs_encoder( + self, + encoder_requests: List[LlmRequest], + resource_manager: Optional[ResourceManager] = None, + ): + """Pack encoder-side inputs for an encoder-decoder forward pass. + + Mirrors the no-cache path used by ``mm_encoder_only`` and the + legacy ``EncoderBuffers`` shape contract: ``encoder_input_ids`` + and ``encoder_position_ids`` are concatenated across requests + into a single ``[sum(encoder_output_len)]`` tensor, with one + non-causal :class:`AttentionMetadata` describing the packed + encoder batch. + + The encoder pass does not touch any KV-cache pool. The cross pool is + only written by the decoder's cross-attention on the first context + step. Self-pool blocks for the decoder are reserved on the next + scheduler iteration when the request transitions to ``CONTEXT_INIT``. + """ + if not encoder_requests: + raise ValueError( + "_prepare_tp_inputs_encoder called with no encoder requests") + + # Feature-driven audio encoders (Whisper) carry a tensor instead of + # encoder token ids; they take a dedicated prep path (which rejects + # mixed feature/token batches). + if any( + getattr(request, "py_encoder_input_features", None) is not None + for request in encoder_requests): + return self._prepare_tp_inputs_encoder_features( + encoder_requests, resource_manager=resource_manager) + + encoder_input_ids: List[int] = [] + encoder_position_ids: List[int] = [] + sequence_lengths: List[int] = [] + request_ids: List[int] = [] + + for request in encoder_requests: + tokens = request.encoder_tokens + if tokens is None: + raise ValueError( + f"Encoder request {request.py_request_id} has no " + "encoder_tokens; encoder_input_token_ids must be wired " + "through executor_request_to_llm_request.") + seq_len = len(tokens) + encoder_input_ids.extend(tokens) + encoder_position_ids.extend( + self._apply_position_id_offset(list(range(seq_len)))) + sequence_lengths.append(seq_len) + request_ids.append(request.py_request_id) + + return self._prepare_encoder_decoder_encoder_inputs( + encoder_input_ids=encoder_input_ids, + encoder_position_ids=encoder_position_ids, + sequence_lengths=sequence_lengths, + request_ids=request_ids, resource_manager=resource_manager, is_dummy=self._warmup_depth > 0, ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e0502b4cd181..c52321250b31 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -9081,11 +9081,10 @@ def _sample_async(self, scheduled_batch, sample_state = self.sampler.sample_async( scheduled_batch, batch_outputs, num_context_logits_prefix_sum) - self._maybe_record_hang_diagnostic_phase( - "sampling_returned", scheduled_batch) + if sample_state is not None: + sample_state.draft_history_update = batch_outputs.get( + 'draft_history_update') return sample_state - self._maybe_record_hang_diagnostic_phase("sampling_skipped", - scheduled_batch) except Exception as e: self._maybe_record_hang_diagnostic_phase("sampling_error", scheduled_batch) @@ -9107,10 +9106,15 @@ def _setup_sampler_step(self, requests: ScheduledRequests): @nvtx_range("_update_requests") def _update_requests(self, - sample_state: SampleState, + sample_state: SampleState | None, resource_manager: Optional[ResourceManager] = None): + if sample_state is None: + return try: self.sampler.update_requests(sample_state, resource_manager) + if sample_state.draft_history_update is not None: + sample_state.draft_history_update.publish() + sample_state.draft_history_update = None self._accumulate_spec_dec_stats(sample_state) except Exception as e: traceback.print_exc() diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index b1f256e560c7..6fcedbea0505 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -42,16 +42,14 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager from ..scheduler import ScheduledRequests -from .beam_search import ( - BeamHistory, - BeamHistoryBuilder, - BeamSearchHandler, - BeamSearchParams, - BeamSearchStep, - BeamSearchStore, - PendingBeamHistories, - finalize_beam, - prepare_beam_search, +from ..standalone_draft_cache import DraftHistoryUpdate +from .beam_search import BeamHistoryBuilder, BeamSearchHandler, finalize_beam, prepare_beam_search +from .finish_reasons import FinishReasonsHandler +from .logprobs import ( + LogProbsState, + LogProbsStateList, + LogProbsStore, + store_logprobs_list_to_request, ) from .finish_reasons import FinishReasonsHandler from .logprobs import LogProbsHandler, LogProbsState, LogProbsStateList, LogProbsStore @@ -114,6 +112,7 @@ class SampleState(Generic[GenericSampleStateTensorsHost, GenericSampleStateTenso host: Optional[GenericSampleStateTensorsHost] = None sampler_event: Optional[SamplerEvent] = None runtime_draft_len: Optional[int] = None + draft_history_update: DraftHistoryUpdate | None = None # Generic bounds not supported, https://github.com/python/typing/issues/548 diff --git a/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py b/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py new file mode 100644 index 000000000000..ce224f3a4649 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Storage and history contracts for manager-owned DSpark/DFlash drafters.""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from .kv_cache_manager_v2 import KVCacheManagerV2 + + +@dataclass(frozen=True) +class StandaloneDraftHistory: + """Committed length and next absolute position, independent of target/scratch state.""" + + valid_length: int + position: int + + def __post_init__(self) -> None: + if type(self.valid_length) is not int or type(self.position) is not int: + raise ValueError("Standalone draft history requires integer length and position") + if self.valid_length < 0 or self.position < self.valid_length: + raise ValueError("Invalid standalone draft history length or position") + + +@dataclass(frozen=True) +class DraftHistoryUpdate: + """One execution's draft history, published after its sampling event completes.""" + + manager: "KVCacheManagerV2" + request_ids: tuple[int, ...] + cache_instances: tuple[object, ...] + values_host: torch.Tensor + + @classmethod + def capture( + cls, + manager: "KVCacheManagerV2", + request_ids: tuple[int, ...], + values: torch.Tensor, + ) -> "DraftHistoryUpdate": + """Queue an independent length/position readback on the execution stream.""" + if values.shape != (len(request_ids), 2): + raise ValueError("Draft history update requires one length/position pair per request") + values_host = torch.empty( + values.shape, dtype=values.dtype, device="cpu", pin_memory=values.is_cuda + ) + values_host.copy_(values, non_blocking=True) + return cls( + manager, + tuple(request_ids), + tuple(manager.kv_cache_map[request_id] for request_id in request_ids), + values_host, + ) + + def publish(self) -> None: + """Publish completed writes without reviving released or replaced requests.""" + for request_id, cache, (length, position) in zip( + self.request_ids, self.cache_instances, self.values_host.tolist() + ): + current = self.manager.kv_cache_map.get(request_id) + if current is cache and current.is_active: + self.manager.set_draft_history(request_id, length, position) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 6b7e6b517876..71b2cad0af3e 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -33,7 +33,9 @@ from ..modules.logits_processor import LogitsProcessor from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID -from ..pyexecutor.resource_manager import BaseResourceManager +from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from ..pyexecutor.resource_manager import BaseResourceManager, ResourceManagerType +from ..pyexecutor.standalone_draft_cache import DraftHistoryUpdate from .accept_stats import maybe_create_recorder from .dflash_attention import ( get_dflash_paged_append, @@ -55,6 +57,7 @@ if TYPE_CHECKING: from ...llmapi.llm_args import DFlashDecodingConfig + from ..pyexecutor.resource_manager import ResourceManager def dflash_context_dtype(draft_model: nn.Module) -> torch.dtype: @@ -599,7 +602,7 @@ def prepare(self): # Update slot mapping for DFlash context buffers worker = getattr(self, "_dflash_worker", None) - if worker is not None and worker._ctx_buf_inited: + if worker is not None and worker._ctx_buf_inited and not worker._has_unified_draft_cache(): current = set(self.request_ids) evicted = {} for rid in list(worker._req_to_slot.keys()): @@ -752,12 +755,21 @@ def __init__( # graph compatible. self._ctx_buf_inited = False self._ctx_len = None + self._ctx_position_offset = None + self._managed_cache_bindings = {} + self._managed_request_ids = () + self._managed_snapshot_slots: Optional[torch.Tensor] = None # Host shadows of _ctx_len and of each request's prompt progress. self._ctx_len_host = None self._req_ctx_pos = {} + # Decode the manager draft pool's device block offsets instead of + # building the table on the host (see _refresh_unified_block_tables). + self._device_draft_tables = False + self._draft_table_checks_left = 0 # Snapshot for rolling back in-place _ctx_len updates when a forward # fails (or after warmup). See _ensure_spec_dec_state_restored. self._saved_ctx_len = None + self._saved_ctx_position_offset = None self._saved_ctx_len_host = None self._saved_req_ctx_pos = None self._ctx_len_restore_pending = False @@ -1004,6 +1016,142 @@ def _argmax_block_tokens( local_block.reshape(-1, shard), out_dtype=torch.int32 ).view(num_gens, block_len) + def get_draft_kv_cache_manager( + self, resource_manager: Optional["ResourceManager"] + ) -> Optional[BaseResourceManager]: + if resource_manager is not None: + manager = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) + if bool(getattr(manager, "draft_layer_ids", ())): + return manager + return super().get_draft_kv_cache_manager(resource_manager) + + def _has_unified_draft_cache(self) -> bool: + return bool(getattr(self._ctx_kv_manager, "draft_layer_ids", ())) + + def _prepare_managed_history(self, request_ids: list[int]) -> list[int]: + """Restore newly resident histories without rewinding overlapping iterations.""" + updates = {self._dummy_slot: 0} + offsets = {self._dummy_slot: 0} + restore_rows = [] + real_request_ids = { + request_id + for request_id in request_ids + if request_id != ATTENTION_DP_DUMMY_REQUEST_ID + and request_id < self._graph_dummy_id_floor + and request_id not in self._ctx_kv_manager._draft_dummy_request_ids + } + for request_id in list(self._req_to_slot): + if request_id not in real_request_ids: + slot = self._req_to_slot.pop(request_id) + updates[slot] = 0 + offsets[slot] = 0 + self._req_ctx_pos.pop(request_id, None) + self._managed_cache_bindings.pop(request_id, None) + self._free_slots.append(slot) + for row, request_id in enumerate(request_ids): + if request_id in real_request_ids: + self._assign_slot(request_id) + slot = self._req_to_slot.get(request_id) + if slot is not None: + cache = self._ctx_kv_manager.kv_cache_map[request_id] + if self._managed_cache_bindings.get(request_id) is cache: + continue + history = self._ctx_kv_manager.get_draft_history(request_id) + updates[slot] = history.valid_length if history is not None else 0 + self._req_ctx_pos[request_id] = history.position if history is not None else 0 + offsets[slot] = self._req_ctx_pos[request_id] - updates[slot] + self._managed_cache_bindings[request_id] = cache + restore_rows.append(row) + self._write_ctx_len(updates) + self._write_slot_values(self._ctx_position_offset, offsets) + self._batch_to_slot[: len(request_ids)].copy_( + torch.tensor( + [self._req_to_slot.get(request_id, self._dummy_slot) for request_id in request_ids], + dtype=torch.long, + device="cpu", + pin_memory=prefer_pinned(), + ), + non_blocking=True, + ) + self._managed_request_ids = tuple( + request_id for request_id in request_ids if request_id in real_request_ids + ) + # Each upload owns fresh host storage. PyTorch's pinned allocator defers + # recycling it until the asynchronous copy completes. + self._managed_snapshot_slots[: len(self._managed_request_ids)].copy_( + torch.tensor( + [self._req_to_slot[request_id] for request_id in self._managed_request_ids], + dtype=torch.long, + device="cpu", + pin_memory=prefer_pinned(), + ), + non_blocking=True, + ) + return restore_rows + + def _gather_managed_context(self, request_ids: list[int], restore_rows: list[int]) -> None: + """Refresh VANILLA's dense staging; manager pages remain authoritative.""" + if self._dflash_attention_backend != "VANILLA": + return + for row in restore_rows: + request_id = request_ids[row] + slot = self._req_to_slot.get(request_id) + if slot is None: + continue + length = self._ctx_len_host[slot] + positions = torch.arange(length, device=self._ctx_k_buf.device) + pages = self._ctx_block_tables[row, positions // self._ctx_page_size].long() + offsets = positions % self._ctx_page_size + for layer_idx, pool in enumerate(self._ctx_kv_buf): + self._ctx_k_buf[slot, layer_idx, :length] = pool[pages, 0, :, offsets, :] + self._ctx_v_buf[slot, layer_idx, :length] = pool[pages, 1, :, offsets, :] + + def prepare_managed_draft_cache( + self, draft_model, spec_metadata, attn_metadata, resource_manager + ) -> None: + """Refresh managed inputs before eager execution, capture, or graph replay.""" + manager = self.get_draft_kv_cache_manager(resource_manager) + if not getattr(manager, "draft_layer_ids", ()): + return + self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata, manager) + spec_metadata._dflash_worker = self + restore_rows = self._prepare_managed_history(spec_metadata.request_ids) + self._refresh_ctx_block_tables( + attn_metadata, attn_metadata.num_seqs, spec_metadata.request_ids + ) + self._gather_managed_context(spec_metadata.request_ids, restore_rows) + + def _write_slot_values(self, target: torch.Tensor, values: dict[int, int]) -> None: + """Scatter a slot-to-value mapping into a device tensor with one upload.""" + if not values: + return + slots, vals = zip(*values.items()) + pinned = prefer_pinned() + target.index_copy_( + 0, + torch.tensor(slots, dtype=torch.long, pin_memory=pinned).to( + target.device, non_blocking=True + ), + torch.tensor(vals, dtype=target.dtype, pin_memory=pinned).to( + target.device, non_blocking=True + ), + ) + + def snapshot_managed_draft_history(self) -> DraftHistoryUpdate | None: + """Queue iteration-owned readback; the executor publishes after completion.""" + if not self._has_unified_draft_cache() or not self._managed_request_ids: + return None + # Preparation, forward, snapshot gathers and the next preparation run + # on the execution stream, so its next upload cannot overtake this read. + slots = self._managed_snapshot_slots[: len(self._managed_request_ids)] + lengths = self._ctx_len[slots] + positions = lengths + self._ctx_position_offset[slots] + return DraftHistoryUpdate.capture( + self._ctx_kv_manager, + self._managed_request_ids, + torch.stack((lengths, positions), dim=1), + ) + def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype, kv_factor=2): """Fail with the arithmetic before allocating the drafter context arena. @@ -1059,13 +1207,21 @@ def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype, kv_facto f"scales it down linearly too." ) - def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor=2): + def _managed_ctx_pool( + self, + draft_kv_cache_manager: BaseResourceManager | None, + L: int, + nkv: int, + hd: int, + dtype: torch.dtype, + kv_factor: int = 2, + ) -> list[torch.Tensor] | None: """Per-layer views of the draft KV cache manager's pool, or None. - The manager already funds a pool sized from the drafter's own config - (K3 at TP8: 5 layers x 2 x 2 kv heads x 64 x 2B = 2560 B/token), and - ``get_buffers(..., "HND")`` hands back exactly the per-layer shape the - private arena uses, so nothing but the ownership of the memory changes. + The manager's buffer accessor supplies the same per-layer HND shape + as the private arena. Unified storage uses ``get_draft_buffers`` and + requires the pool to match the drafter and the manager's page size; + legacy storage uses ``get_buffers`` and may fall back to the arena. A list, not one stacked tensor: only the two consumers that index it by layer need it, and the two managers disagree on what a stack would mean. @@ -1075,19 +1231,33 @@ def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor Note the page count is not the drafter's to interpret: a V2 view spans the whole interleaved pool, not this layer's slice, so only shape[1:] is - checked here and the index space is settled in _init_ctx_block_tables. + checked here. Legacy offsets are settled in _init_ctx_block_offsets; + unified storage supplies draft block indices directly. """ if draft_kv_cache_manager is None: # No separate draft KV cache. Attention DP alone does NOT disable # it: _util.py skips that bail for an external drafter, which is # what DSpark is, so this drafter does get a manager under DP. return None - layers = [draft_kv_cache_manager.get_buffers(i, kv_layout="HND") for i in range(L)] + unified = bool(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) + get_buffers = ( + draft_kv_cache_manager.get_draft_buffers + if unified + else draft_kv_cache_manager.get_buffers + ) + layers = [get_buffers(i, kv_layout="HND") for i in range(L)] base = layers[0] + page_size = draft_kv_cache_manager.tokens_per_block if unified else base.size(-2) # kv_factor 1 is the MLA drafter's SELFKONLY pool: one latent, no V. - expected = (kv_factor, nkv, base.size(-2), hd) + expected = (kv_factor, nkv, page_size, hd) + # A unified pool may store the drafter's K/V in FP8; writes convert. + allowed_dtypes = (dtype, torch.float8_e4m3fn) if unified else (dtype,) for i, layer in enumerate(layers): - if tuple(layer.shape[1:]) != expected or layer.dtype != dtype: + if tuple(layer.shape[1:]) != expected or layer.dtype not in allowed_dtypes: + if unified: + raise ValueError( + "Unified DSpark draft pool does not match the drafter KV layout" + ) # Fall back rather than fail: the private arena is what every # DFlash model shipped with, so an unrecognized pool costs the # old memory footprint, not the run. @@ -1102,10 +1272,8 @@ def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor ) return layers - def _init_ctx_block_tables( - self, draft_kv_cache_manager, pool, num_slots, L, nkv, hd, page_size - ) -> bool: - """Size the per-iteration block table and the offset decode constant. + def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page_size) -> bool: + """Validate the legacy pool layout and resolve its offset decode constant. ``kv_cache_block_offsets`` entries are ``pool_block_index * num_pool_layers * kv_factor`` plus the K/V field index. What that has to @@ -1141,7 +1309,13 @@ def _init_ctx_block_tables( f"rather than guessing block indices." ) return False - max_blocks = mgr.max_blocks_per_seq + logger.info( + f"DFlash: ctx block offsets pool_idx={self._ctx_pool_idx}, " + f"divisor={self._ctx_block_divisor}, view_pages={pool[0].size(0)}" + ) + return True + + def _init_ctx_block_tables(self, num_slots: int, max_blocks: int) -> None: self._ctx_block_tables = torch.zeros( (num_slots, max_blocks), dtype=torch.int32, device="cuda" ) @@ -1149,14 +1323,11 @@ def _init_ctx_block_tables( 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" ) self._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") - logger.info( - f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}, " - f"pool_idx={self._ctx_pool_idx}, divisor={self._ctx_block_divisor}, " - f"view_pages={pool[0].size(0)}" - ) - return True + logger.info(f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}") - def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: + def _refresh_ctx_block_tables( + self, attn_metadata, num_seqs: int, request_ids: list[int] | None = None + ) -> bool: """Decode this iteration's draft block table into the persistent buffer. In place and fixed-shape so a captured graph keeps reading the same @@ -1165,6 +1336,26 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: """ if self._ctx_block_tables is None or num_seqs <= 0: return False + if self._has_unified_draft_cache(): + if self._device_draft_tables and self._refresh_unified_block_tables( + attn_metadata, num_seqs, request_ids + ): + return True + table = self._ctx_kv_manager.get_draft_block_table(request_ids[:num_seqs]) + self._ctx_block_tables[:num_seqs].copy_(table, non_blocking=True) + self._ctx_block_counts[:num_seqs].copy_( + torch.tensor( + [ + self._ctx_kv_manager.get_draft_num_blocks(rid) + for rid in request_ids[:num_seqs] + ], + dtype=torch.long, + device="cpu", + pin_memory=prefer_pinned(), + ), + non_blocking=True, + ) + return True src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None) if src is None: # The table is bound but unfillable. Leaving it at zeros would send @@ -1187,22 +1378,95 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: ) return True - def _lazy_init_ctx_buffers( - self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None - ): - self._validate_draft_attention_backend(draft_model) - # Only TrtllmAttentionMetadata builds draft_kv_cache_block_offsets. - # Without it the table stays all-zero and every request would read and - # write pool block 0 -- silently, at acceptance 1.0. - if draft_kv_cache_manager is not None and not hasattr( - attn_metadata, "draft_kv_cache_block_offsets" + def _init_unified_draft_tables(self, manager) -> None: + """Choose how the unified pool's block table is built each iteration.""" + self._device_draft_tables = False + layer = manager.draft_layer_ids[0] + pool = manager.layer_to_pool_mapping_dict[manager.layer_offsets[layer]] + if ( + getattr(manager, "_helix_cp_size", 1) != 1 + or getattr(manager, "enable_swa_scratch_reuse", False) + or manager._index_scale_ints[pool] != manager.get_layer_page_index_scale(layer) ): logger.warning( - f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; " - "using the private ctx arena." + "DFlash: the draft pool's block offsets cannot be decoded on the device; " + "building draft block tables on the host." ) - draft_kv_cache_manager = None + return + self._ctx_pool_idx = pool + self._ctx_block_divisor = manager.get_layer_kv_factor(layer) + self._device_draft_tables = True + self._draft_table_checks_left = 8 + logger.info( + f"DFlash: draft block tables decoded on the device from pool {pool} " + f"(divisor {self._ctx_block_divisor})" + ) + + def _refresh_unified_block_tables( + self, attn_metadata, num_seqs: int, request_ids: list[int] + ) -> bool: + """Decode the draft pool's rows of the target's device block offsets. + + The target metadata already uploads every pool's offsets for this + batch, so the draft rows need one device divide plus the per-request + block counts, rather than a host table rebuilt from page lists. + Returns False when the host path must build this iteration's table. + """ + src = getattr(attn_metadata, "kv_cache_block_offsets", None) + if src is None or src.size(0) <= self._ctx_pool_idx: + return False + width = min(src.size(-1), self._ctx_block_tables.size(1)) + encoded = src[self._ctx_pool_idx, :num_seqs, 0, :width] + decoded = torch.div(encoded.clamp(min=0), self._ctx_block_divisor, rounding_mode="floor") + self._ctx_block_tables[:num_seqs, :width].copy_(decoded) + cache_map = self._ctx_kv_manager.kv_cache_map + counts = [] + for request_id in request_ids[:num_seqs]: + cache = cache_map.get(request_id) + counts.append(cache.num_blocks if cache is not None else 0) + self._ctx_block_counts[:num_seqs].copy_( + torch.tensor(counts, dtype=torch.long, pin_memory=prefer_pinned()), + non_blocking=True, + ) + if self._draft_table_checks_left > 0 and not torch.cuda.is_current_stream_capturing(): + self._draft_table_checks_left -= 1 + if not self._check_device_block_tables(num_seqs, request_ids, counts): + return False + return True + + def _check_device_block_tables( + self, num_seqs: int, request_ids: list[int], counts: list[int] + ) -> bool: + """Compare device-decoded rows with the manager's page lists (startup only).""" + manager = self._ctx_kv_manager + rows = [] + for row, request_id in enumerate(request_ids[:num_seqs]): + cache = manager.kv_cache_map.get(request_id) + if counts[row] > 0 and cache is not None and cache.is_active: + rows.append(row) + if not rows: + return True + expected = manager.get_batch_cache_indices( + [request_ids[row] for row in rows], manager.draft_layer_ids[0] + ) + device_rows = self._ctx_block_tables[rows].cpu() + for k, row in enumerate(rows): + got = device_rows[k, : counts[row]].tolist() + want = expected[k][: counts[row]] + if any(w >= 0 and g != w for g, w in zip(got, want)): + logger.error( + "DFlash: device-decoded draft block table disagrees with the manager " + f"for request {request_ids[row]}; building draft tables on the host." + ) + self._device_draft_tables = False + return False + return True + def _lazy_init_ctx_buffers( + self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None + ): + self._validate_draft_attention_backend(draft_model) + unified = bool(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) if self._ctx_buf_inited: if self._ctx_kv_manager is draft_kv_cache_manager: return @@ -1270,15 +1534,16 @@ def _lazy_init_ctx_buffers( self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") + self._ctx_position_offset = torch.zeros_like(self._ctx_len) self._ctx_len_host = [0] * num_slots - self._batch_to_slot_buffer = DeviceCopyBuffer( - torch.zeros(max_batch, dtype=torch.long, device="cuda"), - pin_memory=prefer_pinned(), - ) + self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._managed_snapshot_slots = torch.empty_like(self._batch_to_slot) if unified else None self._free_slots = deque(range(max_batch)) self._req_to_slot = {} self._req_ctx_pos = {} + self._managed_cache_bindings = {} + self._managed_request_ids = () # checkpoint's trained block width self._resolved_block_size = getattr(draft_model, "block_size", None) or ( @@ -1320,14 +1585,9 @@ def _lazy_init_ctx_buffers( use_paged = self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS or getattr( draft_model, "_paged_ctx_cache", False ) - self._ctx_paged = use_paged - if use_paged: - # FA4 pages, but on the private arena with its own kernel. - pool = ( - None - if self._dflash_attention_backend == "FA4" - else self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor) - ) + self._ctx_paged = unified or use_paged + if self._ctx_paged: + pool = self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor) # The manager's page size wins when bound to it: its pool is already # carved, so the drafter adopts the geometry rather than imposing one. page_size = self._ctx_page_size if pool is None else pool[0].size(-2) @@ -1355,8 +1615,12 @@ def _lazy_init_ctx_buffers( # Settle the block table before committing to the pool: it is the # last thing that can rule the pool out, and falling back after # taking the pool branch would leave no buffer allocated at all. - if pool is not None and not self._init_ctx_block_tables( - draft_kv_cache_manager, pool, num_slots, L, nkv, hd, page_size + if ( + not unified + and pool is not None + and not self._init_ctx_block_offsets( + draft_kv_cache_manager, pool, L, nkv, hd, page_size + ) ): pool = None if pool is None: @@ -1387,25 +1651,27 @@ def _lazy_init_ctx_buffers( # block table one iteration at a time, so the footprint follows # the sequences served rather than max_batch x max_seq_len. self._ctx_kv_buf = pool - self._ctx_kv_ptrs = torch.tensor( - [layer.data_ptr() for layer in pool], dtype=torch.int64, device="cuda" + max_blocks = ( + draft_kv_cache_manager.draft_max_blocks_per_seq + if unified + else draft_kv_cache_manager.max_blocks_per_seq ) - # The fused preparation stores single-latent context entries - # projected from the captured target features through fc and - # hidden_norm (which every loaded DFlash drafter owns), as the - # eager path does; LiLiCorr drafters keep the eager path. - self._paged_latent_prepare = kv_factor == 1 and not draft_model._uses_lilicorr + self._init_ctx_block_tables(num_slots, max_blocks) + if unified: + self._init_unified_draft_tables(draft_kv_cache_manager) # Only a manager that publishes and matches its own blocks can # hand back a reused prefix's drafter K/V; the private arena and # an unpaired draft pool both start every request at 0. - self._ctx_reuse_addressable = bool( - getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) - ) + if not unified: + self._ctx_reuse_addressable = bool( + getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) + ) self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) - else: # dense private arena (VANILLA, unpaged) - self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) + if not use_paged: + if not unified: + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) kv_shape = (num_slots, L, capacity, nkv, hd) self._ctx_k_buf = torch.zeros(kv_shape, dtype=dtype, device="cuda") # An MLA drafter stores one latent per token; leaving _ctx_v_buf as @@ -1470,14 +1736,21 @@ def clear(slot: int) -> None: if reset and req_id in self._req_to_slot: old_slot = self._req_to_slot.pop(req_id) - clear(old_slot) + self._ctx_len[old_slot] = 0 + self._ctx_len_host[old_slot] = 0 self._free_slots.append(old_slot) if req_id not in self._req_to_slot: if not self._free_slots: + if self._has_unified_draft_cache(): + raise RuntimeError("Unified DSpark has no free request staging slot") return None slot = self._free_slots.popleft() self._req_to_slot[req_id] = slot - clear(slot) + if self._has_unified_draft_cache(): + self._managed_cache_bindings.pop(req_id, None) + self._ctx_position_offset[slot] = 0 + self._ctx_len[slot] = 0 + self._ctx_len_host[slot] = 0 return self._req_to_slot[req_id] def _get_ctx_paged_append(self) -> Callable[..., None]: @@ -1493,10 +1766,9 @@ def _store_context_kv_paged( rows: torch.Tensor, positions: torch.Tensor, ) -> None: - if v is None: - # Single-latent (MLA) cache: one vector per token per layer, so - # there is no K/V interleave for append_paged_kv_cache to do and the - # write is a plain scatter into the page the block table names. + fp8_cache = self._ctx_kv_buf[0].dtype == torch.float8_e4m3fn + if v is None or self._dflash_attention_backend == "VANILLA" or fp8_cache: + # VANILLA, single-latent (MLA) and FP8 caches use direct page writes. table = ( self._ctx_block_tables if self._ctx_block_tables is not None @@ -1506,7 +1778,21 @@ def _store_context_kv_paged( pages = table[rows.long(), positions.long() // page_size].long() offsets = positions.long() % page_size for layer_idx in range(k.size(1)): - self._ctx_kv_buf[layer_idx][pages, 0, 0, offsets] = k[:, layer_idx, 0] + pool = self._ctx_kv_buf[layer_idx] + if fp8_cache: + # Byte views: FP8 index_put support varies across builds. + pool_bytes = pool.view(torch.uint8) + pool_bytes[pages, 0, :, offsets] = ( + k[:, layer_idx].to(pool.dtype).view(torch.uint8) + ) + if v is not None: + pool_bytes[pages, 1, :, offsets] = ( + v[:, layer_idx].to(pool.dtype).view(torch.uint8) + ) + continue + pool[pages, 0, :, offsets] = k[:, layer_idx].to(pool.dtype) + if v is not None: + pool[pages, 1, :, offsets] = v[:, layer_idx].to(pool.dtype) return append_paged_kv_cache = self._get_ctx_paged_append() @@ -1517,12 +1803,13 @@ def _store_context_kv_paged( positions_i32 = positions.to(torch.int32) kv_indices, kv_indptr = self._ctx_paged_index_args() for layer_idx in range(k.size(1)): - append_paged_kv_cache( - append_key=k[:, layer_idx].contiguous(), - append_value=v[:, layer_idx].contiguous(), + pool = self._ctx_kv_buf[layer_idx] + trtllm_gen_ops.append_paged_kv_cache( + append_key=k[:, layer_idx].to(pool.dtype).contiguous(), + append_value=v[:, layer_idx].to(pool.dtype).contiguous(), batch_indices=rows_i32, positions=positions_i32, - paged_kv_cache=self._ctx_kv_buf[layer_idx], + paged_kv_cache=pool, kv_indices=kv_indices, kv_indptr=kv_indptr, kv_last_page_len=self._ctx_kv_last_page_len, @@ -1572,23 +1859,6 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount attn_metadata.kv_lens_cuda[nc:bs].clamp_(min=0) - def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> None: - """Add this step's accepted tokens to the context lengths of ``slots``, on the device. - - Padding, warmup and attention-DP dummy rows share the dummy slot, which no request owns - or frees. Emptying it right after the add keeps it at 0 between steps: those rows attend - no context, and their next context K / V starts at column 0, inside the padding - request's own pages (past them, its page-table row maps page 0, another request's). - The store is to a fixed slot, so a captured step replays it. - """ - self._ctx_len[slots] += num_accepted - self._ctx_len.clamp_(max=self._max_ctx) - self._ctx_len[self._dummy_slot].zero_() - - def _rebuild_batch_to_slot(self, request_ids: List[int]) -> None: - mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] - self._batch_to_slot_buffer.update(torch.tensor(mapping, dtype=torch.long)) - def _write_ctx_len(self, updates: dict[int, int]) -> None: """Apply a slot-to-length mapping to _ctx_len in one async scatter.""" if not updates: @@ -1612,6 +1882,8 @@ def _restore_ctx_len_host(self) -> None: self._ctx_len_host = list(self._saved_ctx_len_host) if self._saved_req_ctx_pos is not None: self._req_ctx_pos = dict(self._saved_req_ctx_pos) + if self._has_unified_draft_cache() and self._saved_ctx_position_offset is not None: + self._ctx_position_offset.copy_(self._saved_ctx_position_offset) def _store_prefill_context( self, @@ -1654,55 +1926,63 @@ def _store_prefill_context( ctx_len_updates = {} offset = 0 num_contexts = attn_metadata.num_contexts - # A context write addresses pages through row i of the block table, so - # it must fit that row's allocation as well as the arena's length. - # One tolist() for the whole loop: the lengths themselves come from - # _ctx_len_host, so the loop body below stays sync-free. - ctx_alloc = ( - (self._ctx_block_counts[:num_contexts] * self._ctx_page_size).tolist() - if self._ctx_block_tables is not None - else None - ) + # Host copies of each chunk's start avoid a device sync per request. + cached_tokens = None + kv_cache_params = getattr(attn_metadata, "kv_cache_params", None) + if self._has_unified_draft_cache() and kv_cache_params is not None: + cached_tokens = getattr(kv_cache_params, "num_cached_tokens_per_seq", None) for i in range(num_contexts): req_id = spec_metadata.request_ids[i] slen = int(attn_metadata._seq_lens[i]) chunk_proj = ctx_proj[offset : offset + slen].detach() chunk_pos = position_ids[offset : offset + slen].long().detach() - if slen == 0: - first_pos = 0 - elif host_pos is not None: - first_pos = int(host_pos[offset]) + if cached_tokens is not None and i < len(cached_tokens): + first_pos = int(cached_tokens[i]) else: - first_pos = chunk_pos[0].item() - - # A chunk continues what we already stored only if it starts exactly - # where the previous one ended. Everything else -- a fresh request, a - # request id reused after completion, a prefill restarted after - # preemption -- has to start from a clean slot. - reset = self._req_ctx_pos.get(req_id) != first_pos - if self._assign_slot(req_id, reset=reset, updates=ctx_len_updates) is None: + first_pos = chunk_pos[0].item() if slen > 0 else 0 + + # Assign slots for new requests or reset reused IDs in private caches. + # Managed preparation already binds fresh histories and records their + # snapshot slots; keep those slots stable through the forward pass. + reset_slot = first_pos == 0 and not self._has_unified_draft_cache() + if self._assign_slot(req_id, reset=reset_slot) is None: logger.warning("DFlash: no free slots, skipping context store") self._req_ctx_pos.pop(req_id, None) offset += slen continue slot = self._req_to_slot[req_id] - cur = ctx_len_updates.get(slot, self._ctx_len_host[slot]) - cap = self._max_ctx if ctx_alloc is None else min(self._max_ctx, ctx_alloc[i]) + if self._has_unified_draft_cache(): + # Restore and prefill keep the host shadow current here. + cur = self._ctx_len_host[slot] + else: + cur = int(self._ctx_len[slot].item()) if cur == 0 and first_pos > 0 and self._ctx_reuse_addressable: # Reuse let the target skip [0, first_pos), so only the tail # reaches the capture buffer. The prefix K/V are still in the # paired pool: precompute_context_kv is per-token in (hidden, p). cur = first_pos - if cur + slen > cap: - # Request-level, like the no-free-slots path above: truncating - # would silently draft from a stale prefix, but killing the - # forward would take every other in-flight request with it. - logger.warning( - f"DFlash: ctx overflow on slot {slot} " - f"({cur} + {slen} > {cap}); skipping its context " - "store, so this request drafts nothing." + if cur + slen > self._max_ctx: + if self._has_unified_draft_cache(): + self._ctx_len[slot] = 0 + self._ctx_len_host[slot] = 0 + self._req_ctx_pos.pop(req_id, None) + self._ctx_position_offset[slot] = first_pos + slen + self._managed_cache_bindings[req_id] = self._ctx_kv_manager.kv_cache_map[req_id] + offset += slen + continue + # Truncating here keeps the OLDEST tokens and drops the newest — + # the arena is append-only, not a sliding window — so the drafter + # would go on drafting from a stale prefix with no error and only + # a quiet acceptance-rate loss to show for it. ctx_len tracks the + # target sequence 1:1, so overflow means max_seq_len disagrees + # with the workload, which is a configuration bug. + raise RuntimeError( + f"DFlash drafter context overflow on slot {slot}: " + f"{cur} + {slen} > max_ctx {self._max_ctx}. The drafter " + f"context tracks the target sequence 1:1, so raise " + f"max_seq_len to cover the sequences being served." ) self._req_to_slot.pop(req_id, None) self._req_ctx_pos.pop(req_id, None) @@ -1720,8 +2000,13 @@ def _store_prefill_context( cache_dtype = ( self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype ) + if cache_dtype == torch.float8_e4m3fn: + # The K/V projection runs in the drafter's weight dtype; + # the store converts its outputs to the cache dtype. + cache_dtype = chunk_proj.dtype chunk_proj_cast = chunk_proj[:actual].to(cache_dtype) - ctx_len_updates[slot] = end + self._ctx_len[slot] = end + self._ctx_len_host[slot] = end # Precompute post-norm/post-RoPE K,V for this prefill chunk # so decode iters can read without re-projecting. # @@ -1747,10 +2032,14 @@ def _store_prefill_context( torch.full((actual,), row, dtype=torch.long, device="cuda"), local_pos, ) - else: # VANILLA DFlash backend (FlashAttention) + if self._ctx_k_buf is not None: self._ctx_k_buf[slot, :, cur:end] = chunk_k.permute(1, 0, 2, 3) if chunk_v is not None: self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) + self._req_ctx_pos[req_id] = first_pos + slen + if self._has_unified_draft_cache(): + self._ctx_position_offset[slot] = self._req_ctx_pos[req_id] - end + self._managed_cache_bindings[req_id] = self._ctx_kv_manager.kv_cache_map[req_id] offset += slen self._write_ctx_len(ctx_len_updates) @@ -1816,16 +2105,29 @@ def _forward_impl( draft_model, ) - # Lazy init buffers and attach worker reference for prepare() draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - self._lazy_init_ctx_buffers( - draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager - ) - spec_metadata._dflash_worker = self - # Before any store: prefill and decode both address pages through it. - # Returning False here means an empty batch -- the missing-offsets case - # raises inside, with the metadata type in the message. - self._refresh_ctx_block_tables(attn_metadata, batch_size) + if bool(getattr(draft_kv_cache_manager, "draft_layer_ids", ())): + if not self._ctx_buf_inited or self._ctx_kv_manager is not draft_kv_cache_manager: + raise RuntimeError("Unified DSpark draft inputs must be prepared before forward") + else: + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) + spec_metadata._dflash_worker = self + # Before any store: prefill and decode both address pages through it. + refreshed = self._refresh_ctx_block_tables( + attn_metadata, batch_size, spec_metadata.request_ids + ) + if self._ctx_block_tables is not None and not refreshed: + # False is normal only while no paged pool is bound (the private + # arena addresses slots directly). Once one IS bound, the counts + # drive the advertised context length, so a missed refresh leaves + # them at zero and every drafter attends an empty context -- no + # error, just acceptance quietly collapsing. + raise RuntimeError( + "DFlash: draft block tables are bound to a paged pool but this " + "iteration's attn_metadata carried no draft_kv_cache_block_offsets." + ) # Save context lengths so both warmup and a failed forward can roll # back the in-place _ctx_len updates made during drafting. @@ -1838,6 +2140,8 @@ def _forward_impl( # during capture aborts the graph itself, and captured ops do not # mutate _ctx_len until replay. self._saved_ctx_len = self._ctx_len.clone() + if self._has_unified_draft_cache(): + self._saved_ctx_position_offset = self._ctx_position_offset.clone() self._saved_ctx_len_host = list(self._ctx_len_host) self._saved_req_ctx_pos = dict(self._req_ctx_pos) self._ctx_len_restore_pending = True @@ -1909,7 +2213,9 @@ def _forward_impl( ) if num_gens > 0: - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): + with self.draft_kv_cache_context( + attn_metadata, None if self._has_unified_draft_cache() else draft_kv_cache_manager + ): hidden_states_out = draft_model.dflash_forward( noise_embedding=inputs["noise_embedding"], query_positions=inputs["query_positions"], @@ -2314,13 +2620,8 @@ def prepare_1st_drafter_inputs( j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") - # _ctx_len is clamped to _max_ctx only AFTER this step's accepted - # tokens are folded in (see the update below), so the running length - # used here has to be clamped on its own -- otherwise a request that - # already sits at the ceiling indexes num_accepted positions past - # any position the sequence can legitimately reach. - ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) - query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) + query_position = (ctx_len_gen + gen_num_accepted.long()).clamp(max=self._max_ctx) + query_position_ids = query_position.unsqueeze(1) + j_block.unsqueeze(0) ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) noise_embed_2d = dflash_noise_block_embedding( @@ -2349,7 +2650,7 @@ def prepare_1st_drafter_inputs( # safety) and a placeholder page resolves to another's. gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) - else: + if self._ctx_block_tables is None or self._ctx_k_buf is not None: col_idx = col_idx.clamp(max=self._max_ctx - 1) # Fixed-size writes for CUDA graph compatibility: @@ -2367,6 +2668,9 @@ def prepare_1st_drafter_inputs( cache_dtype = ( self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype ) + if cache_dtype == torch.float8_e4m3fn: + # Project in the drafter's weight dtype; the store converts. + cache_dtype = proj_flat.dtype k_new, v_new = draft_model.precompute_context_kv( proj_flat.to(cache_dtype), pos_flat ) @@ -2384,12 +2688,20 @@ def prepare_1st_drafter_inputs( else: rows_long = slot_long self._store_context_kv_paged(k_new, v_new, rows_long, col_long) - else: # VANILLA DFlash backend (FlashAttention) + if self._ctx_k_buf is not None: self._ctx_k_buf[slot_long, :, col_long] = k_new if v_new is not None: self._ctx_v_buf[slot_long, :, col_long] = v_new - self._advance_ctx_len(slots, gen_num_accepted_long) + self._ctx_len[slots] += gen_num_accepted_long + if self._has_unified_draft_cache(): + self._ctx_position_offset[slots] += ( + self._ctx_len[slots] - self._max_ctx + ).clamp_min(0) + self._ctx_len.clamp_(max=self._max_ctx) + # Padding and ADP dummy rows share a slot that is never freed; keep its + # context empty. zero_() is graph-capturable, a scalar store is not. + self._ctx_len[self._dummy_slot].zero_() num_ctx_per_req_t = self._ctx_len[slots] if self._ctx_block_tables is not None: @@ -2427,7 +2739,11 @@ def prepare_1st_drafter_inputs( "ctx_v_cache": self._ctx_v_buf, # Slots index the private arena; the manager's block table is keyed # by batch position, so the drafter reads its own rows there. - "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, + "ctx_cache_batch_idx": ( + gen_rows_out + if self._ctx_block_tables is not None and self._ctx_k_buf is None + else slots + ), # Anchor token per gen request (block slot 0): last accepted # token. The dspark Markov chain conditions its first step on it. "first_prev_tokens": bonus, diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index cd38f476a686..fe0377e51dea 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -33,8 +33,8 @@ from ..distributed import AllReduceStrategy from ..modules.linear import Linear, TensorParallelMode from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID -from .dflash import DFlashWorker, VocabParallelDraftTokens, dflash_draft_slot_ids -from .dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT +from ..pyexecutor.resource_manager import ResourceManagerType +from .dflash import DFlashWorker, dflash_draft_slot_ids from .interface import SpecMetadata, SpecWorkerBase from .vocab_parallel import tp_pair_exchange @@ -79,7 +79,7 @@ def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) """Return the number of RoPE entries needed by the DSv4 block drafter. ``start_pos`` is a FRAME index, one above the absolute token position: the - prompt token at position p occupies frame p+1 (``_seed_context_windows``), so a + prompt token at position p occupies frame p+1 (``_seed_context_pages``), so a request served to ``max_ctx`` bootstraps at ``max_ctx + 1``. A verification then accepts up to ``max_draft_len + 1`` tokens and the block drafter indexes ``block_size`` further positions, making the largest index @@ -226,62 +226,6 @@ def prepare(self): ) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - # CUDA-graph-safe path: maintain the request->slot mapping on the host - # (outside the captured region) and mirror it into ``_batch_to_slot`` so the - # captured gen forward can index the rolling windows by tensor. Mirrors - # ``DFlashSpecMetadata.prepare`` (dflash.py:96-113). - worker = getattr(self, "_dspark_worker", None) - if worker is not None and worker._win_inited: - current = set(self.request_ids) - for rid in list(worker._req_to_slot.keys()): - if rid not in current: - slot = worker._req_to_slot.pop(rid) - worker._ctx_len[slot] = 0 - worker._valid_len[slot] = 0 - worker._position_initialized[slot] = False - worker._kv_windows[slot].zero_() - worker._free_slots.append(slot) - # Assign a persistent rolling-window slot to every real generation - # request that never ran a context/seed forward on this worker. In - # disaggregated serving the prompt is prefilled (and the window - # seeded) on the *context* server, so ``_seed_context_windows`` never - # runs on the generation server and ``_req_to_slot`` stays empty; - # without this, all concurrent gen requests fall through to the shared - # scratch row below and corrupt each other's draft window at batch - # size > 1 (GitHub #16767). Context-prefix entries are left to - # ``_seed_context_windows``. CUDA-graph capture dummies and real - # generation requests acquire a persistent slot only once; prepare() - # runs before every replay, so it must never reset a live slot. - # ADP-idle (id 0) and high-ID CUDA-graph padding dummies use scratch. - num_contexts = max(0, len(self.request_ids) - self.num_generations) - for rid in self.request_ids[num_contexts:]: - if ( - rid != ATTENTION_DP_DUMMY_REQUEST_ID - and rid < worker._graph_dummy_id_floor - and rid not in worker._req_to_slot - ): - worker._assign_slot(rid, reset=False) - # Unknown request IDs (synthetic warmup / CUDA-graph padding, ADP idle - # requests, or disagg seed forwards without a real id) map to the - # dedicated throwaway scratch row so they cannot overwrite a live - # request's rolling window (they previously aliased to slot 0). - scratch = worker._scratch_slot - # Reset every iteration. The scratch row has no prefill to seed it - # and no completion to free it, so _advance_generation_state's - # unconditional advance would leave its ABSOLUTE position climbing - # across graph shapes and iterations until it indexes past the RoPE - # table. A real slot is bounded instead by its request's lifetime. - worker._ctx_len[scratch] = 0 - worker._valid_len[scratch] = 0 - worker._position_initialized[scratch] = False - mapping = torch.tensor( - [worker._req_to_slot.get(rid, scratch) for rid in self.request_ids], - dtype=torch.long, - device="cpu", - pin_memory=prefer_pinned(), - ) - worker._batch_to_slot[:num_seqs].copy_(mapping, non_blocking=True) - def is_layer_capture(self, layer_id: int) -> bool: return layer_id in self._capture_layer_set @@ -340,31 +284,30 @@ class DSv4DSparkWorker(SpecWorkerBase): DSpark drafts a whole block of ``block_size`` tokens in one backbone forward (``DSv4DSparkDraftModel.forward``): it projects the captured target-layer hidden states (``main_proj`` + ``main_norm``) into the draft's captured-context - attention, runs the ``num_stages`` DSpark blocks over a rolling captured + attention, runs the ``num_stages`` DSpark blocks over a captured window, refines the per-position logits with the Markov head, and predicts a per-position acceptance confidence used by the verification scheduler. - Unlike DFlash, the draft does NOT use the paged KV cache or mask-token - cross-attention: its attention K/V come from the worker-owned rolling window - of projected captured context (one ``main_kv`` per decode step, per stage). + Attention reads a sliding window of projected captured context directly + from manager-owned KV pages, alongside the current draft block. Acceptance of the previous block goes through the unified :meth:`SpecWorkerBase.sample_and_accept_draft_tokens` (strict target-verify, or rejection sampling for a non-greedy batch), so greedy parity with no-spec is preserved regardless of draft quality. - The rolling window is kept consistent across the whole decode: it is seeded + The context window is kept consistent across the whole decode: it is seeded from the prompt's captured context at prefill and back-filled with the intermediate accepted tokens of a multi-accept step (both via - ``DSv4DSparkDraftModel.write_context_windows``), in addition to the per-step bonus + ``DSv4DSparkDraftModel.write_context_pages``), in addition to the per-step bonus write done by the generation path. These affect draft acceptance rate only, not correctness, which the standard target verify guarantees. Naming: workers are classified by *deployment form*, not by draft backbone (see :class:`DSparkWorker`). This one is form-specific - because it owns a rolling captured-context window and drives the draft + because it tracks captured-context history and drives the draft through attributes only an embedded DeepSeek-V4-Pro draft has -- - ``num_stages``, ``_attn_params``, ``write_context_windows``, - ``write_context_windows_batched`` and ``forward_batched``. A standalone + ``num_stages``, ``_attn_params``, ``write_context_pages`` and + ``forward``. A standalone drafter has none of them and is served by :class:`DSparkWorker`. Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). @@ -380,16 +323,17 @@ def __init__( self.spec_config = spec_config self.mapping = mapping - # Per-slot rolling captured-context KV windows, built lazily on the - # first forward (fixed-size for slot-indexed reads/writes). - self._win_inited = False - self._output_verify_lens: Optional[torch.Tensor] = None - self._attention_warmup_attempted = False - self._kv_windows: Optional[torch.Tensor] = None # [max_batch, num_stages, win, hd] + # Request progress is slot-indexed; managed KV lives only in its pages. self._ctx_len: Optional[torch.Tensor] = None # [max_batch] abs decode position self._valid_len: Optional[torch.Tensor] = None # [max_batch] written window entries - self._position_initialized: Optional[torch.Tensor] = None # [max_batch] bool self._win = 0 + self._draft_kv_manager = None + self._draft_kv_buffers = () + self._draft_block_tables = None + self._draft_capacities = None + self._draft_page_indices: dict[int, tuple[int, ...]] = {} + self._draft_metadata_host_buffers: list[dict] = [] + self._managed_residency = {} # Set in _lazy_init from the RoPE table the drafter will build; None # leaves positions unbounded (direct construction in tests). self._position_cap: Optional[int] = None @@ -398,11 +342,11 @@ def __init__( # source of truth, updated in prepare()/forward(); ``_batch_to_slot`` is the # CUDA mirror (request-order -> slot) read by the CUDA-graph-safe batched # gen path (set on the host in prepare(), so the captured forward indexes - # the rolling windows through a tensor instead of a python dict lookup). + # request progress through a tensor instead of a python dict lookup). self._req_to_slot = {} # request_id -> slot index self._free_slots = deque() # available slot indices self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda - # Index of the throwaway "scratch" window row that absorbs padded / + # Index of the throwaway "scratch" progress row that absorbs padded / # unknown request IDs (set in ``_lazy_init`` to ``max_batch``); it is # never handed out through ``_free_slots``. self._scratch_slot = 0 @@ -423,8 +367,8 @@ def __init__( self.verify_planner = None # The generation draft path is the batched, host-sync-free - # ``_draft_gen_block_batched`` + ``DSv4DSparkDraftModel.forward_batched`` + - # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe + # ``_draft_gen_block_batched`` + ``DSv4DSparkDraftModel.forward`` + + # ``dspark_attention_forward``: it is correct in eager mode AND safe # to capture into the target's CUDA graph (DSpark is a one-engine drafter — # its worker forward runs inside that graph, so the draft path MUST be # capture-safe whenever ``cuda_graph_config`` is set). @@ -501,7 +445,7 @@ def max_draft_len(self) -> int: def _publish_position_ceiling(self, draft_model, attn_metadata, block_size: int) -> None: """Size the drafter's RoPE table from the runtime bound, and bound positions by it. - Runs on EVERY forward, not once behind ``_win_inited``. The first forward can + Refreshes the bound for every batch. The first forward can land on a KV-estimation probe manager whose ``max_seq_len`` is below the real one, and a table pinned to that probe would index out of range for the rest of the process; ``DFlashWorker._lazy_init_ctx_buffers`` re-publishes for that @@ -509,9 +453,8 @@ def _publish_position_ceiling(self, draft_model, attn_metadata, block_size: int) Two bounds exist and neither dominates: ``attn_metadata`` carries the KV manager's ``max_seq_len`` while ``_freqs_cap`` carries the user's, and - ``DummyRequestFactory.create_cuda_graph_warmup_request`` sizes its dummy - request at whichever is larger. Covering both is what keeps warmup off - the end of the table. + ``_create_cuda_graph_warmup_request`` sizes its dummy request at whichever is + larger. Covering both is what keeps warmup off the end of the table. """ inner_model = getattr(draft_model, "dspark_model", None) or draft_model config_cap = int(getattr(inner_model, "_freqs_cap", 0) or 0) @@ -529,7 +472,7 @@ def _publish_position_ceiling(self, draft_model, attn_metadata, block_size: int) draft_model._runtime_position_ceiling = ceiling if inner_model is not draft_model: inner_model._runtime_position_ceiling = ceiling - # Largest absolute position the block drafter may hold: forward_batched + # Largest absolute position the block drafter may hold: forward # gathers ``freqs[start_pos + block_size]`` and the interim back-fill # ``freqs[old + block_size]``, so keep one block of headroom. table_len = int(getattr(inner_model, "_runtime_position_ceiling", 0) or config_cap) @@ -545,20 +488,14 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: self._publish_position_ceiling(draft_model, attn_metadata, block_size) - if not self._win_inited: + if self._ctx_len is None: max_batch = spec_metadata.max_num_requests - self._output_verify_lens = torch.empty(max_batch, dtype=torch.int32, device="cuda") - num_stages = draft_model.num_stages + draft_model.configure_generation_autotuner(max_batch) + self._num_stages = draft_model.num_stages self._win = int(draft_model._attn_params["window_size"]) - head_dim = int(draft_model._attn_params["head_dim"]) - - # Real requests occupy slots ``[0, max_batch)``; one extra "scratch" row - # at index ``max_batch`` absorbs padded / unknown request IDs (CUDA-graph - # padding, ADP idle requests, or disagg seed forwards that arrive without - # a real request id) so they can never overwrite a live request's rolling - # window. Previously such IDs aliased to slot 0 and corrupted whichever - # real request occupied it. The scratch row is never handed out through - # ``_free_slots`` and its contents are throwaway. + self._head_dim = int(draft_model._attn_params["head_dim"]) + + # Padding requests share an extra progress slot, never a real request slot. self._scratch_slot = max_batch num_rows = max_batch + 1 @@ -573,158 +510,175 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len - self._kv_windows = torch.zeros( - (num_rows, num_stages, self._win, head_dim), - dtype=torch.bfloat16, - device="cuda", - ) self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._valid_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") - self._position_initialized = torch.zeros(num_rows, dtype=torch.bool, device="cuda") self._batch_to_slot = torch.full( (max_batch,), self._scratch_slot, dtype=torch.long, device="cuda" ) - if self.return_confidence: - # Neutral fill = large positive logit (sigmoid ~ 1.0): - # unscored slots fail safe by verifying the full block. The - # neutral row sits one past the scratch row and is never used - # as a draft-output destination. - self._neutral_conf_row = num_rows - self._confidence_logits = torch.full( - (num_rows + 1, block_size), - _NEUTRAL_CONFIDENCE_LOGIT, - dtype=torch.float32, - device="cuda", - ) - self._confidence_stamp = torch.zeros( - (num_rows + 1,), dtype=torch.int32, device="cuda" - ) - self._draft_seq_host = 0 - # The executor updates this scalar outside inference mode; - # inference tensors reject that in-place host-side update. - with torch.inference_mode(False): - self._draft_seq_cuda = torch.zeros((), dtype=torch.int32, device="cuda") - self._build_verify_planner(draft_model, block_size) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} - logger.info( - f"DSpark: allocated rolling KV windows " - f"[{num_rows}, {num_stages}, {self._win}, {head_dim}] " - f"({max_batch} request slots + 1 scratch row)" - ) - # Buffer state is complete independently of CuTe DSL prewarming. - # A failed prewarm must not cause the next forward to recreate the - # windows or reset the live slot maps. - self._win_inited = True - - if self._attention_warmup_attempted: - return - - # Prewarm the same self-JIT ops used by production before CUDA graph - # capture. This is best-effort inside the named warmup entry; the ops - # remain able to compile themselves on a later eager first use. - from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE - - if IS_CUTLASS_DSL_AVAILABLE: - from ..custom_ops.dspark_attention_custom_op import ( - is_dsv4_dspark_attention_config_supported, - warmup_fused_dsv4_dspark_attention, - ) - - if is_dsv4_dspark_attention_config_supported( - block_size, - int(draft_model._attn_params["n_heads"]), - int(draft_model._attn_params["head_dim"]), - int(draft_model._attn_params["window_size"]), - ): - warmup_fused_dsv4_dspark_attention( - block_size, - eps=float(draft_model._attn_params["eps"]), - ) - self._attention_warmup_attempted = True - - def _build_verify_planner(self, draft_model, block_size: int) -> None: - """Construct the host-side verify planner (once, at lazy init). - ModelEngine owns and authenticates the deployment cost table. The - worker creates only the policy state and calibration, then receives - that exact runtime object before its first decision. - """ - from .dspark_schedule import DSparkScheduleConfig - from .dspark_verify import DSparkVerifyPlanner - - cfg = DSparkScheduleConfig(block_size=block_size, min_verify_len=1) - - table_path = self.spec_config.confidence_sps_table_path - # ModelEngine owns the actual CUDA graph ladder, so it is the only - # component that may read, authenticate, and construct the deployment - # cost object. Reopening this path here creates a TOCTOU window in - # which graph capture and scheduling can observe different files. - logger.info( - f"DSpark: deferring exact SPS table {table_path} to the validated " - "ModelEngine runtime object" - ) - - # Calibration lives on the draft's confidence head; fall back to a plain - # sigmoid when the head is absent (no confidence to calibrate anyway). - head = resolve_confidence_head(draft_model) - apply_calibration = head.apply_sts if head is not None else None - sts_path = self.spec_config.confidence_sts_path - if sts_path and head is None: - # Refuse rather than degrade: a configured calibration that cannot - # be attached is a wiring bug, never a preference. - raise ValueError( - f"confidence_sts_path is set but no confidence head was found " - f"on {type(draft_model).__name__}; calibration cannot be " - f"applied, so the planner would schedule on raw sigmoid " - f"survivals. This is a model-wiring bug, not a config choice." - ) - if sts_path and head is not None: - # Accepts either key spelling ("sts_temperatures" here, - # "temperatures" in SGLang); the vectors are interchangeable. - temps = load_sts_temperatures_from_path(sts_path) - head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) - logger.info(f"DSpark: loaded STS calibration from {sts_path}") - - self.verify_planner = DSparkVerifyPlanner( - cfg=cfg, - apply_calibration=apply_calibration, - device_windows=bool(self.spec_config.enable_fused_confidence_scheduler), - ) - logger.info( - f"DSpark verify planner: max_verify_len={self.verify_planner.max_verify_len}, " - f"device_windows={self.verify_planner.device_windows}" - ) - - def _assign_slot(self, req_id: int, reset: bool) -> int: - """Get (or refresh) the slot for a request; reset clears its window.""" - if reset and req_id in self._req_to_slot: - old = self._req_to_slot.pop(req_id) - self._ctx_len[old] = 0 - self._valid_len[old] = 0 - self._position_initialized[old] = False - self._kv_windows[old].zero_() - self._free_slots.append(old) + def _assign_slot(self, req_id: int) -> int: + """Get the persistent progress slot for a request.""" if req_id not in self._req_to_slot: if not self._free_slots: raise RuntimeError( - "DSpark has no free rolling-window slots for request " + "DSpark has no free progress slots for request " f"{req_id}; increase max_num_requests" ) slot = self._free_slots.popleft() self._req_to_slot[req_id] = slot self._ctx_len[slot] = 0 self._valid_len[slot] = 0 - self._position_initialized[slot] = False - self._kv_windows[slot].zero_() - if self._confidence_logits is not None: - # A recycled slot still holds the previous occupant's scores; - # reset to neutral so the first scheduling decision verifies - # the full block. - self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] - def _seed_context_windows( + def _release_inactive_slots(self, request_ids: list[int]) -> None: + current = set(request_ids) + for request_id in list(self._req_to_slot): + if request_id not in current: + slot = self._req_to_slot.pop(request_id) + self._managed_residency.pop(request_id, None) + self._ctx_len[slot] = 0 + self._valid_len[slot] = 0 + self._free_slots.append(slot) + + def _is_managed_request(self, request_id: int) -> bool: + return ( + request_id != ATTENTION_DP_DUMMY_REQUEST_ID + and request_id < self._graph_dummy_id_floor + and request_id not in self._draft_kv_manager._draft_dummy_request_ids + ) + + def _bind_managed_history(self, resource_manager) -> None: + manager = ( + resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) + if resource_manager is not None + else None + ) + if not getattr(manager, "draft_layer_ids", ()): + raise ValueError("Embedded DSpark requires manager-owned draft KV pages") + if manager is self._draft_kv_manager: + return + if ( + manager.draft_window_size != self._win + or len(manager.draft_layer_ids) != self._num_stages + or any( + manager.kv_factor_per_layer[manager.layer_offsets[layer_id]] != 1 + or manager.head_dim_per_layer[manager.layer_offsets[layer_id]] != self._head_dim + for layer_id in manager.draft_layer_ids + ) + ): + raise ValueError("Embedded DSpark draft cache does not match its attention dimensions") + buffers = tuple( + manager.get_draft_buffers(stage) for stage in range(len(manager.draft_layer_ids)) + ) + self._draft_kv_manager = manager + self._draft_kv_buffers = buffers + self._managed_residency.clear() + self._draft_page_indices.clear() + self._draft_block_tables = torch.zeros( + (self._batch_to_slot.shape[0], manager.draft_max_blocks_per_seq), + dtype=torch.int32, + device=self._ctx_len.device, + ) + self._draft_capacities = torch.zeros( + self._batch_to_slot.shape[0], dtype=torch.long, device=self._ctx_len.device + ) + + def prepare_managed_draft_cache( + self, draft_model, spec_metadata, attn_metadata, resource_manager + ) -> None: + """Refresh persistent draft inputs before eager execution or graph replay.""" + self._lazy_init(draft_model, spec_metadata, attn_metadata) + self._bind_managed_history(resource_manager) + self._prepare_managed_history(spec_metadata.request_ids, attn_metadata.num_contexts) + + def _prepare_managed_history(self, request_ids: list[int], num_contexts: int) -> None: + """Bind live page tables and initialize progress for newly resident histories.""" + # Startup probes and graph/ADP padding never publish synthetic history. + real_rows = [ + row + for row, request_id in enumerate(request_ids) + if self._is_managed_request(request_id) + ] + real_ids = [request_ids[row] for row in real_rows] + self._release_inactive_slots(real_ids) + batch_indices = self._draft_kv_manager.get_draft_block_indices(real_ids) + mappings = {row: tuple(indices) for row, indices in zip(real_rows, batch_indices)} + updates = [] + for row in sorted(self._draft_page_indices.keys() | mappings.keys()): + previous = self._draft_page_indices.get(row, ()) + current = mappings.get(row, ()) + if current == previous: + continue + # Clear removed mappings as well as uploading new/reassigned pages. + width = max(len(previous), len(current)) + old = previous + (0,) * (width - len(previous)) + new = current + (0,) * (width - len(current)) + first = next((i for i in range(width) if old[i] != new[i]), width) + if first == width: + continue + last = width + while old[last - 1] == new[last - 1]: + last -= 1 + updates.append((row, first, new[first:last])) + self._draft_page_indices = mappings + self._ctx_len[self._scratch_slot] = 0 + self._valid_len[self._scratch_slot] = 0 + batch_slots = [self._scratch_slot] * len(request_ids) + for row in real_rows: + request_id = request_ids[row] + cache = self._draft_kv_manager.kv_cache_map[request_id] + slot = self._assign_slot(request_id) + batch_slots[row] = slot + if self._managed_residency.get(request_id) is cache: + continue + position = cache.history_length + if row >= num_contexts and position == 0: + raise ValueError( + f"Embedded DSpark generation request {request_id} has no committed KV history" + ) + self._ctx_len[slot] = position + self._valid_len[slot] = min(self._win, position) + self._managed_residency[request_id] = cache + # Staging may still be read by an earlier asynchronous H2D. Reuse it + # only after its copy event, independently of overlapped model execution. + count = sum(len(indices) for _, _, indices in updates) + host = next( + (buffers for buffers in self._draft_metadata_host_buffers if buffers["event"].query()), + None, + ) + if host is None: + host = { + "tables": torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()), + "capacities": torch.empty( + self._draft_capacities.shape, dtype=torch.long, pin_memory=prefer_pinned() + ), + "slots": torch.empty( + self._batch_to_slot.shape, dtype=torch.long, pin_memory=prefer_pinned() + ), + "event": torch.cuda.Event(), + } + self._draft_metadata_host_buffers.append(host) + elif host["tables"].numel() < count: + host["tables"] = torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()) + offset = 0 + for row, first, indices in updates: + size = len(indices) + source = host["tables"][offset : offset + size] + source.copy_(torch.tensor(indices, dtype=torch.int32)) + self._draft_block_tables[row, first : first + size].copy_(source, non_blocking=True) + offset += size + host["capacities"].zero_() + for row, request_id in zip(real_rows, real_ids): + host["capacities"][row] = self._draft_kv_manager.kv_cache_map[request_id].capacity + self._draft_capacities.copy_(host["capacities"], non_blocking=True) + host["slots"].fill_(self._scratch_slot) + host["slots"][: len(batch_slots)] = torch.tensor(batch_slots, dtype=torch.long) + self._batch_to_slot.copy_(host["slots"], non_blocking=True) + host["event"].record() + + def _seed_context_pages( self, draft_model, spec_metadata: "DSparkSpecMetadata", @@ -734,9 +688,8 @@ def _seed_context_windows( ) -> None: """Seed context chunks using their absolute positions. - A request can arrive in multiple prefill chunks. Only its first chunk - starts at position zero and resets the persistent rolling window; - continuation chunks append to the same request slot. + A request can arrive in multiple prefill chunks; continuation chunks + append to the same request's pages. """ captured = spec_metadata.get_hidden_states(total_target_tokens) capture_lens = getattr(spec_metadata, "context_capture_lens", None) @@ -755,21 +708,28 @@ def _seed_context_windows( continue req_id = spec_metadata.request_ids[i] - first_position = int(chunk_positions[0].item()) - slot = self._assign_slot(req_id, reset=first_position == 0) - self._ctx_len[slot] = chunk_positions[-1] + 1 - self._position_initialized[slot] = True - - if captured is not None and capture_len: - self._valid_len[slot] = torch.clamp( - self._valid_len[slot] + capture_len, max=self._win + slot = self._req_to_slot.get(req_id, self._scratch_slot) + if captured is None and slot != self._scratch_slot: + raise RuntimeError( + "Embedded DSpark requires captured context to initialize KV history" + ) + if captured is not None: + keep = min(self._win, chunk_len) + # Prefix reuse can publish any complete page in this chunk. + count = chunk_len if self._draft_kv_manager.enable_block_reuse else keep + positions = (chunk_positions[-count:] + 1).unsqueeze(0) + draft_model.write_context_pages( + captured[ + context_offset + chunk_len - count : context_offset + chunk_len + ].unsqueeze(0), + positions, + torch.ones_like(positions, dtype=torch.bool), + self._draft_kv_buffers, + self._draft_block_tables[i : i + 1], + self._draft_capacities[i : i + 1], ) - keep = min(self._win, capture_len) - hidden = captured[context_offset + chunk_len - keep : context_offset + chunk_len] - # A prompt token at absolute position p is stored in frame p+1, - # matching the generation path's start_pos convention. - window_positions = chunk_positions[-keep:] + 1 - draft_model.write_context_windows(hidden, window_positions, self._kv_windows[slot]) + self._valid_len[slot] = torch.clamp(self._valid_len[slot] + keep, max=self._win) + self._ctx_len[slot] = chunk_positions[-1] + 1 context_offset += chunk_len def _advance_generation_state( @@ -785,15 +745,9 @@ def _advance_generation_state( scratch = self._scratch_slot self._ctx_len[scratch].zero_() self._valid_len[scratch].zero_() - self._position_initialized[scratch].zero_() - old = torch.where(self._position_initialized[slots], self._ctx_len[slots], input_positions) + old = torch.where(slots == scratch, input_positions, self._ctx_len[slots]) start_pos = old + num_accepted_tokens - # Bound both positions by the RoPE table they index. Warmup reaches this - # through a captured graph replay, and prepare() cannot tell a synthetic - # warmup request from a real one, so the bound has to live in the graph. - # Only a warmup row climbs -- no completion frees its slot, and it ran 53 - # entries past the table [measured job 3097207] -- while _position_cap sits - # exactly at what a full-length real request reaches. + # Keep dummy graph positions within the RoPE table. if self._position_cap is not None: old = torch.clamp(old, max=self._position_cap) start_pos = torch.clamp(start_pos, max=self._position_cap) @@ -801,7 +755,6 @@ def _advance_generation_state( self._valid_len[slots] = torch.clamp( self._valid_len[slots] + num_accepted_tokens, max=self._win ) - self._position_initialized[slots] = torch.ones_like(slots, dtype=torch.bool) return old, start_pos def _draft_gen_block_batched( @@ -823,7 +776,7 @@ def _draft_gen_block_batched( (``nacc``, the bonus, ``main_hidden``, ``start_pos``, the multi-accept back-fill) are gathered as tensors, slots come from the host-built ``_batch_to_slot`` mirror, and the backbone runs once via - ``DSv4DSparkDraftModel.forward_batched``. Returns the per-position corrected + ``DSv4DSparkDraftModel.forward``. Returns the per-position corrected block logits ``[num_gens, K, vocab]`` (or ``None`` when there is nothing to draft); the worker feeds them to ``SpecWorkerBase.sample_draft_tokens``. Confidence truncation stays disabled — the full block is proposed. When @@ -843,11 +796,7 @@ def _draft_gen_block_batched( return None captured = spec_metadata.get_hidden_states(total_target_tokens) if captured is None: - raise RuntimeError( - "DSpark generation requires captured target hidden states; " - "check target_layer_ids and the target model's capture hooks. " - "Refusing to substitute zero draft tokens for a missing capture buffer." - ) + raise RuntimeError("Embedded DSpark requires captured context to advance KV history") # gen-only graph batches have num_ctx_tokens == 0; mixed eager batches put # the gen tokens after the context tokens. @@ -882,14 +831,10 @@ def _draft_gen_block_batched( bonus_row = torch.where(is_scratch, torch.zeros_like(base), base + gidx) main_hidden = captured.index_select(0, bonus_row) - # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens - # (everything but the bonus) into the rolling window — same frames as the - # eager path (old+1 .. old+nacc-1), with j >= nacc-1 masked out. - # A disaggregated generation worker never sees prompt prefill, so a new - # slot has no absolute decode position. Bootstrap it once from the first - # target input position; locally-prefilled and existing slots keep their - # monotonically advanced position. - input_positions = position_ids.reshape(-1).index_select(0, base).long() + # Back-fill intermediate accepted tokens (everything but the bonus) at + # frames old+1 .. old+nacc-1. Fixed [G, K] tensors keep this capture-safe; + # j >= nacc-1 is masked before any page store. + input_positions = position_ids.reshape(-1)[base].long() old, start_pos = self._advance_generation_state(slots, nacc, input_positions) j = torch.arange(K, device=device) # [K] interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1) # [G, K] @@ -898,19 +843,22 @@ def _draft_gen_block_batched( min=0, max=captured.shape[0] - 1 ) # [G, K] (clamped; invalid entries are masked out anyway) interim_hidden = captured[interim_base] # [G, K, ncap*hidden] - draft_model.write_context_windows_batched( - interim_hidden, interim_pos, slots, interim_valid, self._kv_windows + tables = self._draft_block_tables[num_contexts:batch_size] + capacities = self._draft_capacities[num_contexts:batch_size] + draft_model.write_context_pages( + interim_hidden, interim_pos, interim_valid, self._draft_kv_buffers, tables, capacities ) # Surface the per-position corrected block logits ([num_gens, K, vocab]) # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) # sampling + TP gather + draft_probs scatter, rather than argmaxing here. - _toks, confidence, block_logits = draft_model.forward_batched( + _toks, _num_proposed, block_logits = draft_model.forward( main_hidden, bonus, start_pos, - kv_windows=self._kv_windows, - slots=slots, + kv_pages=self._draft_kv_buffers, + block_tables=tables, + capacities=capacities, valid_len=self._valid_len[slots], temperature=0.0, return_confidence=self.return_confidence, @@ -984,10 +932,9 @@ def _forward_impl( raw_logits = logits K = self.max_draft_len - self._lazy_init(draft_model, spec_metadata, attn_metadata) - # Backref so DSparkSpecMetadata.prepare() can maintain the host slot map - # and mirror it into _batch_to_slot for the CUDA-graph-safe gen path. - spec_metadata._dspark_worker = self + if self._draft_kv_manager is None: + raise RuntimeError("DSpark KV pages must be prepared before worker forward") + self._execute_guided_decoder_if_present(logits) # Target-verify acceptance via the unified SpecWorkerBase entry: it @@ -1001,36 +948,8 @@ def _forward_impl( total_target_tokens = input_ids.shape[0] - # CUDA-graph warmup guard: the warmup forwards (is_cuda_graph set, stream - # NOT yet capturing) run synthetic gen batches that would otherwise advance - # the persistent rolling-window state. Snapshot and restore it so warmup is - # side-effect-free. (During the capture pass itself the stream IS capturing, - # so we skip the save/restore and let the ops be recorded. Capture-time - # mutation stays on dummy/scratch state: a locally-prefilled request - # resets its slot at position zero, while a disaggregated generation-only - # request receives a freshly zeroed slot.) - is_warmup = ( - getattr(spec_metadata, "is_cuda_graph", False) - and not torch.cuda.is_current_stream_capturing() - ) - if is_warmup: - saved_ctx_len = self._ctx_len.clone() - saved_valid_len = self._valid_len.clone() - saved_position_initialized = self._position_initialized.clone() - saved_windows = self._kv_windows.clone() - # The confidence buffer is persistent state too; warmup scores - # must not reach the verification planner. - saved_confidence = ( - None if self._confidence_logits is None else self._confidence_logits.clone() - ) - saved_stamp = None if self._confidence_stamp is None else self._confidence_stamp.clone() - - # Assign / reset window slots for context (prefill) requests and seed each - # request's rolling KV window from its prompt's captured context, so the - # first generation step drafts against real context instead of an all-zero - # window (acceptance-rate only; verified decoding keeps output correct). if num_contexts > 0: - self._seed_context_windows( + self._seed_context_pages( draft_model, spec_metadata, attn_metadata, @@ -1133,17 +1052,7 @@ def _forward_impl( num_accepted_tokens, ) - if is_warmup: - self._ctx_len.copy_(saved_ctx_len) - self._valid_len.copy_(saved_valid_len) - self._position_initialized.copy_(saved_position_initialized) - self._kv_windows.copy_(saved_windows) - if saved_confidence is not None: - self._confidence_logits.copy_(saved_confidence) - if saved_stamp is not None: - self._confidence_stamp.copy_(saved_stamp) - - outputs = { + return { "logits": raw_logits, "new_tokens": accepted_tokens, "new_tokens_lens": num_accepted_tokens, @@ -1192,9 +1101,8 @@ class DSparkWorker(DFlashWorker): subclass. Naming workers by backbone would produce N classes with identical bodies. - Deployment form is the axis the runtime state actually splits on: paged - draft K/V here, a worker-owned rolling window in - :class:`DSv4DSparkWorker`. + Embedded DSpark uses :class:`DSv4DSparkWorker` to manage captured-context + progress and the draft block forward. """ def set_draft_model(self, draft_model) -> None: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index f41d241ce25d..1f9f5747ebb0 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -37,6 +37,8 @@ if TYPE_CHECKING: from ..pyexecutor.guided_decoder import CapturableGuidedDecoder from ..pyexecutor.llm_request import LlmRequest + from ..pyexecutor.resource_manager import ResourceManager + from ..pyexecutor.standalone_draft_cache import DraftHistoryUpdate if IS_FLASHINFER_AVAILABLE: import flashinfer @@ -1692,6 +1694,19 @@ def commit_auxiliary_speculative_states( handler.commit_speculative_states(num_accepted_tokens, state_indices, num_contexts) + def prepare_managed_draft_cache( + self, + draft_model: nn.Module | None, + spec_metadata: "SpecMetadata", + attn_metadata: AttentionMetadata, + resource_manager: "ResourceManager", + ) -> None: + """Stage manager-owned draft history before eager execution or replay.""" + + def snapshot_managed_draft_history(self) -> Optional["DraftHistoryUpdate"]: + """Snapshot this execution's history for publication at completion.""" + return None + def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) if "forward" in cls.__dict__: diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index b2c63bc05f7b..a4d1064aa04d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -756,9 +756,9 @@ def get_spec_worker(spec_config, if spec_dec_mode.is_dflash(): return DFlashWorker(spec_config, mapping, use_separate_draft_kv_cache) # DSpark splits by deployment form, mirroring the draft-model side. The - # embedded DeepSeek-V4-Pro draft needs DSv4DSparkWorker, whose rolling-window - # plumbing reads V4-draft-only attributes (num_stages, write_context_windows, - # forward_batched). A standalone drafter is DFlash lineage and is served by + # embedded DeepSeek-V4-Pro draft needs DSv4DSparkWorker, whose captured-context + # plumbing reads V4-draft-only attributes (num_stages, write_context_pages, + # forward). A standalone drafter is DFlash lineage and is served by # DSparkWorker, which adds only the Markov bias and the shift_label # slot convention on top of DFlashWorker. if spec_dec_mode.is_dspark(): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 3eb903af9973..bbab9e7b77ec 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -137,6 +137,17 @@ class _KVCacheManagerConfigFieldSpec: pool_rebalance: object = None cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + # Symbols added on main that are not yet ported to the C++ backend. + # TODO(kvCacheManagerV2-cpp): port these and replace the fallbacks. + AttnLifeCycle = getattr(_cpp, "AttnLifeCycle", None) + OutOfMemoryError = getattr(_cpp, "OutOfMemoryError", MemoryError) + PageIndexConverter = getattr(_cpp, "PageIndexConverter", None) + ReuseScope = getattr(_cpp, "ReuseScope", ReuseScope) + ScratchDesc = getattr(_cpp, "ScratchDesc", None) + SsmSnapshotIterationStatsDelta = _cpp.SsmSnapshotIterationStatsDelta + SwaScratchReuseConfig = getattr(_cpp, "SwaScratchReuseConfig", None) + UniqueToken = _cpp.UniqueToken + CacheDomain = _cpp.CacheDomain KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ del _KVCacheManagerConfigFieldSpec, _dataclasses @@ -307,3 +318,6 @@ def typed_range(*args: int) -> range: "take_poison", "typed_range", ] + +if _BACKEND == "cpp": + __all__.append("CacheDomain") diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index e418a261853e..9458848c6340 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -32,7 +32,7 @@ from typing import ( Union, ) -from tensorrt_llm.bindings.executor import KvCacheRetentionConfig +from ._utils import HalfOpenRange as HalfOpenRange # From _common.py NDEBUG: Final[int] @@ -75,6 +75,10 @@ def take_poison() -> str | None: def num_live_managers() -> int: """Number of constructed, not-yet-destroyed managers.""" +class CacheDomain(enum.Enum): + STANDALONE_DRAFT = 0 + TARGET = 1 + class CacheTier(enum.IntEnum): GPU_MEM = 0 HOST_MEM = 1 @@ -93,31 +97,27 @@ LifeCycleId = NewType("LifeCycleId", int) LayerGroupId: TypeAlias = LifeCycleId class AttnLifeCycle: - """The attention life cycle, keyed by window, sink tokens, sparsity, and reuse policy.""" + """The attention life cycle, keyed by window, sink-token shape, and cache domain.""" def __init__( self, window_size: int | None, num_sink_blocks: int, - is_sparse: bool = False, - reuse_policy: AttentionReusePolicy = ..., + cache_domain: CacheDomain = CacheDomain.TARGET, ) -> None: ... @staticmethod def make( window_size: int | None, num_sink_tokens: int | None, tokens_per_block: int, - is_sparse: bool = False, - reuse_policy: AttentionReusePolicy = ..., + cache_domain: CacheDomain = CacheDomain.TARGET, ) -> "AttnLifeCycle": ... @property def window_size(self) -> int | None: ... @property def num_sink_blocks(self) -> int: ... @property - def is_sparse(self) -> bool: ... - @property - def reuse_policy(self) -> AttentionReusePolicy: ... + def cache_domain(self) -> CacheDomain: ... def get_stale_range(self, history_length: int, tokens_per_block: int) -> HalfOpenRange: ... CacheLevel = NewType("CacheLevel", int) @@ -276,7 +276,7 @@ class AttentionLayerConfig: buffers: list[BufferConfig] sliding_window_size: int | None = None num_sink_tokens: int | None = None - reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + cache_domain: CacheDomain = CacheDomain.TARGET @property def window_size(self) -> int | None: ... diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py new file mode 100644 index 000000000000..cf0caff78d65 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -0,0 +1,1010 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for KV cache token estimation in KvCacheCreator._get_token_num_for_estimation. + +Guards the ADP (Attention Data Parallelism) cache-block reduction: when +enable_attention_dp is True and tp_size > 1, _create_dummy_context_requests +produces tp_size duplicate requests, but the scheduler distributes them +1-per-rank. Each rank's KV cache therefore only needs capacity for its own +share, not all copies. +""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator +from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + KVCacheManagerV2, + _get_draft_layer_sizes, +) +from tensorrt_llm._utils import torch_dtype_to_binding +from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MultimodalConfig + +pytestmark = pytest.mark.cpu_only + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_mock_request(num_input_tokens, beam_width=1): + """Create a mock request with the fields _get_token_num_for_estimation reads.""" + req = Mock() + req.input_token_ids = list(range(num_input_tokens)) + req.sampling_config.beam_width = beam_width + return req + + +class _TextModel: + def __init__(self, encoder_cache_max_bytes): + self.model_config = ModelConfig( + multimodal_config=MultimodalConfig(encoder_cache_max_bytes=encoder_cache_max_bytes) + ) + + +class _MultimodalModel(MultimodalModelMixin): + def __init__(self, encoder_cache_max_bytes): + self.model_config = ModelConfig( + multimodal_config=MultimodalConfig(encoder_cache_max_bytes=encoder_cache_max_bytes) + ) + + +class _EncoderCacheMultimodalModel(_MultimodalModel): + supports_encoder_cache = True + + +def test_encoder_profiling_uses_full_budget_independent_of_llm_limit( + monkeypatch, +): + class _InputProcessor: + def __init__(self): + self.calls = [] + + def get_mm_max_tokens_per_item(self, max_num_encoder_tokens=None): + del max_num_encoder_tokens + return {"image": 4096} + + def get_dummy_mm_data( + self, + *, + max_num_encoder_tokens, + mm_counts, + dtype, + ): + self.calls.append((max_num_encoder_tokens, mm_counts, dtype)) + item_count = mm_counts["image"] + return {"image": {"item_count": item_count}} + + class _Model(MultimodalModelMixin): + dtype = torch.float16 + mm_encoder = object() + + def __init__(self): + self.forwarded_item_counts = [] + self.last_output = None + + def encode_multimodal_inputs(self, multimodal_params): + count = multimodal_params[0].multimodal_data["image"]["item_count"] + self.forwarded_item_counts.append(count) + self.last_output = torch.arange(count * 4).reshape(count, 4) + return self.last_output + + input_processor = _InputProcessor() + model = _Model() + model.model_config = SimpleNamespace(pretrained_config=SimpleNamespace(vocab_size=128)) + creator = object.__new__(KvCacheCreator) + creator._model_engine = SimpleNamespace( + model=model, + input_processor=input_processor, + encoder_batch_size=4, + encoder_max_num_tokens=8192, + use_mrope=False, + ) + creator._profiling_stage_data = {"enable_mm_reqs": True} + creator._mapping = SimpleNamespace(enable_attention_dp=False) + creator._max_num_tokens = 18 + creator._max_beam_width = 1 + + # The LLM dummy remains text-only and is not allowed to shrink the + # independent encoder profiling budget. + requests = creator._create_dummy_context_requests(input_seq_len=18) + assert sum(len(request.input_token_ids) for request in requests) == 18 + assert input_processor.calls == [] + assert all(getattr(request, "py_multimodal_data", None) is None for request in requests) + + creator._dummy_encoder_inputs = creator._create_dummy_encoder_inputs() + assert input_processor.calls == [(8192, {"image": 2}, torch.float16)] + + monkeypatch.setattr(MultimodalParams, "to_device", lambda self, *args, **kwargs: self) + retained_output = creator._encode_dummy_inputs() + assert model.forwarded_item_counts == [2] + assert retained_output.data_ptr() != model.last_output.data_ptr() + assert creator._dummy_encoder_inputs == [] + + +def _make_creator( + tokens_per_block, + dummy_reqs, + enable_attention_dp, + tp_size, + batch_size=1, + model_max_seq_len=1, + max_cuda_graph_batch_size=1, + layer_types=None, + sliding_window=None, + use_sliding_window=None, + max_attention_window=None, +): + """Build a minimal KvCacheCreator (bypasses __init__) wired up for + _get_token_num_for_estimation only.""" + c = object.__new__(KvCacheCreator) + + c._tokens_per_block = tokens_per_block + c._net_max_seq_len = 2048 + c._speculative_config = None + c._dummy_reqs = dummy_reqs + + c._mapping = Mock(enable_attention_dp=enable_attention_dp, tp_size=tp_size, cp_config={}) + + c._llm_args = Mock(disable_overlap_scheduler=True) + + pretrained = SimpleNamespace( + layer_types=layer_types, + num_hidden_layers=(len(layer_types) if isinstance(layer_types, (list, tuple)) else None), + sliding_window=sliding_window, + use_sliding_window=use_sliding_window, + ) + + model_config = Mock() + model_config.pretrained_config = pretrained + + c._model_engine = Mock( + batch_size=batch_size, + max_seq_len=model_max_seq_len, + _max_cuda_graph_batch_size=max_cuda_graph_batch_size, + ) + c._model_engine.model = Mock(model_config=model_config) + + c._kv_cache_config = Mock( + free_gpu_memory_fraction=0.9, max_attention_window=max_attention_window + ) + + # _get_token_num_for_estimation gates pool-group scaling on V2 (the + # split-pool layout that motivates the scaling). Mamba hybrid uses + # MambaHybridCacheManager and must NOT scale. These tests target V2 + # behavior; production sets this in __init__ via + # _get_model_kv_cache_manager_cls(), which we bypass here. + c._kv_cache_manager_cls = KVCacheManagerV2 + c._is_kv_cache_manager_v2 = True + + return c + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _no_gpu(): + """Stub out CUDA memory queries and per-token KV size so the test runs on + any machine and the memory cap never constrains the result. + + _get_kv_size_per_token now returns a CacheCost; using slope=1 + zero + intercept keeps the legacy ``budget // bytes_per_token`` arithmetic + untouched downstream.""" + huge = 100 * (1 << 30) + with ( + patch("torch.cuda.mem_get_info", return_value=(huge, huge)), + patch.object(KvCacheCreator, "_get_kv_size_per_token", return_value=CacheCost(slope=1)), + ): + yield + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_cache_cost_tokens_for_budget_not_int_division(): + """Regression: V2 KV cache capacity must derive max_tokens via + ``CacheCost.tokens_for_budget``, not ``budget // CacheCost``.""" + budget = 10_000 + cost = CacheCost(slope=512, intercept=128) + assert cost.tokens_for_budget(budget) == (budget - 128) // 512 + with pytest.raises(TypeError): + _ = budget // cost + + +def test_adp_reduces_blocks_to_per_rank_share(): + """With ADP + tp_size duplicated requests the result must equal a single + rank's share, not the sum across all duplicates.""" + tpb = 64 + tp = 4 + n_in = 128 # ceil((128+1)/64) = 3 blocks per request + + baseline = _make_creator(tpb, [_make_mock_request(n_in)], enable_attention_dp=False, tp_size=1) + adp = _make_creator( + tpb, [_make_mock_request(n_in) for _ in range(tp)], enable_attention_dp=True, tp_size=tp + ) + + assert adp._get_token_num_for_estimation() == baseline._get_token_num_for_estimation() + + +def test_without_adp_all_blocks_counted(): + """Without ADP every request's blocks contribute to the total.""" + tpb = 64 + n_in = 128 # 3 blocks each + n_reqs = 4 + + c = _make_creator( + tpb, [_make_mock_request(n_in) for _ in range(n_reqs)], enable_attention_dp=False, tp_size=1 + ) + + # 4 reqs * 3 blocks * 64 tokens/block = 768 + assert c._get_token_num_for_estimation() == n_reqs * 3 * tpb + + +@pytest.mark.parametrize("tp_size", [2, 4, 8]) +def test_adp_various_tp_sizes(tp_size): + """ADP division must hold for several representative tp_size values.""" + tpb = 64 + n_in = 128 # 3 blocks per request + + c = _make_creator( + tpb, + [_make_mock_request(n_in) for _ in range(tp_size)], + enable_attention_dp=True, + tp_size=tp_size, + ) + + total = tp_size * 3 + expected_blocks = (total + tp_size - 1) // tp_size + assert c._get_token_num_for_estimation() == expected_blocks * tpb + + +def test_regression_without_fix_would_overcount(): + """If the ADP ceil-division fix were removed, the returned + value would be tp_size times too large. This test guards that fix.""" + tpb = 64 + tp = 4 + n_in = 128 + + c = _make_creator( + tpb, [_make_mock_request(n_in) for _ in range(tp)], enable_attention_dp=True, tp_size=tp + ) + + result = c._get_token_num_for_estimation() + + correct = 3 * tpb # 192 (per-rank share) + wrong = tp * 3 * tpb # 768 (all duplicates summed) + assert result == correct + assert result != wrong + + +@pytest.mark.parametrize( + ("model_cls", "encoder_cache_max_bytes", "expected_reserve"), + [ + (_TextModel, 64, 0), + (_MultimodalModel, 0, 0), + (_MultimodalModel, 64, 0), + (_EncoderCacheMultimodalModel, 0, 0), + (_EncoderCacheMultimodalModel, 64, 64), + ], +) +def test_kv_cache_estimation_reserves_multimodal_encoder_cache( + model_cls, + encoder_cache_max_bytes, + expected_reserve, +): + creator = object.__new__(KvCacheCreator) + creator._model_engine = SimpleNamespace(model=model_cls(encoder_cache_max_bytes)) + + assert creator._get_multimodal_encoder_memory_reserve() == expected_reserve + + +def test_reserve_adds_only_unprofiled_output_capacity(): + creator = object.__new__(KvCacheCreator) + creator._model_engine = SimpleNamespace( + mm_encoder_output_budget_bytes=512, + model=_MultimodalModel(0), + ) + assert creator._get_multimodal_encoder_memory_reserve(profiled_output_bytes=400) == 112 + + +# --------------------------------------------------------------------------- +# VSWA hybrid attention pool-group scaling (Gemma4 hybrid MMMU Pro hang fix) +# --------------------------------------------------------------------------- +# +# KVCacheManagerV2 creates one pool group per distinct attention-window size. +# The quota derived from max_tokens is split proportionally across pool +# groups, so each pool ends up with roughly num_cache_blocks / num_pool_groups +# blocks. A single long-context request then overflows the full-attention +# pool and the scheduler livelocks on suspend/retry. The fix scales +# num_cache_blocks by the number of distinct attention-window sizes inferred +# effective per-layer sliding windows on the pretrained config (preferred) or +# an explicit ``max_attention_window`` list on kv_cache_config (fallback). + + +def test_uniform_layer_types_no_scaling(): + """All-sliding or all-full layers stay a single pool group.""" + tpb = 32 + max_seq_len = 4096 + uniform = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=["sliding_attention"] * 26, + ) + # num_pool_groups = 1 -> behaviour unchanged from legacy ADP-only case. + baseline = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + ) + assert uniform._get_token_num_for_estimation() == baseline._get_token_num_for_estimation() + + +def test_get_layer_attention_window_honors_max_window_layers(): + config = SimpleNamespace( + use_sliding_window=True, + sliding_window=4096, + max_window_layers=2, + ) + + assert [get_layer_attention_window(config, layer_idx) for layer_idx in range(4)] == [ + None, + None, + 4096, + 4096, + ] + + +def test_gemma4_hybrid_scales_by_num_pool_groups(): + """Gemma4 hybrid attention (mixed sliding/full layers) must scale the + estimated block count by the number of distinct layer types. Otherwise + the per-pool quota is too small to hold a single max_seq_len request, + which is the MMMU Pro livelock reproducer.""" + tpb = 32 + max_seq_len = 12288 + layer_types = ["sliding_attention"] * 28 + ["full_attention"] * 7 + # Mixed sliding/full attention must include the model's actual window size; + # Gemma4-E2B uses a 512-token sliding window. + sliding_window = 512 + assert len(set(layer_types)) == 2 + + hybrid = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1), _make_mock_request(1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=layer_types, + sliding_window=sliding_window, + ) + uniform = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1), _make_mock_request(1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=["full_attention"] * 35, + ) + + hybrid_tokens = hybrid._get_token_num_for_estimation() + uniform_tokens = uniform._get_token_num_for_estimation() + assert hybrid_tokens == 2 * uniform_tokens, ( + f"Expected 2x scaling for 2 pool groups, got " + f"hybrid={hybrid_tokens}, uniform={uniform_tokens}" + ) + + +def test_hybrid_linear_attention_scales_by_num_pool_groups(): + """Hybrid linear/attention V2 managers retain both estimation pools.""" + tpb = 32 + max_seq_len = 4096 + layer_types = ["linear_attention"] * 30 + ["full_attention"] * 10 + + hybrid = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=layer_types, + ) + uniform = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=["full_attention"] * 40, + ) + + assert hybrid._get_token_num_for_estimation() == 2 * uniform._get_token_num_for_estimation() + + +@pytest.mark.parametrize( + ("sliding_window", "use_sliding_window"), + [ + ([512, 1024], None), + (None, True), + ], + ids=["multiple_window_sizes", "missing_window"], +) +def test_v2_pool_estimation_falls_back_for_unsupported_window_metadata( + sliding_window, + use_sliding_window, +): + tpb = 32 + max_seq_len = 4096 + hybrid = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=["sliding_attention", "full_attention"], + sliding_window=sliding_window, + use_sliding_window=use_sliding_window, + ) + uniform = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=["full_attention", "full_attention"], + ) + + assert hybrid._get_token_num_for_estimation() == 2 * uniform._get_token_num_for_estimation() + + +def test_vswa_max_attention_window_fallback_scales(): + """When layer_types is absent but kv_cache_config.max_attention_window is + a heterogeneous list (VSWA), we still scale by the number of distinct + windows.""" + tpb = 32 + max_seq_len = 12288 + max_attention_window = [1024] * 28 + [max_seq_len] * 7 + assert len(set(max_attention_window)) == 2 + + c = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=None, + max_attention_window=max_attention_window, + ) + uniform = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + ) + assert c._get_token_num_for_estimation() == 2 * uniform._get_token_num_for_estimation() + + +def test_pool_scaling_prevents_mmmu_pro_underestimation(): + """Regression: with max_seq_len=12288 and max_num_tokens=12288 (MMMU Pro + config), hybrid estimation must produce enough capacity to hold one full + max_seq_len request per pool, with max_util_for_resume headroom.""" + tpb = 32 + max_seq_len = 12288 + layer_types = ["sliding_attention"] * 28 + ["full_attention"] * 7 + # Mixed sliding/full attention must include the model's actual window size; + # Gemma4-E2B uses a 512-token sliding window. + sliding_window = 512 + + c = _make_creator( + tpb, + [_make_mock_request(max_seq_len - 1), _make_mock_request(1)], + enable_attention_dp=False, + tp_size=1, + model_max_seq_len=max_seq_len, + max_cuda_graph_batch_size=4, + layer_types=layer_types, + sliding_window=sliding_window, + ) + + total_tokens = c._get_token_num_for_estimation() + per_pool_tokens = total_tokens // 2 # 2 pool groups + # Must be enough to hold a max_seq_len request per pool. + assert per_pool_tokens >= max_seq_len + + +def test_v2_cache_size_per_token_models_generation_swa_cost(): + class FakeModelConfig: + quant_config = None + pretrained_config = SimpleNamespace( + hidden_size=32, + num_attention_heads=4, + num_key_value_heads=2, + ) + + def get_num_attention_layers(self): + return 3 + + mapping = Mock(enable_attention_dp=False, tp_size=1) + mapping.pp_layers.return_value = [0, 1, 2] + + no_scratch_size_per_token = CacheCost.from_raw( + KVCacheManagerV2.get_cache_size_per_token( + FakeModelConfig(), + mapping, + tokens_per_block=64, + max_seq_len=4096, + max_batch_size=3, + kv_cache_config=KvCacheConfig(max_attention_window=[2048, 2048, 4096]), + ) + ) + scratch_size_per_token = CacheCost.from_raw( + KVCacheManagerV2.get_cache_size_per_token( + FakeModelConfig(), + mapping, + tokens_per_block=64, + max_seq_len=4096, + max_batch_size=3, + kv_cache_config=KvCacheConfig(max_attention_window=[2048, 2048, 4096]), + enable_swa_scratch_reuse=True, + ) + ) + + # Per layer: K+V * kv_heads * head_dim * bf16 bytes = 2 * 2 * 8 * 2. + expected = CacheCost(slope=64, intercept=3 * 2 * 2048 * 64) + assert no_scratch_size_per_token == expected + assert scratch_size_per_token == expected + + +def test_creator_uses_v2_affine_cache_cost(): + class FakeV2Manager(KVCacheManagerV2): + @staticmethod + def get_cache_size_per_token(model_config, mapping, **kwargs): + return 20, 21 + + creator = object.__new__(KvCacheCreator) + creator._mapping = Mock() + creator._tokens_per_block = 64 + creator._max_seq_len = 1024 + creator._max_batch_size = 3 + creator._kv_cache_config = KvCacheConfig() + creator._speculative_config = None + + cost = creator._per_manager_cache_cost(FakeV2Manager, Mock()) + + assert cost == CacheCost(slope=20, intercept=21) + + +def test_v2_quota_from_max_tokens_models_context_swa_scratch(): + manager = object.__new__(KVCacheManagerV2) + manager._has_cp_helix = False + manager.num_local_layers = 3 + manager.pp_layers = [0, 1, 2] + manager.max_attention_window_vec = [128, 128, None] + manager.tokens_per_block = 64 + manager.max_batch_size = 4 + manager.max_num_tokens = 1000 + manager.get_layer_bytes_per_token = lambda local_layer_idx, data_role: [10, 10, 20][ + local_layer_idx + ] + + max_tokens = 1200 + + manager.enable_swa_scratch_reuse = False + no_scratch_quota = manager._get_quota_from_max_tokens(max_tokens) + assert no_scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 20 + 4 * 2 * 128 * 10) + assert manager._get_max_tokens_from_quota(no_scratch_quota) == max_tokens + + manager.enable_swa_scratch_reuse = True + scratch_quota = manager._get_quota_from_max_tokens(max_tokens) + assert scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 10 + 4 * 2 * 128 * 10) + assert manager._get_max_tokens_from_quota(scratch_quota) == max_tokens + + +# --------------------------------------------------------------------------- +# KVCacheManagerV2 clamp_max_seq_len_for_mem float-to-int cast regression +# --------------------------------------------------------------------------- +# +# ``KVCacheManagerV2.get_num_available_tokens`` returns +# ``clamp_max_seq_len_for_mem(...) - extra_tokens``, where the clamp helper +# does a floating-point memory-budget division and returns a float. When +# ``max_seq_len > max_num_tokens`` the ``__init__`` body assigned the float +# directly to ``self.max_seq_len``, which then propagated into +# ``_util.py::_create_dummy_context_requests``'s ``torch.randint(size=(...))`` +# call and crashed on a float size tuple. +# +# resource_manager.py:1820 now casts ``int(max_num_tokens)`` at the assign +# site. The tests below lock in that invariant: (a) the replay test +# demonstrates the fix applied to the scenario that previously crashed, +# and (b) the propagation test shows that downstream arithmetic in +# ``_util.py:610-620`` stays int-safe. + + +def test_kv_cache_manager_v2_clamp_casts_float_to_int(): + """Replay of the V2 clamp block with a float clamp result. + + Reproduces the exact arithmetic from ``resource_manager.py`` + ``KVCacheManagerV2.__init__`` lines 1813-1825. Pre-fix, the ``if`` + branch stored a float on ``self.max_seq_len``. Post-fix, it casts + to int. A broken fix (e.g. dropping the cast) will make this test + fail on the ``isinstance(..., int)`` assertion.""" + # Scenario pulled from the MMMU Pro 26B + s=131K crash log: + # clamp_max_seq_len_for_mem returned 60160.0 when requested + # 131072 tokens of max_seq_len. + initial_max_seq_len = 131072 + # Simulate the float return from clamp_max_seq_len_for_mem. + clamp_float_result = 60160.0 + assert isinstance(clamp_float_result, float) + # Match the production code block (resource_manager.py:1816-1825). + self_max_seq_len = initial_max_seq_len + if self_max_seq_len > clamp_float_result: + self_max_seq_len = int(clamp_float_result) + + assert isinstance(self_max_seq_len, int), ( + "KVCacheManagerV2 must cast the clamp result to int — float " + "propagates into torch.randint(size=...) and crashes." + ) + assert self_max_seq_len == 60160 + + +def test_kv_cache_manager_v2_float_max_seq_len_would_crash_torch_randint(): + """Pre-fix behaviour: a float max_seq_len propagating into + torch.randint(size=(...,)) raises. This test documents WHY the cast + is necessary — if the cast is dropped, the following code crashes.""" + + float_seq_len = 60160.0 + with pytest.raises((TypeError, RuntimeError)): + # torch.randint only accepts int-valued size tuples. + torch.randint(low=0, high=32000, size=(float_seq_len,)) + + +def test_kv_manager_int_max_seq_len_stays_int_through_util_expression(): + """The downstream expression in ``_util.py:615-620`` (which builds the + dummy context request input_seq_len) must stay int when the V2 + manager's max_seq_len is int. Covers the full propagation chain + from the root-cause fix to the downstream consumer.""" + net_max_seq_len = 131072 # from ModelEngine + creator_max_seq_len = 131072 # KvCacheCreator.self._max_seq_len + kv_manager_max_seq_len = 60160 # post-fix int (was 60160.0 pre-fix) + + # Replay _util.py:616-619 arithmetic. + input_seq_len = max( + 1, + net_max_seq_len - 1 - (creator_max_seq_len - kv_manager_max_seq_len), + ) + assert isinstance(input_seq_len, int) + assert input_seq_len >= 1 + + # Sanity: the same expression with a float would yield float (the + # original bug), ensuring this test would notice regression. + buggy_input = max( + 1, + net_max_seq_len - 1 - (creator_max_seq_len - float(kv_manager_max_seq_len)), + ) + assert isinstance(buggy_input, float), ( + "Sanity check: the pre-fix float path must still be demonstrable " + "in the test, otherwise this regression guard is hollow." + ) + + +# --------------------------------------------------------------------------- +# _create_kv_cache_manager: MLA branch must forward max_num_tokens +# --------------------------------------------------------------------------- + + +def test_mla_branch_forwards_max_num_tokens_to_manager() -> None: + """The is_mla branch of ``_create_kv_cache_manager`` must pass + ``max_num_tokens`` to the manager, like the generic branch does. + + Regression guard: when the argument is dropped, + ``DeepseekV4CacheManager._max_num_tokens`` stays ``None`` and the + profiling-phase context extra quota is sized by the full ``max_tokens`` + estimate instead of the runtime chunk size (observed as a 27.59 GiB vs + 11.92 GiB temp-quota inflation on DeepSeek-V4-Pro, raising peak memory + and OOM risk during KV cache estimation). + """ + + from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager + + captured_kwargs = {} + + class _RecordingManager: + def __init__(self, *args, **kwargs) -> None: + captured_kwargs.update(kwargs) + + # Minimal MLA pretrained config: is_mla() keys on kv_lora_rank and + # qk_rope_head_dim being set. + pretrained = SimpleNamespace( + kv_lora_rank=512, + qk_rope_head_dim=64, + hidden_size=1024, + num_attention_heads=8, + num_key_value_heads=8, + num_hidden_layers=2, + vocab_size=32000, + ) + model_config = Mock() + model_config.pretrained_config = pretrained + model_config.quant_config = None + + _create_kv_cache_manager( + model_engine=None, + kv_cache_manager_cls=_RecordingManager, + mapping=Mock(), + kv_cache_config=KvCacheConfig(), + tokens_per_block=32, + max_seq_len=2048, + max_batch_size=8, + spec_config=None, + sparse_attention_config=None, + max_num_tokens=333, + max_beam_width=1, + kv_connector_manager=None, + model_config=model_config, + dtype=torch.bfloat16, + is_draft=False, + ) + + assert captured_kwargs.get("max_num_tokens") == 333, ( + "is_mla branch dropped max_num_tokens: DeepseekV4CacheManager then " + "falls back to _max_num_tokens=None and over-sizes the estimation " + "temp quota." + ) + + +def test_estimation_temporarily_uses_inferred_pool_sizing() -> None: + pool_ratio = [0.2, 0.3, 0.5] + avg_seq_len = 128 + max_seq_len = 4096 + user_max_tokens = 1024 + estimation_max_tokens = 256 + kv_cache_config = KvCacheConfig( + max_tokens=user_max_tokens, + pool_ratio=pool_ratio, + avg_seq_len=avg_seq_len, + ) + model_engine = Mock() + model_engine.model.model_config.attn_backend = "TRTLLM" + # Explicit False: try_prepare_estimation skips estimation for + # encoder-decoder models, and a bare Mock attribute is truthy. + model_engine.model.model_config.is_encoder_decoder = False + # A bare Mock would auto-create the attribute; real engines set it to + # None unless the model opted into MM item scheduling. + model_engine.mm_encoder_output_budget_bytes = None + llm_args = Mock(cache_transceiver_config=None) + + with patch.object( + KvCacheCreator, + "_get_model_kv_cache_manager_cls", + return_value=KVCacheManagerV2, + ): + creator = KvCacheCreator( + model_engine=model_engine, + draft_model_engine=None, + mapping=Mock(cp_config={}), + net_max_seq_len=max_seq_len, + kv_connector_manager=None, + max_num_tokens=256, + max_beam_width=1, + tokens_per_block=128, + max_seq_len=max_seq_len, + max_batch_size=8, + kv_cache_config=kv_cache_config, + llm_args=llm_args, + speculative_config=None, + sparse_attention_config=None, + profiling_stage_data=None, + is_disagg=False, + ) + + with ( + patch.object( + creator, + "_get_token_num_for_estimation", + return_value=estimation_max_tokens, + ), + patch.object(creator, "_cal_max_memory", return_value=512), + patch.object(torch.cuda, "mem_get_info", return_value=(768, 1024)), + patch.object(torch.cuda, "memory_stats", return_value={"allocated_bytes.all.current": 128}), + patch.object(torch.cuda, "empty_cache"), + patch.object(torch.cuda, "reset_peak_memory_stats"), + # This test exercises inferred pool sizing, not the backend workspace reserve; the mock + # model_config would otherwise walk into backend resolution and the MLA byte-cost path. + # Neutralize it so no reserve applies. + patch( + "tensorrt_llm._torch.pyexecutor._util.get_attention_workspace_bytes_per_token", + return_value=0, + ), + ): + assert creator.try_prepare_estimation() + assert kv_cache_config.max_tokens == estimation_max_tokens + assert kv_cache_config.pool_ratio is None + assert kv_cache_config.avg_seq_len == max_seq_len + + creator.configure_kv_cache_capacity() + + assert kv_cache_config.max_tokens == user_max_tokens + assert kv_cache_config.pool_ratio == pool_ratio + assert kv_cache_config.avg_seq_len == avg_seq_len + + +@pytest.mark.parametrize( + ("estimating_kv_cache", "expected_avg_seq_len"), + [(True, 2045), (False, 2055)], +) +def test_manager_estimation_clamps_only_temporary_avg_seq_len( + estimating_kv_cache, + expected_avg_seq_len, +) -> None: + import torch + + from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager + + captured_configs = [] + + class _RecordingKVCacheManagerV2(KVCacheManagerV2): + def __init__(self, kv_cache_config, _kv_cache_type, **kwargs) -> None: + captured_configs.append(kv_cache_config) + self.max_seq_len = kwargs["max_seq_len"] + + pretrained = SimpleNamespace( + hidden_size=1024, + num_attention_heads=8, + num_key_value_heads=8, + num_hidden_layers=2, + vocab_size=32000, + ) + model_config = Mock() + model_config.pretrained_config = pretrained + model_config.quant_config = None + kv_cache_config = KvCacheConfig( + max_tokens=2048, + avg_seq_len=2055, + ) + + _create_kv_cache_manager( + model_engine=None, + kv_cache_manager_cls=_RecordingKVCacheManagerV2, + mapping=Mock(), + kv_cache_config=kv_cache_config, + tokens_per_block=32, + max_seq_len=2045, + max_batch_size=4, + spec_config=None, + sparse_attention_config=None, + max_num_tokens=2048, + max_beam_width=1, + kv_connector_manager=None, + estimating_kv_cache=estimating_kv_cache, + model_config=model_config, + dtype=torch.bfloat16, + is_draft=False, + ) + + assert captured_configs[0].avg_seq_len == expected_avg_seq_len + assert kv_cache_config.avg_seq_len == 2055 + + +def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: + creator = object.__new__(KvCacheCreator) + target_pool_ratio = [0.32, 0.68] + creator._kv_cache_config = KvCacheConfig( + pool_ratio=target_pool_ratio, + max_attention_window=None, + ) + creator._max_seq_len = 9472 + creator._tokens_per_block = 32 + creator._max_batch_size = 1024 + creator._max_num_tokens = 9472 + creator._max_beam_width = 1 + creator._kv_connector_manager = None + creator._skip_est = False + creator._execution_stream = None + creator._is_disagg = False + creator._mapping = Mock() + creator._speculative_config = Mock() + + effective_draft_config = Mock() + effective_draft_config.pretrained_config.torch_dtype = "bfloat16" + effective_draft_config.sparse_attention_config = None + + with ( + patch.object(creator, "_get_num_draft_layers", return_value=1), + patch.object(creator, "_get_one_model_draft_layer_mask", return_value=[True]), + patch.object( + creator, + "_get_effective_draft_config", + return_value=effective_draft_config, + ), + patch.object(creator, "_enable_kv_cache_stats", return_value=False), + patch.object( + creator, + "_validate_or_fallback_kv_cache_manager_v2", + return_value=KVCacheManagerV2, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util._derive_draft_max_attention_window", + return_value=None, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util.get_kv_cache_manager_cls", + return_value=KVCacheManagerV2, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", + return_value=Mock(), + ) as create_manager, + ): + creator._create_one_model_draft_kv_cache_manager(creator._max_seq_len) + + draft_config = create_manager.call_args.kwargs["kv_cache_config"] + assert draft_config.pool_ratio == [1.0] + assert creator._kv_cache_config.pool_ratio == target_pool_ratio + + +@pytest.mark.parametrize( + "attention_backend,dtype,bytes_per_layer", + [("TRTLLM", torch.float8_e4m3fn, 2048), ("VANILLA", torch.bfloat16, 4096)], +) +def test_standalone_draft_kv_dtype_follows_the_drafter_attention( + attention_backend, dtype, bytes_per_layer +) -> None: + """The TRTLLM drafter attention stores the drafter's history in FP8; VANILLA keeps BF16.""" + creator = object.__new__(KvCacheCreator) + draft = SimpleNamespace( + fc=SimpleNamespace(weight=torch.empty(1, dtype=torch.bfloat16)), + dflash_attention_backend=attention_backend, + config=SimpleNamespace(max_position_embeddings=4096), + ) + creator._model_engine = SimpleNamespace(model=SimpleNamespace(draft_model=draft)) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace(is_dspark=lambda: True), + draft_is_embedded_in_target=False, + attention_backend="AUTO", + max_draft_len=3, + ) + # Drafter shapes remain unsharded under attention DP. + creator._draft_config = SimpleNamespace( + pretrained_config=SimpleNamespace( + num_attention_heads=64, + num_key_value_heads=16, + head_dim=64, + hidden_size=4096, + num_hidden_layers=5, + ) + ) + creator._mapping = SimpleNamespace( + has_cp_helix=lambda: False, enable_attention_dp=True, tp_size=16 + ) + + kwargs = creator._get_unified_draft_cache_kwargs() + + assert kwargs["draft_dtype"] == torch_dtype_to_binding(dtype) + assert _get_draft_layer_sizes(**kwargs) == (bytes_per_layer,) * 5 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py index e77b14140dd5..31355f231135 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py @@ -87,80 +87,118 @@ def test_trtllm_backend_builds_private_paged_context_cache(monkeypatch): ] -def test_generation_step_keeps_the_dummy_slot_context_empty(): - """Padding and ADP dummy rows must not accumulate drafter context. - - They all map to the shared dummy slot, which no request ever frees. If the - generation step advanced its length like a real slot's, the dummies' drafter - context would grow every step until it reached ``_max_ctx``. - """ - max_draft_len = 3 - num_slots = 3 # two request slots plus the dummy slot - hidden, num_layers, num_kv_heads, head_dim = 8, 1, 1, 4 - max_ctx = 64 - config = DFlashDecodingConfig( - max_draft_len=max_draft_len, attention_backend="TRTLLM", mask_token_id=0 - ) - worker = DFlashWorker(config, Mapping()) - worker._dummy_slot = num_slots - 1 - worker._max_ctx = max_ctx - worker._compute_block_size = max_draft_len + 1 - worker._ctx_len = torch.tensor([5, 0, 0], dtype=torch.long, device="cuda") - # Gen request 0 owns slot 0; gen request 1 is padding on the dummy slot. - worker._batch_to_slot_buffer = DeviceCopyBuffer( - torch.tensor([0, worker._dummy_slot], dtype=torch.long, device="cuda") - ) - buf_shape = (num_slots, num_layers, max_ctx + max_draft_len + 1, num_kv_heads, head_dim) - worker._ctx_k_buf = torch.zeros(buf_shape, device="cuda") - worker._ctx_v_buf = torch.zeros(buf_shape, device="cuda") - - num_gens = 2 - tokens_per_req = max_draft_len + 1 - num_target_tokens = num_gens * tokens_per_req - - def precompute_context_kv(projected, positions): - rows = projected.shape[0] - k = torch.ones(rows, num_layers, num_kv_heads, head_dim, device="cuda") - return k, k.clone() - - embed_tokens = torch.nn.Embedding(16, hidden).cuda() - draft_model = SimpleNamespace( - fc=object(), - hidden_norm=object(), - project_target_hidden=lambda hs: hs, - precompute_context_kv=precompute_context_kv, - draft_model_full=SimpleNamespace(model=SimpleNamespace(embed_tokens=embed_tokens)), - ) - captured = torch.randn(num_target_tokens, hidden, device="cuda") - spec_metadata = SimpleNamespace( - hidden_size=hidden, - runtime_draft_len=max_draft_len, - get_hidden_states=lambda num_tokens: captured[:num_tokens], - ) - attn_metadata = SimpleNamespace( - num_contexts=0, - num_seqs=num_gens, - num_ctx_tokens=0, - _seq_lens_cuda=torch.zeros(num_gens, dtype=torch.int32, device="cuda"), - _seq_lens=torch.zeros(num_gens, dtype=torch.int32), - ) - accepted_tokens = torch.arange(num_gens * tokens_per_req, dtype=torch.long, device="cuda").view( - num_gens, tokens_per_req +def _make_worker(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", + lambda **kwargs: None, ) - num_accepted_tokens = torch.tensor([2, 3], dtype=torch.int32, device="cuda") - - for step in range(2): - inputs = worker.prepare_1st_drafter_inputs( - input_ids=torch.zeros(num_target_tokens, dtype=torch.long, device="cuda"), - position_ids=torch.zeros(num_target_tokens, dtype=torch.long, device="cuda"), - hidden_states=torch.zeros(num_target_tokens, hidden, device="cuda"), - accepted_tokens=accepted_tokens, - num_accepted_tokens=num_accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_model=draft_model, - total_target_tokens=num_target_tokens, - ) - real_len = 5 + 2 * (step + 1) - assert worker._ctx_len.tolist() == [real_len, 0, 0] - assert inputs["num_ctx_per_req"].tolist() == [real_len, 0] + config = DFlashDecodingConfig(max_draft_len=7, attention_backend="TRTLLM") + return DFlashWorker(config, Mapping()) + + +class _FakeDraftCacheManager: + """Unified manager whose draft layer 61 lives in pool 1 with two KV planes.""" + + draft_layer_ids = [61] + layer_offsets = {61: 7} + layer_to_pool_mapping_dict = {7: 1} + _index_scale_ints = [1, 1] + enable_swa_scratch_reuse = False + + def __init__(self, pages, helix_cp_size=1): + self.pages = pages + self._helix_cp_size = helix_cp_size + self.kv_cache_map = { + request_id: SimpleNamespace(num_blocks=len(request_pages), is_active=True) + for request_id, request_pages in pages.items() + } + + def get_layer_page_index_scale(self, layer_idx): + return 1 + + def get_layer_kv_factor(self, layer_idx): + return 2 + + def get_batch_cache_indices(self, request_ids, layer_idx): + return [list(self.pages[request_id]) for request_id in request_ids] + + def block_offsets(self, request_ids, width): + """The target's [pool, batch, plane, block] offsets: K at page * 2, V at page * 2 + 1.""" + offsets = torch.full((2, len(request_ids), 2, width), -1, dtype=torch.int32) + for row, request_id in enumerate(request_ids): + pages = torch.tensor(self.pages[request_id], dtype=torch.int32) + offsets[0, row, :, : len(pages)] = 1000 # the target pool's own pages + offsets[1, row, 0, : len(pages)] = pages * 2 + offsets[1, row, 1, : len(pages)] = pages * 2 + 1 + return offsets.cuda() + + +def _bind_device_tables(worker, manager): + worker._ctx_kv_manager = manager + worker._init_ctx_block_tables(num_slots=3, max_blocks=4) + worker._init_unified_draft_tables(manager) + + +def test_unified_draft_tables_are_decoded_from_target_block_offsets(monkeypatch): + worker = _make_worker(monkeypatch) + manager = _FakeDraftCacheManager({11: [5, 9, 2], 12: [7]}) + _bind_device_tables(worker, manager) + assert worker._device_draft_tables + assert (worker._ctx_pool_idx, worker._ctx_block_divisor) == (1, 2) + + attn_metadata = SimpleNamespace(kv_cache_block_offsets=manager.block_offsets([11, 12], 4)) + assert worker._refresh_unified_block_tables(attn_metadata, 2, [11, 12]) + + tables = worker._ctx_block_tables.cpu() + assert tables[0, :3].tolist() == [5, 9, 2] + assert tables[1, :1].tolist() == [7] + assert worker._ctx_block_counts[:2].tolist() == [3, 1] + # The startup check compared the rows with the manager's page lists. + assert worker._device_draft_tables + assert worker._draft_table_checks_left == 7 + + +def test_unified_draft_tables_fall_back_to_host_on_mismatch(monkeypatch): + worker = _make_worker(monkeypatch) + manager = _FakeDraftCacheManager({11: [5, 9, 2], 12: [7]}) + _bind_device_tables(worker, manager) + attn_metadata = SimpleNamespace(kv_cache_block_offsets=manager.block_offsets([11, 12], 4)) + manager.pages[12] = [8] + + assert not worker._refresh_unified_block_tables(attn_metadata, 2, [11, 12]) + assert not worker._device_draft_tables + + +def test_unified_draft_tables_stay_on_host_for_helix_managers(monkeypatch): + worker = _make_worker(monkeypatch) + _bind_device_tables(worker, _FakeDraftCacheManager({}, helix_cp_size=2)) + + assert not worker._device_draft_tables + + +def test_fp8_draft_cache_stores_context_kv_by_direct_page_writes(monkeypatch): + worker = _make_worker(monkeypatch) + num_layers, nkv, page_size, hd = 2, 2, 32, 64 + worker._ctx_kv_buf = [ + torch.zeros(4, 2, nkv, page_size, hd, dtype=torch.float8_e4m3fn, device="cuda") + for _ in range(num_layers) + ] + worker._ctx_block_tables = torch.tensor([[3, 1], [0, 2]], dtype=torch.int32, device="cuda") + worker._ctx_page_size = page_size + rows = torch.tensor([0, 0, 1], device="cuda") + positions = torch.tensor([5, 40, 33], device="cuda") + k = torch.randn(3, num_layers, nkv, hd, dtype=torch.bfloat16, device="cuda") + v = torch.randn_like(k) + + worker._store_context_kv_paged(k, v, rows, positions) + + pages = torch.tensor([3, 1, 2], device="cuda") + offsets = positions % page_size + for layer, pool in enumerate(worker._ctx_kv_buf): + stored = pool.view(torch.uint8) + expected_k = k[:, layer].to(pool.dtype).view(torch.uint8) + expected_v = v[:, layer].to(pool.dtype).view(torch.uint8) + assert torch.equal(stored[pages, 0, :, offsets], expected_k) + assert torch.equal(stored[pages, 1, :, offsets], expected_v) + stored[pages, :, :, offsets] = 0 + assert not stored.any(), "the store wrote outside the addressed slots" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index 40a841f3656f..098e42770acf 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -12,13 +12,7 @@ # 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. -"""Unit tests for the DSpark captured-context attention primitives (CPU). - -The two primitives are validated against fully independent computations: -``dspark_sparse_attn`` vs (a) a per-element loop reference of the kernel formula -and (b) ``torch.nn.functional.scaled_dot_product_attention`` for the no-sink, -all-valid case; ``get_dspark_topk_idxs`` vs the reference index formula. -""" +"""DSpark attention math and managed-page addressing tests.""" import types from unittest.mock import Mock @@ -31,12 +25,12 @@ from tensorrt_llm._torch.models.modeling_dspark import ( DSv4DSparkDraftModel, apply_dspark_rotary, - dspark_attention_forward, - dspark_sparse_attn, - get_dspark_topk_idxs, precompute_dspark_freqs_cis, ) +__extra_import_path__ = [".."] + + # The captured-context attention primitives were folded into # modeling_dspark; keep the historical alias so monkeypatch targets # below read unchanged. @@ -86,32 +80,353 @@ def test_rmsnorm_rope_skips_frequency_view_on_unsupported_arch(monkeypatch): def test_batched_attention_rejects_mismatched_window_size(): - unused_tensor = torch.empty(0) - with pytest.raises(ValueError, match="does not match window_size"): - dspark_attention.dspark_attention_forward_batched( - torch.empty(1, 1, 1), - torch.empty(1, 1, 1), - torch.zeros(1, dtype=torch.long), - torch.empty(1, 2, 1), - torch.zeros(1, dtype=torch.long), - wq_a=unused_tensor, - q_norm_w=unused_tensor, - wq_b=unused_tensor, - wkv=unused_tensor, - kv_norm_w=unused_tensor, - wo_a=unused_tensor, - wo_b=unused_tensor, - attn_sink=unused_tensor, - n_heads=1, - head_dim=1, - rope_head_dim=0, - n_groups=1, - o_lora_rank=1, - window_size=3, - eps=1e-6, - softmax_scale=1.0, - freqs_cis=unused_tensor, - ) + import inspect + + forward = modeling_dspark.dspark_attention_forward + args = { + name: None + for name, param in inspect.signature(forward).parameters.items() + if param.default is inspect.Parameter.empty + } + args.update(window_size=64, rope_head_dim=64) + with pytest.raises(ValueError, match="128-token window"): + forward(**args) + + +@pytest.mark.parametrize( + "value,tp_size,enable_attention_dp,expected", + [ + (None, 1, False, False), + (None, 2, True, False), + (None, 2, False, True), + ("fp8", 2, False, True), + ("bf16", 1, False, True), + ], +) +def test_use_dspark_bf16_linear(monkeypatch, value, tp_size, enable_attention_dp, expected): + if value is None: + monkeypatch.delenv("TRTLLM_DSPARK_LINEAR_BACKEND", raising=False) + else: + monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", value) + + mapping = types.SimpleNamespace( + tp_size=tp_size, + enable_attention_dp=enable_attention_dp, + ) + assert modeling_dspark._use_dspark_bf16_linear(mapping) is expected + + +def test_use_dspark_bf16_linear_rejects_invalid_backend(monkeypatch): + monkeypatch.setenv("TRTLLM_DSPARK_LINEAR_BACKEND", "auto") + mapping = types.SimpleNamespace(tp_size=1, enable_attention_dp=False) + + with pytest.raises(ValueError, match="must be 'fp8' or 'bf16'"): + modeling_dspark._use_dspark_bf16_linear(mapping) + + +def test_fused_qkv_projection_groups_rows_without_copy() -> None: + batch, block, hidden = 2, 3, 4 + q_rank, kv_rank = 6, 2 + x = torch.arange(batch * block * hidden, dtype=torch.float32).view(batch, block, hidden) + main_x = -torch.arange(batch * hidden, dtype=torch.float32).view(batch, 1, hidden) + weight = torch.arange((q_rank + kv_rank) * hidden, dtype=torch.float32).view( + q_rank + kv_rank, hidden + ) + projected_inputs: list[torch.Tensor] = [] + + def qkv_proj(rows: torch.Tensor) -> torch.Tensor: + projected_inputs.append(rows.clone()) + return F.linear(rows, weight) + + q_latent, kv, main_kv = _qkv_projections( + x, + main_x, + wq_a=None, + wkv=None, + q_norm_w=torch.empty(q_rank), + qkv_proj=qkv_proj, + ) + + assert len(projected_inputs) == 1 + expected_rows = torch.cat((main_x.view(-1, hidden), x.view(-1, hidden)), dim=0) + torch.testing.assert_close(projected_inputs[0], expected_rows) + + expected_main = F.linear(main_x, weight) + expected_block = F.linear(x, weight) + torch.testing.assert_close(q_latent, expected_block[..., :q_rank]) + torch.testing.assert_close(kv, expected_block[..., q_rank:]) + torch.testing.assert_close(main_kv, expected_main[..., q_rank:]) + + for tensor in (q_latent, kv, main_kv): + flat = tensor.view(-1, tensor.shape[-1]) + assert flat.stride(-1) == 1 + assert flat.untyped_storage().data_ptr() == tensor.untyped_storage().data_ptr() + + +def test_fused_fp8_input_producers_feed_prequantized_linears(monkeypatch) -> None: + import tensorrt_llm._torch.models.modeling_dspark as dspark_attention + + batch, block, hidden = 2, 3, 4 + q_rank, kv_rank, q_out = 8, 2, 12 + x = torch.randn(batch, block, hidden) + main_x = torch.randn(batch, 1, hidden) + concat_input = (object(), object()) + norm_input = (object(), object()) + projection_inputs = [] + prevalidated_calls = [] + + def prevalidated_projection(input, projection): + prevalidated_calls.append((input, projection)) + return projection(input) + + monkeypatch.setattr(dspark_attention, "_prequantized_fp8_projection", prevalidated_projection) + + def concat_quant(first, second): + assert first.shape == (batch, hidden) + assert second.shape == (batch * block, hidden) + return concat_input + + def qkv_proj(input): + projection_inputs.append(input) + return torch.randn(batch * (block + 1), q_rank + kv_rank) + + monkeypatch.setattr(dspark_attention, "cute_dsl_dspark_concat_fp8_quant", concat_quant) + q_latent, _, _ = _qkv_projections( + x, + main_x, + wq_a=None, + wkv=None, + q_norm_w=torch.empty(q_rank), + qkv_proj=qkv_proj, + use_fused_fp8_input=True, + ) + assert projection_inputs == [concat_input] + assert prevalidated_calls == [(concat_input, qkv_proj)] + + def rmsnorm_quant(input, weight, eps): + assert input is q_latent + assert weight.shape == (q_rank,) + assert eps == 1e-6 + return norm_input + + def q_b_proj(input): + projection_inputs.append(input) + return torch.randn(batch * block, q_out) + + monkeypatch.setattr(dspark_attention, "cute_dsl_dspark_rmsnorm_fp8_quant", rmsnorm_quant) + q = dspark_attention._q_projection(q_latent, torch.ones(q_rank), 1e-6, q_b_proj, True) + assert projection_inputs[-1] == norm_input + assert prevalidated_calls[-1] == (norm_input, q_b_proj) + assert q.shape == (batch, block, q_out) + + +def test_output_projection_uses_fused_fp8_path(monkeypatch): + import tensorrt_llm._torch.models.modeling_dspark as dspark_attention + + batch, block, groups, rank, group_dim, hidden = 2, 3, 2, 128, 16, 24 + o = torch.randn(batch, block, groups, group_dim, dtype=torch.bfloat16) + o_a_proj = torch.empty(groups, rank, group_dim, dtype=torch.float8_e4m3fn) + o_a_scale = object() + calls = [] + prevalidated_calls = [] + + def prevalidated_projection(input, projection): + prevalidated_calls.append((input, projection)) + return projection(input) + + monkeypatch.setattr(dspark_attention, "_prequantized_fp8_projection", prevalidated_projection) + + class FakeOutputLinear: + has_fp8_block_scales = True + use_cute_dsl_blockscaling_mm = True + disable_deep_gemm = False + in_features = groups * rank + weight = torch.empty(hidden, groups * rank, dtype=torch.float8_e4m3fn) + weight_scale = torch.empty(1, dtype=torch.uint8) + + def __call__(self, input): + calls.append(input) + return torch.zeros(batch * block, hidden, dtype=torch.bfloat16) + + def fake_quantize(input): + assert input.shape == (batch * block, groups, group_dim) + return input.permute(1, 0, 2).contiguous().to(torch.float8_e4m3fn), object() + + def fake_bmm_quantize( + input, + weight, + input_scale, + weight_scale, + output, + packed_scale, + use_tvm_ffi=True, + dspark_max_batch=0, + dspark_rows_per_request=0, + ): + assert input.shape == (groups, batch * block, group_dim) + assert weight is o_a_proj + assert input_scale is not None + assert weight_scale is o_a_scale + assert output.shape == (groups, batch * block, rank) + assert packed_scale.dtype == torch.uint8 + assert use_tvm_ffi is True + assert dspark_max_batch == 0 + assert dspark_rows_per_request == 0 + output.zero_() + packed_scale.zero_() + + monkeypatch.setattr(dspark_attention, "get_sm_version", lambda: 107) + monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_INTERNAL_AVAILABLE", True) + monkeypatch.setattr( + torch.ops.trtllm, + "fp8_batched_quantize_1x128_permute102", + fake_quantize, + ) + monkeypatch.setattr( + torch.ops.trtllm, + "cute_dsl_dspark_fp8_bmm_quantize_rubin_out", + fake_bmm_quantize, + ) + + output = _output_projection( + o, + wo_a=None, + wo_b=FakeOutputLinear(), + n_groups=groups, + o_lora_rank=rank, + o_a_proj=o_a_proj, + o_a_scale=o_a_scale, + o_a_dequant=None, + ) + + assert output.shape == (batch, block, hidden) + assert len(calls) == 1 + assert len(prevalidated_calls) == 1 + assert prevalidated_calls[0][0] is calls[0] + fp8_output, packed_scale = calls[0] + assert fp8_output.shape == (batch * block, groups * rank) + assert fp8_output.dtype == torch.float8_e4m3fn + assert packed_scale.dtype == torch.uint8 + + +def test_fp8_output_projection_uses_prepared_attention_quantized_output(monkeypatch): + import tensorrt_llm._torch.models.modeling_dspark as dspark_attention + + batch, block, hidden = 1, 5, 4 + heads, head_dim, rope_dim, groups, rank, window = 2, 64, 64, 2, 2, 128 + x = torch.zeros(batch, block, hidden, dtype=torch.bfloat16) + main_x = torch.zeros(batch, 1, hidden, dtype=torch.bfloat16) + start_pos = torch.tensor([3], dtype=torch.long) + block_tables = torch.tensor([[0]], dtype=torch.int32) + capacities = torch.tensor([window], dtype=torch.long) + valid_len = torch.tensor([3], dtype=torch.long) + kv_cache = torch.zeros(batch, window, head_dim, dtype=torch.bfloat16) + q_latent = torch.zeros(batch, block, 2, dtype=torch.bfloat16) + block_kv = torch.zeros(batch, block, head_dim, dtype=torch.bfloat16) + main_kv = torch.zeros(batch, 1, head_dim, dtype=torch.bfloat16) + q_projected = torch.zeros(batch, block, heads * head_dim, dtype=torch.bfloat16) + prepared_fp8 = object() + prepared_scale = object() + expected_output = torch.zeros(batch, block, hidden, dtype=torch.bfloat16) + calls = [] + + monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_AVAILABLE", True) + monkeypatch.setattr(dspark_attention, "_get_dspark_arch_str", lambda: "sm_107") + monkeypatch.setattr( + dspark_attention, + "_qkv_projections", + lambda *args, **kwargs: (q_latent, block_kv, main_kv), + ) + monkeypatch.setattr( + dspark_attention, + "_q_projection", + lambda *args, **kwargs: q_projected, + ) + monkeypatch.setattr( + dspark_attention, + "_rmsnorm_rope_batched", + lambda tensor, *args, **kwargs: tensor, + ) + monkeypatch.setattr( + dspark_attention, + "cute_dsl_dspark_rmsnorm_rope_page_write", + lambda *args: None, + ) + monkeypatch.setattr( + dspark_attention, + "cute_dsl_dspark_rmsnorm_rope_draft_block", + lambda *args: torch.zeros(batch, 8, head_dim, dtype=torch.bfloat16), + ) + + def prepared_attention_fp8(*args): + calls.append("prepared_fp8_attention") + assert args[-1] == groups + return prepared_fp8, prepared_scale + + def output_projection(grouped_o_fp8, grouped_o_scale, **kwargs): + calls.append("fp8_output_consumer") + assert grouped_o_fp8 is prepared_fp8 + assert grouped_o_scale is prepared_scale + assert kwargs["batch"] == batch + assert kwargs["block"] == block + return expected_output + + monkeypatch.setattr( + torch.ops.trtllm, + "fused_dsv4_dspark_attention_fp8", + prepared_attention_fp8, + raising=False, + ) + monkeypatch.setattr( + torch.ops.trtllm, + "fused_dsv4_dspark_attention", + lambda *args: pytest.fail("BF16 prepared attention should not run"), + raising=False, + ) + monkeypatch.setattr( + dspark_attention, + "_fused_fp8_output_projection_from_quantized", + output_projection, + ) + monkeypatch.setattr( + dspark_attention, + "_output_projection", + lambda *args, **kwargs: pytest.fail("standalone output quantization should not run"), + ) + + actual = dspark_attention.dspark_attention_forward( + x, + main_x, + start_pos, + kv_cache, + block_tables, + capacities, + valid_len, + wq_a=None, + q_norm_w=torch.ones(2, dtype=torch.bfloat16), + wq_b=object(), + wkv=None, + kv_norm_w=torch.ones(head_dim, dtype=torch.bfloat16), + wo_a=None, + wo_b=object(), + attn_sink=torch.zeros(heads, dtype=torch.float32), + n_heads=heads, + head_dim=head_dim, + rope_head_dim=rope_dim, + n_groups=groups, + o_lora_rank=rank, + window_size=window, + eps=1e-6, + softmax_scale=head_dim**-0.5, + freqs_cis=torch.ones(16, rope_dim // 2, dtype=torch.complex64), + qkv_proj=object(), + o_a_proj=torch.empty(groups, rank, heads * head_dim // groups, dtype=torch.float8_e4m3fn), + o_a_scale=torch.ones(1), + use_fused_output_projection=True, + ) + + assert actual is expected_output + assert calls == ["prepared_fp8_attention", "fp8_output_consumer"] def test_rope_table_is_cached_once_per_device(): @@ -180,6 +495,478 @@ def fake_decoder_layer_init( assert captured["disable_post_moe_fusion"] is True +def test_context_wkv_autotune_bucket_contract(): + from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( + _gen_dspark_context_wkv_m_profiles, + _map_dspark_context_wkv_m_profile, + ) + + assert [_map_dspark_context_wkv_m_profile(m) for m in (1, 32, 128, 129, 511, 1024)] == [ + 128, + 128, + 128, + 256, + 512, + 1024, + ] + assert _gen_dspark_context_wkv_m_profiles(1) == (128,) + assert _gen_dspark_context_wkv_m_profiles(128) == (128,) + assert _gen_dspark_context_wkv_m_profiles(129) == (128, 256) + assert _gen_dspark_context_wkv_m_profiles(1024) == (128, 256, 512, 1024) + + +def test_generation_autotune_bucket_contract_uses_dspark_rows(): + from tensorrt_llm._torch.autotuner import AutoTuner + from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( + _get_dspark_generation_fp8_bmm_quant_tuning_config, + _get_dspark_generation_mxfp8_tuning_config, + ) + + tuner = AutoTuner() + block_config = _get_dspark_generation_mxfp8_tuning_config(5, 32) + profiles = tuner._optimization_profiles(block_config, [torch.empty(5, 8)]) + assert [profile.get_opt_shapes() for profile in profiles] == [ + ((5, 8),), + ((10, 8),), + ((20, 8),), + ((40, 8),), + ((80, 8),), + ((160, 8),), + ] + block_spec = block_config.dynamic_tensor_specs[0] + assert [block_spec.map_to_tuning_buckets(m) for m in (5, 15, 25, 160)] == [ + 5, + 20, + 40, + 160, + ] + + qkv_config = _get_dspark_generation_mxfp8_tuning_config(6, 32) + profiles = tuner._optimization_profiles(qkv_config, [torch.empty(6, 8)]) + assert [profile.get_opt_shapes()[0][0] for profile in profiles] == [ + 6, + 12, + 24, + 48, + 96, + 192, + ] + + oa_config = _get_dspark_generation_fp8_bmm_quant_tuning_config(5, 32) + oa_spec = oa_config.dynamic_tensor_specs[0] + assert oa_spec.gen_tuning_buckets == ( + 5, + 6, + 10, + 12, + 20, + 24, + 40, + 48, + 80, + 96, + 160, + 192, + ) + assert [oa_spec.map_to_tuning_buckets(m) for m in (5, 6, 15, 18, 160, 192)] == [ + 5, + 6, + 20, + 20, + 160, + 192, + ] + + +def test_model_configures_projection_owned_generation_buckets(): + class Projection: + def __call__(self, x): + return x + + qkv, q_b, o_b = Projection(), Projection(), Projection() + model = types.SimpleNamespace( + _use_bf16_linear=False, + block_size=5, + mtp_layers=[ + types.SimpleNamespace(_dspark_attn={"qkv_proj": qkv, "wq_b": q_b, "wo_b": o_b}) + ], + ) + + DSv4DSparkDraftModel.configure_generation_autotuner(model, max_batch=32) + + assert ( + qkv._dspark_generation_tuning_config.dynamic_tensor_specs[0].gen_tuning_buckets[-1] == 192 + ) + assert ( + q_b._dspark_generation_tuning_config.dynamic_tensor_specs[0].gen_tuning_buckets[-1] == 160 + ) + assert o_b._dspark_generation_tuning_config is q_b._dspark_generation_tuning_config + assert o_b._dspark_generation_max_batch == 32 + assert o_b._dspark_generation_rows_per_request == 5 + + +def test_stack_fp8_projection_output_rows_preserves_stage_order(monkeypatch): + rows, output_rows, hidden = 256, 128, 128 + scale_cols = hidden // 32 + weights = [torch.full((rows, hidden), stage, dtype=torch.float8_e4m3fn) for stage in range(2)] + unswizzled = [ + torch.arange(rows * scale_cols).view(rows, scale_cols) + stage * 10000 for stage in range(2) + ] + projections = [ + types.SimpleNamespace(weight=weight, weight_scale=object()) for weight in weights + ] + + def fake_unswizzle(scale, scale_rows, cols, scaling_vector_size): + assert (scale_rows, cols, scaling_vector_size) == (rows, hidden, 32) + stage = next( + stage + for stage, projection in enumerate(projections) + if scale is projection.weight_scale + ) + return unswizzled[stage] + + def fake_swizzle(scale, scale_rows, cols, scaling_vector_size): + assert (scale_rows, cols, scaling_vector_size) == ( + len(projections) * output_rows, + hidden, + 32, + ) + return scale.flatten() + + monkeypatch.setattr(modeling_dspark, "unswizzle_sf", fake_unswizzle) + monkeypatch.setattr(modeling_dspark, "swizzle_sf", fake_swizzle) + + stacked_weight, stacked_scale = modeling_dspark._stack_fp8_projection_output_rows( + projections, output_rows + ) + + assert stacked_weight.shape == (len(projections) * output_rows, hidden) + torch.testing.assert_close( + stacked_weight, + torch.cat([weight[-output_rows:] for weight in weights]), + ) + torch.testing.assert_close( + stacked_scale, + torch.cat([scale[-output_rows:] for scale in unswizzled]).flatten(), + ) + + +@pytest.mark.parametrize( + "shape,output_rows", + [ + ((256, 128), 64), + ((192, 128), 128), + ((256, 96), 128), + ], +) +def test_stack_fp8_projection_output_rows_rejects_unaligned_geometry(shape, output_rows): + projection = types.SimpleNamespace( + weight=torch.empty(shape, dtype=torch.float8_e4m3fn), + weight_scale=torch.empty(0), + ) + + with pytest.raises(ValueError, match="128"): + modeling_dspark._stack_fp8_projection_output_rows([projection], output_rows) + + +def test_bind_attn_projections_reuses_loaded_mla_modules(monkeypatch): + import tensorrt_llm._torch.models.modeling_dspark as dspark_attention + + qkv_proj = Mock() + qkv_proj.in_features = 4 + q_b_proj = Mock() + q_b_proj.in_features = 6 + o_b_proj = Mock() + o_b_proj.has_fp8_block_scales = True + o_b_proj.use_cute_dsl_blockscaling_mm = True + o_b_proj.disable_deep_gemm = False + o_b_proj.in_features = 16 * 128 + o_b_proj.weight = torch.empty(8, 16 * 128, dtype=torch.float8_e4m3fn) + o_b_proj.weight_scale = torch.empty(1, dtype=torch.uint8) + o_a_proj = torch.empty(16, 128, 8, dtype=torch.float8_e4m3fn) + o_a_scale = torch.empty(16, 1, 1) + o_a_dequant = object() + q_norm = torch.ones(6) + kv_norm = torch.ones(8) + sink = torch.zeros(128) + attn = types.SimpleNamespace( + num_heads_tp=128, + n_local_groups=16, + q_a_layernorm=types.SimpleNamespace(weight=q_norm), + q_b_proj=q_b_proj, + kv_a_layernorm=types.SimpleNamespace(weight=kv_norm), + o_b_proj=o_b_proj, + mqa=types.SimpleNamespace(attn_sink=sink), + kv_a_proj_with_mqa=qkv_proj, + o_a_proj=o_a_proj, + o_a_proj_scale=o_a_scale, + o_a_proj_dequant=o_a_dequant, + ) + stage = types.SimpleNamespace(self_attn=attn) + model = types.SimpleNamespace( + _attn_params={ + "n_heads": 128, + "n_groups": 16, + "head_dim": 8, + "eps": 1e-6, + }, + mtp_layers=[stage], + ) + monkeypatch.setattr(dspark_attention, "get_sm_version", lambda: 107) + monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_INTERNAL_AVAILABLE", True) + + DSv4DSparkDraftModel._bind_attn_projections(model) + + bound = stage._dspark_attn + assert bound["wq_a"] is None + assert bound["wkv"] is None + assert bound["wo_a"] is None + assert bound["qkv_proj"] is qkv_proj + assert bound["wq_b"] is q_b_proj + assert bound["wo_b"] is o_b_proj + assert model._context_wkv_fp8 is None + assert bound["o_a_proj"] is o_a_proj + assert bound["o_a_scale"] is o_a_scale + assert bound["o_a_dequant"] is o_a_dequant + assert bound["use_fused_output_projection"] is True + assert bound["use_fused_qkv_input_projection"] is False + assert bound["use_fused_q_input_projection"] is False + assert bound["q_norm_w"] is q_norm + assert bound["kv_norm_w"] is kv_norm + assert bound["attn_sink"] is sink + + attn.o_a_proj = torch.empty(16, 128, 8, dtype=torch.bfloat16) + DSv4DSparkDraftModel._bind_attn_projections(model) + assert stage._dspark_attn["use_fused_output_projection"] is False + + context_wkv = (object(), object()) + stack_projections = Mock(return_value=context_wkv) + monkeypatch.setattr( + modeling_dspark, + "_can_use_fused_fp8_input_projection", + lambda projection, _in_features: projection is qkv_proj, + ) + monkeypatch.setattr( + modeling_dspark, + "_stack_fp8_projection_output_rows", + stack_projections, + ) + from tensorrt_llm._torch.custom_ops import dspark_rmsnorm_rope_custom_op as rmsnorm_rope_op + + monkeypatch.setattr( + rmsnorm_rope_op, + "precompile_dspark_fp8_input_fusions", + Mock(), + ) + + DSv4DSparkDraftModel._bind_attn_projections(model) + + stack_projections.assert_called_once_with([qkv_proj], 8) + assert model._context_wkv_fp8 is context_wkv + + +@pytest.mark.parametrize( + "tp_size,enable_attention_dp,use_bf16", + [ + (1, False, False), + (2, True, False), + (2, False, True), + ], +) +def test_post_load_weights_selects_dspark_linear_backend( + monkeypatch, tp_size, enable_attention_dp, use_bf16 +): + model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) + torch.nn.Module.__init__(model) + monkeypatch.delenv("TRTLLM_DSPARK_LINEAR_BACKEND", raising=False) + mapping = types.SimpleNamespace( + tp_size=tp_size, + enable_attention_dp=enable_attention_dp, + ) + model._use_bf16_linear = modeling_dspark._use_dspark_bf16_linear(mapping) + model.mtp_layers = [] + fp8_bind = Mock() + bf16_bind = Mock() + monkeypatch.setattr(DSv4DSparkDraftModel, "_bind_attn_projections", fp8_bind) + monkeypatch.setattr(DSv4DSparkDraftModel, "_bind_bf16_attn_projections", bf16_bind) + weights = {} + + model.post_load_weights() + model.post_load_weights(weights) + model.post_load_weights() + + if use_bf16: + fp8_bind.assert_not_called() + bf16_bind.assert_called_once_with(weights) + else: + fp8_bind.assert_called_once_with() + bf16_bind.assert_not_called() + + +def test_post_load_weights_selects_rubin_cute_backend_before_transform(monkeypatch): + model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) + torch.nn.Module.__init__(model) + model._use_bf16_linear = False + + projection = Linear.__new__(Linear) + torch.nn.Module.__init__(projection) + projection.use_cute_dsl_blockscaling_mm = False + projection.transform_weights = Mock() + stage = types.SimpleNamespace(self_attn=types.SimpleNamespace(modules=lambda: [projection])) + model.mtp_layers = [stage] + model.modules = lambda: [projection] + + bind = Mock() + monkeypatch.setattr(modeling_dspark, "get_sm_version", lambda: 107) + monkeypatch.setattr(DSv4DSparkDraftModel, "_bind_attn_projections", bind) + + model.post_load_weights({}) + + assert projection.use_cute_dsl_blockscaling_mm is True + projection.transform_weights.assert_called_once_with() + bind.assert_called_once_with() + + +def test_bind_bf16_attn_projections_disables_every_fp8_projection(): + model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) + torch.nn.Module.__init__(model) + model._context_wkv_fp8 = object() + stage = types.SimpleNamespace(input_layernorm=types.SimpleNamespace(weight=torch.ones(1))) + model.mtp_layers = [stage] + prefix = "mtp.0.attn." + shapes = { + "wq_a": (2, 4), + "wq_b": (6, 2), + "wkv": (3, 4), + "wo_a": (8, 3), + "wo_b": (4, 8), + } + weights = { + f"{prefix}q_norm.weight": torch.ones(2), + f"{prefix}kv_norm.weight": torch.ones(3), + f"{prefix}attn_sink": torch.zeros(6), + } + for index, (name, shape) in enumerate(shapes.items(), start=1): + weights[f"{prefix}{name}.weight"] = torch.full(shape, index, dtype=torch.float8_e4m3fn) + weights[f"{prefix}{name}.scale"] = torch.full((1, 1), index + 0.5) + + model._bind_bf16_attn_projections(weights) + + bound = stage._dspark_attn + for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): + expected = model._block_dequant( + weights[f"{prefix}{name}.weight"], weights[f"{prefix}{name}.scale"] + ) + assert bound[name].dtype == torch.bfloat16 + torch.testing.assert_close(bound[name], expected) + assert bound["qkv_proj"] is None + assert bound["o_a_proj"] is None + assert bound["o_a_scale"] is None + assert bound["o_a_dequant"] is None + assert bound["use_fused_output_projection"] is False + assert bound["use_fused_qkv_input_projection"] is False + assert bound["use_fused_q_input_projection"] is False + assert model._context_wkv_fp8 is None + + +def test_context_wkv_fp8_quantizes_once_per_page_write(monkeypatch): + num_stages, hidden, head_dim, window = 3, 4, 2, 8 + model = DSv4DSparkDraftModel.__new__(DSv4DSparkDraftModel) + torch.nn.Module.__init__(model) + model._attn_params = { + "head_dim": head_dim, + "rope_head_dim": 0, + "window_size": window, + "eps": 1e-6, + } + weights = [ + torch.arange(head_dim * hidden, dtype=torch.float32).view(head_dim, hidden) + stage_id + for stage_id in range(num_stages) + ] + qkv_projections = [ + Mock(side_effect=AssertionError("full q_a+wkv projection is forbidden")) + for _ in range(num_stages) + ] + stages = [ + types.SimpleNamespace( + _dspark_attn={ + "qkv_proj": qkv_projections[stage_id], + "kv_norm_w": torch.ones(head_dim), + "wkv": weights[stage_id], + }, + ) + for stage_id in range(num_stages) + ] + stages[0].main_proj = Mock(side_effect=lambda x: x) + stages[0].main_norm = Mock(side_effect=lambda x: x) + model.mtp_layers = stages + model._dspark_freqs_table = lambda device: torch.empty( + 32, 1, dtype=torch.complex64, device=device + ) + model._context_wkv_fp8 = (torch.cat(weights), object()) + + quant_calls = [] + gemm_calls = [] + input_scale = object() + + def fake_quantize(rows): + quant_calls.append(rows) + return rows, input_scale + + def fake_gemm(rows, weight, scale, weight_scale): + assert scale is input_scale + gemm_calls.append((rows, weight, weight_scale)) + return F.linear(rows, weight) + + monkeypatch.setattr( + torch.ops.trtllm, + "fp8_quantize_1x128_packed_ue8m0", + fake_quantize, + ) + monkeypatch.setattr( + torch.ops.trtllm, + "cute_dsl_dspark_context_wkv_gemm_rubin", + fake_gemm, + ) + monkeypatch.setattr(modeling_dspark, "_rmsnorm", lambda value, *args: value) + monkeypatch.setattr(modeling_dspark, "_rope_last_dims_batched", lambda value, *args: value) + + # Capture the paged-store boundary; projection math runs on CPU. + writes = [] + monkeypatch.setattr( + modeling_dspark, + "write_dspark_context", + lambda values, *args: writes.append((values.clone(), args)), + ) + main_hidden = torch.arange(2 * 2 * hidden, dtype=torch.float32).view(2, 2, hidden) + positions = torch.tensor([[1, 32], [3, 4]]) + mask = torch.tensor([[True, False], [True, True]]) + pools = tuple(torch.zeros(4, 1, 1, window, head_dim) for _ in stages) + tables = torch.tensor([[2], [1]], dtype=torch.int32) + capacities = torch.tensor([window, window]) + model.write_context_pages(main_hidden, positions, mask, pools, tables, capacities) + + assert len(quant_calls) == len(gemm_calls) == 1 + assert len(writes) == num_stages + for stage_id, weight in enumerate(weights): + values, args = writes[stage_id] + torch.testing.assert_close(values, F.linear(main_hidden, weight)) + assert args[0] is positions and args[1] is mask + assert args[2].data_ptr() == pools[stage_id].data_ptr() + assert args[3] is tables and args[4] is capacities + assert all(projection.call_count == 0 for projection in qkv_projections) + + expected = [values for values, _ in writes] + writes.clear() + model._context_wkv_fp8 = None + for stage in stages: + stage._dspark_attn["qkv_proj"] = None + model.write_context_pages(main_hidden, positions, mask, pools, tables, capacities) + for (actual, _), wanted in zip(writes, expected): + torch.testing.assert_close(actual, wanted) + assert len(quant_calls) == len(gemm_calls) == 1 + + @pytest.mark.parametrize("enable_fused_hc", [True, False]) def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): """The draft stage must use the inherited fused-HC rollback setting.""" @@ -253,9 +1040,13 @@ def call(*args, **kwargs): stage, h, torch.randn(num_requests, hidden_size), - 1, + torch.ones(num_requests, dtype=torch.long), torch.empty(0), torch.zeros(num_requests, block_size, dtype=torch.long), + torch.empty(1, 128, 1), + torch.zeros(num_requests, 1, dtype=torch.int32), + torch.full((num_requests,), 128, dtype=torch.long), + torch.ones(num_requests, dtype=torch.long), ) assert actual is final_h @@ -304,123 +1095,6 @@ def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) -def _loop_reference(q, kv, attn_sink, topk_idxs, scale): - """Obvious, slow per-(b,m,h) reference of the exact kernel math.""" - b, m, h, d = q.shape - out = torch.zeros(b, m, h, d, dtype=torch.float32) - qf, kvf, sink = q.float(), kv.float(), attn_sink.float() - for bi in range(b): - for mi in range(m): - idxs = topk_idxs[bi, mi].tolist() - for hi in range(h): - scores, vecs = [], [] - for j in idxs: - if j < 0: - continue - k = kvf[bi, j] - scores.append(torch.dot(qf[bi, mi, hi], k) * scale) - vecs.append(k) - if not scores: - continue - s = torch.stack(scores) - smax = s.max() - p = torch.exp(s - smax) - denom = p.sum() + torch.exp(sink[hi] - smax) - num = (p.unsqueeze(-1) * torch.stack(vecs)).sum(0) - out[bi, mi, hi] = num / denom - return out - - -@pytest.mark.parametrize("seed", [0, 1, 2]) -def test_sparse_attn_matches_loop_reference(seed): - torch.manual_seed(seed) - b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, n, d) - attn_sink = torch.randn(h) - idx = torch.stack( - [torch.stack([torch.randperm(n)[:topk] for _ in range(m)]) for _ in range(b)] - ).int() - scale = d**-0.5 - got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() - ref = _loop_reference(q, kv, attn_sink, idx, scale) - torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) - - -def test_sparse_attn_no_sink_matches_sdpa(): - """With a -inf sink (no sink mass) and all-valid contiguous indices, the - primitive must equal standard scaled-dot-product attention over the gathered - KV — an independent implementation.""" - torch.manual_seed(0) - b, m, h, d, topk = 2, 4, 5, 16, 9 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, topk, d) # n == topk, attend to all - attn_sink = torch.full((h,), float("-inf")) - idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() - scale = d**-0.5 - got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() - - # SDPA: q [b,h,m,d], k/v [b,h,topk,d] (broadcast the shared KV over heads). - qh = q.permute(0, 2, 1, 3) - kvh = kv.unsqueeze(1).expand(b, h, topk, d) - ref = F.scaled_dot_product_attention(qh, kvh, kvh, scale=scale).permute(0, 2, 1, 3) - torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) - - -def test_sparse_attn_sink_reduces_mass(): - """A finite sink must strictly shrink the attention output magnitude vs an - infinitely-negative (disabled) sink, because it adds denominator mass only.""" - torch.manual_seed(0) - b, m, h, d, topk = 1, 2, 2, 16, 6 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, topk, d) - idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() - scale = d**-0.5 - no_sink = dspark_sparse_attn(q, kv, torch.full((h,), float("-inf")), idx, scale) - with_sink = dspark_sparse_attn(q, kv, torch.zeros(h), idx, scale) - assert with_sink.abs().sum() < no_sink.abs().sum() - - -def test_sparse_attn_masked_indices_excluded(): - """An index of -1 must be excluded exactly (equiv. to dropping that column).""" - torch.manual_seed(0) - b, m, h, d = 1, 1, 2, 16 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, 5, d) - sink = torch.full((h,), float("-inf")) - scale = d**-0.5 - full = torch.tensor([[[0, 1, 2, 3]]]).int() - masked = torch.tensor([[[0, 1, 2, -1]]]).int() - drop3 = torch.tensor([[[0, 1, 2]]]).int() - got_masked = dspark_sparse_attn(q, kv, sink, masked, scale) - got_drop = dspark_sparse_attn(q, kv, sink, drop3, scale) - torch.testing.assert_close(got_masked, got_drop, rtol=1e-5, atol=1e-5) - # And masking genuinely changes the result vs attending to position 3. - got_full = dspark_sparse_attn(q, kv, sink, full, scale) - assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) - - -@pytest.mark.parametrize( - "start_pos,window,block", [(1, 128, 5), (3, 128, 5), (10, 4, 5), (200, 128, 6)] -) -def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): - bsz = 3 - got = get_dspark_topk_idxs(window, bsz, block, start_pos) - # Reference formula (DeepSpec get_dspark_topk_idxs). - ctx = torch.arange(min(window, start_pos + 1)) - blk = window + torch.arange(block) - ref_row = torch.cat([ctx, blk]).int() - assert got.shape == (bsz, block, ref_row.numel()) - for bi in range(bsz): - for mi in range(block): - torch.testing.assert_close(got[bi, mi], ref_row) - - -def test_get_dspark_topk_idxs_requires_generation(): - with pytest.raises(AssertionError): - get_dspark_topk_idxs(128, 1, 5, 0) - - @pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) @@ -448,83 +1122,3 @@ def test_apply_rotary_inverse_roundtrip(ndim): fc = precompute_dspark_freqs_cis(rd, s) roundtrip = apply_dspark_rotary(apply_dspark_rotary(x, fc), fc, inverse=True) torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) - - -def _make_attn_inputs(seed=0): - """Small synthetic DSpark attention inputs/weights (CPU bf16).""" - torch.manual_seed(seed) - dim, n_heads, head_dim, rd = 12, 4, 8, 4 - q_lora, o_lora, n_groups = 6, 5, 2 - window, block, start_pos = 8, 3, 5 - b = 2 - g = dict( - dim=dim, - n_heads=n_heads, - head_dim=head_dim, - rope_head_dim=rd, - q_lora=q_lora, - o_lora=o_lora, - n_groups=n_groups, - window=window, - block=block, - start_pos=start_pos, - b=b, - eps=1e-6, - softmax_scale=head_dim**-0.5, - ) - bf = torch.bfloat16 - g["x"] = torch.randn(b, block, dim, dtype=bf) - g["main_x"] = torch.randn(b, 1, dim, dtype=bf) - g["kv_cache0"] = torch.randn(b, window, head_dim, dtype=bf) - g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 - g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 - g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 - g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 - g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 - g["q_norm"] = torch.ones(q_lora) - g["kv_norm"] = torch.ones(head_dim) - g["attn_sink"] = torch.randn(n_heads) - g["freqs"] = precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) - return g - - -def _run(g): - return dspark_attention_forward( - g["x"], - g["main_x"], - g["start_pos"], - g["kv_cache0"], - wq_a=g["wq_a"], - q_norm_w=g["q_norm"], - wq_b=g["wq_b"], - wkv=g["wkv"], - kv_norm_w=g["kv_norm"], - wo_a=g["wo_a"], - wo_b=g["wo_b"], - attn_sink=g["attn_sink"], - n_heads=g["n_heads"], - head_dim=g["head_dim"], - rope_head_dim=g["rope_head_dim"], - n_groups=g["n_groups"], - o_lora_rank=g["o_lora"], - window_size=g["window"], - eps=g["eps"], - softmax_scale=g["softmax_scale"], - freqs_cis=g["freqs"], - ) - - -def test_attention_forward_shape_and_determinism(): - g = _make_attn_inputs() - o = _run(g) - assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) - assert torch.isfinite(o.float()).all() - torch.testing.assert_close(o, _run(g)) # deterministic - - -def test_attention_forward_does_not_mutate_kv_cache(): - """The rolling window write must be functional (cache cloned, not mutated).""" - g = _make_attn_inputs() - before = g["kv_cache0"].clone() - _run(g) - torch.testing.assert_close(g["kv_cache0"], before) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index 95fa963affed..2e47dc49014f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -12,17 +12,7 @@ # 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-graph-safety unit tests for the DSpark batched draft-attention path. - -The load-bearing invariant for the DSpark batched draft path (the default, used -whenever ``cuda_graph_config`` is set, since the one-engine drafter is captured in -the target's graph) is that the batched, sync-free primitives are **numerically -identical, per request**, to the validated scalar path — only the host-int -``start_pos`` and the per-request window indexing are tensorized. These tests assert -that equivalence on CPU (so they run in pre-merge CI without a GPU), plus a -GPU-gated capture+replay smoke test that proves the batched attention is actually -graph-capturable. -""" +"""Per-request rotary equivalence for batched DSpark graph inputs.""" import pytest import torch @@ -30,240 +20,19 @@ from tensorrt_llm._torch.models.modeling_dspark import ( apply_dspark_rotary, apply_dspark_rotary_batched, - dspark_attention_forward, - dspark_attention_forward_batched, - get_dspark_topk_idxs, - get_dspark_topk_idxs_batched, precompute_dspark_freqs_cis, ) -def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): - """Per-request DSpark attention inputs/weights (shared weights, distinct pos). - - Mirrors ``test_dspark_attention._make_attn_inputs`` but builds ``G`` requests - each with its own ``start_pos`` (small => partial context, large => full - rolling window) and its own pre-seeded window, so the batched-vs-scalar - comparison exercises the per-request RoPE gather + windowed context read. - """ - torch.manual_seed(seed) - dim, n_heads, head_dim, rd = 12, 4, 8, 4 - q_lora, o_lora, n_groups = 6, 5, 2 - window, block = 8, 3 - G = len(start_positions) - bf = torch.bfloat16 - # A single fixed RoPE table covering every request's positions (both paths - # index/gather the same values, so freqs are identical across paths). - maxlen = max(start_positions) + 1 + block + 4 - g = dict( - dim=dim, - n_heads=n_heads, - head_dim=head_dim, - rope_head_dim=rd, - q_lora=q_lora, - o_lora=o_lora, - n_groups=n_groups, - window=window, - block=block, - G=G, - start_positions=list(start_positions), - eps=1e-6, - softmax_scale=head_dim**-0.5, - x=torch.randn(G, block, dim, dtype=bf), - main_x=torch.randn(G, 1, dim, dtype=bf), - # Distinct, non-zero seeded window per request to exercise context reads. - kv_cache=torch.randn(G, window, head_dim, dtype=bf), - wq_a=torch.randn(q_lora, dim, dtype=bf) * 0.1, - wq_b=torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1, - wkv=torch.randn(head_dim, dim, dtype=bf) * 0.1, - wo_a=torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1, - wo_b=torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1, - q_norm=torch.ones(q_lora), - kv_norm=torch.ones(head_dim), - attn_sink=torch.randn(n_heads), - freqs=precompute_dspark_freqs_cis(rd, maxlen), - ) - return g - - -def _attn_kwargs(g): - return dict( - wq_a=g["wq_a"], - q_norm_w=g["q_norm"], - wq_b=g["wq_b"], - wkv=g["wkv"], - kv_norm_w=g["kv_norm"], - wo_a=g["wo_a"], - wo_b=g["wo_b"], - attn_sink=g["attn_sink"], - n_heads=g["n_heads"], - head_dim=g["head_dim"], - rope_head_dim=g["rope_head_dim"], - n_groups=g["n_groups"], - o_lora_rank=g["o_lora"], - window_size=g["window"], - eps=g["eps"], - softmax_scale=g["softmax_scale"], - freqs_cis=g["freqs"], - ) - - -def _scalar_reference(g): - """Run the validated scalar attention once per request and stack the outputs.""" - outs = [] - for i, sp in enumerate(g["start_positions"]): - out_i = dspark_attention_forward( - g["x"][i : i + 1], - g["main_x"][i : i + 1], - int(sp), - g["kv_cache"][i : i + 1].clone(), - **_attn_kwargs(g), - ) - outs.append(out_i) - return torch.cat(outs, dim=0) - - -def _batched(g, persist=False): - G = g["G"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long) - slots = torch.arange(G, dtype=torch.long) - return dspark_attention_forward_batched( - g["x"], - g["main_x"], - start_pos, - g["kv_cache"].clone(), - slots, - persist=persist, - **_attn_kwargs(g), - ) - - -@pytest.mark.parametrize("seed", [0, 1, 2]) -def test_batched_attention_matches_scalar_per_request(seed): - """Batched attention == per-request scalar attention at distinct start_pos. - - This is the invariant that lets the batched path replace the per-request loop - under CUDA graphs without changing draft quality / greedy parity. - """ - g = _make_batched_inputs(seed=seed) - ref = _scalar_reference(g) - got = _batched(g) - assert tuple(got.shape) == (g["G"], g["block"], g["dim"]) - torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) - - -def test_batched_attention_persist_writes_through_window(): - """persist=True writes main_kv into the shared window at start_pos%window.""" - g = _make_batched_inputs(seed=3) - G, win = g["G"], g["window"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long) - slots = torch.arange(G, dtype=torch.long) - cache = g["kv_cache"].clone() - before = cache.clone() - dspark_attention_forward_batched( - g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) - ) - # Exactly the start_pos%win row of each request changed. - for i, sp in enumerate(g["start_positions"]): - changed = (cache[i] != before[i]).any(dim=-1) - expected = torch.zeros(win, dtype=torch.bool) - expected[sp % win] = True - assert torch.equal(changed, expected), f"req {i}: wrong window row written" - - -def test_batched_attention_no_persist_keeps_window(): - """persist=False must not mutate the caller's window (functional).""" - g = _make_batched_inputs(seed=4) - before = g["kv_cache"].clone() - _batched(g, persist=False) - torch.testing.assert_close(g["kv_cache"], before) - - -@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) -def test_batched_topk_matches_scalar(start_positions): - """Fixed-size masked batched topk == scalar topk per request (valid slots).""" - window, block = 8, 3 - start_pos = torch.tensor(start_positions, dtype=torch.long) - batched = get_dspark_topk_idxs_batched(window, block, start_pos) - assert tuple(batched.shape) == (len(start_positions), block, window + block) - for i, sp in enumerate(start_positions): - scalar = get_dspark_topk_idxs(window, 1, block, int(sp))[0] # [block, topk_i] - # The batched row drops the -1-masked context slots to recover the scalar - # (variable-width) index list; both must then be identical. - for m in range(block): - valid = batched[i, m][batched[i, m] >= 0] - torch.testing.assert_close(valid, scalar[m].to(valid.dtype)) - - -def test_batched_topk_respects_partial_window_valid_len(): - """Only actually written circular slots are visible for partial seeds.""" - window, block = 8, 3 - start_pos = torch.tensor([3, 10, 20], dtype=torch.long) - valid_len = torch.tensor([0, 3, 8], dtype=torch.long) - batched = get_dspark_topk_idxs_batched(window, block, start_pos, valid_len) - assert tuple(batched.shape) == (3, block, window + block) - - for i, (pos, count) in enumerate(zip(start_pos.tolist(), valid_len.tolist())): - expected_context = sorted((pos - age) % window for age in range(count)) - row = batched[i, 0] - context = row[(row >= 0) & (row < window)].tolist() - assert context == expected_context - block_indices = row[row >= window].tolist() - assert block_indices == list(range(window, window + block)) - - @pytest.mark.parametrize("ndim", [3, 4]) def test_batched_rotary_matches_scalar_per_row(ndim): - """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row.""" torch.manual_seed(0) G, s, h, rd = 3, 4, 2, 8 x = torch.randn(G, s, h, rd) if ndim == 4 else torch.randn(G, s, rd) table = precompute_dspark_freqs_cis(rd, 64) - # Per-row absolute start positions -> per-row freq windows. starts = [1, 9, 30] - per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) # [G, s, rd/2] + per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) got = apply_dspark_rotary_batched(x, per_row) for i, sp in enumerate(starts): ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) torch.testing.assert_close(got[i : i + 1], ref_i) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") -def test_batched_attention_cuda_graph_capture_replay(): - """The batched attention captures + replays and matches eager output. - - Proves the path is free of capture-illegal ops (host syncs, dynamic shapes). - """ - g = _make_batched_inputs(seed=0) - dev = "cuda" - G = g["G"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) - slots = torch.arange(G, dtype=torch.long, device=dev) - # Static input tensors the graph reads/writes. - x = g["x"].to(dev) - main_x = g["main_x"].to(dev) - cache = g["kv_cache"].to(dev) - kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} - - def run(persist): - return dspark_attention_forward_batched( - x, main_x, start_pos, cache, slots, persist=persist, **kw - ) - - eager = run(persist=False) - - # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so - # the comparison isn't perturbed by the window write-through. - s = torch.cuda.Stream() - s.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s): - for _ in range(3): - run(persist=False) - torch.cuda.current_stream().wait_stream(s) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = run(persist=False) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 40786627b200..1d846c73ed6a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -16,7 +16,7 @@ Covers the framework-side logic that does NOT need the full draft model: ``DSparkSpecMetadata`` hidden-state capture (incl. the mHC hc-mean reduction) -and ``DSv4DSparkWorker`` slot / rolling-KV-window management. The end-to-end block +and ``DSv4DSparkWorker`` progress-slot / managed-page management. The end-to-end block draft and acceptance path is covered by the DSpark test in ``integration/defs/accuracy/test_llm_api_pytorch.py``. """ @@ -125,16 +125,69 @@ def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): ) +def _bind_manager(worker, request_ids=()): + from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType + + manager = types.SimpleNamespace( + draft_layer_ids=tuple(range(worker._num_stages)), + draft_window_size=worker._win, + kv_factor_per_layer=[1] * worker._num_stages, + head_dim_per_layer=[worker._head_dim] * worker._num_stages, + layer_offsets=list(range(worker._num_stages)), + draft_max_blocks_per_seq=8192, + tokens_per_block=128, + _draft_dummy_request_ids=set(), + enable_block_reuse=False, + kv_cache_map={}, + ) + buffers = tuple( + torch.zeros(65, 1, 1, 128, worker._head_dim, device="cuda", dtype=torch.bfloat16) + for _ in range(worker._num_stages) + ) + manager.get_draft_buffers = lambda stage: buffers[stage] + manager.get_draft_block_indices = lambda ids: [manager.kv_cache_map[rid].indices for rid in ids] + resource_manager = types.SimpleNamespace( + get_resource_manager=lambda kind: manager + if kind == ResourceManagerType.KV_CACHE_MANAGER + else None + ) + worker._bind_managed_history(resource_manager) + for i, rid in enumerate(request_ids): + manager.kv_cache_map[rid] = types.SimpleNamespace( + history_length=1, capacity=4096, indices=[i + 1] * 32 + ) + return manager + + +def _prepare(worker, meta, num_contexts=0): + if worker._draft_kv_manager is None: + _bind_manager(worker) + manager = worker._draft_kv_manager + for rid in meta.request_ids: + if worker._is_managed_request(rid) and rid not in manager.kv_cache_map: + manager.kv_cache_map[rid] = types.SimpleNamespace( + history_length=1, capacity=4096, indices=[len(manager.kv_cache_map) + 1] * 32 + ) + # A previously assigned slot represents already initialized resident history. + for rid in worker._req_to_slot: + if rid in manager.kv_cache_map: + worker._managed_residency.setdefault(rid, manager.kv_cache_map[rid]) + meta.prepare() + worker._prepare_managed_history(meta.request_ids, num_contexts) + + def test_worker_lazy_init_window_buffers(): worker = _make_worker() dm = _fake_draft_model(num_stages=3, window_size=128, head_dim=64) meta = _make_metadata(max_num_requests=8) worker._lazy_init(dm, meta) + manager = _bind_manager(worker) # max_batch (8) request slots + 1 scratch row for padded / unknown IDs. - assert worker._kv_windows.shape == (9, 3, 128, 64) + assert worker._draft_block_tables.shape == (8, manager.draft_max_blocks_per_seq) + assert worker._draft_kv_buffers[0] is manager.get_draft_buffers(0) assert worker._ctx_len.shape == (9,) assert worker._valid_len.shape == (9,) - assert worker._position_initialized.shape == (9,) + assert worker._draft_capacities.shape == (8,) assert worker._scratch_slot == 8 # Dummy-id floor separates real request ids from CUDA-graph padding ids. assert worker._graph_dummy_id_floor == (1 << 64) - 1 - worker.max_draft_len @@ -145,9 +198,9 @@ def test_worker_lazy_init_window_buffers(): assert worker._batch_to_slot.device.type == "cuda" assert worker._batch_to_slot.tolist() == [worker._scratch_slot] * 8 # idempotent - buf_id = id(worker._kv_windows) + buf_id = worker._draft_block_tables.data_ptr() worker._lazy_init(dm, meta) - assert id(worker._kv_windows) == buf_id + assert worker._draft_block_tables.data_ptr() == buf_id def test_worker_rejects_mismatched_block_size(): @@ -163,41 +216,39 @@ def test_worker_slot_assignment_and_reset(): worker = _make_worker() worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=4)) - s0 = worker._assign_slot(100, reset=False) - s1 = worker._assign_slot(101, reset=False) + s0 = worker._assign_slot(100) + s1 = worker._assign_slot(101) assert s0 != s1 # same request id -> same slot (no reset) - assert worker._assign_slot(100, reset=False) == s0 + assert worker._assign_slot(100) == s0 # mark a position, then reset -> slot freed + window/pos cleared worker._ctx_len[s0] = 42 worker._valid_len[s0] = 8 - worker._position_initialized[s0] = True - worker._kv_windows[s0].fill_(1.0) - s0b = worker._assign_slot(100, reset=True) + worker._release_inactive_slots([101]) + s0b = worker._assign_slot(100) assert int(worker._ctx_len[s0b]) == 0 assert int(worker._valid_len[s0b]) == 0 - assert not bool(worker._position_initialized[s0b]) - assert float(worker._kv_windows[s0b].abs().sum()) == 0.0 + assert worker._req_to_slot[101] == s1 def test_worker_slot_exhaustion_preserves_live_request(): worker = _make_worker() worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=1)) - slot = worker._assign_slot(100, reset=False) + slot = worker._assign_slot(100) worker._ctx_len[slot] = 42 - worker._kv_windows[slot].fill_(1.0) + worker._valid_len[slot] = 1 - with pytest.raises(RuntimeError, match="no free rolling-window slots"): - worker._assign_slot(101, reset=False) + with pytest.raises(RuntimeError, match="no free progress slots"): + worker._assign_slot(101) assert worker._req_to_slot == {100: slot} assert int(worker._ctx_len[slot]) == 42 - assert torch.all(worker._kv_windows[slot] == 1.0) + assert int(worker._valid_len[slot]) == 1 -def test_seed_context_windows_preserves_state_across_prefill_chunks(): +def test_seed_context_pages_preserves_state_across_prefill_chunks(): class DraftModel: num_stages = 1 block_size = 5 @@ -212,9 +263,21 @@ class DraftModel: def __init__(self): self.written_positions = [] - def write_context_windows(self, hidden, positions, windows): + def configure_generation_autotuner(self, _max_batch): + pass + + def write_context_pages(self, hidden, positions, mask, pools, tables, capacities): + from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context + self.written_positions.append(positions.clone()) - windows.add_(1) + write_dspark_context( + torch.ones(*positions.shape, 4, device="cuda", dtype=torch.bfloat16), + positions, + mask, + pools[0][:, 0, 0], + tables, + capacities, + ) worker = _make_worker() draft_model = DraftModel() @@ -226,31 +289,35 @@ def write_context_windows(self, hidden, positions, windows): ), ) worker._lazy_init(draft_model, metadata) + manager = _bind_manager(worker, [100]) + manager.kv_cache_map[100].history_length = 0 + worker._prepare_managed_history([100], 1) first_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[3]) - worker._seed_context_windows( + worker._seed_context_pages( draft_model, metadata, first_chunk, torch.tensor([[0, 1, 2]], device="cuda"), 3 ) slot = worker._req_to_slot[100] assert int(worker._ctx_len[slot]) == 3 assert int(worker._valid_len[slot]) == 3 - assert bool(worker._position_initialized[slot]) + assert worker._managed_residency[100] is manager.kv_cache_map[100] metadata.get_hidden_states = lambda _num_tokens: torch.zeros( 2, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 ) second_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[2]) - worker._seed_context_windows( + worker._seed_context_pages( draft_model, metadata, second_chunk, torch.tensor([[3, 4]], device="cuda"), 2 ) assert int(worker._ctx_len[slot]) == 5 assert int(worker._valid_len[slot]) == 5 assert [positions.tolist() for positions in draft_model.written_positions] == [ - [1, 2, 3], - [4, 5], + [[1, 2, 3]], + [[4, 5]], ] - assert torch.all(worker._kv_windows[slot] == 2.0) + assert torch.all(worker._draft_kv_buffers[0][1, 0, 0, :5] == 1.0) + assert torch.count_nonzero(worker._draft_kv_buffers[0][1, 0, 0, 5:]) == 0 @pytest.mark.parametrize( @@ -324,11 +391,11 @@ def test_prepare_builds_batch_to_slot_on_batched_path(): meta._dspark_worker = worker # Assign slots for two requests (as the prefill path would). - sa = worker._assign_slot(100, reset=True) - sb = worker._assign_slot(101, reset=True) + sa = worker._assign_slot(100) + sb = worker._assign_slot(101) meta.request_ids = [101, 100] - meta.prepare() + _prepare(worker, meta) # Mirror reflects request-order -> slot. assert worker._batch_to_slot[:2].tolist() == [sb, sa] @@ -340,20 +407,19 @@ def test_prepare_frees_stale_slots_on_batched_path(): worker._lazy_init(_fake_draft_model(), meta) meta._dspark_worker = worker - sa = worker._assign_slot(100, reset=True) - worker._assign_slot(101, reset=True) + sa = worker._assign_slot(100) + worker._assign_slot(101) worker._ctx_len[sa] = 17 worker._valid_len[sa] = 8 - worker._position_initialized[sa] = True # Only request 101 survives; 100's slot must be freed + cleared. meta.request_ids = [101] - meta.prepare() + _prepare(worker, meta) assert 100 not in worker._req_to_slot assert sa in worker._free_slots assert int(worker._ctx_len[sa]) == 0 assert int(worker._valid_len[sa]) == 0 - assert not bool(worker._position_initialized[sa]) + assert 100 not in worker._managed_residency def test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(): @@ -369,14 +435,16 @@ def test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(): meta._dspark_worker = worker # A live request takes the first free slot (0) and populates its window. - s_real = worker._assign_slot(100, reset=True) + s_real = worker._assign_slot(100) assert s_real == 0 worker._ctx_len[s_real] = 17 - worker._kv_windows[s_real].fill_(1.0) + worker._valid_len[s_real] = 5 + manager = _bind_manager(worker) + manager._draft_dummy_request_ids.add(999) # Batch contains the live request plus an unknown id (e.g. graph padding). meta.request_ids = [100, 999] - meta.prepare() + _prepare(worker, meta) # The unknown id maps to the scratch row, not to slot 0. assert worker._batch_to_slot[:2].tolist() == [s_real, worker._scratch_slot] @@ -386,7 +454,7 @@ def test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(): assert list(worker._free_slots) == [1, 2, 3] # The live request's rolling window and position are untouched. assert int(worker._ctx_len[s_real]) == 17 - assert torch.all(worker._kv_windows[s_real] == 1.0) + assert int(worker._valid_len[s_real]) == 5 def test_prepare_assigns_slots_to_disagg_generation_requests(): @@ -394,7 +462,7 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): Regression for GitHub #16767: on the disaggregated generation server the prompt is prefilled (and the DSpark window seeded) on the *context* server, - so ``_seed_context_windows`` never runs here and ``_req_to_slot`` stays empty. + so ``_seed_context_pages`` never runs here and ``_req_to_slot`` stays empty. ``prepare()`` must therefore assign each real generation request its own rolling-window slot instead of collapsing them all onto the shared scratch row (which corrupts drafts and collapses accept length at batch size > 1). @@ -407,7 +475,7 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): # All-generation batch (num_contexts == 0), no prior seeding. meta.request_ids = [1000, 1001] meta.num_generations = 2 - meta.prepare() + _prepare(worker, meta) s0 = worker._req_to_slot[1000] s1 = worker._req_to_slot[1001] @@ -416,7 +484,7 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): assert worker._batch_to_slot[:2].tolist() == [s0, s1] # Stable across steps: the same ids keep their slots (no churn / reassignment). - meta.prepare() + _prepare(worker, meta) assert worker._req_to_slot[1000] == s0 assert worker._req_to_slot[1001] == s1 assert worker._batch_to_slot[:2].tolist() == [s0, s1] @@ -435,7 +503,7 @@ def test_prepare_keeps_dummy_generation_requests_on_scratch_row(): graph_dummy = CUDA_GRAPH_DUMMY_REQUEST_ID - worker.max_draft_len meta.request_ids = [1000, ATTENTION_DP_DUMMY_REQUEST_ID, graph_dummy] meta.num_generations = 3 - meta.prepare() + _prepare(worker, meta) s_real = worker._req_to_slot[1000] assert s_real != worker._scratch_slot @@ -459,17 +527,15 @@ def test_prepare_keeps_small_real_request_slots_across_steps(): meta._dspark_worker = worker meta.request_ids = [1, 2] meta.num_generations = 2 - meta.prepare() + _prepare(worker, meta) slots = worker._batch_to_slot[:2] worker._ctx_len[slots] = torch.tensor([7, 8], device="cuda") worker._valid_len[slots] = torch.tensor([1, 2], device="cuda") - worker._position_initialized[slots] = True - meta.prepare() + _prepare(worker, meta) assert worker._ctx_len[slots].tolist() == [7, 8] assert worker._valid_len[slots].tolist() == [1, 2] - assert worker._position_initialized[slots].tolist() == [True, True] def test_prepare_keeps_cuda_graph_slots_across_replays(): @@ -483,20 +549,18 @@ def test_prepare_keeps_cuda_graph_slots_across_replays(): meta.request_ids = [11, 12, 13] meta.num_generations = 3 meta.is_cuda_graph = True - meta.prepare() + _prepare(worker, meta) slots = worker._batch_to_slot[:3] assert len(set(slots.tolist())) == 3 assert all(slot != worker._scratch_slot for slot in slots.tolist()) worker._ctx_len[slots] = torch.tensor([7, 8, 9], device="cuda") worker._valid_len[slots] = torch.tensor([1, 2, 3], device="cuda") - worker._position_initialized[slots] = True - meta.prepare() + _prepare(worker, meta) slots = worker._batch_to_slot[:3] assert worker._ctx_len[slots].tolist() == [7, 8, 9] assert worker._valid_len[slots].tolist() == [1, 2, 3] - assert worker._position_initialized[slots].tolist() == [True, True, True] class _RecordingDraftModel: @@ -513,10 +577,13 @@ class _RecordingDraftModel: def __init__(self): self.forward_calls = [] - def write_context_windows_batched(self, *args): + def configure_generation_autotuner(self, _max_batch): + pass + + def write_context_pages(self, *args): pass - def forward_batched(self, main_hidden, bonus, start_pos, **kwargs): + def forward(self, main_hidden, bonus, start_pos, **kwargs): self.forward_calls.append( { "main_hidden": main_hidden.clone(), @@ -541,13 +608,11 @@ def test_generation_state_cuda_graph_bootstrap_and_replay(): with torch.cuda.graph(graph): worker._advance_generation_state(slots, num_accepted, input_positions) - worker._ctx_len[slots] = 0 + worker._ctx_len[slots] = input_positions worker._valid_len[slots] = 0 - worker._position_initialized[slots] = False graph.replay() assert worker._ctx_len[slots].tolist() == [4017, 89] assert worker._valid_len[slots].tolist() == [1, 2] - assert worker._position_initialized[slots].tolist() == [True, True] num_accepted.copy_(torch.tensor([2, 3], device="cuda")) graph.replay() @@ -594,8 +659,13 @@ def test_disagg_position_bootstrap_uses_actual_positions_and_target_width(): verify_lens=None, ) worker._lazy_init(draft_model, metadata) + manager = _bind_manager(worker, [1000, 1001]) + manager.kv_cache_map[1000].history_length = 4016 + manager.kv_cache_map[1001].history_length = 87 + worker._prepare_managed_history([1000, 1001], 0) + worker._valid_len.zero_() - slots = [worker._assign_slot(1000, reset=False), worker._assign_slot(1001, reset=False)] + slots = [worker._assign_slot(1000), worker._assign_slot(1001)] worker._batch_to_slot[:2] = torch.tensor(slots, device="cuda") captured = torch.stack( @@ -628,7 +698,6 @@ def test_disagg_position_bootstrap_uses_actual_positions_and_target_width(): assert worker._ctx_len[slots].tolist() == [4017, 88] assert first_call["valid_len"].tolist() == [1, 1] assert worker._valid_len[slots].tolist() == [1, 1] - assert worker._position_initialized[slots].tolist() == [True, True] # Existing slots retain their state. The normal K+1 target layout must still # select each accepted bonus hidden from its own packed request row. @@ -710,7 +779,7 @@ def fake_accept(logits, am, sm): monkeypatch.setattr(worker, "sample_and_accept_draft_tokens", fake_accept) # Context-window seeding is covered by its own test; stub it out here. - monkeypatch.setattr(worker, "_seed_context_windows", lambda *a, **k: None) + monkeypatch.setattr(worker, "_seed_context_pages", lambda *a, **k: None) # The gen-block helper now returns the corrected block logits [num_gens,K,vocab]. gen_logits = torch.randn(num_gens, K, vocab, device="cuda") @@ -744,6 +813,8 @@ def fake_sample_draft_tokens(gl, sm, bs, *, num_contexts): hidden = torch.zeros(batch_size, HIDDEN, device="cuda", dtype=torch.bfloat16) logits = torch.zeros(batch_size, vocab, device="cuda") + worker._lazy_init(dm, meta) + _prepare(worker, meta, num_contexts) out = worker.forward(input_ids, position_ids, hidden, logits, attn_metadata, meta, dm) # Acceptance went through the unified entry with the right metadata objects. @@ -810,7 +881,7 @@ def test_forward_guided_batch_masks_and_advances_matcher_per_step(monkeypatch): monkeypatch.setattr( worker, "sample_and_accept_draft_tokens", lambda *a, **k: (accepted, num_accepted) ) - monkeypatch.setattr(worker, "_seed_context_windows", lambda *a, **k: None) + monkeypatch.setattr(worker, "_seed_context_pages", lambda *a, **k: None) gen_logits = torch.randn(num_gens, K, vocab, device="cuda") monkeypatch.setattr(worker, "_draft_gen_block_batched", lambda *a, **k: gen_logits) @@ -873,6 +944,8 @@ def fake_sample_draft_tokens(gl, sm, bs, *, draft_step): hidden = torch.zeros(batch_size, HIDDEN, device="cuda", dtype=torch.bfloat16) logits = torch.zeros(batch_size, vocab, device="cuda") + worker._lazy_init(dm, meta) + _prepare(worker, meta, num_contexts) out = worker.forward(input_ids, position_ids, hidden, logits, attn_metadata, meta, dm) # Matcher advanced once per draft position, in order 0..K-1. @@ -943,8 +1016,7 @@ def test_worker_and_draft_kv_follow_the_draft_form(embedded, worker_cls, uses_se worker = get_spec_worker(spec_config, None, Mapping()) assert type(worker) is worker_cls - # The embedded draft owns a rolling window and never reads the paged draft - # KV cache; the standalone one is DFlash lineage and does. + # Embedded drafts share the target manager; standalone drafts can use a separate manager. assert should_use_separate_draft_kv_cache(spec_config) is uses_separate_draft_kv @@ -1036,8 +1108,7 @@ def test_position_bound_holds_across_cuda_graph_replays(): with torch.cuda.graph(graph): worker._advance_generation_state(slots, num_accepted, input_positions) - worker._ctx_len[slots] = 0 - worker._position_initialized[slots] = False + worker._ctx_len[slots] = input_positions for _ in range(5): graph.replay() assert worker._ctx_len[0].item() <= cap @@ -1096,89 +1167,18 @@ def test_prepare_resets_the_scratch_row_every_iteration(): # State a previous padding / ADP-idle forward left behind. worker._ctx_len[scratch] = 4096 worker._valid_len[scratch] = 8 - worker._position_initialized[scratch] = True - live = worker._assign_slot(100, reset=True) + live = worker._assign_slot(100) worker._ctx_len[live] = 17 worker._valid_len[live] = 5 - worker._position_initialized[live] = True meta.request_ids = [100, CUDA_GRAPH_DUMMY_REQUEST_ID] - meta.prepare() + _prepare(worker, meta) assert int(worker._ctx_len[scratch]) == 0 assert int(worker._valid_len[scratch]) == 0 - assert not bool(worker._position_initialized[scratch]) + assert worker._draft_capacities[1].item() == 0 # A live request's own state is bounded by its lifetime and must survive. assert int(worker._ctx_len[live]) == 17 assert int(worker._valid_len[live]) == 5 - assert bool(worker._position_initialized[live]) - - -def test_ragged_draft_gathers_each_requests_hidden_state(): - worker = _make_worker() - draft_model = _RecordingDraftModel() - metadata = types.SimpleNamespace( - max_num_requests=2, - qo_indptr=torch.tensor([0, 4, 6], dtype=torch.int32, device="cuda"), - verify_lens=torch.tensor([4, 2], dtype=torch.int32, device="cuda"), - ) - worker._lazy_init(draft_model, metadata) - worker._batch_to_slot[:2] = torch.tensor([0, 1], device="cuda") - captured = torch.stack( - [torch.full((HIDDEN * NCAP,), float(index), device="cuda") for index in range(6)] - ) - metadata.get_hidden_states = lambda _num_tokens: captured - accepted = torch.zeros((2, 6), dtype=torch.int32, device="cuda") - - worker._draft_gen_block_batched( - draft_model, - metadata, - types.SimpleNamespace(num_ctx_tokens=0), - accepted, - torch.ones(2, dtype=torch.int32, device="cuda"), - num_contexts=0, - batch_size=2, - total_target_tokens=6, - position_ids=torch.arange(6, device="cuda").unsqueeze(0), - ) - - assert draft_model.forward_calls[-1]["main_hidden"][:, 0].tolist() == [0.0, 4.0] - - -def test_confidence_rows_are_slot_indexed_and_fail_open(): - worker = _make_worker(enable_confidence_scheduling=True) - worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) - - buffer = worker.staged_confidence_buffer() - assert buffer is not None - neutral = worker.confidence_row_for(12345) - assert neutral == worker._neutral_conf_row - assert neutral > worker._scratch_slot - assert torch.sigmoid(buffer[neutral]).min().item() > 0.999 - - worker._req_to_slot = {100: 5, 101: 2, 102: 7} - for slot in range(worker._scratch_slot): - buffer[slot] = float(slot) - rows = [worker.confidence_row_for(request_id) for request_id in (100, 101, 102)] - assert rows == [5, 2, 7] - assert buffer[rows, 0].tolist() == [5.0, 2.0, 7.0] - - -def test_recycled_confidence_slot_is_reset_to_neutral(): - worker = _make_worker(enable_confidence_scheduling=True) - worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) - worker._confidence_logits[: worker._scratch_slot] = -20.0 - - slot = worker._assign_slot(100, reset=False) - - assert torch.sigmoid(worker._confidence_logits[slot]).min().item() > 0.999 - - -def test_confidence_state_is_absent_when_scheduling_is_disabled(): - worker = _make_worker(enable_confidence_scheduling=False) - worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) - - assert worker.staged_confidence_buffer() is None - assert worker.confidence_stamp_buffer() is None - assert worker.verify_planner is None + assert worker._managed_residency[100] is worker._draft_kv_manager.kv_cache_map[100] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 77093bec8ba6..336ab0b22e20 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -60,6 +60,13 @@ def _has_absorbed_mla_kernel() -> bool: reason="paged absorbed-MLA decode needs an SM100-family GPU", ) +# The TRTLLM-Gen paged attention the DFlash drafter uses ships for the same +# SM100 family. +needs_trtllm_gen = pytest.mark.skipif( + not _has_absorbed_mla_kernel(), + reason="TRTLLM-Gen paged block decode needs an SM100-family GPU", +) + # --------------------------------------------------------------------------- # Reference: line-for-line port of DeepSpec VanillaMarkov # (deepspec/modeling/dspark/markov_head.py) at temperature 0. @@ -827,6 +834,86 @@ def test_plain_dflash_block_decode_matches_full_attention_reference(): assert diff < 0.02, f"plain DFlash parity failed: max abs diff {diff}" +def _run_paged_block_decode(drafter, captured, noise_embed, cache_dtype, page_size=32): + """Block decode through the TRTLLM paged cache, one request per history. + + Pages are shuffled across requests, so only the block tables decide which + history each request reads. + """ + dev = "cuda" + blk = TINY["block_size"] + lengths = [ctx.shape[0] for ctx in captured] + history = [] + for ctx in captured: + proj = drafter.project_target_hidden(ctx.to(dev, torch.bfloat16)) + history.append(drafter.precompute_context_kv(proj, torch.arange(ctx.shape[0], device=dev))) + # [N, layers, kv heads, head dim]; the drafter sizes its fused K/V lazily. + num_layers, nkv, hd = history[0][0].shape[1:] + pages_per_req = max((n + blk + page_size - 1) // page_size for n in lengths) + num_pages = len(lengths) * pages_per_req + order = torch.randperm(num_pages, generator=torch.Generator().manual_seed(3)) + page_table = order.view(len(lengths), pages_per_req).to(dev, torch.int32) + caches = [ + torch.zeros(num_pages, 2, nkv, page_size, hd, dtype=torch.bfloat16, device=dev) + for _ in range(num_layers) + ] + for row, (k, v) in enumerate(history): + positions = torch.arange(lengths[row], device=dev) + pages = page_table[row, positions // page_size].long() + offsets = positions % page_size + for layer, cache in enumerate(caches): + cache[pages, 0, :, offsets] = k[:, layer] + cache[pages, 1, :, offsets] = v[:, layer] + caches = [cache.to(cache_dtype) for cache in caches] + out = drafter.dflash_forward( + noise_embedding=torch.stack(noise_embed).to(dev, torch.bfloat16), + query_positions=torch.stack([torch.arange(n, n + blk) for n in lengths]).to(dev), + num_ctx_per_req=torch.tensor(lengths, device=dev), + ctx_k_cache=None, + ctx_v_cache=None, + ctx_cache_batch_idx=torch.arange(len(lengths), device=dev), + ctx_kv_cache=caches, + ctx_page_table=page_table, + ) + return out.float().cpu().view(len(lengths), blk, -1) + + +@needs_trtllm_gen +@pytest.mark.parametrize( + "cache_dtype,tolerance", + [(torch.bfloat16, 0.02), (torch.float8_e4m3fn, 0.05)], + ids=["bf16", "fp8"], +) +def test_trtllm_paged_block_decode_matches_full_attention_reference(cache_dtype, tolerance): + """The TRTLLM non-causal block attention over a paged BF16 or FP8 cache + matches the eager reference for requests with different history lengths. + + Each block query runs as its own one-token decode request over its + parent's pages, so a wrong table, length or block write shows up here. + An FP8 cache also takes the block's K/V and queries in FP8. + """ + weights = _tiny_weights() + drafter = _build_drafter(False, weights, dflash_attention_backend="TRTLLM") + g = torch.Generator().manual_seed(44) + lengths = (CTX_LEN, 70) + captured = [ + torch.randn(n, TINY["hidden_size"] * NUM_CAPTURE, generator=g) * 0.5 for n in lengths + ] + noise_embed = [ + torch.randn(TINY["block_size"], TINY["hidden_size"], generator=g) * 0.5 for _ in lengths + ] + + out = _run_paged_block_decode(drafter, captured, noise_embed, cache_dtype) + + for row in range(len(lengths)): + expected = _reference_block_decode( + weights, captured[row].to(torch.bfloat16), noise_embed[row].to(torch.bfloat16), False + ) + expected = expected.float() + rel_err = ((out[row] - expected).norm() / expected.norm()).item() + assert rel_err < tolerance, f"request {row}: relative error {rel_err}" + + # --------------------------------------------------------------------------- # MLA-backboned DSpark drafter (Inferact/Kimi-K3-DSpark shape). # diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 0c9468c48a75..878a70531840 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -1,18 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU correctness tests for the DSpark CuteDSL attention op.""" +"""GPU correctness tests for paged DSpark CuteDSL attention and preparation.""" + +from unittest.mock import Mock import pytest import torch +import torch.nn.functional as F -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.models.modeling_dspark import ( - _rope_last_dims_batched, - dspark_sparse_attn, - get_dspark_topk_idxs_batched, - precompute_dspark_freqs_cis, +from tensorrt_llm._torch.custom_ops import dspark_attention_custom_op as attention_op +from tensorrt_llm._torch.custom_ops import dspark_rmsnorm_rope_custom_op as preparation_op +from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_INTERNAL_AVAILABLE, ) +from tensorrt_llm._torch.models import modeling_dspark as dspark from tensorrt_llm._utils import get_sm_version _ROPE_DIM = 64 @@ -25,669 +28,1044 @@ ) +@pytest.fixture(params=[False, True], ids=["native", "unsupported_arch_fallback"]) +def force_fallback(request, monkeypatch): + if request.param: + monkeypatch.setattr(preparation_op, "_get_dspark_arch_str", lambda: None) + monkeypatch.setattr(dspark, "_get_dspark_arch_str", lambda: None) + + def unsupported_kernel(*args, **kwargs): + pytest.fail("Unsupported architecture entered a fused DSpark kernel") + + monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope", unsupported_kernel) + monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope_page_write", unsupported_kernel) + monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope_draft_block", unsupported_kernel) + monkeypatch.setattr(torch.ops.trtllm, "fused_dsv4_dspark_attention", unsupported_kernel) + assert preparation_op._get_dspark_arch_str() is None + return request.param + + def _make_inputs( - seed: int = 0, - batch: int = 2, - block: int = 6, - start_pos_values=None, - cache_pages: int | None = None, + seed=0, + batch=2, + block=6, + heads=128, + page_size=32, + positions=None, + valid_lengths=None, + capacities=None, + extra_pages=0, ): torch.manual_seed(seed) device = torch.device("cuda") - if start_pos_values is not None: - batch = len(start_pos_values) - heads, head_dim, window = 128, 512, 128 - q = torch.randn(batch, block, heads, head_dim, device=device, dtype=torch.bfloat16) - main_kv = torch.randn(batch, head_dim, device=device, dtype=torch.bfloat16) - block_kv = torch.randn(batch, block, head_dim, device=device, dtype=torch.bfloat16) - - # The worker cache window is a strided view. Exercise that layout contract. - if cache_pages is None: - cache_pages = max(4, batch + 1) - cache_storage = torch.randn( - cache_pages, 3, window, head_dim, device=device, dtype=torch.bfloat16 - ) - kv_cache = cache_storage[:, 1] - assert not kv_cache.is_contiguous() - # Slots deliberately differ from decode positions so a kernel that derives - # the window-validity mask from anything but start_pos fails the tests. - slots = torch.arange(batch - 1, -1, -1, device=device, dtype=torch.int32) - if start_pos_values is None: - start_pos_values = [199 * i + 1 for i in range(batch)] - start_pos = torch.tensor(start_pos_values, device=device, dtype=torch.int32) - sink = torch.randn(heads, device=device, dtype=torch.float32) - # Per-request RoPE phases match the runtime indexing contract. - table = precompute_dspark_freqs_cis(_ROPE_DIM, int(start_pos.max()) + block + 2, device=device) - blk_freqs = table[start_pos.long().unsqueeze(1) + 1 + torch.arange(block, device=device)] - inverse_rope_freqs = torch.view_as_real(blk_freqs).contiguous() - return q, main_kv, block_kv, kv_cache, slots, start_pos, sink, blk_freqs, inverse_rope_freqs - - -def _legacy_valid_len(start_pos): - return (start_pos.long() + 1).clamp(max=128) - - -def _prepare_attention_inputs(main_kv, block_kv, kv_cache, slots, start_pos): - kv_cache[slots.long(), (start_pos % kv_cache.shape[1]).long()] = main_kv - draft_block = block_kv.new_zeros((block_kv.shape[0], 8, block_kv.shape[2])) - draft_block[:, : block_kv.shape[1]].copy_(block_kv) - return draft_block, slots.to(torch.int32), start_pos.to(torch.int32) - - -def _with_row_padding(x: torch.Tensor, padding: int = 128) -> torch.Tensor: - storage = torch.empty((*x.shape[:-1], x.shape[-1] + padding), device=x.device, dtype=x.dtype) - padded = storage[..., : x.shape[-1]] - padded.copy_(x) - return padded + if positions is None: + positions = [1 + (37 * i) % 386 for i in range(batch)] + batch = len(positions) + positions = torch.tensor(positions, device=device, dtype=torch.int64) + lengths = ( + positions.clamp(max=128) + if valid_lengths is None + else torch.tensor(valid_lengths, device=device, dtype=torch.int64) + ) + capacities = torch.tensor( + [512] * batch if capacities is None else capacities, device=device, dtype=torch.int64 + ) + pages_per_request = 512 // page_size + num_pages = batch * pages_per_request + extra_pages + # Both the pool and table rows have padding, as in managed cache views. + storage = torch.randn(num_pages, 3, page_size, 512, device=device, dtype=torch.bfloat16) + table_storage = torch.full((batch, pages_per_request + 1), -1, device=device, dtype=torch.int32) + tables = table_storage[:, :pages_per_request] + tables.copy_( + torch.randperm(num_pages, device=device, dtype=torch.int32)[ + : batch * pages_per_request + ].reshape(batch, pages_per_request) + ) + tables[capacities == 0] = -1 + freqs = dspark.precompute_dspark_freqs_cis(_ROPE_DIM, 1024, device=device) + block_freqs = freqs[positions[:, None] + 1 + torch.arange(block, device=device)] + return dict( + q=torch.randn(batch, block, heads, 512, device=device, dtype=torch.bfloat16), + draft_block=torch.randn(batch, 8, 512, device=device, dtype=torch.bfloat16), + kv_pages=storage[:, 1], + block_tables=tables, + positions=positions, + valid_lengths=lengths, + capacities=capacities, + attn_sink=torch.randn(heads, device=device), + inverse_rope_freqs=torch.view_as_real(block_freqs).contiguous(), + softmax_scale=512**-0.5, + ) def _reference( q, - main_kv, - block_kv, - kv_cache, - slots, - start_pos, - valid_len, - sink, - blk_freqs, + draft_block, + kv_pages, + block_tables, + positions, + valid_lengths, + capacities, + attn_sink, + inverse_rope_freqs, softmax_scale, ): - cache = kv_cache.clone() - window = cache.shape[1] - cache[slots.long(), (start_pos % window).long()] = main_kv - kv_full = torch.cat([cache[slots.long()], block_kv], dim=1) - topk = get_dspark_topk_idxs_batched(window, q.shape[1], start_pos.long(), valid_len) - o = dspark_sparse_attn(q, kv_full, sink, topk, softmax_scale) - return _rope_last_dims_batched(o, _ROPE_DIM, blk_freqs, inverse=True), cache + """Gather absolute context positions and compute sink attention in FP32.""" + tables = block_tables.cpu() + outputs = [] + for row, (end, length, capacity) in enumerate( + zip(positions.tolist(), valid_lengths.tolist(), capacities.tolist()) + ): + tokens = list(range(max(0, end - min(128, length)), min(end, capacity))) + if tokens: + pages = [int(tables[row, token // kv_pages.shape[1]]) for token in tokens] + offsets = [token % kv_pages.shape[1] for token in tokens] + context = kv_pages[pages, offsets] + kv = torch.cat((context, draft_block[row, : q.shape[1]])) + else: + kv = draft_block[row, : q.shape[1]] + scores = torch.einsum("shd,td->sht", q[row].float(), kv.float()) * softmax_scale + sink = attn_sink[None, :, None].expand(q.shape[1], -1, 1) + probs = torch.cat((scores, sink), dim=-1).softmax(-1)[..., :-1] + outputs.append(torch.einsum("sht,td->shd", probs, kv.float()).to(q.dtype)) + return dspark._rope_last_dims_batched( + torch.stack(outputs), _ROPE_DIM, torch.view_as_complex(inverse_rope_freqs), inverse=True + ) + + +def _norm_rope(x, weight, freqs): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6) + normalized = (normalized * weight.float()).to(x.dtype) + return dspark._rope_last_dims_batched(normalized, _ROPE_DIM, freqs) + + +def _write_reference(cache, main_kv, inputs): + cache = cache.clone() + for row, (end, capacity) in enumerate( + zip(inputs["positions"].tolist(), inputs["capacities"].tolist()) + ): + token = end - 1 + if 0 <= token < capacity: + page = int(inputs["block_tables"][row, token // cache.shape[1]]) + cache[page, token % cache.shape[1]] = main_kv[row, 0] + return cache @pytest.mark.parametrize( - ("invalid_case", "reason"), + "invalid_case,reason", [ - ("q_dtype", "q dtype must be BF16"), - ("block_size", "draft block size must be 5 or 6"), - ("valid_len_dtype", "valid_len must use INT64"), - ("freqs_shape", "inverse_rope_freqs shape must be"), + ("q_dtype", "must use BF16"), + ("block_size", "requires block 5/6"), + ("heads", "64/128 heads"), + ("page_size", "16/32/64/128/256-token pages"), + ("draft_shape", "draft KV must have shape"), + ("draft_layout", "must be contiguous"), + ("table_dtype", "contiguous INT32 rows"), + ("table_layout", "contiguous INT32 rows"), + ("positions", "contiguous INT64 batch vectors"), + ("valid_lengths", "contiguous INT64 batch vectors"), + ("capacities", "contiguous INT64 batch vectors"), + ("sink_dtype", "sink must be FP32"), + ("freqs_shape", "inverse RoPE frequencies must be FP32"), ], ) -def test_fused_dsv4_dspark_attention_support_gate_logs_rejection(monkeypatch, invalid_case, reason): - import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op - - q, _, _, kv_cache, _, start_pos, sink, _, freqs = _make_inputs() - inputs = [q, kv_cache, _legacy_valid_len(start_pos), sink, freqs] +def test_fused_dsv4_dspark_attention_rejects_invalid_inputs_before_launch( + monkeypatch, invalid_case, reason +): + inputs = _make_inputs() if invalid_case == "q_dtype": - inputs[0] = q.float() + inputs["q"] = inputs["q"].float() elif invalid_case == "block_size": - inputs[0] = q[:, :4].contiguous() - inputs[4] = freqs[:, :4].contiguous() - elif invalid_case == "valid_len_dtype": - inputs[2] = inputs[2].to(torch.int32) + inputs["q"] = inputs["q"][:, :4].contiguous() + elif invalid_case == "heads": + inputs["q"] = inputs["q"][:, :, :24].contiguous() + elif invalid_case == "page_size": + inputs["kv_pages"] = inputs["kv_pages"][:, :24] + elif invalid_case == "draft_shape": + inputs["draft_block"] = inputs["draft_block"][:, :7].contiguous() + elif invalid_case == "draft_layout": + inputs["draft_block"] = F.pad(inputs["draft_block"], (0, 128))[..., :512] + elif invalid_case == "table_dtype": + inputs["block_tables"] = inputs["block_tables"].long() + elif invalid_case == "table_layout": + inputs["block_tables"] = inputs["block_tables"][:, ::2] + elif invalid_case in ("positions", "valid_lengths", "capacities"): + inputs[invalid_case] = inputs[invalid_case].int() + elif invalid_case == "sink_dtype": + inputs["attn_sink"] = inputs["attn_sink"].bfloat16() else: - inputs[4] = freqs[:, :, :-1].contiguous() + inputs["inverse_rope_freqs"] = inputs["inverse_rope_freqs"][:, :, :-1].contiguous() - messages = [] - monkeypatch.setattr( - dspark_attention_op.logger, - "debug_once", - lambda *message, key: messages.append((" ".join(map(str, message)), key)), - ) - - assert not dspark_attention_op.is_fused_dsv4_dspark_attention_supported(*inputs) - assert len(messages) == 1 - assert reason in messages[0][0] - assert messages[0][1][0] == "fused_dsv4_dspark_attention_unsupported" + def unexpected_launch(*args): + pytest.fail("invalid inputs reached compilation or a kernel launch") + key = (6, 128, 32, attention_op._get_dspark_arch_str(), None) + monkeypatch.setattr(attention_op, "_dspark_attention_kernel_cache", {key: unexpected_launch}) + monkeypatch.setattr(attention_op, "_compile_dspark_attention", unexpected_launch) + with pytest.raises(ValueError, match=reason): + attention_op.fused_dsv4_dspark_attention(**inputs) -@pytest.mark.parametrize("invalid_case", ("draft_block_shape", "draft_block_layout", "index_dtype")) -def test_fused_dsv4_dspark_attention_rejects_invalid_inputs_before_launch( - monkeypatch, invalid_case -): - import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op - - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs() - valid_len = _legacy_valid_len(start_pos) - draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( - main_kv, block_kv, kv_cache, slots, start_pos - ) - if invalid_case == "draft_block_shape": - draft_block = draft_block[:, :-1].contiguous() - elif invalid_case == "draft_block_layout": - draft_block = _with_row_padding(draft_block) - assert not draft_block.is_contiguous() - else: - slots_i32 = slots_i32.long() - monkeypatch.setattr( - dspark_attention_op, - "_run_dspark_attention", - lambda *args: pytest.fail("invalid inputs reached the kernel launch"), - ) - with pytest.raises(ValueError, match="requires contiguous supported DSV4 DSpark tensors"): - dspark_attention_op.fused_dsv4_dspark_attention( - q, - draft_block, - kv_cache, - slots_i32, - cache_seqs, - valid_len, - sink, - freqs, - q.shape[-1] ** -0.5, - ) +@pytest.mark.parametrize("heads", (64, 128)) +@pytest.mark.parametrize("block", (5, 6)) +@pytest.mark.parametrize("page_size", (16, 32, 64, 128, 256)) +def test_fused_dsv4_dspark_attention_matches_reference(heads, block, page_size): + inputs = _make_inputs( + 29, + heads=heads, + block=block, + page_size=page_size, + positions=[0, 1, 32, 33, 127, 128, 129, 257, 390, 99], + valid_lengths=[0, 1, 3, 32, 127, 128, 17, 128, 3, 128], + capacities=[512] * 9 + [0], + ) + original_cache = inputs["kv_pages"].clone() + expected = _reference(**inputs) + actual = attention_op.fused_dsv4_dspark_attention(**inputs) + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) + torch.testing.assert_close(inputs["kv_pages"], original_cache, rtol=0, atol=0) -@pytest.mark.parametrize("block", (5, 6)) -@pytest.mark.parametrize( - ("start_pos_values", "valid_len_values"), - [ - # Full windows with slots != start_pos: catches the window-validity - # mask being fed anything but the absolute decode position. - ([257, 390], [128, 128]), - # Partially filled windows: only rows 0..start_pos are attended. - ([5, 100], [6, 101]), - # Bootstrapped positions with short physical suffixes exercise wraparound. - ([257, 390], [3, 5]), - ], - ids=["full_window", "partial_window", "wrapped_suffix"], -) -def test_fused_dsv4_dspark_attention_matches_reference(block, start_pos_values, valid_len_values): - from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( - fused_dsv4_dspark_attention, +@pytest.mark.parametrize("heads,page_size", [(64, 32), (128, 64), (128, 256)]) +def test_fused_dsv4_dspark_attention_cuda_graph_replay(heads, page_size): + inputs = _make_inputs( + 3, heads=heads, page_size=page_size, positions=[257, 5, 390], capacities=[512, 512, 0] ) + main_x = torch.randn(3, 1, 512, device="cuda", dtype=torch.bfloat16) + block_x = torch.randn(3, 6, 512, device="cuda", dtype=torch.bfloat16) + weight = torch.ones(512, device="cuda", dtype=torch.bfloat16) + freqs = dspark.precompute_dspark_freqs_cis(_ROPE_DIM, 1024, device="cuda") + main_freqs = torch.view_as_real(freqs[inputs["positions"]]).contiguous() + block_freqs = inputs["inverse_rope_freqs"].view(-1, 32, 2) + + def forward(): + preparation_op.cute_dsl_dspark_rmsnorm_rope_page_write( + main_x, + weight, + main_freqs, + inputs["kv_pages"], + inputs["block_tables"], + inputs["positions"], + inputs["capacities"], + 1e-6, + ) + draft = preparation_op.cute_dsl_dspark_rmsnorm_rope_draft_block( + block_x, weight, block_freqs, 1e-6 + ) + return attention_op.fused_dsv4_dspark_attention(**(inputs | {"draft_block": draft})), draft - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, blk_freqs, inverse_rope_freqs = ( - _make_inputs(29, block=block, start_pos_values=start_pos_values) + # Eager execution warms all preparation and attention specializations. + forward() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured, draft = forward() + for end in (127, 128, 129, 255): + inputs["positions"][:2].copy_(torch.tensor([end, end + 1], device="cuda")) + inputs["valid_lengths"].copy_(torch.tensor([128, 3, 128], device="cuda")) + inputs["block_tables"][:2].copy_(inputs["block_tables"][:2].roll(1, dims=1)) + main_x.normal_() + block_x.normal_() + main_freqs.copy_(torch.view_as_real(freqs[inputs["positions"]])) + block_phases = freqs[inputs["positions"][:, None] + 1 + torch.arange(6, device="cuda")] + inputs["inverse_rope_freqs"].copy_(torch.view_as_real(block_phases)) + expected_main = _norm_rope(main_x, weight, freqs[inputs["positions"]][:, None]) + expected_block = _norm_rope(block_x, weight, block_phases) + expected_cache = _write_reference(inputs["kv_pages"], expected_main, inputs) + expected_draft = F.pad(expected_block, (0, 0, 0, 2)) + expected = _reference( + **(inputs | {"kv_pages": expected_cache, "draft_block": expected_draft}) + ) + graph.replay() + torch.testing.assert_close(captured, expected, rtol=5e-2, atol=3e-2) + torch.testing.assert_close(inputs["kv_pages"], expected_cache, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(draft, expected_draft, rtol=2e-2, atol=2e-2) + assert torch.count_nonzero(draft[:, 6:]) == 0 + + +@pytest.mark.parametrize("heads", (64, 128)) +@pytest.mark.parametrize("page_size", (32, 256)) +@pytest.mark.parametrize("block", (3, 5, 6)) +def test_dspark_attention_forward_matches_reference( + heads, page_size, block, force_fallback, monkeypatch +): + inputs = _make_inputs( + 17, + block=block, + heads=heads, + page_size=page_size, + positions=[1, 33, 129, 390, 0], + valid_lengths=[1, 3, 128, 3, 0], + capacities=[512, 512, 512, 512, 0], ) - valid_len = torch.tensor(valid_len_values, device=q.device, dtype=torch.long) - expected, expected_cache = _reference( - q, - main_kv, - block_kv, - kv_cache, - slots, - start_pos, - valid_len, - sink, - blk_freqs, - q.shape[-1] ** -0.5, + batch, hidden, rank, groups, o_rank = 5, 64, 64, 8, 32 + + def weight(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.02 + + x, main_x = weight(batch, block, hidden), weight(batch, 1, hidden) + kwargs = dict( + wq_a=weight(rank, hidden), + q_norm_w=torch.ones(rank, device="cuda", dtype=x.dtype), + wq_b=weight(heads * 512, rank), + wkv=weight(512, hidden), + kv_norm_w=torch.ones(512, device="cuda", dtype=x.dtype), + wo_a=weight(groups * o_rank, heads * 512 // groups), + wo_b=weight(hidden, groups * o_rank), + attn_sink=inputs["attn_sink"], + n_heads=heads, + head_dim=512, + rope_head_dim=_ROPE_DIM, + n_groups=groups, + o_lora_rank=o_rank, + window_size=128, + eps=1e-6, + softmax_scale=inputs["softmax_scale"], + freqs_cis=dspark.precompute_dspark_freqs_cis(_ROPE_DIM, 1024, device="cuda"), ) - draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( - main_kv, block_kv, kv_cache, slots, start_pos + positions = inputs["positions"] + block_freqs = torch.view_as_complex(inputs["inverse_rope_freqs"]) + q = F.linear( + dspark._rmsnorm(F.linear(x, kwargs["wq_a"]), kwargs["q_norm_w"], 1e-6), kwargs["wq_b"] + ).unflatten(-1, (heads, 512)) + q = _norm_rope(q, torch.ones(512, device="cuda", dtype=x.dtype), block_freqs) + main = _norm_rope( + F.linear(main_x, kwargs["wkv"]), + kwargs["kv_norm_w"], + kwargs["freqs_cis"][positions][:, None], ) - - actual = fused_dsv4_dspark_attention( - q, - draft_block, - kv_cache, - slots_i32, - cache_seqs, - valid_len, - sink, - inverse_rope_freqs, - q.shape[-1] ** -0.5, + draft = _norm_rope(F.linear(x, kwargs["wkv"]), kwargs["kv_norm_w"], block_freqs) + expected_cache = _write_reference(inputs["kv_pages"], main, inputs) + output = _reference(**(inputs | {"q": q, "draft_block": draft, "kv_pages": expected_cache})) + output = torch.einsum( + "bsgd,grd->bsgr", + output.reshape(batch, block, groups, -1), + kwargs["wo_a"].view(groups, o_rank, -1), ) - - torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) - torch.testing.assert_close(kv_cache, expected_cache, rtol=0, atol=0) + expected = F.linear(output.flatten(2), kwargs["wo_b"]) + fused_calls = [] + for module, name in ( + (dspark, "cute_dsl_dspark_rmsnorm_rope_page_write"), + (dspark, "cute_dsl_dspark_rmsnorm_rope_draft_block"), + (torch.ops.trtllm, "fused_dsv4_dspark_attention"), + ): + call = Mock(wraps=getattr(module, name)) + monkeypatch.setattr(module, name, call) + fused_calls.append(call) + actual = dspark.dspark_attention_forward( + x, + main_x, + positions, + inputs["kv_pages"], + inputs["block_tables"], + inputs["capacities"], + inputs["valid_lengths"], + **kwargs, + ) + for call in fused_calls: + if not force_fallback and block in (5, 6): + call.assert_called_once() + else: + call.assert_not_called() + torch.testing.assert_close(actual, expected, rtol=8e-2, atol=1e-2) + torch.testing.assert_close(inputs["kv_pages"], expected_cache, rtol=2e-2, atol=2e-2) -def test_fused_dsv4_dspark_attention_cuda_graph_replay(): - from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( - fused_dsv4_dspark_attention, - warmup_fused_dsv4_dspark_attention, - ) - from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( - cute_dsl_dspark_rmsnorm_rope, - cute_dsl_dspark_rmsnorm_rope_cache_write, - cute_dsl_dspark_rmsnorm_rope_draft_block, - ) - - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, blk_freqs, inverse_rope_freqs = ( - _make_inputs(3, start_pos_values=[257, 5, 390]) - ) - batch, block, dim = block_kv.shape - scale = dim**-0.5 - main_x = main_kv.unsqueeze(1) - block_x = block_kv - weight = torch.ones(dim, device=q.device, dtype=q.dtype) - main_freqs = torch.zeros(batch, _ROPE_DIM // 2, 2, device=q.device) - block_freqs = torch.zeros(batch * block, _ROPE_DIM // 2, 2, device=q.device) - main_freqs[..., 0] = 1 - block_freqs[..., 0] = 1 - slots = slots.long() - start_pos = start_pos.long() - valid_len = _legacy_valid_len(start_pos) - - # The named best-effort prewarm calls the same self-JIT ops used below, - # so graph capture sees only hot compile-cache entries. - warmup_fused_dsv4_dspark_attention(block, 1e-6) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - slots_i32, cache_seqs = cute_dsl_dspark_rmsnorm_rope_cache_write( - main_x, weight, main_freqs, kv_cache, slots, start_pos, 1e-6 - ) - draft_block = cute_dsl_dspark_rmsnorm_rope_draft_block(block_x, weight, block_freqs, 1e-6) - captured = fused_dsv4_dspark_attention( - q, - draft_block, - kv_cache, - slots_i32, - cache_seqs, - valid_len, - sink, - inverse_rope_freqs, - scale, - ) +def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(monkeypatch): + assert [attention_op._get_dspark_arch_str(sm) for sm in (100, 103, 107)] == [ + "sm_100", + "sm_103", + "sm_107", + ] + assert all(attention_op._get_dspark_arch_str(sm) is None for sm in (90, 101, 109)) + cache = {} + monkeypatch.setattr(attention_op, "_dspark_attention_kernel_cache", cache) + for batch in (1, 3, 32): + inputs = _make_inputs(17 + batch, batch=batch, extra_pages=batch + 41) + if batch == 3: + inputs["kv_pages"] = inputs["kv_pages"].clone() + inputs["block_tables"] = inputs["block_tables"].clone() + inputs["softmax_scale"] *= 0.5 if batch % 2 == 0 else 1 + expected = _reference(**inputs) + actual = attention_op.fused_dsv4_dspark_attention(**inputs) + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) + assert len(cache) == 1 + if batch == 1: - main_x.copy_(torch.randn_like(main_x)) - block_x.copy_(torch.randn_like(block_x)) - valid_len.copy_(torch.tensor([3, 2, 7], device=q.device)) - expected_main = cute_dsl_dspark_rmsnorm_rope( - main_x, weight, main_freqs, 1, _ROPE_DIM, 1e-6, True, True, False - ).squeeze(1) - expected_block = cute_dsl_dspark_rmsnorm_rope( - block_x, weight, block_freqs, 1, _ROPE_DIM, 1e-6, True, True, False - ) - expected, expected_cache = _reference( - q, - expected_main, - expected_block, - kv_cache, - slots, - start_pos, - valid_len, - sink, - blk_freqs, - scale, - ) - graph.replay() + def unexpected_compile(*args): + pytest.fail("runtime shape or scale change caused recompilation") - torch.testing.assert_close(captured, expected, rtol=5e-2, atol=3e-2) - torch.testing.assert_close(kv_cache, expected_cache, rtol=0, atol=0) - torch.testing.assert_close(draft_block[:, :block], expected_block, rtol=0, atol=0) - torch.testing.assert_close( - draft_block[:, block:], torch.zeros_like(draft_block[:, block:]), rtol=0, atol=0 - ) + monkeypatch.setattr(attention_op, "_compile_dspark_attention", unexpected_compile) -@pytest.mark.parametrize("persist", (False, True)) -def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, persist): - import tensorrt_llm._torch.models.modeling_dspark as dspark_attention +def test_compile_without_real_specimens_and_cached_wrapper_avoids_views(monkeypatch): + arch_str = attention_op._get_dspark_arch_str() - torch.manual_seed(17) - device = torch.device("cuda") - dtype = torch.bfloat16 - batch, block, hidden = 2, 6, 64 - heads, head_dim, rope_dim = 128, 512, 64 - q_rank, groups, o_rank, window = 1024, 8, 32, 128 - - def scaled_randn(*shape): - return torch.randn(*shape, device=device, dtype=dtype) * 0.02 - - x = torch.randn(batch, block, hidden, device=device, dtype=dtype) * 0.1 - main_x = torch.randn(batch, 1, hidden, device=device, dtype=dtype) * 0.1 - start_pos = torch.tensor([5, 390], device=device, dtype=torch.long) - slots = torch.tensor([1, 0], device=device, dtype=torch.long) - valid_len = torch.tensor([6, 3], device=device, dtype=torch.long) - kwargs = { - "wq_a": scaled_randn(q_rank, hidden), - "q_norm_w": torch.ones(q_rank, device=device, dtype=dtype), - "wq_b": scaled_randn(heads * head_dim, q_rank), - "wkv": scaled_randn(head_dim, hidden), - "kv_norm_w": torch.ones(head_dim, device=device, dtype=dtype), - "wo_a": scaled_randn(groups * o_rank, heads * head_dim // groups), - "wo_b": scaled_randn(hidden, groups * o_rank), - "attn_sink": torch.randn(heads, device=device, dtype=torch.float32) * 0.1, - "n_heads": heads, - "head_dim": head_dim, - "rope_head_dim": rope_dim, - "n_groups": groups, - "o_lora_rank": o_rank, - "window_size": window, - "eps": 1e-6, - "softmax_scale": head_dim**-0.5, - "freqs_cis": dspark_attention.precompute_dspark_freqs_cis(rope_dim, 512, device=device), - "persist": persist, - } - cache_storage = torch.randn(3, 3, window, head_dim, device=device, dtype=dtype) * 0.1 - op_cache = cache_storage[:, 1] - fallback_cache = op_cache.clone() - calls = { - "fused_attention": 0, - "cache_write": 0, - "draft_block": 0, - "rmsnorm_rope": 0, - } - op_attention = dspark_attention.fused_dsv4_dspark_attention - op_cache_write = dspark_attention.cute_dsl_dspark_rmsnorm_rope_cache_write - op_draft_block = dspark_attention.cute_dsl_dspark_rmsnorm_rope_draft_block - op_rmsnorm_rope = dspark_attention.cute_dsl_dspark_rmsnorm_rope - - def counted_attention(*args): - calls["fused_attention"] += 1 - return op_attention(*args) - - def counted_cache_write(*args): - calls["cache_write"] += 1 - return op_cache_write(*args) - - def counted_draft_block(*args): - calls["draft_block"] += 1 - return op_draft_block(*args) - - def counted_rmsnorm_rope(*args): - calls["rmsnorm_rope"] += 1 - return op_rmsnorm_rope(*args) + def unexpected_compile_op(*args, **kwargs): + pytest.fail("DSpark compilation used a real tensor specimen") with monkeypatch.context() as patch: - patch.setattr(dspark_attention, "fused_dsv4_dspark_attention", counted_attention) - patch.setattr( - dspark_attention, - "cute_dsl_dspark_rmsnorm_rope_cache_write", - counted_cache_write, - ) - patch.setattr( - dspark_attention, - "cute_dsl_dspark_rmsnorm_rope_draft_block", - counted_draft_block, - ) - patch.setattr(dspark_attention, "cute_dsl_dspark_rmsnorm_rope", counted_rmsnorm_rope) - actual = dspark_attention.dspark_attention_forward_batched( - x, main_x, start_pos, op_cache, slots, valid_len, **kwargs - ) + patch.setattr(torch, "empty", unexpected_compile_op) + patch.setattr(torch, "empty_like", unexpected_compile_op) + patch.setattr(attention_op.cute.runtime, "from_dlpack", unexpected_compile_op) + compiled = attention_op._compile_dspark_attention(5, arch_str, 32, None, 128) + monkeypatch.setattr( + attention_op, "_dspark_attention_kernel_cache", {(5, 128, 32, arch_str, None): compiled} + ) + inputs = _make_inputs(41, block=5, positions=[257, 9]) - # Main K/V and block K/V reuse their existing RMSNorm/RoPE launches for - # physical preparation; only q_a and per-head q use the generic op. - assert calls == { - "fused_attention": 1, - "cache_write": 1, - "draft_block": 1, - "rmsnorm_rope": 2, - } + def unexpected_host_op(*args, **kwargs): + pytest.fail("cached DSpark host wrapper used a Python tensor conversion or view") - with monkeypatch.context() as patch: - patch.setattr( - dspark_attention, - "is_fused_dsv4_dspark_attention_supported", - lambda *args: False, - ) - patch.setattr( - dspark_attention, - "is_fused_dspark_rmsnorm_rope_supported", - lambda *args: False, - ) - expected = dspark_attention.dspark_attention_forward_batched( - x, main_x, start_pos, fallback_cache, slots, valid_len, **kwargs - ) + monkeypatch.setattr(attention_op.cute.runtime, "from_dlpack", unexpected_host_op) + monkeypatch.setattr(torch.cuda, "current_stream", unexpected_host_op) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", unexpected_host_op) + for method in ("permute", "unsqueeze", "reshape"): + monkeypatch.setattr(torch.Tensor, method, unexpected_host_op) + output = attention_op.fused_dsv4_dspark_attention(**inputs) + assert output.shape == inputs["q"].shape - torch.testing.assert_close(actual, expected, rtol=8e-2, atol=1e-2) - torch.testing.assert_close(op_cache, fallback_cache, rtol=2e-2, atol=2e-2) +def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch): + inputs = _make_inputs(43, block=5) + monkeypatch.setattr(attention_op, "_dspark_attention_kernel_cache", {}) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) -def test_warmup_is_best_effort(monkeypatch): - import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op + def unexpected_compile(*args, **kwargs): + pytest.fail("compiler or allocation was called before the capture guard") - warnings = [] - monkeypatch.setattr(dspark_attention_op, "_get_dspark_arch_str", lambda: "sm_100") - monkeypatch.setattr( - dspark_attention_op, - "cute_dsl_dspark_rmsnorm_rope_cache_write", - lambda *args: (_ for _ in ()).throw(RuntimeError("synthetic compile failure")), - ) - monkeypatch.setattr(dspark_attention_op.logger, "warning", warnings.append) + monkeypatch.setattr(attention_op, "_compile_dspark_attention", unexpected_compile) + monkeypatch.setattr(torch, "empty_like", unexpected_compile) + with pytest.raises(RuntimeError, match="must run eagerly before graph capture"): + attention_op.fused_dsv4_dspark_attention(**inputs) - dspark_attention_op.warmup_fused_dsv4_dspark_attention(5, 1e-6) - assert len(warnings) == 1 - assert "self-JIT on first use" in warnings[0] - assert "synthetic compile failure" in warnings[0] +def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch): + preparation_op._compile_dspark_rmsnorm_rope_page_write.cache_clear() + preparation_op._compile_dspark_rmsnorm_rope_draft_block.cache_clear() + device = torch.cuda.current_device() + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + def unexpected_compile(*args, **kwargs): + pytest.fail("preparation compiler was called during CUDA graph capture") -def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(): - from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( - _dspark_attention_kernel_cache, - _get_dspark_arch_str, - fused_dsv4_dspark_attention, - is_dsv4_dspark_attention_config_supported, - ) + monkeypatch.setattr(preparation_op.cute, "compile", unexpected_compile) + with pytest.raises(RuntimeError, match="page-write must be warmed up"): + preparation_op._compile_dspark_rmsnorm_rope_page_write(32, 1e-6, device) + with pytest.raises(RuntimeError, match="draft-block must be warmed up"): + preparation_op._compile_dspark_rmsnorm_rope_draft_block(5, 1e-6, device) - assert [_get_dspark_arch_str(sm) for sm in (100, 103)] == [ - "sm_100", - "sm_103", - ] - assert _get_dspark_arch_str(101) is None - assert _get_dspark_arch_str(109) is None - assert _get_dspark_arch_str() in ("sm_100", "sm_103") - assert is_dsv4_dspark_attention_config_supported(5, 128, 512, 128) - assert not is_dsv4_dspark_attention_config_supported(4, 128, 512, 128) - assert not is_dsv4_dspark_attention_config_supported(5, 24, 512, 128) - - _dspark_attention_kernel_cache.clear() - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs( - 11, start_pos_values=[300, 4, 250], cache_pages=40 - ) - scale = q.shape[-1] ** -0.5 - valid_len = _legacy_valid_len(start_pos) - - expected, expected_cache = _reference( - q, - main_kv, - block_kv, - kv_cache, - slots, - start_pos, - valid_len, - sink, - torch.view_as_complex(freqs), - scale, - ) - # A missing key self-JITs through the production op. - draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( - main_kv, block_kv, kv_cache, slots, start_pos - ) - actual = fused_dsv4_dspark_attention( - q, - draft_block, - kv_cache, - slots_i32, - cache_seqs, - valid_len, - sink, - freqs, - scale, - ) - torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) - torch.testing.assert_close(kv_cache, expected_cache, rtol=0, atol=0) - assert len(_dspark_attention_kernel_cache) == 1 - # Every runtime batch, page count, page stride, and softmax scale is a hot - # hit on the same compiled object without compiler calls or runtime padding. +def test_preparation_self_jit_covers_dynamic_batches(): + preparation_op._compile_dspark_rmsnorm_rope_page_write.cache_clear() + preparation_op._compile_dspark_rmsnorm_rope_draft_block.cache_clear() + weight = torch.ones(512, device="cuda", dtype=torch.bfloat16) for batch in (1, 3, 32): - values = [5 + (37 * i) % 386 for i in range(batch)] - args = _make_inputs(17 + batch, start_pos_values=values, cache_pages=batch + 41) - q_b, main_b, block_b, cache_b, slots_b, pos_b, sink_b, blk_freqs_b, freqs_b = args - if batch == 3: - cache_b = cache_b.clone() - assert cache_b.is_contiguous() - runtime_scale = scale if batch % 2 else scale * 0.5 - valid_len_b = _legacy_valid_len(pos_b) - expected, expected_cache = _reference( - q_b, - main_b, - block_b, - cache_b, - slots_b, - pos_b, - valid_len_b, - sink_b, - blk_freqs_b, - runtime_scale, + inputs = _make_inputs(31 + batch, batch=batch) + main = torch.randn(batch, 1, 512, device="cuda", dtype=torch.bfloat16) + block = inputs["draft_block"][:, :6].contiguous() + freqs = dspark.precompute_dspark_freqs_cis(_ROPE_DIM, 1024, device="cuda") + main_freqs = freqs[inputs["positions"]][:, None] + block_freqs = torch.view_as_complex(inputs["inverse_rope_freqs"]) + expected_cache = _write_reference( + inputs["kv_pages"], _norm_rope(main, weight, main_freqs), inputs ) - draft_block_b, slots_i32_b, cache_seqs_b = _prepare_attention_inputs( - main_b, block_b, cache_b, slots_b, pos_b + expected_block = _norm_rope(block, weight, block_freqs) + preparation_op.cute_dsl_dspark_rmsnorm_rope_page_write( + main, + weight, + torch.view_as_real(main_freqs).view(batch, 32, 2), + inputs["kv_pages"], + inputs["block_tables"], + inputs["positions"], + inputs["capacities"], + 1e-6, ) - actual = fused_dsv4_dspark_attention( - q_b, - draft_block_b, - cache_b, - slots_i32_b, - cache_seqs_b, - valid_len_b, - sink_b, - freqs_b, - runtime_scale, + draft = preparation_op.cute_dsl_dspark_rmsnorm_rope_draft_block( + block, weight, inputs["inverse_rope_freqs"].view(-1, 32, 2), 1e-6 ) - torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) - torch.testing.assert_close(cache_b, expected_cache, rtol=0, atol=0) - assert len(_dspark_attention_kernel_cache) == 1 + torch.testing.assert_close(inputs["kv_pages"], expected_cache, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(draft[:, :6], expected_block, rtol=2e-2, atol=2e-2) + assert torch.count_nonzero(draft[:, 6:]) == 0 + assert preparation_op._compile_dspark_rmsnorm_rope_page_write.cache_info().misses == 1 + assert preparation_op._compile_dspark_rmsnorm_rope_draft_block.cache_info().misses == 1 + + +@pytest.mark.parametrize("page_size", (128, 256)) +def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, page_size): + g = _make_attn_inputs(seed=17, device="cuda", page_size=page_size) + for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): + g[name].mul_(0.2) + for name in ("x", "main_x", "kv_cache0"): + g[name].mul_(0.1) + before = g["kv_cache0"].clone() + actual = _run(g) + actual_cache = g["kv_cache0"].clone() + g["kv_cache0"].copy_(before) + with monkeypatch.context() as patch: + patch.setattr(torch.ops.trtllm, "fused_dsv4_dspark_attention", _reference) + patch.setattr(dspark, "is_fused_dspark_rmsnorm_rope_supported", lambda *a: False) + expected = _run(g) + torch.testing.assert_close(actual, expected, rtol=8e-2, atol=1e-2) + torch.testing.assert_close(actual_cache, g["kv_cache0"], rtol=2e-2, atol=2e-2) -def test_compile_without_real_specimens_and_cached_wrapper_avoids_views(monkeypatch): - import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op +def test_warmup_propagates_compile_failure(monkeypatch): + inputs = _make_inputs(page_size=256) + monkeypatch.setattr(attention_op, "_dspark_attention_kernel_cache", {}) - dspark_attention_op._dspark_attention_kernel_cache.clear() - arch_str = dspark_attention_op._get_dspark_arch_str() - assert arch_str is not None + def fail_compile(*args): + raise RuntimeError("synthetic compile failure") - def unexpected_compile_op(*args, **kwargs): - del args, kwargs - pytest.fail("DSpark compilation used a real tensor specimen") + monkeypatch.setattr(attention_op, "_compile_dspark_attention", fail_compile) + with pytest.raises(RuntimeError, match="synthetic compile failure"): + attention_op.fused_dsv4_dspark_attention(**inputs) - # Compile before runtime inputs exist. Compile-only pointers and the fake - # output anchor must not allocate or wrap a PyTorch tensor. - with monkeypatch.context() as patch: - patch.setattr(torch, "empty", unexpected_compile_op) - patch.setattr(torch, "empty_like", unexpected_compile_op) - patch.setattr(dspark_attention_op.cute.runtime, "from_dlpack", unexpected_compile_op) - compiled = dspark_attention_op._compile_dspark_attention(5, arch_str) - dspark_attention_op._dspark_attention_kernel_cache[(5, arch_str)] = compiled - args = _make_inputs(41, block=5, start_pos_values=[257, 9]) - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = args - draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( - main_kv, block_kv, kv_cache, slots, start_pos +def _attention_case(seed=0, start_positions=(40, 40), lengths=None, sink=None): + inputs = _make_inputs( + seed, block=5, page_size=256, positions=start_positions, valid_lengths=lengths ) - valid_len = _legacy_valid_len(start_pos) - scale = q.shape[-1] ** -0.5 + if sink is not None: + inputs["attn_sink"].fill_(sink) + inputs["inverse_rope_freqs"].zero_() + inputs["inverse_rope_freqs"][..., 0] = 1 + output = attention_op.fused_dsv4_dspark_attention(**inputs) + histories = [] + pages, tables = inputs["kv_pages"], inputs["block_tables"] + for row, (end, length) in enumerate(zip(start_positions, inputs["valid_lengths"].tolist())): + logical = torch.arange(end - length, end, device="cuda") + history = pages[tables[row, logical // pages.shape[1]].long(), logical % pages.shape[1]] + histories.append(torch.cat((history, inputs["draft_block"][row, :5]))) + return output, inputs["q"], histories, inputs["attn_sink"] + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_sparse_attn_matches_loop_reference(seed): + output, q, histories, sink = _attention_case(seed) + expected = torch.empty_like(output) + for row, kv in enumerate(histories): + for token in range(q.shape[1]): + scores = q[row, token].float() @ kv.float().T * 512**-0.5 + weights = torch.softmax(torch.cat((scores, sink[:, None]), dim=-1), dim=-1) + expected[row, token] = (weights[:, :-1] @ kv.float()).to(output.dtype) + torch.testing.assert_close(output, expected, rtol=5e-2, atol=3e-2) + + +def test_sparse_attn_no_sink_matches_sdpa(): + output, q, histories, _ = _attention_case(sink=float("-inf")) + kv = torch.stack(histories).unsqueeze(1).expand(-1, 128, -1, -1) + expected = F.scaled_dot_product_attention(q.transpose(1, 2), kv, kv).transpose(1, 2) + torch.testing.assert_close(output, expected, rtol=5e-2, atol=3e-2) + + +def test_sparse_attn_sink_reduces_mass(): + no_sink, *_ = _attention_case(start_positions=(4,), sink=float("-inf")) + with_sink, *_ = _attention_case(start_positions=(4,), sink=0.0) + assert with_sink.abs().sum() < no_sink.abs().sum() + + +def test_sparse_attn_masked_indices_excluded(): + masked, q, histories, _ = _attention_case(start_positions=(5,), lengths=[3], sink=float("-inf")) + kv = histories[0][None, None].expand(1, 128, -1, -1) + expected = F.scaled_dot_product_attention(q.transpose(1, 2), kv, kv).transpose(1, 2) + torch.testing.assert_close(masked, expected, rtol=5e-2, atol=3e-2) + full, *_ = _attention_case(start_positions=(5,), lengths=[4], sink=float("-inf")) + assert not torch.allclose(full, masked, rtol=1e-3, atol=1e-3) - def unexpected_host_op(*args, **kwargs): - del args, kwargs - pytest.fail("cached DSpark host wrapper used a Python tensor conversion or view") - - monkeypatch.setattr(dspark_attention_op.cute.runtime, "from_dlpack", unexpected_host_op) - monkeypatch.setattr(torch.cuda, "current_stream", unexpected_host_op) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", unexpected_host_op) - monkeypatch.setattr(torch.Tensor, "permute", unexpected_host_op) - monkeypatch.setattr(torch.Tensor, "unsqueeze", unexpected_host_op) - monkeypatch.setattr(torch.Tensor, "reshape", unexpected_host_op) - output = dspark_attention_op._run_dspark_attention( - q, - draft_block, - kv_cache, - slots_i32, - cache_seqs, - valid_len, - sink, - freqs, - scale, +@pytest.mark.parametrize( + "start_pos,page_size,block", [(1, 128, 5), (3, 64, 5), (10, 16, 5), (200, 256, 6)] +) +def test_page_mapping_matches_reference(start_pos, page_size, block): + from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context + + inputs = _make_inputs(block=block, positions=[start_pos] * 3, page_size=page_size) + values = inputs["draft_block"][:, :1].contiguous() + expected = _write_reference(inputs["kv_pages"], values, inputs) + write_dspark_context( + values, + inputs["positions"][:, None], + torch.ones(3, 1, device="cuda", dtype=torch.bool), + inputs["kv_pages"], + inputs["block_tables"], + inputs["capacities"], + ) + torch.testing.assert_close(inputs["kv_pages"], expected, rtol=0, atol=0) + + +def test_page_mapping_masks_non_generation_rows(): + from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context + + values = torch.randn(1, 1, 16, device="cuda") + pool = torch.randn(1, 128, 16, device="cuda") + before = pool.clone() + write_dspark_context( + values, + torch.zeros(1, 1, device="cuda", dtype=torch.long), + torch.ones(1, 1, device="cuda", dtype=torch.bool), + pool, + torch.zeros(1, 1, device="cuda", dtype=torch.int32), + torch.zeros(1, device="cuda", dtype=torch.long), ) - assert output.shape == q.shape + torch.testing.assert_close(pool, before, rtol=0, atol=0) -def test_attention_cache_miss_rejects_cuda_graph_capture(monkeypatch): - import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op +def _make_attn_inputs(seed=0, device="cuda", page_size=256): + """Small synthetic DSpark attention inputs/weights (CPU bf16).""" + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 128, 512, 64 + q_lora, o_lora, n_groups = 64, 8, 8 + window, block, start_pos = 128, 5, 200 + b = 2 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + start_pos=start_pos, + b=b, + eps=1e-6, + softmax_scale=head_dim**-0.5, + ) + bf = torch.bfloat16 + g["x"] = torch.randn(b, block, dim, dtype=bf) + g["main_x"] = torch.randn(b, 1, dim, dtype=bf) + width = (start_pos + page_size - 1) // page_size + g["kv_cache0"] = torch.randn(b * width + 1, page_size, head_dim, dtype=bf) + g["tables"] = torch.arange(b * width, 0, -1, dtype=torch.int32).reshape(b, width) + g["capacities"] = torch.full((b,), width * page_size, dtype=torch.long) + g["valid_len"] = torch.full((b,), min(start_pos, window), dtype=torch.long) + g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 + g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 + g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 + g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 + g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 + g["q_norm"] = torch.ones(q_lora, dtype=bf) + g["kv_norm"] = torch.ones(head_dim, dtype=bf) + g["attn_sink"] = torch.randn(n_heads) + g["freqs"] = dspark.precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) + g["start_pos"] = torch.full((b,), start_pos, dtype=torch.long) + return {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in g.items()} + + +def _run(g): + return dspark.dspark_attention_forward( + g["x"], + g["main_x"], + g["start_pos"], + g["kv_cache0"], + g["tables"], + g["capacities"], + g["valid_len"], + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) - dspark_attention_op._dspark_attention_kernel_cache.clear() - q, main_kv, block_kv, kv_cache, slots, start_pos, sink, _, freqs = _make_inputs(43, block=5) - draft_block, slots_i32, cache_seqs = _prepare_attention_inputs( - main_kv, block_kv, kv_cache, slots, start_pos + +def test_attention_forward_shape_and_determinism(): + g = _make_attn_inputs() + o = _run(g) + assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) + assert torch.isfinite(o.float()).all() + torch.testing.assert_close(o, _run(g)) # deterministic + + +def test_attention_forward_updates_only_mapped_kv(): + g = _make_attn_inputs() + before = g["kv_cache0"].clone() + _run(g) + changed = (g["kv_cache0"] != before).any(dim=-1) + expected = torch.zeros_like(changed) + for row, position in enumerate(g["start_pos"].tolist()): + logical = position - 1 + page_size = before.shape[1] + expected[int(g["tables"][row, logical // page_size]), logical % page_size] = True + assert torch.equal(changed, expected) + + +def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): + g = _make_attn_inputs(seed) + batch = len(start_positions) + page_size = g["kv_cache0"].shape[1] + width = (max(start_positions) + page_size - 1) // page_size + g["x"] = g["x"][:1].repeat(batch, 1, 1) + g["main_x"] = g["main_x"][:1].repeat(batch, 1, 1) + g["kv_cache0"] = torch.randn( + batch * width + 1, page_size, g["head_dim"], device="cuda", dtype=torch.bfloat16 ) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) - monkeypatch.setattr( - dspark_attention_op, - "_compile_dspark_attention", - lambda *args: pytest.fail("compiler was called during CUDA graph capture"), + g["tables"] = torch.arange(batch * width, 0, -1, device="cuda", dtype=torch.int32).reshape( + batch, width ) - monkeypatch.setattr( - torch, - "empty_like", - lambda *args, **kwargs: pytest.fail("output allocated before the capture guard"), - ) - - with pytest.raises(RuntimeError, match="must be warmed up before CUDA graph capture"): - dspark_attention_op._run_dspark_attention( - q, - draft_block, - kv_cache, - slots_i32, - cache_seqs, - _legacy_valid_len(start_pos), - sink, - freqs, - q.shape[-1] ** -0.5, + g["start_pos"] = torch.tensor(start_positions, device="cuda", dtype=torch.long) + g["capacities"] = torch.full((batch,), width * page_size, device="cuda", dtype=torch.long) + g["valid_len"] = g["start_pos"].clamp(max=128) + g["b"] = batch + return g + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_batched_attention_matches_scalar_per_request(seed): + g = _make_batched_inputs(seed) + before = g["kv_cache0"].clone() + actual = _run(g) + outputs = [] + for row in range(g["b"]): + single = dict(g) + for key in ("x", "main_x", "tables", "start_pos", "capacities", "valid_len"): + single[key] = g[key][row : row + 1].contiguous() + single["kv_cache0"] = before.clone() + outputs.append(_run(single)) + torch.testing.assert_close(actual, torch.cat(outputs), rtol=2e-2, atol=2e-2) + + +def test_batched_attention_writes_through_pages(): + g = _make_batched_inputs(seed=3) + before = g["kv_cache0"].clone() + _run(g) + changed = (g["kv_cache0"] != before).any(dim=-1) + expected = torch.zeros_like(changed) + for row, pos in enumerate(g["start_pos"].tolist()): + page_size = before.shape[1] + expected[int(g["tables"][row, (pos - 1) // page_size]), (pos - 1) % page_size] = True + assert torch.equal(changed, expected) + + +def test_batched_attention_dummy_rows_keep_pages(): + g = _make_batched_inputs(seed=4) + g["capacities"].zero_() + before = g["kv_cache0"].clone() + _run(g) + torch.testing.assert_close(g["kv_cache0"], before) + + +@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) +def test_batched_page_writes_match_per_request(start_positions): + from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context + + g = _make_batched_inputs(start_positions=start_positions) + values = torch.randn(g["b"], 1, g["head_dim"], device="cuda", dtype=torch.bfloat16) + mask = torch.ones(g["b"], 1, device="cuda", dtype=torch.bool) + expected = g["kv_cache0"].clone() + for row in range(g["b"]): + write_dspark_context( + values[row : row + 1], + g["start_pos"][row : row + 1, None], + mask[row : row + 1], + expected, + g["tables"][row : row + 1], + g["capacities"][row : row + 1], ) + write_dspark_context( + values, g["start_pos"][:, None], mask, g["kv_cache0"], g["tables"], g["capacities"] + ) + torch.testing.assert_close(g["kv_cache0"], expected, rtol=0, atol=0) -def test_preparation_cache_misses_reject_cuda_graph_capture(monkeypatch): - import tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op as preparation_op +def test_batched_attention_respects_partial_window_valid_len(): + actual, q, histories, sink = _attention_case( + start_positions=(3, 10, 200), lengths=[0, 3, 128], sink=0.0 + ) + for row, kv in enumerate(histories): + scores = torch.einsum("mhd,nd->mhn", q[row].float(), kv.float()) * 512**-0.5 + scores = torch.cat((scores, sink[None, :, None].expand(q.shape[1], -1, 1)), dim=-1) + probs = torch.softmax(scores, dim=-1)[..., :-1] + expected = torch.einsum("mhn,nd->mhd", probs, kv.float()).to(q.dtype) + torch.testing.assert_close(actual[row], expected, rtol=5e-2, atol=3e-2) + + +def test_batched_attention_cuda_graph_capture_replay(force_fallback): + g = _make_batched_inputs(seed=0) + eager = _run(g) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _run(g) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = _run(g) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, eager, rtol=2e-2, atol=2e-2) - preparation_op._compile_dspark_rmsnorm_rope_cache_write.cache_clear() - preparation_op._compile_dspark_rmsnorm_rope_draft_block.cache_clear() - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) - monkeypatch.setattr( - preparation_op.cute, - "compile", - lambda *args: pytest.fail("preparation compiler was called during CUDA graph capture"), + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="The fused DSpark FP8 output projection is a Rubin specialization", +) +def test_fused_dsv4_dspark_attention_fp8_matches_separate_quant_and_replays_graph(): + from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( + fused_dsv4_dspark_attention, + fused_dsv4_dspark_attention_fp8, ) - with pytest.raises(RuntimeError, match="cache-write must be warmed up"): - preparation_op._compile_dspark_rmsnorm_rope_cache_write(1e-6) - with pytest.raises(RuntimeError, match="draft-block must be warmed up"): - preparation_op._compile_dspark_rmsnorm_rope_draft_block(5, 1e-6) + groups = 16 + inputs = _make_inputs(47, page_size=256, positions=[257, 390], valid_lengths=[3, 5]) + q = inputs["q"] + block = q.shape[1] + + def separate_quantize(): + bf16 = fused_dsv4_dspark_attention(**inputs) + rows = q.shape[0] * block + grouped = bf16.reshape(rows, groups, -1).permute(1, 0, 2).contiguous() + quant_blocks = grouped.view(groups, rows, -1, 128) + amax = quant_blocks.abs().amax(dim=-1).float().clamp_min(1e-10) + quant_scale = 448.0 / amax + quantized = (quant_blocks.float() * quant_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + scale_rows = (rows + 3) // 4 * 4 + scales = torch.zeros( + groups, + quant_blocks.shape[2], + scale_rows, + device=q.device, + dtype=torch.float32, + ) + scales[..., :rows] = (1.0 / quant_scale).transpose(1, 2) + return quantized.view_as(grouped), scales + + def assert_quantized_close(actual_fp8, actual_scale, expected_fp8, expected_scale): + rows = q.shape[0] * block + torch.testing.assert_close( + actual_scale[..., :rows], expected_scale[..., :rows], rtol=1e-6, atol=0 + ) + actual = actual_fp8.float() + expected = expected_fp8.float() + mismatch_rate = torch.count_nonzero(actual != expected) / actual.numel() + assert mismatch_rate <= 1e-3 + # CuTe and Torch can choose neighboring E4M3 values at a rounding tie. + torch.testing.assert_close(actual, expected, rtol=0.125, atol=0.015625) + actual_fp8, actual_scale = fused_dsv4_dspark_attention_fp8(**inputs, output_groups=groups) + expected_fp8, expected_scale = separate_quantize() + assert_quantized_close(actual_fp8, actual_scale, expected_fp8, expected_scale) -def test_preparation_self_jit_covers_dynamic_batches(): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_fp8, captured_scale = fused_dsv4_dspark_attention_fp8( + **inputs, output_groups=groups + ) + q.copy_(torch.randn_like(q)) + inputs["valid_lengths"].copy_(torch.tensor([7, 2], device=q.device)) + graph.replay() + expected_fp8, expected_scale = separate_quantize() + assert_quantized_close(captured_fp8, captured_scale, expected_fp8, expected_scale) + + +@pytest.mark.skipif( + get_sm_version() != 107 or not IS_CUTLASS_DSL_INTERNAL_AVAILABLE, + reason="The fused DSpark FP8 input producers require SM107 Rubin CuTe DSL", +) +def test_dspark_fp8_input_producers_match_reference_and_reuse_dynamic_m_compile(): from tensorrt_llm._torch.custom_ops.dspark_rmsnorm_rope_custom_op import ( - _compile_dspark_rmsnorm_rope_cache_write, - _compile_dspark_rmsnorm_rope_draft_block, - cute_dsl_dspark_rmsnorm_rope_cache_write, - cute_dsl_dspark_rmsnorm_rope_draft_block, + _compile_dspark_concat_fp8_quant, + _compile_dspark_rmsnorm_fp8_quant, + cute_dsl_dspark_concat_fp8_quant, + cute_dsl_dspark_rmsnorm_fp8_quant, + precompile_dspark_fp8_input_fusions, ) + from tensorrt_llm._torch.models.modeling_dspark import _rmsnorm + + torch.manual_seed(53) + hidden_dim = 1536 + eps = 1e-6 + weight = torch.randn(hidden_dim, device="cuda", dtype=torch.bfloat16) + _compile_dspark_concat_fp8_quant.cache_clear() + _compile_dspark_rmsnorm_fp8_quant.cache_clear() + + def active_scales(scale: torch.Tensor, rows: int) -> torch.Tensor: + offsets = [] + quant_blocks = hidden_dim // 128 + for row in range(rows): + for quant_block in range(quant_blocks): + base = ( + (row // 128) * quant_blocks * 512 + + quant_block * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + ) + offsets.extend(range(base, base + 4)) + return scale[torch.tensor(offsets, device=scale.device, dtype=torch.long)] + + missing = torch.randn(2, hidden_dim, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="not precompiled"): + cute_dsl_dspark_concat_fp8_quant(missing[:1], missing[1:]) + with pytest.raises(RuntimeError, match="not precompiled"): + cute_dsl_dspark_rmsnorm_fp8_quant(missing, weight, eps) + + _compile_dspark_concat_fp8_quant.cache_clear() + _compile_dspark_rmsnorm_fp8_quant.cache_clear() + precompile_dspark_fp8_input_fusions(hidden_dim, hidden_dim, eps) + assert _compile_dspark_concat_fp8_quant.cache_info().misses == 1 + assert _compile_dspark_rmsnorm_fp8_quant.cache_info().misses == 1 + + for first_rows, second_rows in ((1, 1), (3, 10), (64, 65)): + first_storage = torch.randn( + first_rows, hidden_dim + 128, device="cuda", dtype=torch.bfloat16 + ) + second_storage = torch.randn( + second_rows, hidden_dim + 256, device="cuda", dtype=torch.bfloat16 + ) + first = first_storage[:, :hidden_dim] + second = second_storage[:, :hidden_dim] + if first_rows > 1: + assert not first.is_contiguous() + if second_rows > 1: + assert not second.is_contiguous() + + expected_input = torch.cat((first, second), dim=0) + expected_fp8, expected_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + expected_input + ) + actual_fp8, actual_scale = cute_dsl_dspark_concat_fp8_quant(first, second) + torch.testing.assert_close(actual_fp8.float(), expected_fp8.float(), rtol=0, atol=0) + torch.testing.assert_close( + active_scales(actual_scale, first_rows + second_rows), + active_scales(expected_scale, first_rows + second_rows), + rtol=0, + atol=0, + ) - _compile_dspark_rmsnorm_rope_cache_write.cache_clear() - _compile_dspark_rmsnorm_rope_draft_block.cache_clear() - q, main_kv, block_kv, kv_cache, slots, start_pos, _, _, _ = _make_inputs( - 23, start_pos_values=[300, 4, 250], cache_pages=40 - ) - weight = torch.ones(512, device=q.device, dtype=q.dtype) - main_freqs = torch.zeros(q.shape[0], _ROPE_DIM // 2, 2, device=q.device) - block_freqs = torch.zeros(q.shape[0] * q.shape[1], _ROPE_DIM // 2, 2, device=q.device) - main_freqs[..., 0] = 1 - block_freqs[..., 0] = 1 + norm_storage = torch.randn( + first_rows + second_rows, + hidden_dim + 128, + device="cuda", + dtype=torch.bfloat16, + ) + norm_input = norm_storage[:, :hidden_dim] + expected_norm = _rmsnorm(norm_input, weight, eps) + expected_fp8, expected_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + expected_norm + ) + actual_fp8, actual_scale = cute_dsl_dspark_rmsnorm_fp8_quant(norm_input, weight, eps) + torch.testing.assert_close(actual_fp8.float(), expected_fp8.float(), rtol=0, atol=0) + torch.testing.assert_close( + active_scales(actual_scale, first_rows + second_rows), + active_scales(expected_scale, first_rows + second_rows), + rtol=0, + atol=0, + ) - for batch in (1, 3, 32): - args = _make_inputs(31 + batch, batch=batch, cache_pages=40) - q_b, main_b, block_b, cache_b, slots_b, pos_b, _, _, _ = args - main_freqs_b = torch.zeros(batch, _ROPE_DIM // 2, 2, device=q.device) - block_freqs_b = torch.zeros(batch * q_b.shape[1], _ROPE_DIM // 2, 2, device=q.device) - main_freqs_b[..., 0] = 1 - block_freqs_b[..., 0] = 1 - slots_i32, cache_seqs = cute_dsl_dspark_rmsnorm_rope_cache_write( - main_b.unsqueeze(1), - weight, - main_freqs_b, - cache_b, - slots_b.long(), - pos_b.long(), - 1e-6, + # Runtime row counts and row strides reuse the two initialization artifacts. + assert _compile_dspark_concat_fp8_quant.cache_info().misses == 1 + assert _compile_dspark_rmsnorm_fp8_quant.cache_info().misses == 1 + + +@pytest.mark.skipif( + get_sm_version() != 107 or not IS_CUTLASS_DSL_INTERNAL_AVAILABLE, + reason="The fused DSpark output projection requires SM107 Rubin CuTe DSL", +) +def test_dspark_fused_fp8_output_projection_matches_separate_and_reuses_cache(): + from _torch.helpers import per_block_cast_to_fp8, per_block_cast_to_fp8_e8m0 + + from tensorrt_llm._torch.autotuner import autotune + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLFp8RubinBmmQuantRunner, + CuteDSLMXFP8RubinLinear, + ) + from tensorrt_llm._torch.models.modeling_dspark import _output_projection + from tensorrt_llm.quantization.utils import fp8_utils + + torch.manual_seed(41) + batch, block, groups = 2, 6, 8 + rank, group_dim, hidden = 128, 512, 128 + rows, flat_n = batch * block, groups * rank + o = ( + torch.randn( + batch, + block, + groups, + group_dim, + device="cuda", + dtype=torch.bfloat16, ) - draft_block = cute_dsl_dspark_rmsnorm_rope_draft_block(block_b, weight, block_freqs_b, 1e-6) - torch.testing.assert_close(slots_i32, slots_b, rtol=0, atol=0) - torch.testing.assert_close(cache_seqs, pos_b, rtol=0, atol=0) - assert draft_block.shape == (batch, 8, 512) - assert torch.count_nonzero(draft_block[:, q_b.shape[1] :]) == 0 - - assert _compile_dspark_rmsnorm_rope_cache_write.cache_info().misses == 1 - assert _compile_dspark_rmsnorm_rope_draft_block.cache_info().misses == 1 + / group_dim + ) + o_a = torch.randn(groups, rank, group_dim, device="cuda", dtype=torch.bfloat16) / group_dim + o_a_fp8 = torch.empty_like(o_a, dtype=torch.float8_e4m3fn) + o_a_scale = torch.empty( + groups, rank // 128, group_dim // 128, device="cuda", dtype=torch.float32 + ) + for group in range(groups): + o_a_fp8[group], o_a_scale[group] = per_block_cast_to_fp8(o_a[group]) + + wo_b = torch.randn(hidden, flat_n, device="cuda", dtype=torch.bfloat16) / flat_n + wo_b_fp8, wo_b_scale = per_block_cast_to_fp8_e8m0(wo_b) + wo_b_scale = fp8_utils.transform_k128_scales_to_cutedsl_mxfp8_layout( + wo_b_scale, mn=hidden, k=flat_n + ) + + class OutputLinear: + has_fp8_block_scales = True + use_cute_dsl_blockscaling_mm = True + disable_deep_gemm = False + in_features = flat_n + weight = wo_b_fp8 + weight_scale = wo_b_scale + + def __call__(self, input): + if isinstance(input, tuple): + activation, scale = input + else: + activation, scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(input) + return torch.ops.trtllm.cute_dsl_mxfp8_gemm_rubin( + activation, self.weight, scale, self.weight_scale + ) + + output_linear = OutputLinear() + args = dict( + wo_a=None, + wo_b=output_linear, + n_groups=groups, + o_lora_rank=rank, + o_a_proj=o_a_fp8, + o_a_scale=o_a_scale, + o_a_dequant=None, + use_fused_output_projection=True, + ) + with autotune(skip_dynamic_tuning_buckets=True): + actual = _output_projection(o, **args) + cache_sizes = ( + len(CuteDSLFp8RubinBmmQuantRunner.kernel_cache), + len(CuteDSLMXFP8RubinLinear.kernel_cache), + ) + replay = _output_projection(o, **args) + assert cache_sizes[0] > 0 and cache_sizes[1] > 0 + assert cache_sizes == ( + len(CuteDSLFp8RubinBmmQuantRunner.kernel_cache), + len(CuteDSLMXFP8RubinLinear.kernel_cache), + ) + torch.testing.assert_close(replay, actual, rtol=0, atol=0) + + grouped_o = o.reshape(rows, groups, group_dim) + grouped_o_fp8, grouped_o_scale = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102( + grouped_o + ) + assert grouped_o_fp8.shape == (groups, rows, group_dim) + o_lora = torch.empty(groups, rows, rank, device="cuda", dtype=torch.bfloat16) + torch.ops.trtllm.cute_dsl_fp8_bmm_rubin( + grouped_o_fp8, o_a_fp8, grouped_o_scale, o_a_scale, o_lora + ) + expected = output_linear(o_lora.transpose(0, 1).reshape(rows, flat_n)) + expected = expected.view(batch, block, hidden) + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=3e-2) diff --git a/tests/unittest/disaggregated/region/test_aux.py b/tests/unittest/disaggregated/region/test_aux.py index 7c2528fbf81c..2571b5869084 100644 --- a/tests/unittest/disaggregated/region/test_aux.py +++ b/tests/unittest/disaggregated/region/test_aux.py @@ -7,6 +7,8 @@ AuxBuffer, AuxBufferMeta, AuxSlot, + _decode_draft_history, + _encode_draft_history, build_aux_transfer_layout, ) @@ -185,3 +187,23 @@ def test_aux_transfer_layout_ctx_no_spec_gen_sa(): np.testing.assert_array_equal(layout.dst_item_sizes, gen.meta.item_sizes[expected_keep]) np.testing.assert_array_equal(layout.src_base_ptrs, ctx.meta.ptrs[expected_keep]) np.testing.assert_array_equal(layout.dst_base_ptrs, gen.meta.ptrs[expected_keep]) + + +@pytest.mark.parametrize("dtype", ["torch.bfloat16", "torch.float8_e4m3fn"]) +def test_draft_history_record_round_trips_the_storage_dtype(dtype): + """Prefill and generation servers compare the decoded layout, dtype included.""" + history = { + "valid_length": 89, + "position": 89, + "layout": { + "num_layers": 5, + "num_kv_heads": 16, + "head_dim": 64, + "dtype": dtype, + "attention_backend": "VANILLA", + "kv_factor": 2, + "window_size": None, + }, + } + + assert _decode_draft_history(_encode_draft_history(history)) == history diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec6ef42ad0f7..b8d1f656bc2d 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -314,6 +314,7 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l impl=impl, pp_layers=[0, 1], num_kv_heads_per_layer=[1, 1], + cp_as_tp_layer_ids=frozenset(), get_disagg_role_mapper_kinds=lambda: role_mapper_kinds, get_disagg_role_layouts=lambda: {}, ) From a5329b6d200bfc99e6ea224a63baa0c5d1e51c2b Mon Sep 17 00:00:00 2001 From: Lizhi Zhou Date: Thu, 15 Oct 2026 23:51:08 -0700 Subject: [PATCH 05/16] [None][perf] Split PARALLEL autotuning across Helix CP ranks Signed-off-by: Lizhi Zhou Signed-off-by: Mingyang Hao --- tensorrt_llm/_torch/autotuner.py | 12 ++++++++++++ tests/unittest/_torch/misc/test_autotuner.py | 17 +++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index f55006af4213..e47c6fa71b19 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -2140,6 +2140,18 @@ def _maybe_parallelize_tactics( # Each rank of the tuning group only tunes its share of tactics; # _merge_cache_data then keeps the fastest across the group. group_size, group_index = self._tuning_group() + # Sharding a candidate list shorter than the group size by + # idx % group_size leaves the ranks with + # group_index >= len(all_valid_tactics) with an EMPTY shard. Those + # ranks then find no valid (runner, tactic) pair and emit a + # spurious "No valid runner/tactic was found" warning, while also + # diverging in tuning wall-time from the ranks that did get work + # (a per-rank desync hazard for the following collective). + # Parallelizing such a short list saves almost nothing, so have + # every rank profile the full list instead; the subsequent + # _merge_cache_data allgather still keeps the min-time winner. + if len(all_valid_tactics) < group_size: + return all_valid_tactics return [ tactic for idx, tactic in enumerate(all_valid_tactics) if idx % group_size == group_index diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 7b7d44422ff6..dbfeb4807253 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -1791,6 +1791,23 @@ def test_parallel_tactics_split_skips_non_helix_cp_ranks(): tactics, DistributedTuningStrategy.PARALLEL) == tactics +def test_parallel_tactics_short_list_not_split_across_helix_cp_ranks(): + """A tactic list shorter than the Helix TP*CP group is not sharded.""" + from tensorrt_llm.mapping import CpType + + tactics = list(range(3)) + for rank in range(4): + tuner = autotuner.AutoTuner() + mapping = Mapping(world_size=4, + rank=rank, + tp_size=1, + cp_size=4, + cp_config={"cp_type": CpType.HELIX}) + tuner.setup_distributed_state(mapping, dist=object()) + assert tuner._maybe_parallelize_tactics( + tactics, DistributedTuningStrategy.PARALLEL) == tactics + + def test_profiling_cache_enum_tactic_roundtrip(tmp_path): # An enum tactic must survive save -> load: repr("") isn't # ast.literal_eval-parsable, so the cache serializes tactic.value instead. From b91aec3d9a08f955a6239ba5effe7d9891810932 Mon Sep 17 00:00:00 2001 From: Xin Guan Date: Mon, 19 Oct 2026 04:19:20 -0700 Subject: [PATCH 06/16] [None][fix] Helix all-to-all: fence shared memory before the async-proxy bulk copies Signed-off-by: Xin Guan Signed-off-by: Mingyang Hao --- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 11 +- tensorrt_llm/_torch/attention/mla.py | 45 ++++- .../_torch/models/modeling_kimi_linear.py | 12 ++ .../attention/multi_gpu/test_mla_helix.py | 185 ++++++++++++++++++ 4 files changed, 248 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index 48e1d6ff552c..9f9ed7fac6fd 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -474,7 +474,10 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) fifoEntry128ByteIndexBase += singleProtoTransfer128ByteCount; // ensure that we can over-write shmem in next iteration - // (it must be fully read by all threads when doing S2G above) + // (it must be fully read by all threads when doing S2G above). + // The next iteration refills shmem with a bulk copy, which must + // also be ordered after the sanitize and protoPack stores above. + tensorrt_llm::common::fence_view_async_shared(); __syncwarp(); } if (fifoEntry128ByteIndexBase > 0) @@ -539,9 +542,9 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) LL128Proto::protoUnpack( shmem, tail, singlePacked128ByteCount, fifoEntry128ByteIndexBase, loaded128ByteCount, laneId); - // protoUnpack restores the words displaced by the LL128 flags through the generic proxy. Make - // those writes visible to the async proxy before cp.async.bulk reads the same buffer for S2G; - // otherwise the store can copy a flag (the FIFO step) in place of a restored data word. + // protoUnpack restores the u64 under each line's flag with ordinary + // stores, and the bulk store below reads shared memory through the + // async proxy: without this fence it can still read the flag value. tensorrt_llm::common::fence_view_async_shared(); __syncwarp(); diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index 328bb37accbf..4568814542a7 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -824,8 +824,15 @@ def _attn_forward_gen( v: torch.Tensor, position_ids: Optional[torch.Tensor], attn_metadata: AttentionMetadata, + helix_v_up: bool = False, **kwargs, ): + """Run generation attention; under Helix CP also exchange and combine partials. + + With ``helix_v_up`` the V up-projection of every head is applied to this + rank's partial outputs before the all-to-all, so the result holds + ``num_heads_tp_cp * v_head_dim`` values per token instead of latents. + """ if self.mapping.has_cp_helix(): # partial_o: [num_tokens, num_heads_tp * kv_lora_rank] # softmax_stats: [num_tokens, num_heads_tp, 2] @@ -867,12 +874,16 @@ def _attn_forward_gen( seq_start=attn_metadata.num_contexts, num_seqs=attn_metadata.num_generations, ) + value_dim = kv_lora_rank + if helix_v_up: + partial_o = self._helix_project_v_partials(partial_o) + value_dim = self.v_head_dim return _helix_post_process( partial_o, softmax_stats, self.mapping, self.num_heads_tp_cp, - kv_lora_rank, + value_dim, self.aux_stream, self.ln_events, zero_kv_mask=zero_kv_mask, @@ -1561,6 +1572,30 @@ def forward_context( q, compressed_kv, k_pe, position_ids, attn_metadata, output, latent_cache ) + def _helix_v_up_enabled(self) -> bool: + """Whether Helix generation projects V before the CP all-to-all.""" + return ( + self.mapping.has_cp_helix() and getattr(self, "helix_v_b_proj_full", None) is not None + ) + + def _helix_project_v_partials(self, partial_o: torch.Tensor) -> torch.Tensor: + """Apply every head's V up-projection to this CP rank's partial outputs. + + The CP combine is a per-head weighted sum of the ranks' partial outputs, + so it commutes with the per-head linear V projection. Projecting first + shrinks each head's all-to-all payload from ``kv_lora_rank`` to + ``v_head_dim`` values. ``helix_v_b_proj_full`` is the BF16 + ``[num_heads_tp, v_head_dim, kv_lora_rank]`` weight of all heads. + """ + num_tokens = partial_o.shape[0] + partial_v = partial_o.new_empty((num_tokens, self.num_heads_tp, self.v_head_dim)) + torch.ops.trtllm.bmm_out( + partial_o.view(num_tokens, self.num_heads_tp, self.kv_lora_rank).transpose(0, 1), + self.helix_v_b_proj_full.transpose(1, 2), + partial_v.transpose(0, 1), + ) + return partial_v.view(num_tokens, self.num_heads_tp * self.v_head_dim) + def _bmm_bf16_out( self, a, b_no_transpose, b_transposed, output, locality_domain_weight_shards=None ): @@ -1761,6 +1796,7 @@ def _mla_gen_rope(): else: raise RuntimeError("MLA absorption requires k_b_proj_trans") + helix_v_up = self._helix_v_up_enabled() attn_out_latent = self._attn_forward_gen( self.mqa, fused_q, @@ -1768,6 +1804,7 @@ def _mla_gen_rope(): None, position_ids, attn_metadata, + helix_v_up=helix_v_up, attention_input_type=AttentionInputType.generation_only, out_scale=self.out_scale, latent_cache=latent_cache, # kvcache and k_pe @@ -1782,6 +1819,12 @@ def _mla_gen_rope(): ) fused_q = None + if helix_v_up: + # The V up-projection already ran before the all-to-all. + assert attn_out_latent.shape == (q.shape[0], self.num_heads_tp_cp * self.v_head_dim) + output.view(attn_out_latent.shape).copy_(attn_out_latent) + 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] diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 8ad22f849c4b..81443f6c18e6 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -135,6 +135,11 @@ from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently +# Helix MLA generation applies the V up-projection of all heads before the CP +# all-to-all (see ``MLA._helix_project_v_partials``), which needs every head's +# BF16 V weight on each rank. +_K3_HELIX_V_UP_BEFORE_A2A = os.environ.get("TRTLLM_HELIX_V_UP_BEFORE_A2A", "0") == "1" + # A/B escape hatch: restore nn.Linear for the K3 latent MoE projections # instead of the min-latency fused GEMM op (read once at import). _K3_DISABLE_MIN_LATENCY_LATENT_PROJ = ( @@ -2979,6 +2984,11 @@ def _load_mla_fp8_kv_b(attention: nn.Module, weight: torch.Tensor, scale: torch. ) attention.k_b_proj_trans.data.copy_(key_weight.transpose(1, 2)) attention.k_b_proj_trans_scale.data.copy_(key_scale.transpose(1, 2)) + if _K3_HELIX_V_UP_BEFORE_A2A and attention.mapping.has_cp_helix(): + expanded = value_scale.repeat_interleave(_FP8_BLOCK, -2).repeat_interleave( + _FP8_BLOCK, -1 + ) + attention.helix_v_b_proj_full = (value_weight.float() * expanded).to(torch.bfloat16) value_weight = _helix_cp_v_b_shard( value_weight, num_heads_tp_cp=attention.num_heads_tp_cp, @@ -3296,6 +3306,8 @@ def load_param(name: str, param: torch.nn.Parameter) -> None: cp_rank=mla_mixer.mapping.cp_rank, ) mla_mixer.v_b_proj = nn.Parameter(v_weight, requires_grad=False) + if _K3_HELIX_V_UP_BEFORE_A2A and mla_mixer.mapping.has_cp_helix(): + mla_mixer.helix_v_b_proj_full = param[h * n :].view(h, v, kv) return if name.endswith(".A_log") and src.numel() != param.numel(): # The checkpoint pads A_log from [num_heads] to [head_dim] diff --git a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py index 741b2511bbb2..e4a6e397ab49 100644 --- a/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py +++ b/tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import dataclasses import pickle import sys import weakref @@ -35,6 +36,7 @@ split_inputs_for_rank, ) from mpi4py import MPI +from mpi4py.futures import MPIPoolExecutor from utils.util import skip_pre_blackwell from tensorrt_llm._torch.attention.backends.interface import ( @@ -94,6 +96,8 @@ class Scenario: # a lot of precision and we're using bf16 here. atol: float = 1e-1 rtol: float = 5e-2 + # Apply the V up-projection of all heads before the CP all-to-all. + helix_v_up: bool = False @property def max_position_embeddings(self) -> int: @@ -358,9 +362,12 @@ def _run_mla_distributed( dtype=scenario.dtype, config=config, ).cuda() + full_v_b_proj = weights["v_b_proj"] copy_weights_for_cp(weights, "o_proj.weight", 1, rank, world_size) copy_weights_for_cp(weights, "v_b_proj", 0, rank, world_size) mla.load_state_dict(weights) + if scenario.helix_v_up: + mla.helix_v_b_proj_full = full_v_b_proj.to(device="cuda", dtype=scenario.dtype) # Set up KVCacheManager and attn_metadata for distributed. kv_cache_manager, attn_metadata = setup_kv_and_metadata( @@ -570,6 +577,18 @@ def test_mla_helix_distributed( run_helix_test(_full_test_multi_gpu, scenario, comms_medium) +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") +@skip_pre_blackwell +@pytest.mark.parametrize( + "scenario", + [dataclasses.replace(s, helix_v_up=True) for s in test_scenarios], + ids=lambda x: f"scenario: {x}", +) +@pytest.mark.parametrize("comms_medium", ["nccl", "fifo_v2"]) +def test_mla_helix_v_up_before_alltoall(scenario: Scenario, comms_medium: str): + run_helix_test(_full_test_multi_gpu, scenario, comms_medium) + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs 2 GPUs to run this test") @skip_pre_blackwell def test_mla_helix_fifo_v2_unaligned_softmax_stats(): @@ -577,3 +596,169 @@ def test_mla_helix_fifo_v2_unaligned_softmax_stats(): # row, exercising the float2 fallback used by Kimi K3 over CP32. scenario = Scenario(num_heads=6, num_kv_heads=6, batch=1, ctx_len=64) run_helix_test(_full_test_multi_gpu, scenario, "fifo_v2") + + +@skip_pre_blackwell +@pytest.mark.parametrize("num_tokens", [1, 37, 224]) +def test_helix_v_up_commutes_with_combine(num_tokens: int): + """Projecting V before the CP exchange equals projecting the combined latent. + + Emulates the fifo_v2 exchange of ``cp_size`` ranks on one GPU and runs the + real combine kernel on both the latent partials (current path, then the + rank's V shard) and the V-projected partials (``_helix_project_v_partials``). + """ + torch.manual_seed(0) + cp_size, num_heads, kv_lora_rank, v_head_dim = 4, 16, 512, 128 + heads_per_rank = num_heads // cp_size + v_full = torch.randn(num_heads, v_head_dim, kv_lora_rank, device="cuda", dtype=torch.bfloat16) + v_full = v_full / kv_lora_rank**0.5 + mla_like = type("MlaLike", (), {})() + mla_like.num_heads_tp = num_heads + mla_like.kv_lora_rank = kv_lora_rank + mla_like.v_head_dim = v_head_dim + mla_like.helix_v_b_proj_full = v_full + + # Partial outputs and (max, sum) softmax stats of every CP rank for all heads. + partials = [ + torch.randn(num_tokens, num_heads * kv_lora_rank, device="cuda", dtype=torch.bfloat16) + for _ in range(cp_size) + ] + stats = [] + for _ in range(cp_size): + s = torch.empty(num_tokens, num_heads, 2, device="cuda", dtype=torch.float32) + s[..., 0].uniform_(-2.0, 2.0) + s[..., 1].uniform_(0.5, 4.0) + stats.append(s) + projected = [MLA._helix_project_v_partials(mla_like, p) for p in partials] + + def gather(per_rank, rank, value_dim): + # What rank ``rank`` holds after the exchange: its head chunk from every rank. + chunks = [ + t.view(num_tokens, cp_size, heads_per_rank * value_dim)[:, rank] for t in per_rank + ] + return torch.stack(chunks, dim=1).view(num_tokens, cp_size, heads_per_rank, value_dim) + + for rank in range(cp_size): + gathered_stats = torch.stack( + [s.view(num_tokens, cp_size, heads_per_rank * 2)[:, rank] for s in stats], dim=1 + ).view(num_tokens, cp_size, heads_per_rank, 2) + latent = torch.ops.trtllm.helix_post_process_native( + gather(partials, rank, kv_lora_rank).contiguous(), gathered_stats.contiguous(), 1.0, 1 + ).view(num_tokens, heads_per_rank, kv_lora_rank) + v_shard = v_full[rank * heads_per_rank : (rank + 1) * heads_per_rank] + reference = torch.empty( + num_tokens, heads_per_rank, v_head_dim, device="cuda", dtype=torch.bfloat16 + ) + torch.ops.trtllm.bmm_out( + latent.transpose(0, 1), v_shard.transpose(1, 2), reference.transpose(0, 1) + ) + result = torch.ops.trtllm.helix_post_process_native( + gather(projected, rank, v_head_dim).contiguous(), gathered_stats.contiguous(), 1.0, 1 + ).view(num_tokens, heads_per_rank, v_head_dim) + torch.testing.assert_close(result.float(), reference.float(), rtol=2e-2, atol=2e-2) + + +def _alltoall_after_v_up_gemm(rank: int, world_size: int, num_replays: int) -> int: + """Count wrong (token, source rank) rows of the fifo_v2 exchange of V-projected partials. + + The V-up projection GEMM and the all-to-all are captured in one CUDA graph, so the + all-to-all launches while the GEMM still runs on the SMs (programmatic dependent + launch). Every replay checks the received partials and softmax stats bit for bit. + """ + from tensorrt_llm._torch.distributed.ops import HelixAllToAllNative + + mapping = Mapping( + world_size=world_size, + rank=rank, + cp_size=world_size, + cp_config={"cp_type": CpType.HELIX, "use_nccl_for_alltoall": False, "fifo_version": 2}, + ) + helix = HelixAllToAllNative.get(mapping) + comm = MPI.COMM_WORLD + heads_per_rank, kv_lora_rank, v_head_dim = 24, 512, 128 + num_heads = heads_per_rank * world_size + torch.manual_seed(rank) + mla_like = type("MlaLike", (), {})() + mla_like.num_heads_tp = num_heads + mla_like.kv_lora_rank = kv_lora_rank + mla_like.v_head_dim = v_head_dim + mla_like.helix_v_b_proj_full = ( + torch.randn(num_heads, v_head_dim, kv_lora_rank, device="cuda") / kv_lora_rank**0.5 + ).to(torch.bfloat16) + + num_bad = 0 + for num_tokens in (224, 448): + latent = torch.empty( + num_tokens, num_heads * kv_lora_rank, device="cuda", dtype=torch.bfloat16 + ) + stats = torch.empty(num_tokens, num_heads, 2, device="cuda") + zero_kv_mask = torch.empty(num_tokens, device="cuda", dtype=torch.bool) + + def fill_inputs(): + latent.copy_(torch.randn_like(latent, dtype=torch.float32)) + stats.copy_(torch.randn_like(stats)) + zero_kv_mask.copy_(torch.rand(num_tokens, device="cuda") < 0.25) + + def exchange(): + partial_v = MLA._helix_project_v_partials(mla_like, latent) + partial_v_out, stats_out = helix.alltoall_native( + partial_v.view(num_tokens, world_size, heads_per_rank * v_head_dim), + stats.view(num_tokens, world_size, heads_per_rank * 2), + zero_kv_mask, + ) + return partial_v, partial_v_out, stats_out + + fill_inputs() + exchange() + torch.cuda.synchronize() + comm.Barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + partial_v, partial_v_out, stats_out = exchange() + comm.Barrier() + for _ in range(num_replays): + fill_inputs() + comm.Barrier() + graph.replay() + torch.cuda.synchronize() + # What this rank must hold: its head chunk from every rank, with the rows + # a source rank owns no KV for replaced by the no-op contribution. + sent = comm.allgather((partial_v.cpu(), stats.cpu(), zero_kv_mask.cpu())) + expected_v = torch.stack( + [p.view(num_tokens, world_size, -1)[:, rank] for p, _, _ in sent], dim=1 + ) + expected_stats = torch.stack( + [s.view(num_tokens, world_size, heads_per_rank, 2)[:, rank] for _, s, _ in sent], + dim=1, + ) + for src, (_, _, mask) in enumerate(sent): + expected_v[mask, src] = 0 + expected_stats[mask, src, :, 0] = float("-inf") + expected_stats[mask, src, :, 1] = 0 + bad_rows = (partial_v_out.cpu() != expected_v).any(-1) | ( + stats_out.cpu().view_as(expected_stats) != expected_stats + ).flatten(2).any(-1) + num_bad += int(bad_rows.sum()) + del graph + return num_bad + + +@skip_pre_blackwell +@pytest.mark.parametrize("world_size", [2, 4]) +def test_helix_alltoall_after_v_up_gemm_in_cuda_graph(world_size: int): + """The exchange stays exact when it overlaps the V-up projection GEMM. + + The receiver restores each 128-byte line's flagged u64 with ordinary shared-memory + stores and then writes the line out with a bulk copy, which only sees those stores + after a proxy fence. The overlap with the GEMM makes a missing fence show up as + flag values in the received partials or softmax stats. + """ + if torch.cuda.device_count() < world_size: + pytest.skip(f"needs {world_size} GPUs to run this test") + with MPIPoolExecutor(max_workers=world_size) as executor: + results = executor.map( + helix_utils.run_single_rank, + *zip(*[(_alltoall_after_v_up_gemm, world_size, 20)] * world_size), + ) + num_bad = list(results) + assert num_bad == [0] * world_size, f"wrong (token, source) rows per rank: {num_bad}" From 80ecd72a48d1d140d04d9c5ec70635b96ee24ad8 Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Wed, 21 Oct 2026 20:52:29 -0700 Subject: [PATCH 07/16] [None][perf] FP4 MLA packed decode: shorter tile critical path and producer pipelining Signed-off-by: Mingyang Hao --- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 1621 ++++++++++------- 1 file changed, 936 insertions(+), 685 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 1c0b92059305..301ef2e2148a 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -377,13 +377,19 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_PV_SFB_SCORE_SLOT_COLS = SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT * SMEM_P4_PV_SFB_COLS_PER_KBLOCK SMEM_P4_P_ARCHIVE_TMEM_COLS = SMEM_P4_SCORE_TMEM_COLS SMEM_P4_P_ARCHIVE_SLOT_STRIDE = SMEM_P4_QK_PIPELINE_SLOT_STRIDE -SMEM_P4_PV_VSF_N_TILE_COLS = SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_PV_SFB_SCORE_SLOT_COLS SMEM_P4_SCALE_FEED_STAGE_COLS = 28 SMEM_P4_SCALE_FEED_TMEM_COLS = ( SMEM_P4_O_ACC_TMEM_COLS[1], SMEM_P4_O_ACC_TMEM_COLS[1] + SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_SCALE_FEED_STAGE_COLS, ) SMEM_P4_QK_ACC_COL_OFFSET = SMEM_P4_SCORE_TMEM_COLS[0] +# P scale factors per score bank, outside the score columns so the next QK into +# the bank cannot overwrite a pending PV's scales. +SMEM_P4_PSF_TMEM_COL_OFFSET = SMEM_P4_SCALE_FEED_TMEM_COLS[1] +if SMEM_P4_PSF_TMEM_COL_OFFSET + SMEM_P4_TMEM_BANKS * SMEM_P4_P_SCALE_FEED_STAGE_COLS > ( + NUM_TMEM_ALLOC_COLS +): + raise ValueError("P scale-factor banks do not fit in the TMEM allocation") SMEM_P4_O_ACC_COL_OFFSET = SMEM_P4_O_ACC_TMEM_COLS[0] SMEM_P4_SCALE_A_COL_OFFSET = SMEM_P4_SCALE_FEED_TMEM_COLS[0] SMEM_P4_QK_COMPACT_SFB_S2T_BYTE_OFFSET = 0 @@ -408,7 +414,6 @@ def _initial_float_from_argv(option: str, default: float) -> float: ) SMEM_P4_QK_SFA_FEED_COLS = SMEM_P4_QK_RESIDENT_SF_COLS SMEM_P4_SCALE_B_COL_OFFSET = (SMEM_P4_SCALE_A_COL_OFFSET + SMEM_P4_QK_SFA_FEED_COLS + 15) // 16 * 16 -SMEM_P4_P_SFA_BANK_REL_COL_OFFSET = 32 SMEM_P4_PV_VSF_TAIL_COL_OFFSET = SMEM_P4_SCALE_B_COL_OFFSET SMEM_P4_RESIDENT_Q_SF_COL_OFFSET = SMEM_P4_SCALE_FEED_TMEM_COLS[0] SMEM_P4_QK_SFA_COL_OFFSET = SMEM_P4_RESIDENT_Q_SF_COL_OFFSET @@ -417,8 +422,6 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS = ( SMEM_P4_TMEM_BANKS * SMEM_P4_SCORE_SLOTS_PER_KV_TILE * SMEM_P4_SCORE_SLOT_ROW_STATE_STRIDE ) -SMEM_P4_ROW_SCALE_RING_FLOATS = SMEM_P4_TMEM_BANKS * SMEM_P4_BMM1_M -SMEM_P4_FINAL_ANCHOR_ROWSUM_FLOATS = SMEM_P4_BMM1_M SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS = SMEM_P4_BMM1_M SMEM_P4_TMEM_SCORE_PIPELINE_STAGES = SMEM_P4_QK_PIPELINE_SLOTS SMEM_P4_MBAR_PARITY_PHASES = 2 @@ -443,6 +446,9 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_PAIR_P_HALF_READY_MBAR_OFFSET + SMEM_P4_P_HALF_COUNT * SMEM_P4_QK_PIPELINE_SLOTS ) SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET = SMEM_P4_PAIR_ROWSUM_DONE_MBAR_OFFSET +# Per score bank: all producer warps of both CTAs have read the bank's scores, so +# the MMA warp may already issue the QK two tiles ahead into it. +SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET = SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET SMEM_P4_QK15_STAGE1_MBAR_OFFSET = ( SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET + SMEM_P4_QK_PIPELINE_SLOTS ) @@ -450,6 +456,8 @@ def _initial_float_from_argv(option: str, default: float) -> float: SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + SMEM_P4_QK_PIPELINE_SLOTS ) SMEM_P4_V_PIPELINE_STAGES = 3 +# Poll period of the K/V load warps while their ring slot is still in use. +SMEM_P4_LOAD_WAIT_SLEEP_NS = 128 SMEM_P4_PV_PAGE_B_MMA_LEADING_BYTES = 0 SMEM_P4_PV_PAGE_B_MMA_STRIDE_BYTES = 512 SMEM_P4_PV_PAGE_B_MMA_LAYOUT = 4 @@ -513,6 +521,16 @@ class KvCache3DLayout: TMA_QK_WARP_ID = 9 TMA_V_WARP_ID = 10 ROWMETA_WARP_ID = 11 +# QK and PV B scale factors use disjoint TMEM columns: a ring holding one QK data stage of K scales +# (4 kblocks) and one PV N tile of V scales (4 kblocks). Each issuer reuses its +# columns in issue order. +PV_ISSUE_WARP_ID = ROWMETA_WARP_ID +SMEM_P4_KSF_RING_COL_OFFSET = SMEM_P4_SCALE_B_COL_OFFSET +SMEM_P4_KSF_RING_COLS = SMEM_P4_QK_DATA_STAGE_KBLOCKS * SMEM_P4_QK_SFB_COLS_PER_KBLOCK +SMEM_P4_VSF_TMEM_COL_OFFSET = SMEM_P4_KSF_RING_COL_OFFSET + SMEM_P4_KSF_RING_COLS +SMEM_P4_VSF_TMEM_COLS = SMEM_P4_PV_KBLOCKS_PER_LI * SMEM_P4_PV_SFB_COLS_PER_KBLOCK +if SMEM_P4_VSF_TMEM_COL_OFFSET + SMEM_P4_VSF_TMEM_COLS > SMEM_P4_PSF_TMEM_COL_OFFSET: + raise ValueError("K/V scale-factor rings overlap the P scale-factor banks") SMEM_P4_ALU_REGISTER_BUDGET = 224 SMEM_P4_AUX_REGISTER_BUDGET = 56 A_STAGE_BYTES = MMA_TILER_MNK_PER_CTA[0] * MMA_TILER_MNK[2] // 2 @@ -752,9 +770,9 @@ def _refresh_sf_vec_config(sf_vec_size: int) -> None: + SMEM_P4_P4_OWNER_SMEM_SF_MAILBOX_WORDS * 4 + SMEM_P4_V_PIPELINE_STAGES * sfa_tma_stage_bytes * pv_sfb_n_chunks * SMEM_P4_N_OUT_TILES + SMEM_P4_STAGE_ROW_STATE_SMEM_FLOATS * 4 - + SMEM_P4_ROW_SCALE_RING_FLOATS * 4 - + SMEM_P4_FINAL_ANCHOR_ROWSUM_FLOATS * 4 + SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS * 4 + # Row maxima of the pipelined leading tiles, by tile parity. + + 2 * SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS * 4 + SMEM_P4_PAGE_ID_PLAN_BYTES + SMEM_P4_RUNTIME_SCALE_PAIR_BYTES ) @@ -902,7 +920,8 @@ def _qk_compact_ring_stage_byte_offset( TMEM_BAR_ID = 1 -TMEM_BAR_THREADS = 32 * (1 + SMEM_P4_TOTAL_PRODUCER_WARPS) +# TMEM allocation is published to the producers and the PV issue warp. +TMEM_BAR_THREADS = 32 * (2 + SMEM_P4_TOTAL_PRODUCER_WARPS) O_STORE_BAR_ID = 2 SMEM_P4_P4_SLOT0_BAR_ID = 6 SMEM_P4_P4_SLOT1_BAR_ID = 7 @@ -1860,6 +1879,10 @@ def _load_runtime_t336_qk_tile( use_consecutive_page_pair: ctm.Constexpr = False, use_ksf_gather4: ctm.Constexpr = False, ) -> None: + # Steady tiles: CTA1 loads its own K half (_load_runtime_t336_qk_tile_rank1). + rank1_self_issued: ctm.Constexpr = ( + not initial_tile and not split_prefix_owner and not qk15_final + ) qk_bulk_tma_mbar = qk_tma_mbar if cutlass.const_expr(not initial_tile): qk_bulk_tma_mbar = ctm.Array(qk_bulk_ready_mbar_ptr + qk_slot_idx, shape=1) @@ -1924,7 +1947,7 @@ def _load_runtime_t336_qk_tile( kv_tile_idx, page_size, 0, - issue_rank1=not split_prefix_owner, + issue_rank1=not split_prefix_owner and not rank1_self_issued, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, ) @@ -1954,7 +1977,7 @@ def _load_runtime_t336_qk_tile( kv_tile_idx, page_size, 1, - issue_rank1=not split_prefix_owner, + issue_rank1=not split_prefix_owner and not rank1_self_issued, rank1_first=qk15_final, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, @@ -1980,12 +2003,72 @@ def _load_runtime_t336_qk_tile( page_size, SMEM_P4_QK_FULL_DATA_STAGES, issue_rank0=True, - issue_rank1=not initial_tile, + issue_rank1=not initial_tile and not rank1_self_issued, use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, ) +@cute.jit +def _load_runtime_t336_qk_tile_rank1( + mPageTable_pl: cute.Tensor, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_full_mbar_ptr, + qk_bulk_ready_mbar_ptr, + sB, + sSFB, + physical_page0: ctm.Int32, + physical_page1: ctm.Int32, + kv_tile_idx: ctm.Int32, + qk_slot_idx: ctm.Int32, + page_size: ctm.Constexpr, +) -> None: + """CTA1's half of a steady K tile: its own K rows, completing on the leader's barriers.""" + leader_bulk_mbar = _as_shared_cta( + _mapa_shared_cluster(ctm.Array(qk_bulk_ready_mbar_ptr + qk_slot_idx, shape=1), ctm.Int32(0)) + ) + leader_full_mbar = _as_shared_cta( + _mapa_shared_cluster(ctm.Array(qk_full_mbar_ptr, shape=1), ctm.Int32(0)) + ) + if prims.elect_sync(): + sB_slot = sB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_K_COMPACT_SLOT_BYTES // 4)) + sSFB_slot = sSFB.subview(qk_slot_idx * ctm.Int32(SMEM_P4_KSF_COMPACT_SLOT_BYTES)) + for qk_stage_idx in ctm.range_constexpr(SMEM_P4_QK_FULL_DATA_STAGES + 1): + _load_qk_konly_kblock_stage( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + leader_full_mbar + if qk_stage_idx == SMEM_P4_QK_FULL_DATA_STAGES + else leader_bulk_mbar, + sB_slot, + sSFB_slot, + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(1), + ctm.Int32(0), + ctm.Int32(0), + ctm.Int32(0), + physical_page0, + physical_page1, + kv_tile_idx, + qk_stage_idx * SMEM_P4_QK_DATA_STAGE_KBLOCKS, + 0, + qk_stage_idx, + page_size, + 0, + issue_ksf=False, + manage_mbarrier=False, + ) + + @cute.jit def _issue_runtime_t336_qk_prefix_rank1( mPageTable_pl: cute.Tensor, @@ -2185,6 +2268,19 @@ def _fence_async_shared_cta() -> None: prims.fence_proxy(kind=prims.Proxy.ASYNC_SHARED, space=prims.SharedSpace.shared_cta) +@cute.jit +def _sleep_until_producer_slot_free(producer) -> ctm.Boolean: + """Poll a load pipeline slot with a sleep between acquire attempts.""" + acquired = producer.try_acquire() + while not acquired: + cute_inline_ptx( + "nanosleep.u32 {$r0};", + read_only_args=[ctm.Int32(SMEM_P4_LOAD_WAIT_SLEEP_NS)], + ) + acquired = producer.try_acquire() + return acquired + + @cute.jit def _permute_raw_trtllm_q_tail_in_smem(sQ, local_warp_idx: ctm.Int32, lane_idx: ctm.Int32) -> None: tail_stage_word_offset: ctm.Constexpr = ( @@ -2359,6 +2455,15 @@ def _softmax_exp2_pair_packed_f16x2( ) +@cute.jit +def _add_f16x2(a: ctm.Int32, b: ctm.Int32) -> ctm.Int32: + return cute_inline_ptx( + "add.rn.f16x2 {$w0}, {$r0}, {$r1};", + write_only_types=[ctm.Int32], + read_only_args=[a, b], + ) + + @cute.jit def _unpack_f16x2_to_f32_pair(packed: ctm.Int32) -> tuple: """Widen one packed probability pair for the denominator only.""" @@ -2506,7 +2611,9 @@ def _pack_p4_16_scores_from_group_packed_f16x2( elem_base: ctm.Constexpr, ) -> tuple: words = ctm.Array(ctm.Int32, 2, space=ctm.AddressSpace.rmem) - p_sum_pair = (ctm.Float32(0.0), ctm.Float32(0.0)) + # One FP16x2 partial sum per 8-element word; P <= 6, so a 16-element + # group sum stays far inside FP16 range. + word_sums = ctm.Array(ctm.Int32, 2, space=ctm.AddressSpace.rmem) for word_part in ctm.range_constexpr(2): packed_pair0 = ctm.Int32(0) packed_pair1 = ctm.Int32(0) @@ -2528,11 +2635,10 @@ def _pack_p4_16_scores_from_group_packed_f16x2( score0 = _mask_invalid_score(score0, global_k_idx0, valid_k) score1 = _mask_invalid_score(score1, global_k_idx1, valid_k) p_h2 = _softmax_exp2_pair_packed_f16x2(score0, score1, softmax_scale_log2, p_bias) - p0_f32, p1_f32 = _unpack_f16x2_to_f32_pair(p_h2) - if cutlass.const_expr(word_part == 0 and pair_idx == 0): - p_sum_pair = (p0_f32, p1_f32) + if cutlass.const_expr(pair_idx == 0): + word_sums[word_part] = p_h2 else: - p_sum_pair = nvvm_add_packed_f32x2(p_sum_pair, (p0_f32, p1_f32)) + word_sums[word_part] = _add_f16x2(word_sums[word_part], p_h2) if cutlass.const_expr(pair_idx == 0): packed_pair0 = p_h2 elif cutlass.const_expr(pair_idx == 1): @@ -2544,7 +2650,8 @@ def _pack_p4_16_scores_from_group_packed_f16x2( words[word_part] = _pack_e2m1x8_from_f16x2( packed_pair0, packed_pair1, packed_pair2, packed_pair3 ) - p_sum = p_sum_pair[0] + p_sum_pair[1] + p_sum0, p_sum1 = _unpack_f16x2_to_f32_pair(_add_f16x2(word_sums[0], word_sums[1])) + p_sum = p_sum0 + p_sum1 return (words[0], words[1], p_sum) @@ -3216,19 +3323,15 @@ def _issue_smem_p4_qk_ksf_stage_from_desc( else QK_TAIL_DATA_STAGE_K_DIM // SMEM_P4_QK_MMA_KBLOCK_DIM ) for local_kblock_idx in ctm.range_constexpr(stage_kblocks): - global_kblock_idx: ctm.Constexpr = ( - qk_stage_idx * SMEM_P4_QK_DATA_STAGE_KBLOCKS + local_kblock_idx - ) for scale_chunk_idx in ctm.range_constexpr(SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK): logical_local_scale_word_idx: ctm.Constexpr = ( local_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK + scale_chunk_idx ) physical_local_scale_word_idx: ctm.Constexpr = logical_local_scale_word_idx - global_scale_word_idx: ctm.Constexpr = ( - global_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK + scale_chunk_idx - ) + # One data stage at a time lives in the K scale ring. sfb_tmem_ptr = ctm.inttoptr( - sfb_tmem_addr_base + ctm.Int32(global_scale_word_idx * sfb_cols_per_scale_chunk), + sfb_tmem_addr_base + + ctm.Int32(logical_local_scale_word_idx * sfb_cols_per_scale_chunk), 6, ctm.Int32, ) @@ -3249,13 +3352,6 @@ def _issue_smem_p4_qk_ksf_stage_from_desc( ) -@cute.jit -def _complete_smem_p4_qk_ksf_store() -> None: - prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) - cute.arch.fence_view_async_tmem_store() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - - @cute.jit def _issue_smem_p4_qk_ksf_range_from_base( sKSF, @@ -3286,132 +3382,14 @@ def _issue_smem_p4_qk_ksf_range_from_base( @cute.jit -def _stage_smem_p4_qk_ksf_range_from_base( - sKSF, - base_col_id: ctm.Int32, - base_row_id: ctm.Int32, - sfb_col_offset: ctm.Constexpr, - sfb_stage_base: ctm.Int32, - qk_stage_start: ctm.Constexpr, - qk_stage_count: ctm.Constexpr, - sksf_is_compact_slot_base: ctm.Constexpr = False, - page_major_ksf: ctm.Constexpr = False, -) -> None: - _issue_smem_p4_qk_ksf_range_from_base( - sKSF, - base_col_id, - base_row_id, - sfb_col_offset, - sfb_stage_base, - qk_stage_start, - qk_stage_count, - sksf_is_compact_slot_base, - page_major_ksf, - ) - _complete_smem_p4_qk_ksf_store() - - -@cute.jit -def _stage_smem_p4_qk_full_ksf_from_base( - sKSF, +def _issue_smem_p4_pv_vsf_ntile_from_base( + sVSF, base_col_id: ctm.Int32, base_row_id: ctm.Int32, - sfb_col_offset: ctm.Constexpr, - sfb_stage_base: ctm.Constexpr, - page_major_ksf: ctm.Constexpr = False, -) -> None: - if cutlass.const_expr(SMEM_P4_QK_SFB_WARPX2_MODE == 2): - sfb_s2t_shape, sfb_s2t_multicast = prims.S2TCopyMode.S2T_64x128b_WARPX2_02_13 - else: - sfb_s2t_shape, sfb_s2t_multicast = prims.S2TCopyMode.S2T_64x128b_WARPX2_01_23 - sfb_tmem_addr_base = base_row_id << ctm.Int32(16) | base_col_id + ctm.Int32(sfb_col_offset) - sfb_cols_per_scale_chunk: ctm.Constexpr = _compact_sfb_mma_cols_per_scale_chunk( - SMEM_P4_CTA_GROUP_M, SMEM_P4_QK_SCORE_SLOT_N - ) - full_root_byte_offset: ctm.Constexpr = _qk_compact_ring_stage_byte_offset( - sfb_stage_base, QK_SFB_STAGE_BYTES, SMEM_P4_KSF_COMPACT_SLOT_BYTES - ) - tail_byte_offset: ctm.Constexpr = _qk_compact_ring_stage_byte_offset( - sfb_stage_base + SMEM_P4_QK_FULL_DATA_STAGES, - QK_SFB_STAGE_BYTES, - SMEM_P4_KSF_COMPACT_SLOT_BYTES, - ) - full_s2t_stride_bytes: ctm.Constexpr = ( - SMEM_P4_QK_SFB_S2T_STRIDE_BYTES * SMEM_P4_QK_DATA_STAGE_SF_CHUNKS - if page_major_ksf - else SMEM_P4_QK_SFB_S2T_STRIDE_BYTES - ) - tail_s2t_stride_bytes: ctm.Constexpr = ( - SMEM_P4_QK_SFB_S2T_STRIDE_BYTES * SMEM_P4_QK_K_TAIL_DATA_STAGE_SF_CHUNKS - if page_major_ksf - else SMEM_P4_QK_SFB_S2T_STRIDE_BYTES - ) - desc_ksf_s2t_full_root = prims.Tcgen05SmemDesc.build( - sKSF.subview(full_root_byte_offset), - leading_byte_offset=16, - stride_byte_offset=full_s2t_stride_bytes, - layout=0, - ) - desc_ksf_s2t_tail = prims.Tcgen05SmemDesc.build( - sKSF.subview(tail_byte_offset), - leading_byte_offset=16, - stride_byte_offset=tail_s2t_stride_bytes, - layout=0, - ) - for qk_stage_idx in ctm.range_constexpr(SMEM_P4_QK_DATA_STAGES): - if cutlass.const_expr(qk_stage_idx < SMEM_P4_QK_FULL_DATA_STAGES): - desc_ksf_s2t = desc_ksf_s2t_full_root.advance_start_address( - qk_stage_idx * QK_SFB_STAGE_BYTES - ) - else: - desc_ksf_s2t = desc_ksf_s2t_tail - stage_kblocks: ctm.Constexpr = ( - SMEM_P4_QK_DATA_STAGE_KBLOCKS - if qk_stage_idx < SMEM_P4_QK_FULL_DATA_STAGES - else QK_TAIL_DATA_STAGE_K_DIM // SMEM_P4_QK_MMA_KBLOCK_DIM - ) - for local_kblock_idx in ctm.range_constexpr(stage_kblocks): - global_kblock_idx: ctm.Constexpr = ( - qk_stage_idx * SMEM_P4_QK_DATA_STAGE_KBLOCKS + local_kblock_idx - ) - for scale_chunk_idx in ctm.range_constexpr(SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK): - logical_local_scale_word_idx: ctm.Constexpr = ( - local_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK + scale_chunk_idx - ) - physical_local_scale_word_idx: ctm.Constexpr = logical_local_scale_word_idx - global_scale_word_idx: ctm.Constexpr = ( - global_kblock_idx * SMEM_P4_QK_MMA_SCALE_CHUNKS_PER_KBLOCK + scale_chunk_idx - ) - sfb_tmem_ptr = ctm.inttoptr( - sfb_tmem_addr_base - + ctm.Int32(global_scale_word_idx * sfb_cols_per_scale_chunk), - 6, - ctm.Int32, - ) - source_scale_word_stride: ctm.Constexpr = ( - 1 if page_major_ksf else SMEM_P4_QK_SFB_N_CHUNKS - ) - source_byte_offset: ctm.Constexpr = ( - physical_local_scale_word_idx * source_scale_word_stride * (512 // 16) - + SMEM_P4_QK_COMPACT_SFB_S2T_BYTE_OFFSET - ) - if prims.elect_sync(): - prims.tcgen05_cp( - sfb_s2t_shape, - sfb_tmem_ptr, - desc_ksf_s2t + ctm.Int64(source_byte_offset), - group=prims.CTAGroup.CTA_2, - multicast=sfb_s2t_multicast, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) - cute.arch.fence_view_async_tmem_store() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - - -@cute.jit -def _stage_smem_p4_pv_vsf_bank_from_base( - sVSF, base_col_id: ctm.Int32, base_row_id: ctm.Int32, v_stage: ctm.Int32 + v_stage: ctm.Int32, + n_tile_idx: ctm.Constexpr, ) -> None: + """Copy one PV N tile's V scales (all four kblocks) into the V scale columns.""" sfb_s2t_shape, sfb_s2t_multicast = prims.S2TCopyMode.S2T_64x128b_WARPX2_01_23 sVSF_stage = sVSF.subview(ctm.Int32(SMEM_P4_V_SFB_TMA_STAGE_BYTES) * v_stage) desc_vsf_s2t = prims.Tcgen05SmemDesc.build( @@ -3421,30 +3399,26 @@ def _stage_smem_p4_pv_vsf_bank_from_base( layout=0, ) vsf_tmem_addr_base = base_row_id << ctm.Int32(16) | base_col_id + ctm.Int32( - SMEM_P4_PV_VSF_TAIL_COL_OFFSET + SMEM_P4_VSF_TMEM_COL_OFFSET ) total_n_chunks: ctm.Constexpr = SMEM_P4_PV_SFB_N_CHUNKS * SMEM_P4_N_OUT_TILES - for n_tile_idx in ctm.range_constexpr(SMEM_P4_N_OUT_TILES): - for global_kblock_idx in ctm.range_constexpr(SMEM_P4_PV_KBLOCKS_PER_LI): - vsf_col_delta: ctm.Constexpr = ( - n_tile_idx * SMEM_P4_PV_VSF_N_TILE_COLS - + global_kblock_idx * SMEM_P4_PV_SFB_COLS_PER_KBLOCK + for global_kblock_idx in ctm.range_constexpr(SMEM_P4_PV_KBLOCKS_PER_LI): + vsf_tmem_ptr = ctm.inttoptr( + vsf_tmem_addr_base + ctm.Int32(global_kblock_idx * SMEM_P4_PV_SFB_COLS_PER_KBLOCK), + 6, + ctm.Int32, + ) + source_byte_offset: ctm.Constexpr = ( + global_kblock_idx * total_n_chunks + n_tile_idx * SMEM_P4_PV_SFB_N_CHUNKS + ) * (512 // 16) + if prims.elect_sync(): + prims.tcgen05_cp( + sfb_s2t_shape, + vsf_tmem_ptr, + desc_vsf_s2t + ctm.Int64(source_byte_offset), + group=prims.CTAGroup.CTA_2, + multicast=sfb_s2t_multicast, ) - vsf_tmem_addr = vsf_tmem_addr_base + ctm.Int32(vsf_col_delta) - vsf_tmem_ptr = ctm.inttoptr(vsf_tmem_addr, 6, ctm.Int32) - source_byte_offset: ctm.Constexpr = ( - global_kblock_idx * total_n_chunks + n_tile_idx * SMEM_P4_PV_SFB_N_CHUNKS - ) * (512 // 16) - if prims.elect_sync(): - prims.tcgen05_cp( - sfb_s2t_shape, - vsf_tmem_ptr, - desc_vsf_s2t + ctm.Int64(source_byte_offset), - group=prims.CTAGroup.CTA_2, - multicast=sfb_s2t_multicast, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) - cute.arch.fence_view_async_tmem_store() @cute.jit @@ -3500,7 +3474,7 @@ def _issue_smem_p4_qk_stage_burst_from_base( qk_kblocks=stage_kblocks, acc_col_offset=acc_col_offset, sfa_col_offset=sfa_col_offset_base, - sfb_col_offset=sfb_col_offset_base + stage_kblock_start * qk_sfb_scale_stride_cols, + sfb_col_offset=sfb_col_offset_base, stage_sfa_from_smem=False, stage_sfb_from_smem=False, use_split_operand_stages=True, @@ -3617,21 +3591,6 @@ def _store_smem_p4_n256_rowsum_part( sRowSumParts[row_state_offset + local_row] = value -@cute.jit -def _load_smem_p4_n256_rowsum_part( - sRowSumParts, - local_row: ctm.Int32, - bank_idx: ctm.Int32, - score_half_idx: ctm.Constexpr, - col_band_idx: ctm.Constexpr, -) -> ctm.Float32: - slot_idx = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + ctm.Int32(score_half_idx) - row_state_offset = slot_idx * ctm.Int32(SMEM_P4_SCORE_SLOT_ROW_STATE_STRIDE) + ctm.Int32( - col_band_idx * SMEM_P4_BMM1_M - ) - return sRowSumParts[row_state_offset + local_row] - - @cute.jit def _smem_p4_qk_score_group_col( score_slot_idx: ctm.Constexpr, col_band: ctm.Int32, local_group_idx: ctm.Constexpr @@ -4376,8 +4335,8 @@ def _make_p4_n256_psf_bank_to_tmem_desc( ) p_sfa_col_id = ( base_col_id - + ctm.Int32(SMEM_P4_P_SFA_BANK_REL_COL_OFFSET) - + bank_idx * ctm.Int32(SMEM_P4_QK_PIPELINE_SLOT_STRIDE) + + ctm.Int32(SMEM_P4_PSF_TMEM_COL_OFFSET) + + bank_idx * ctm.Int32(SMEM_P4_P_SCALE_FEED_STAGE_COLS) ) p_sfa_tmem_addr = base_row_id << ctm.Int32(16) | p_sfa_col_id p_sfa_tmem = ctm.inttoptr(p_sfa_tmem_addr, 6, ctm.Int32) @@ -4432,20 +4391,6 @@ def _p4_n256_rowsum_from_packed_psf_word(sf_word: ctm.Int32, packed_group_sums) ) -@cute.jit -def _load_smem_p4_n256_quadrant_rowsum( - sRowSumParts, - sP, - local_row: ctm.Int32, - bank_idx: ctm.Int32, - score_half_idx: ctm.Constexpr, - col_band_idx: ctm.Constexpr, -) -> ctm.Float32: - return _load_smem_p4_n256_rowsum_part( - sRowSumParts, local_row, bank_idx, score_half_idx, col_band_idx - ) - - @cute.jit def _issue_smem_p4_pv_score_slot_from_base( sP, @@ -4493,9 +4438,7 @@ def _issue_smem_p4_pv_score_slot_from_base( if cutlass.const_expr(clear_accumulator_override >= 0): clear_pv_accumulator = ctm.Boolean(clear_accumulator_override != 0) sfb_col_offset: ctm.Constexpr = ( - SMEM_P4_PV_VSF_TAIL_COL_OFFSET - + n_tile_idx * SMEM_P4_PV_VSF_N_TILE_COLS - + physical_global_kblock_idx * SMEM_P4_PV_SFB_COLS_PER_KBLOCK + SMEM_P4_VSF_TMEM_COL_OFFSET + physical_global_kblock_idx * SMEM_P4_PV_SFB_COLS_PER_KBLOCK ) p_stage_slot_col_offset = p_bank_col_offset + ctm.Int32( physical_p_half_idx * SMEM_P4_P_TMEM_SLOT_STRIDE @@ -4741,7 +4684,6 @@ def _runtime_issue_pv_tile_slot( sVSF, tmem_base_col_id: ctm.Int32, tmem_base_row_id: ctm.Int32, - qk_full_mbar, pair_overlap_sync_mbars, p4_smem_slot0_mbars, pv_full_mbar, @@ -4752,17 +4694,11 @@ def _runtime_issue_pv_tile_slot( initial_tile: ctm.Boolean, ) -> None: v_stage = stream_li_idx % ctm.Int32(SMEM_P4_V_PIPELINE_STAGES) - qk_phase = ( - stream_li_idx - // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - while not prims.mbarrier_try_wait_parity( - qk_full_mbar.subview(bank_idx), qk_phase, time_limit=10000000 - ): - pass + # QK of this tile is complete: its producers read the scores before arriving + # on P_SOURCE_READY below. Waiting on qk_full here instead could miss a phase, + # because the QK two tiles ahead already targets the same bank. prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - _stage_smem_p4_pv_vsf_bank_from_base(sVSF, tmem_base_col_id, tmem_base_row_id, v_stage) + _issue_smem_p4_pv_vsf_ntile_from_base(sVSF, tmem_base_col_id, tmem_base_row_id, v_stage, 0) prefetched_psf_desc, prefetched_psf_tmem = _make_p4_n256_psf_bank_to_tmem_desc( sPSF, tmem_base_col_id, tmem_base_row_id, bank_idx ) @@ -4776,88 +4712,50 @@ def _runtime_issue_pv_tile_slot( ) while not prims.mbarrier_try_wait_parity(p_source_ready_mbar, pair_phase, time_limit=10000000): pass - _fence_async_shared_cta() _issue_p4_n256_psf_bank_to_tmem_from_desc(prefetched_psf_desc, prefetched_psf_tmem) prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) p_archive_col_offset = ctm.Int32(SMEM_P4_P_ARCHIVE_TMEM_COLS[0]) + bank_idx * ctm.Int32( SMEM_P4_P_ARCHIVE_SLOT_STRIDE ) p_half0_smem_stage = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) - p_half1_smem_stage = p_half0_smem_stage + ctm.Int32(1) - p_sfa_col_offset = ctm.Int32(SMEM_P4_P_SFA_BANK_REL_COL_OFFSET) + bank_idx * ctm.Int32( - SMEM_P4_QK_PIPELINE_SLOT_STRIDE - ) - for local_kblock_idx in ctm.range_constexpr(SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT): - for n_tile in cutlass.range_constexpr(SMEM_P4_N_OUT_TILES): - a_collector_op: ctm.Constexpr = ( - prims.Tcgen05MMACollectorOp.FILL - if n_tile == 0 - else prims.Tcgen05MMACollectorOp.LASTUSE - ) - _issue_smem_p4_pv_score_slot_from_base( - sP, - sV, - sPSF, - sVSF, - tmem_base_col_id, - tmem_base_row_id, - p_half0_smem_stage, - initial_tile, - v_stage=v_stage, - acc_col_offset=SMEM_P4_O_ACC_COL_OFFSETS[n_tile], - n_tile_idx=n_tile, - p_tmem_col_offset=p_archive_col_offset, - bank_idx=bank_idx, - p_half_idx=0, - local_kblock_idx=local_kblock_idx, - global_kblock_idx=local_kblock_idx, - pv_kblocks=1, - p_sfa_col_offset=p_sfa_col_offset, - a_collector_op=a_collector_op, - ) - for local_kblock_idx in ctm.range_constexpr(SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT): - global_kblock_idx: ctm.Constexpr = SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT + local_kblock_idx - for n_tile in cutlass.range_constexpr(SMEM_P4_N_OUT_TILES): - a_collector_op: ctm.Constexpr = ( - prims.Tcgen05MMACollectorOp.FILL - if n_tile == 0 - else prims.Tcgen05MMACollectorOp.LASTUSE - ) - _issue_smem_p4_pv_score_slot_from_base( - sP, - sV, - sPSF, - sVSF, - tmem_base_col_id, - tmem_base_row_id, - p_half1_smem_stage, - initial_tile, - v_stage=v_stage, - acc_col_offset=SMEM_P4_O_ACC_COL_OFFSETS[n_tile], - n_tile_idx=n_tile, - p_tmem_col_offset=p_archive_col_offset, - bank_idx=bank_idx, - p_half_idx=1, - local_kblock_idx=local_kblock_idx, - global_kblock_idx=global_kblock_idx, - pv_kblocks=1, - p_sfa_col_offset=p_sfa_col_offset, - a_collector_op=a_collector_op, + p_sfa_col_offset = ctm.Int32(SMEM_P4_PSF_TMEM_COL_OFFSET) + bank_idx * ctm.Int32( + SMEM_P4_P_SCALE_FEED_STAGE_COLS + ) + # N-tile major: the V scale columns hold one N tile at a time, refilled in + # issue order after the previous N tile's MMAs. + for n_tile in cutlass.range_constexpr(SMEM_P4_N_OUT_TILES): + if cutlass.const_expr(n_tile > 0): + _issue_smem_p4_pv_vsf_ntile_from_base( + sVSF, tmem_base_col_id, tmem_base_row_id, v_stage, n_tile ) + for p_half_idx in cutlass.range_constexpr(SMEM_P4_SCORE_SLOTS_PER_KV_TILE): + for local_kblock_idx in ctm.range_constexpr(SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT): + _issue_smem_p4_pv_score_slot_from_base( + sP, + sV, + sPSF, + sVSF, + tmem_base_col_id, + tmem_base_row_id, + p_half0_smem_stage + ctm.Int32(p_half_idx), + initial_tile, + v_stage=v_stage, + acc_col_offset=SMEM_P4_O_ACC_COL_OFFSETS[n_tile], + n_tile_idx=n_tile, + p_tmem_col_offset=p_archive_col_offset, + bank_idx=bank_idx, + p_half_idx=p_half_idx, + local_kblock_idx=local_kblock_idx, + global_kblock_idx=p_half_idx * SMEM_P4_PV_KBLOCKS_PER_SCORE_SLOT + + local_kblock_idx, + pv_kblocks=1, + p_sfa_col_offset=p_sfa_col_offset, + ) if stream_li_idx > ctm.Int32(0) and stream_li_idx + ctm.Int32(1) == stream_li_total: if prims.elect_sync(): _cross_arrive_smem_p4_cta2_mbar( pair_overlap_sync_mbars.subview(SMEM_P4_PAIR_FINAL_PV_ISSUED_MBAR_OFFSET + bank_idx) ) - p4_phase = ( - stream_li_idx - // ctm.Int32(SMEM_P4_P4_PIPELINE_STAGES) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - while not prims.mbarrier_try_wait_parity( - p4_smem_slot0_mbars.subview(bank_idx), p4_phase, time_limit=10000000 - ): - pass if prims.elect_sync(): prims.tcgen05_commit( pv_full_mbar.subview(bank_idx), @@ -4884,7 +4782,6 @@ def _runtime_issue_pv_tile_dispatch( sVSF, tmem_base_col_id: ctm.Int32, tmem_base_row_id: ctm.Int32, - qk_full_mbar, pair_overlap_sync_mbars, p4_smem_slot0_mbars, pv_full_mbar, @@ -4901,7 +4798,6 @@ def _runtime_issue_pv_tile_dispatch( sVSF, tmem_base_col_id, tmem_base_row_id, - qk_full_mbar, pair_overlap_sync_mbars, p4_smem_slot0_mbars, pv_full_mbar, @@ -4924,7 +4820,6 @@ def _runtime_refill_qk_tile_state( qk_bulk_ready_mbars, qk_full_mbar, pair_overlap_sync_mbars, - tmem_score_slot0_producer, qk_smem_consumer, qk_next_li_idx: ctm.Int32, qk_bulk_slot_idx: ctm.Int32, @@ -4934,38 +4829,28 @@ def _runtime_refill_qk_tile_state( split_qk15_stage1: ctm.Boolean, page_major_ksf: ctm.Constexpr, ) -> None: + # The score bank's previous tile was read before its scores were consumed, + # which this warp waited on before issuing this QK. prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - score_slot0_handle = tmem_score_slot0_producer.acquire_and_advance() while not prims.mbarrier_try_wait_parity( qk_bulk_ready_mbars.subview(qk_bulk_slot_idx), qk_bulk_phase, time_limit=10000000 ): pass - _fence_async_shared_cta() sK_slot = sB.subview(qk_bulk_slot_idx * ctm.Int32(SMEM_P4_K_COMPACT_SLOT_BYTES // 4)) sKSF_slot = sSFB.subview(qk_bulk_slot_idx * ctm.Int32(SMEM_P4_KSF_COMPACT_SLOT_BYTES)) desc_ksf_s2t_root = _make_smem_p4_qk_ksf_stage_desc( sKSF_slot, ctm.Int32(0), 0, True, page_major_ksf ) + # Each data stage copies its K scales into the ring right before its MMAs; the + # issue order keeps a stage's copy behind the previous stage's MMAs. _issue_smem_p4_qk_ksf_stage_from_desc( desc_ksf_s2t_root, tmem_base_col_id, tmem_base_row_id, - SMEM_P4_SCALE_B_COL_OFFSET, + SMEM_P4_KSF_RING_COL_OFFSET, 0, page_major_ksf, ) - if not split_qk15_stage1: - desc_ksf_s2t_stage1 = desc_ksf_s2t_root.advance_start_address(QK_SFB_STAGE_BYTES) - _issue_smem_p4_qk_ksf_stage_from_desc( - desc_ksf_s2t_stage1, - tmem_base_col_id, - tmem_base_row_id, - SMEM_P4_SCALE_B_COL_OFFSET, - 1, - page_major_ksf, - ) - _complete_smem_p4_qk_ksf_store() _issue_smem_p4_qk_stage_burst_from_base( sA, sK_slot, @@ -4978,7 +4863,7 @@ def _runtime_refill_qk_tile_state( qk_kblocks=SMEM_P4_QK_DATA_STAGE_KBLOCKS, acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, - sfb_col_offset_base=SMEM_P4_SCALE_B_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, use_resident_q_stages=True, b_stage_base=ctm.Int32(0), fence_after=False, @@ -4988,18 +4873,14 @@ def _runtime_refill_qk_tile_state( qk15_stage1_ready_mbar = pair_overlap_sync_mbars.subview(SMEM_P4_QK15_STAGE1_MBAR_OFFSET) while not prims.mbarrier_try_wait_parity(qk15_stage1_ready_mbar, 0, time_limit=10000000): pass - _fence_async_shared_cta() - _stage_smem_p4_qk_ksf_range_from_base( - sKSF_slot, - tmem_base_col_id, - tmem_base_row_id, - SMEM_P4_SCALE_B_COL_OFFSET, - ctm.Int32(0), - 1, - 1, - True, - page_major_ksf, - ) + _issue_smem_p4_qk_ksf_stage_from_desc( + desc_ksf_s2t_root.advance_start_address(QK_SFB_STAGE_BYTES), + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + 1, + page_major_ksf, + ) _issue_smem_p4_qk_stage_burst_from_base( sA, sK_slot, @@ -5012,7 +4893,7 @@ def _runtime_refill_qk_tile_state( qk_kblocks=SMEM_P4_QK_DATA_STAGE_KBLOCKS, acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, - sfb_col_offset_base=SMEM_P4_SCALE_B_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, use_resident_q_stages=True, b_stage_base=ctm.Int32(0), fence_after=False, @@ -5020,12 +4901,11 @@ def _runtime_refill_qk_tile_state( ) qk_handle = qk_smem_consumer.wait() qk_smem_consumer.advance() - _fence_async_shared_cta() - _stage_smem_p4_qk_ksf_range_from_base( + _issue_smem_p4_qk_ksf_range_from_base( sKSF_slot, tmem_base_col_id, tmem_base_row_id, - SMEM_P4_SCALE_B_COL_OFFSET, + SMEM_P4_KSF_RING_COL_OFFSET, ctm.Int32(0), SMEM_P4_QK_FULL_DATA_STAGES, 1, @@ -5044,7 +4924,7 @@ def _runtime_refill_qk_tile_state( qk_kblocks=QK_TAIL_DATA_STAGE_K_DIM // SMEM_P4_QK_MMA_KBLOCK_DIM, acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, - sfb_col_offset_base=SMEM_P4_SCALE_B_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, use_resident_q_stages=True, b_stage_base=ctm.Int32(0), fence_after=True, @@ -5056,7 +4936,6 @@ def _runtime_refill_qk_tile_state( multicast_mask=ctm.Int16(3), group=prims.CTAGroup.CTA_2, ) - score_slot0_handle.commit() qk_handle.release() @@ -5071,7 +4950,6 @@ def _runtime_refill_qk_tile_dispatch( qk_bulk_ready_mbars, qk_full_mbar, pair_overlap_sync_mbars, - tmem_score_slot0_producer, qk_smem_consumer, qk_next_li_idx: ctm.Int32, qk_next_stage_idx: ctm.Int32, @@ -5090,7 +4968,6 @@ def _runtime_refill_qk_tile_dispatch( qk_bulk_ready_mbars, qk_full_mbar, pair_overlap_sync_mbars, - tmem_score_slot0_producer, qk_smem_consumer, qk_next_li_idx, qk_next_stage_idx, @@ -5102,6 +4979,53 @@ def _runtime_refill_qk_tile_dispatch( ) +@cute.jit +def _runtime_refill_qk_tile_at( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id: ctm.Int32, + tmem_base_row_id: ctm.Int32, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_li_idx: ctm.Int32, + stream_li_total: ctm.Int32, +) -> None: + """Issue QK for KV tile ``qk_li_idx`` from its K ring slot into its score bank.""" + qk_stage_idx = qk_li_idx % ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS) + qk_consumer_phase = ( + qk_li_idx + // ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + # Tile 0 is loaded on the pipeline barrier, so slot 0's bulk barrier starts at tile 3. + qk_bulk_phase = qk_consumer_phase ^ ctm.Int32(qk_stage_idx == ctm.Int32(0)) + split_qk15_stage1 = ctm.Boolean(stream_li_total == ctm.Int32(16)) & ctm.Boolean( + qk_li_idx == ctm.Int32(15) + ) + _runtime_refill_qk_tile_dispatch( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id, + tmem_base_row_id, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_li_idx, + qk_stage_idx, + qk_bulk_phase, + qk_consumer_phase, + split_qk15_stage1, + False, + ) + + @cute.jit def _runtime_t336_prepare_p_source( sAtomicRunningRowMax, @@ -5109,9 +5033,8 @@ def _runtime_t336_prepare_p_source( sPSF, qk_full_mbar, pair_rowmax_read_done_mbar, - direct_p_ready_mbar, + leader_score_consumed_mbar, p4_smem_slot0_producer, - tmem_score_slot0_consumer, warp_idx: ctm.Int32, tidx: ctm.Int32, p4_tmem_base_col_id: ctm.Int32, @@ -5153,8 +5076,6 @@ def _runtime_t336_prepare_p_source( p4_tmem_base_row_id, qk_score_chunk_col_offset, ) - if is_compute_warp: - tmem_score_slot0_consumer.wait() score_values, group_max_values, group_valid_values = _load_p4_n256_score_half_from_tmem( sAtomicRunningRowMax, qk_full_mbar.subview(bank_idx), @@ -5170,6 +5091,10 @@ def _runtime_t336_prepare_p_source( producer_warp_base=producer_warp_base, score_half_idx=score_half_idx, ) + # The scores are in registers now: hand the bank back to the MMA warp. + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_score_consumed_mbar) packed_p4_words, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( score_values, group_max_values, @@ -5231,8 +5156,8 @@ def _runtime_t336_prepare_p_source( p_stage_base=p_stage_base, score_half_idx=score_half_idx, ) - if prims.elect_sync(): - prims.mbarrier_arrive(direct_p_ready_mbar) + if is_compute_warp: + cute.nvgpu.cfence() owned_psf_word, prequant_row_sum = _finalize_p4_n256_score_half_psf_source( sPSF, group_max_values, @@ -5624,20 +5549,402 @@ def _runtime_t336_producer_tile( @cute.jit -def _runtime_t336_producer_tile_v23( - sAtomicRunningRowMax, - sP, - sPSF, - sScoreSlotRowMax, - sRowScaleRing, +def _issue_pipelined_score_loads( + qk_full_mbar, stream_li_idx: ctm.Int32, warp_idx, base_col_id, base_row_id +): + """Wait for tile ``stream_li_idx``'s scores and issue their TMEM loads without waiting on them.""" + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + qk_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + score_tmem_addresses = _prepare_p4_n256_score_half_tmem_addresses( + warp_idx, + base_col_id, + base_row_id, + ctm.Int32(SMEM_P4_QK_ACC_COL_OFFSET) + + bank_idx * ctm.Int32(SMEM_P4_QK_PIPELINE_SLOT_STRIDE), + ) + while not prims.mbarrier_try_wait_parity( + qk_full_mbar.subview(bank_idx), qk_phase, time_limit=10000000 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + pending = [] + for local_group_idx in ctm.range_constexpr(4): + tmem = ctm.inttoptr(score_tmem_addresses[local_group_idx], 6, ctm.Float32) + pending.append(_tcgen05_ld_red_32x32b_x16_max_f32(tmem)) + return pending + + +@cute.jit +def _finish_pipelined_score_loads( + pending, + sPipelinedRowMax, + leader_pair_overlap_sync_mbars, + stream_li_idx: ctm.Int32, + warp_idx: ctm.Int32, + tidx: ctm.Int32, +): + """Wait for the score loads, publish the tile row max by tile parity and release the score bank.""" + producer_warp_base = (warp_idx >> ctm.Int32(2)) << ctm.Int32(2) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + score_values = [] + group_max_values = [] + tile_row_max = ctm.Float32(-ctm.Float32.inf) + for local_group_idx in ctm.range_constexpr(4): + regs = pending[local_group_idx] + for elem in ctm.range_constexpr(SF_VEC_SIZE): + score_values.append(regs[elem].bitcast(ctm.Float32)) + group_max = regs[SF_VEC_SIZE].bitcast(ctm.Float32) + group_max_values.append(group_max) + tile_row_max = cute.arch.fmax(tile_row_max, group_max) + parity = stream_li_idx & ctm.Int32(1) + _publish_p4_n256_atomic_rowmax( + sPipelinedRowMax.subview(parity * ctm.Int32(SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS)), + tile_row_max, + warp_idx, + tidx, + producer_warp_base, + ) + cute.nvgpu.cfence() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar( + leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + + stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + ) + ) + return ( + ctm.Vector.from_elements(tuple(score_values), dtype=ctm.Float32), + ctm.Vector.from_elements(tuple(group_max_values), dtype=ctm.Float32), + ) + + +@cute.jit +def _prepack_pipelined_tile( + score_values, + group_max_values, + warp_idx, + softmax_scale_log2, + stream_li_idx, + valid_k, + use_mixed_imlp: ctm.Constexpr, +): + """Pack an unmasked tile's FP4 P words and group sums ahead of its row-max exchange. + + Returns 16 scalars (8 P words, 4 group maxima, 4 group sums): the tile loop carries them to + the next iteration, and the DSL cannot rebind a list inside a dynamic loop. + """ + producer_warp_base = (warp_idx >> ctm.Int32(2)) << ctm.Int32(2) + score_half_idx = warp_idx >> ctm.Int32(2) + group_valid_values = ctm.Vector.from_elements( + tuple((ctm.Boolean(True) for _ in range(4))), dtype=ctm.Boolean + ) + packed, group_p_sums = _prepack_p4_n256_score_half_before_rowmax( + score_values, + group_max_values, + group_valid_values, + warp_idx, + softmax_scale_log2, + kv_tile_idx=stream_li_idx, + valid_k=valid_k, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + mask_valid_k=False, + use_mixed_imlp=use_mixed_imlp, + ) + return ( + packed[0], + packed[1], + packed[2], + packed[3], + packed[4], + packed[5], + packed[6], + packed[7], + group_max_values[0], + group_max_values[1], + group_max_values[2], + group_max_values[3], + group_p_sums[0], + group_p_sums[1], + group_p_sums[2], + group_p_sums[3], + ) + + +@cute.jit +def _finalize_pipelined_tile( + sPipelinedRowMax, + sP, + sPSF, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + carried, + valid_k: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + running_row_max: ctm.Float32, + running_row_anchor: ctm.Float32, + running_part_sum: ctm.Float32, + stream_li_idx: ctm.Int32, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + """Finish tile ``stream_li_idx`` from its prepacked P words. + + Same steps as ``_runtime_t336_prepare_p_source`` from the dual-half barrier on followed by the + tail of ``_runtime_t336_producer_tile_v23``: row anchor, P and PSF stores, lazy O rescale and + the P_SOURCE_READY arrive. + """ + bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + p_stage_base = bank_idx * ctm.Int32(SMEM_P4_SCORE_SLOTS_PER_KV_TILE) + score_half_idx = warp_idx >> ctm.Int32(2) + producer_warp_base = score_half_idx << ctm.Int32(2) + is_compute_warp = score_half_idx == ctm.Int32(0) + has_previous_tile = stream_li_idx > ctm.Int32(0) + packed_p4_words = ctm.Vector.from_elements(tuple(carried[0:8]), dtype=ctm.Int32) + group_max_values = ctm.Vector.from_elements(tuple(carried[8:12]), dtype=ctm.Float32) + group_p_sums = ctm.Vector.from_elements(tuple(carried[12:16]), dtype=ctm.Float32) + pair_rowmax_read_done_mbar = pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_mbar = leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + leader_pair_p_source_ready_dsmem_mbar = leader_pair_overlap_sync_mbars_dsmem.subview( + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx + ) + if is_compute_warp: + p4_smem_slot0_producer.acquire() + prims.barrier( + barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, + number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, + ) + parity = stream_li_idx & ctm.Int32(1) + new_row_max = cute.arch.fmax( + running_row_max, + _load_p4_n256_atomic_running_rowmax( + sPipelinedRowMax.subview(parity * ctm.Int32(SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS)), + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + ), + ) + next_running_row_anchor, lane_rebase, anchor_delta_log2 = _select_smem_p4_lazy_anchor_runtime( + new_row_max, + running_row_anchor, + softmax_scale_log2, + has_previous_tile, + ) + if is_compute_warp: + cute.nvgpu.cfence() + if prims.elect_sync(): + prims.mbarrier_arrive(pair_rowmax_read_done_mbar) + _store_prepacked_p4_n256_score_half_to_direct_a( + sP, + packed_p4_words, + warp_idx, + tidx, + producer_warp_base=producer_warp_base, + p_stage_base=p_stage_base, + score_half_idx=score_half_idx, + ) + if is_compute_warp: + cute.nvgpu.cfence() + owned_psf_word, row_sum = _finalize_p4_n256_score_half_psf_source( + sPSF, + group_max_values, + group_p_sums, + next_running_row_anchor, + warp_idx, + tidx, + softmax_scale_log2, + pv_psf_rescale, + kv_tile_idx=stream_li_idx, + valid_k=valid_k, + mask_valid_k=False, + producer_warp_base=producer_warp_base, + score_half_idx=score_half_idx, + bank_idx=bank_idx, + use_mixed_imlp=use_mixed_imlp, + ) + current_old_scale = ctm.Float32(1.0) + warp_rebase = cute.arch.vote_any_sync(lane_rebase) + if has_previous_tile: + common_source_ready_pred = prims.elect_sync() & ~warp_rebase + if common_source_ready_pred: + prims.mbarrier_arrive( + leader_pair_p_source_ready_dsmem_mbar, + count=1, + scope=prims.MemScope.CLUSTER, + relaxed=True, + ) + if warp_rebase: + pv_done_li_idx = stream_li_idx - ctm.Int32(1) + pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + pv_done_phase = ( + pv_done_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pv_full_mbar.subview(pv_done_slot), pv_done_phase, time_limit=10000000 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + current_old_scale = _compute_smem_p4_lazy_rebase_scale(lane_rebase, anchor_delta_log2) + _rescale_smem_p4_resident_o_in_tmem(corr_o_acc_tmem, current_old_scale) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + elif prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + if is_compute_warp: + p4_smem_slot0_producer.advance() + next_running_part_sum = running_part_sum * current_old_scale + row_sum + return (new_row_max, next_running_row_anchor, next_running_part_sum) + + +@cute.jit +def _run_pipelined_leading_tiles( + sAtomicRunningRowMax, + sPipelinedRowMax, + sP, + sPSF, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx: ctm.Int32, + tidx: ctm.Int32, + base_col_id: ctm.Int32, + base_row_id: ctm.Int32, + valid_k: ctm.Int32, + softmax_scale_log2: ctm.Float32, + pv_psf_rescale: ctm.Float32, + pipelined_tiles: ctm.Int32, + use_mixed_imlp: ctm.Constexpr, +) -> tuple: + """Run producer tiles [0, pipelined_tiles), which need no causal mask, pipelined by one tile. + + Tile t + 1's score loads are in flight while tile t is finished, so the TMEM load latency + and the row-max exchange of one tile overlap the P stores of the previous one. + """ + pending = _issue_pipelined_score_loads( + qk_full_mbar, ctm.Int32(0), warp_idx, base_col_id, base_row_id + ) + scores, group_max = _finish_pipelined_score_loads( + pending, sPipelinedRowMax, leader_pair_overlap_sync_mbars, ctm.Int32(0), warp_idx, tidx + ) + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15) = ( + _prepack_pipelined_tile( + scores, group_max, warp_idx, softmax_scale_log2, ctm.Int32(0), valid_k, use_mixed_imlp + ) + ) + running_row_max = ctm.Float32(-ctm.Float32.inf) + running_row_anchor = ctm.Float32(-ctm.Float32.inf) + running_part_sum = ctm.Float32(0.0) + for stream_li_idx in cutlass.range(0, pipelined_tiles - ctm.Int32(1), 1, unroll=1): + next_li_idx = stream_li_idx + ctm.Int32(1) + next_pending = _issue_pipelined_score_loads( + qk_full_mbar, next_li_idx, warp_idx, base_col_id, base_row_id + ) + running_row_max, running_row_anchor, running_part_sum = _finalize_pipelined_tile( + sPipelinedRowMax, + sP, + sPSF, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15), + valid_k, + softmax_scale_log2, + pv_psf_rescale, + running_row_max, + running_row_anchor, + running_part_sum, + stream_li_idx, + use_mixed_imlp, + ) + next_scores, next_group_max = _finish_pipelined_score_loads( + next_pending, + sPipelinedRowMax, + leader_pair_overlap_sync_mbars, + next_li_idx, + warp_idx, + tidx, + ) + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15) = ( + _prepack_pipelined_tile( + next_scores, + next_group_max, + warp_idx, + softmax_scale_log2, + next_li_idx, + valid_k, + use_mixed_imlp, + ) + ) + running_row_max, running_row_anchor, running_part_sum = _finalize_pipelined_tile( + sPipelinedRowMax, + sP, + sPSF, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15), + valid_k, + softmax_scale_log2, + pv_psf_rescale, + running_row_max, + running_row_anchor, + running_part_sum, + pipelined_tiles - ctm.Int32(1), + use_mixed_imlp, + ) + # The per-tile loop that follows reads the running row max from sAtomicRunningRowMax. + _publish_p4_n256_atomic_rowmax( + sAtomicRunningRowMax, + running_row_max, + warp_idx, + tidx, + (warp_idx >> ctm.Int32(2)) << ctm.Int32(2), + ) + return (running_row_max, running_row_anchor, running_part_sum) + + +@cute.jit +def _runtime_t336_producer_tile_v23( + sAtomicRunningRowMax, + sP, + sPSF, qk_full_mbar, pair_overlap_sync_mbars, leader_pair_overlap_sync_mbars, leader_pair_overlap_sync_mbars_dsmem, p4_smem_slot0_producer, - p4_smem_slot0_mbar_ptr, - tmem_score_slot0_consumer, - tmem_score_slot0_mbar_ptr, pv_full_mbar, corr_o_acc_tmem, warp_idx: ctm.Int32, @@ -5651,6 +5958,7 @@ def _runtime_t336_producer_tile_v23( running_row_anchor: ctm.Float32, carried_dq1024: ctm.Int32, carried_old_scale: ctm.Float32, + running_part_sum: ctm.Float32, stream_li_idx: ctm.Int32, has_previous_tile: ctm.Boolean, tail_needs_mask: ctm.Boolean, @@ -5665,24 +5973,19 @@ def _runtime_t336_producer_tile_v23( // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) ) - rowsum_full_mbar = pair_overlap_sync_mbars.subview(SMEM_P4_PAIR_P4_READY_MBAR_OFFSET + bank_idx) pair_rowmax_read_done_mbar = pair_overlap_sync_mbars.subview( SMEM_P4_PAIR_ROWMAX_READ_DONE_MBAR_OFFSET + bank_idx ) - direct_p_ready_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET + bank_idx - ) leader_pair_p_source_ready_mbar = leader_pair_overlap_sync_mbars.subview( SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx ) leader_pair_p_source_ready_dsmem_mbar = leader_pair_overlap_sync_mbars_dsmem.subview( SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + bank_idx ) + leader_score_consumed_mbar = leader_pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + bank_idx + ) score_half_idx = warp_idx >> ctm.Int32(2) - producer_warp_base = score_half_idx << ctm.Int32(2) - producer_local_warp_idx = warp_idx & ctm.Int32(3) - producer_row_band = producer_local_warp_idx & ctm.Int32(1) - producer_col_band = producer_local_warp_idx >> ctm.Int32(1) is_compute_warp = score_half_idx == ctm.Int32(0) next_running_row_max = running_row_max @@ -5706,9 +6009,8 @@ def _runtime_t336_producer_tile_v23( sPSF, qk_full_mbar, pair_rowmax_read_done_mbar, - direct_p_ready_mbar, + leader_score_consumed_mbar, p4_smem_slot0_producer, - tmem_score_slot0_consumer, warp_idx, tidx, p4_tmem_base_col_id, @@ -5735,9 +6037,8 @@ def _runtime_t336_producer_tile_v23( sPSF, qk_full_mbar, pair_rowmax_read_done_mbar, - direct_p_ready_mbar, + leader_score_consumed_mbar, p4_smem_slot0_producer, - tmem_score_slot0_consumer, warp_idx, tidx, p4_tmem_base_col_id, @@ -5764,9 +6065,8 @@ def _runtime_t336_producer_tile_v23( sPSF, qk_full_mbar, pair_rowmax_read_done_mbar, - direct_p_ready_mbar, + leader_score_consumed_mbar, p4_smem_slot0_producer, - tmem_score_slot0_consumer, warp_idx, tidx, p4_tmem_base_col_id, @@ -5790,84 +6090,34 @@ def _runtime_t336_producer_tile_v23( leader_pair_p_source_ready_dsmem_mbar, count=1, scope=prims.MemScope.CLUSTER, - relaxed=True, - ) - if warp_rebase: - pv_done_li_idx = stream_li_idx - ctm.Int32(1) - pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - pv_done_phase = ( - pv_done_li_idx - // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - while not prims.mbarrier_try_wait_parity( - pv_full_mbar.subview(pv_done_slot), pv_done_phase, time_limit=10000000 - ): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - current_old_scale = _compute_smem_p4_lazy_rebase_scale(lane_rebase, anchor_delta_log2) - _rescale_smem_p4_resident_o_in_tmem(corr_o_acc_tmem, current_old_scale) - if prims.elect_sync(): - _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) - elif prims.elect_sync(): - _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) - if is_compute_warp: - if cutlass.const_expr(not is_final_tile): - while not prims.mbarrier_try_wait_parity( - direct_p_ready_mbar, qk_phase, time_limit=10000000 - ): - pass - cute.arch.mbarrier_arrive( - p4_smem_slot0_mbar_ptr + p4_smem_slot0_producer.current_handle().index, - cutlass.Int32(0), - ) - p4_smem_slot0_producer.advance() - cute.arch.mbarrier_arrive( - tmem_score_slot0_mbar_ptr - + SMEM_P4_TMEM_SCORE_PIPELINE_STAGES - + tmem_score_slot0_consumer.current_handle().index, - cutlass.Int32(0), - ) - tmem_score_slot0_consumer.advance() - if is_compute_warp: - prims.barrier( - barrier_id=SMEM_P4_P4_SLOT0_BAR_ID, - number_of_threads=SMEM_P4_SCORE_SLOT_PRODUCER_THREADS, - ) - else: - prims.barrier( - barrier_id=SMEM_P4_P4_SLOT1_BAR_ID, - number_of_threads=SMEM_P4_CORRECTION_THREADS, - ) - rowsum_empty_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + bank_idx - ) - if stream_li_idx >= ctm.Int32(SMEM_P4_TMEM_BANKS): - previous_rowsum_phase = ( - (stream_li_idx - ctm.Int32(SMEM_P4_TMEM_BANKS)) - // ctm.Int32(SMEM_P4_TMEM_BANKS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - while not prims.mbarrier_try_wait_parity( - rowsum_empty_mbar, previous_rowsum_phase, time_limit=10000000 - ): - pass - _store_smem_p4_n256_rowsum_part( - sScoreSlotRowMax, - warp_idx, - tidx, - producer_warp_base=producer_warp_base, - bank_idx=bank_idx, - score_half_idx=score_half_idx, - value=row_sum, - ) - if is_compute_warp & (producer_col_band == ctm.Int32(0)): - row_state_local_row = producer_row_band * ctm.Int32(32) + (tidx & ctm.Int32(31)) - sRowScaleRing[bank_idx * ctm.Int32(SMEM_P4_BMM1_M) + row_state_local_row] = ( - current_old_scale - ) - if prims.elect_sync(): - prims.mbarrier_arrive(rowsum_full_mbar) + relaxed=True, + ) + if warp_rebase: + pv_done_li_idx = stream_li_idx - ctm.Int32(1) + pv_done_slot = pv_done_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + pv_done_phase = ( + pv_done_li_idx + // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pv_full_mbar.subview(pv_done_slot), pv_done_phase, time_limit=10000000 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + current_old_scale = _compute_smem_p4_lazy_rebase_scale(lane_rebase, anchor_delta_log2) + _rescale_smem_p4_resident_o_in_tmem(corr_o_acc_tmem, current_old_scale) + if prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + elif prims.elect_sync(): + _arrive_smem_p4_mapped_leader_mbar(leader_pair_p_source_ready_mbar) + # P_SOURCE_READY is the only P-side signal the MMA warp needs: it implies + # this tile's P/PSF are in smem and its score bank has been read. + if is_compute_warp: + p4_smem_slot0_producer.advance() + # Each warp keeps the denominator of its own 64-column slice in a register, + # rebased with the row anchor; the slices are combined once after the loop. + next_running_part_sum = running_part_sum * current_old_scale + row_sum if not is_compute_warp: while not prims.mbarrier_try_wait_parity( pair_rowmax_read_done_mbar, qk_phase, time_limit=10000000 @@ -5878,6 +6128,7 @@ def _runtime_t336_producer_tile_v23( next_running_row_anchor, next_carried_dq1024, next_carried_old_scale, + next_running_part_sum, ) @@ -6159,12 +6410,6 @@ def _run_mla_decode_body( space=ctm.AddressSpace.smem, alignment=8, ) - tmem_score_slot0_mbars = ctm.Array( - ctm.Int64, - SMEM_P4_TMEM_SCORE_PIPELINE_STAGES * 2, - space=ctm.AddressSpace.smem, - alignment=8, - ) p4_smem_slot0_mbars = ctm.Array( ctm.Int64, SMEM_P4_P4_CONTROL_MBAR_STAGES * 2, @@ -6234,21 +6479,16 @@ def _run_mla_decode_body( space=ctm.AddressSpace.smem, alignment=128, ) - sRowScaleRing = ctm.Array( - ctm.Float32, - SMEM_P4_ROW_SCALE_RING_FLOATS, - space=ctm.AddressSpace.smem, - alignment=128, - ) - sFinalAnchorRowSum = ctm.Array( - ctm.Float32, - SMEM_P4_FINAL_ANCHOR_ROWSUM_FLOATS, + sAtomicRunningRowMax = ctm.Array( + ctm.Uint32, + SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS, space=ctm.AddressSpace.smem, alignment=128, ) - sAtomicRunningRowMax = ctm.Array( + # Pipelined producer tiles publish their row maxima by tile parity. + sPipelinedRowMax = ctm.Array( ctm.Uint32, - SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS, + 2 * SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS, space=ctm.AddressSpace.smem, alignment=128, ) @@ -6280,9 +6520,6 @@ def _run_mla_decode_body( v_smem_mbar_ptr = cute.make_ptr( ctm.Int64, v_smem_mbars.data_ptr().toint(), cute.AddressSpace.smem ) - tmem_score_slot0_mbar_ptr = cute.make_ptr( - ctm.Int64, tmem_score_slot0_mbars.data_ptr().toint(), cute.AddressSpace.smem - ) p4_smem_slot0_mbar_ptr = cute.make_ptr( ctm.Int64, p4_smem_slot0_mbars.data_ptr().toint(), cute.AddressSpace.smem ) @@ -6319,15 +6556,6 @@ def _run_mla_decode_body( defer_sync=True, name="fp4_v_smem", ).make_participants() - tmem_score_slot0_producer, tmem_score_slot0_consumer = pipeline.PipelineUmmaAsync.create( - num_stages=SMEM_P4_TMEM_SCORE_PIPELINE_STAGES, - producer_group=umma_group, - consumer_group=p4_stage_group, - barrier_storage=tmem_score_slot0_mbar_ptr, - cta_layout_vmnk=cta_layout_vmnk, - defer_sync=True, - name="fp4_tmem_score_slot0", - ).make_participants() p4_smem_slot0_producer, _ = pipeline.PipelineAsyncUmma.create( num_stages=SMEM_P4_P4_PIPELINE_STAGES, producer_group=p4_stage_group, @@ -6370,13 +6598,13 @@ def _run_mla_decode_body( ) prims.mbarrier_init( pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET + pair_slot_idx + SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + pair_slot_idx ), - SMEM_P4_TOTAL_PRODUCER_WARPS, + SMEM_P4_TOTAL_PRODUCER_WARPS * CLUSTER_SHAPE_MNK[0], ) prims.mbarrier_init( pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_P_SOURCE_READY_MBAR_OFFSET + pair_slot_idx + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + pair_slot_idx ), SMEM_P4_TOTAL_PRODUCER_WARPS * CLUSTER_SHAPE_MNK[0], ) @@ -6534,7 +6762,8 @@ def _run_mla_decode_body( qk_steady_end = ctm.Int32(15) for kv_tile_idx in cutlass.range(3, qk_steady_end, 1, unroll=1): qk_handle = qk_smem_producer.acquire_and_advance( - expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] + _sleep_until_producer_slot_free(qk_smem_producer), + expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0], ) qk_tma_mbar = ctm.Array(qk_handle.barrier, shape=1) if is_leader_cta: @@ -6571,6 +6800,32 @@ def _run_mla_decode_body( use_consecutive_page_pair=use_consecutive_page_pair, use_ksf_gather4=use_ksf_gather4, ) + else: + if cutlass.const_expr(use_smem_page_plan): + qk_rank1_page0, qk_rank1_page1 = _load_staged_page_native_tile_pair( + sPageIdPlan, kv_tile_idx + ) + else: + qk_rank1_page0, qk_rank1_page1 = _load_global_page_native_tile_pair( + mPageTable_pl, kv_tile_idx, csr_page_begin, csr_page_count + ) + _load_runtime_t336_qk_tile_rank1( + mPageTable_pl, + tma_k_page_ptr, + tma_k_tail_page_ptr, + tma_ksf_ptr, + tma_ksf_stage_ptr, + tma_ksf_tail_stage_ptr, + qk_handle.barrier, + qk_bulk_ready_mbar_ptr, + sB, + sSFB, + qk_rank1_page0, + qk_rank1_page1, + kv_tile_idx, + kv_tile_idx % ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS), + page_size, + ) if qk15_final_enabled: qk15_handle = qk_smem_producer.acquire_and_advance( expected_tx=SMEM_P4_QK_KONLY_TAIL_STAGE_BYTES * CLUSTER_SHAPE_MNK[0] @@ -6621,7 +6876,9 @@ def _run_mla_decode_body( sPageIdPlan, ctm.Int32(0) ) for kv_tile_idx in cutlass.range(0, kv_tiles, 1, unroll=1): - v_handle = v_smem_producer.acquire_and_advance() + v_handle = v_smem_producer.acquire_and_advance( + _sleep_until_producer_slot_free(v_smem_producer) + ) v_tma_mbar = ctm.Array(v_handle.barrier, shape=1) v_stage = kv_tile_idx % ctm.Int32(SMEM_P4_V_PIPELINE_STAGES) _load_v_tile_stage( @@ -6657,7 +6914,8 @@ def _run_mla_decode_body( sPageIdPlan, kv_tile_idx + ctm.Int32(1) ) v_smem_producer.tail() - if warp_idx == ROWMETA_WARP_ID: + if warp_idx == PV_ISSUE_WARP_ID: + prims.barrier(barrier_id=TMEM_BAR_ID, number_of_threads=TMEM_BAR_THREADS) if is_leader_cta: qk_prefix_end = min(kv_tiles, ctm.Int32(3)) for qk_prefix_li_idx in cutlass.range(1, qk_prefix_end, 1, unroll=1): @@ -6690,63 +6948,26 @@ def _run_mla_decode_body( qk_prefix_li_idx, page_size, ) - rowmeta_lane = tidx & ctm.Int32(31) - running_anchor_row_sum0 = ctm.Float32(0.0) - running_anchor_row_sum1 = ctm.Float32(0.0) - for rowmeta_li_idx in cutlass.range(0, kv_tiles, 1, unroll=1): - rowmeta_bank_idx = rowmeta_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) - rowmeta_phase = ( - rowmeta_li_idx - // ctm.Int32(SMEM_P4_TMEM_BANKS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - rowsum_full_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_P4_READY_MBAR_OFFSET + rowmeta_bank_idx - ) - while not prims.mbarrier_try_wait_parity( - rowsum_full_mbar, rowmeta_phase, time_limit=10000000 - ): - pass - row0 = rowmeta_lane - h0c0_row0 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row0, rowmeta_bank_idx, 0, 0 - ) - h0c1_row0 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row0, rowmeta_bank_idx, 0, 1 - ) - h1c0_row0 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row0, rowmeta_bank_idx, 1, 0 - ) - h1c1_row0 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row0, rowmeta_bank_idx, 1, 1 - ) - stage_row_sum0 = h0c0_row0 + h1c0_row0 + (h0c1_row0 + h1c1_row0) - old_scale0 = sRowScaleRing[ctm.Int32(rowmeta_bank_idx * SMEM_P4_BMM1_M) + row0] - running_anchor_row_sum0 = running_anchor_row_sum0 * old_scale0 + stage_row_sum0 - row1 = rowmeta_lane + ctm.Int32(32) - h0c0_row1 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row1, rowmeta_bank_idx, 0, 0 - ) - h0c1_row1 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row1, rowmeta_bank_idx, 0, 1 - ) - h1c0_row1 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row1, rowmeta_bank_idx, 1, 0 - ) - h1c1_row1 = _load_smem_p4_n256_quadrant_rowsum( - sScoreSlotRowMax, sP, row1, rowmeta_bank_idx, 1, 1 - ) - stage_row_sum1 = h0c0_row1 + h1c0_row1 + (h0c1_row1 + h1c1_row1) - old_scale1 = sRowScaleRing[ctm.Int32(rowmeta_bank_idx * SMEM_P4_BMM1_M) + row1] - running_anchor_row_sum1 = running_anchor_row_sum1 * old_scale1 + stage_row_sum1 - if rowmeta_li_idx + ctm.Int32(1) == kv_tiles: - sFinalAnchorRowSum[row0] = running_anchor_row_sum0 - sFinalAnchorRowSum[row1] = running_anchor_row_sum1 - if prims.elect_sync(): - rowsum_empty_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + rowmeta_bank_idx + pv_tmem_raw_addr = tmem_holding_buf.load() + pv_tmem_base_col_id = pv_tmem_raw_addr & ctm.Int32(65535) + pv_tmem_base_row_id = pv_tmem_raw_addr >> ctm.Int32(16) + for pv_li_idx in cutlass.range(0, kv_tiles, 1, unroll=1): + v_smem_consumer.wait_and_advance() + _runtime_issue_pv_tile_dispatch( + sP, + sV, + sPSF, + sVSF, + pv_tmem_base_col_id, + pv_tmem_base_row_id, + pair_overlap_sync_mbars, + p4_smem_slot0_mbars, + pv_full_mbar, + v_smem_mbars, + pv_li_idx, + kv_tiles, + ctm.Boolean(pv_li_idx == ctm.Int32(0)), ) - prims.mbarrier_arrive(rowsum_empty_mbar) if warp_idx == MMA_WARP_ID: if is_leader_cta: tmem_raw_addr = tmem_holding_buf.load() @@ -6780,19 +7001,31 @@ def _run_mla_decode_body( stream_li_total = kv_tiles qk0_handle = qk_smem_consumer.wait() qk_smem_consumer.advance() - qk0_score_handle = tmem_score_slot0_producer.acquire_and_advance() - for qk_burst_start in ctm.range_constexpr( - 0, QK_MMA_KBLOCKS, SMEM_P4_QK_SCALE_BURST_KBLOCKS - ): - _fence_async_shared_cta() - _stage_smem_p4_qk_full_ksf_from_base( - sSFB, - tmem_base_col_id, - tmem_base_row_id, - sfb_col_offset=SMEM_P4_SCALE_B_COL_OFFSET, - sfb_stage_base=0, - page_major_ksf=False, - ) + # Tile 0 sits in K slot 0; its stages go through the K scale ring + # like every refill. + qk0_desc_ksf = _make_smem_p4_qk_ksf_stage_desc(sSFB, ctm.Int32(0), 0, True, False) + for qk0_stage in ctm.range_constexpr(SMEM_P4_QK_DATA_STAGES): + if cutlass.const_expr(qk0_stage < SMEM_P4_QK_FULL_DATA_STAGES): + _issue_smem_p4_qk_ksf_stage_from_desc( + qk0_desc_ksf.advance_start_address(qk0_stage * QK_SFB_STAGE_BYTES), + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + qk0_stage, + False, + ) + else: + _issue_smem_p4_qk_ksf_range_from_base( + sSFB, + tmem_base_col_id, + tmem_base_row_id, + SMEM_P4_KSF_RING_COL_OFFSET, + ctm.Int32(0), + SMEM_P4_QK_FULL_DATA_STAGES, + 1, + True, + False, + ) _issue_smem_p4_qk_stage_burst_from_base( sA, sB, @@ -6801,13 +7034,17 @@ def _run_mla_decode_body( tmem_base_col_id, tmem_base_row_id, score_slot_idx=0, - qk_kblock_start=qk_burst_start, - qk_kblocks=SMEM_P4_QK_SCALE_BURST_KBLOCKS, + qk_kblock_start=qk0_stage * SMEM_P4_QK_DATA_STAGE_KBLOCKS, + qk_kblocks=SMEM_P4_QK_DATA_STAGE_KBLOCKS + if qk0_stage < SMEM_P4_QK_FULL_DATA_STAGES + else QK_TAIL_DATA_STAGE_K_DIM // SMEM_P4_QK_MMA_KBLOCK_DIM, acc_col_offset=SMEM_P4_QK_ACC_COL_OFFSET, sfa_col_offset_base=SMEM_P4_QK_SFA_COL_OFFSET, - sfb_col_offset_base=SMEM_P4_SCALE_B_COL_OFFSET, + sfb_col_offset_base=SMEM_P4_KSF_RING_COL_OFFSET, use_resident_q_stages=True, - b_stage_base=0, + b_stage_base=ctm.Int32(0), + fence_after=qk0_stage == SMEM_P4_QK_DATA_STAGES - 1, + b_smem_is_compact_slot_base=True, ) if prims.elect_sync(): prims.tcgen05_commit( @@ -6815,26 +7052,45 @@ def _run_mla_decode_body( multicast_mask=ctm.Int16(3), group=prims.CTAGroup.CTA_2, ) - qk0_score_handle.commit() qk0_handle.release() - for stream_li_idx in cutlass.range(0, stream_li_total, 1, unroll=1): - qk_next_li_idx = stream_li_idx + ctm.Int32(1) - qk_next_stage_idx = qk_next_li_idx % ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS) - qk_consumer_phase = ( - qk_next_li_idx - // ctm.Int32(SMEM_P4_QK_SMEM_PIPELINE_SLOTS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + # QK runs two tiles ahead of PV: QK(t) refills score bank t % 2 as + # soon as the producers have read tile t - 2 from it instead of + # waiting behind PV(t - 2), which needs that tile's P. + for qk_ahead_li_idx in cutlass.range( + 1, min(stream_li_total, ctm.Int32(2)), 1, unroll=1 + ): + _runtime_refill_qk_tile_at( + sA, + sB, + sSFA, + sSFB, + tmem_base_col_id, + tmem_base_row_id, + qk_bulk_ready_mbars, + qk_full_mbar, + pair_overlap_sync_mbars, + qk_smem_consumer, + qk_ahead_li_idx, + stream_li_total, ) - qk_stage_wrapped = ctm.Int32(qk_next_stage_idx == ctm.Int32(0)) - v_stage = ( - qk_next_stage_idx + ctm.Int32(SMEM_P4_V_PIPELINE_STAGES - 1) - ) % ctm.Int32(SMEM_P4_V_PIPELINE_STAGES) - v_phase = qk_consumer_phase ^ qk_stage_wrapped - if qk_next_li_idx < stream_li_total: - split_qk15_stage1 = ctm.Boolean( - stream_li_total == ctm.Int32(16) - ) & ctm.Boolean(qk_next_li_idx == ctm.Int32(15)) - _runtime_refill_qk_tile_dispatch( + for stream_li_idx in cutlass.range(0, stream_li_total, 1, unroll=1): + qk_ahead_li_idx = stream_li_idx + ctm.Int32(2) + if qk_ahead_li_idx < stream_li_total: + score_bank_idx = stream_li_idx % ctm.Int32(SMEM_P4_TMEM_BANKS) + score_consumed_phase = ( + stream_li_idx + // ctm.Int32(SMEM_P4_TMEM_BANKS) + % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) + ) + while not prims.mbarrier_try_wait_parity( + pair_overlap_sync_mbars.subview( + SMEM_P4_PAIR_SCORE_CONSUMED_MBAR_OFFSET + score_bank_idx + ), + score_consumed_phase, + time_limit=10000000, + ): + pass + _runtime_refill_qk_tile_at( sA, sB, sSFA, @@ -6844,37 +7100,13 @@ def _run_mla_decode_body( qk_bulk_ready_mbars, qk_full_mbar, pair_overlap_sync_mbars, - tmem_score_slot0_producer, qk_smem_consumer, - qk_next_li_idx, - qk_next_stage_idx, - v_phase, - qk_consumer_phase, - split_qk15_stage1, - False, + qk_ahead_li_idx, + stream_li_total, ) - v_smem_consumer.wait_and_advance() - _fence_async_shared_cta() - _runtime_issue_pv_tile_dispatch( - sP, - sV, - sPSF, - sVSF, - tmem_base_col_id, - tmem_base_row_id, - qk_full_mbar, - pair_overlap_sync_mbars, - p4_smem_slot0_mbars, - pv_full_mbar, - v_smem_mbars, - stream_li_idx, - stream_li_total, - ctm.Boolean(stream_li_idx == ctm.Int32(0)), - ) q_res_handle0.release() q_res_handle1.release() q_res_handle2.release() - tmem_score_slot0_producer.tail() if warp_idx == MMA_WARP_ID and (not is_leader_cta): while not prims.mbarrier_try_wait_parity(qsf_owner_ready_mbar, 0, time_limit=10000000): pass @@ -6894,8 +7126,11 @@ def _run_mla_decode_body( sAtomicRunningRowMax[atomic_rowmax_row] = _float_to_ordered_u32_for_atomic_max( ctm.Float32(-ctm.Float32.inf) ) + if warp_idx < ctm.Int32(2 * SMEM_P4_ATOMIC_RUNNING_ROWMAX_SMEM_ROWS // 32): + sPipelinedRowMax[warp_idx * ctm.Int32(32) + (tidx & ctm.Int32(31))] = ( + _float_to_ordered_u32_for_atomic_max(ctm.Float32(-ctm.Float32.inf)) + ) prims.barrier(barrier_id=TMEM_BAR_ID, number_of_threads=TMEM_BAR_THREADS) - is_compute_warp = warp_idx < ctm.Int32(SMEM_P4_CORRECTION_WARP_ID_BEGIN) is_correction_warp = warp_idx >= ctm.Int32(SMEM_P4_CORRECTION_WARP_ID_BEGIN) if warp_idx < ctm.Int32(SMEM_P4_COMPUTE_WARPS): qsf_tmem_raw_addr = tmem_holding_buf.load() @@ -6984,40 +7219,36 @@ def _run_mla_decode_body( leader_pair_overlap_sync_mbars_dsmem = leader_pair_overlap_sync_mbars carried_dq1024 = ctm.Int32(0) carried_old_scale = ctm.Float32(1.0) + running_part_sum = ctm.Float32(0.0) stream_li_idx = ctm.Int32(0) stream_li_remaining = stream_li_total - while stream_li_remaining > ctm.Int32(1): - producer_needs_mask = ctm.Boolean(False) - if cutlass.const_expr(folded_query_heads > 0): - # The MTP window may straddle a 256-token tile boundary, so - # the penultimate tile can also need per-row causal masking. - producer_needs_mask = folded_first_valid_k < ( - (stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE) - ) + if cutlass.const_expr(folded_query_heads == 0): + # Leading tiles that need no per-row mask run pipelined; the final tile always goes + # through the per-tile loop below. + pipelined_tiles = stream_li_total - ctm.Int32(1) if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): - # Rows of one tile differ by at most query_len tokens, so only - # tiles past the earliest row bound need per-row masking. The - # bound is re-read from smem each tile: kept in a register it - # was spilled at the producers' register ceiling and reloaded - # from local memory here instead. - producer_needs_mask = _load_smem_i32( - sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) - ) < ((stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE)) - running_row_max, running_row_anchor, carried_dq1024, carried_old_scale = ( - _runtime_t336_producer_tile_v23( + pipelined_tiles = ctm.min( + pipelined_tiles, + _load_smem_i32( + sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) + ) + // ctm.Int32(KV_TILE), + ) + if pipelined_tiles >= ctm.Int32(2): + ( + running_row_max, + running_row_anchor, + running_part_sum, + ) = _run_pipelined_leading_tiles( sAtomicRunningRowMax, + sPipelinedRowMax, sP, sPSF, - sScoreSlotRowMax, - sRowScaleRing, qk_full_mbar, pair_overlap_sync_mbars, leader_pair_overlap_sync_mbars, leader_pair_overlap_sync_mbars_dsmem, p4_smem_slot0_producer, - p4_smem_slot0_mbar_ptr, - tmem_score_slot0_consumer, - tmem_score_slot0_mbar_ptr, pv_full_mbar, corr_o_acc_tmem, warp_idx, @@ -7027,46 +7258,43 @@ def _run_mla_decode_body( producer_valid_k, softmax_scale_log2, pv_psf_rescale, - running_row_max, - running_row_anchor, - carried_dq1024, - carried_old_scale, - stream_li_idx, - stream_li_idx > ctm.Int32(0), - producer_needs_mask, - False, + pipelined_tiles, use_mixed_imlp, - folded_query_heads > 0 or (packed_query_heads > 0 and query_len_per_seq > 1), ) - ) - stream_li_idx = stream_li_idx + ctm.Int32(1) - stream_li_remaining = stream_li_remaining - ctm.Int32(1) - final_tile_k_end = (stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE) - tail_needs_mask = valid_k_for_l < final_tile_k_end - if cutlass.const_expr(folded_query_heads > 0): - tail_needs_mask = folded_first_valid_k < final_tile_k_end - if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): - tail_needs_mask = ( - _load_smem_i32( - sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) + stream_li_idx = pipelined_tiles + stream_li_remaining = stream_li_total - pipelined_tiles + while stream_li_remaining > ctm.Int32(1): + producer_needs_mask = ctm.Boolean(False) + if cutlass.const_expr(folded_query_heads > 0): + # The MTP window may straddle a 256-token tile boundary, so + # the penultimate tile can also need per-row causal masking. + producer_needs_mask = folded_first_valid_k < ( + (stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE) ) - < final_tile_k_end - ) - running_row_max, running_row_anchor, carried_dq1024, carried_old_scale = ( - _runtime_t336_producer_tile_v23( + if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): + # Rows of one tile differ by at most query_len tokens, so only + # tiles past the earliest row bound need per-row masking. The + # bound is re-read from smem each tile: kept in a register it + # was spilled at the producers' register ceiling and reloaded + # from local memory here instead. + producer_needs_mask = _load_smem_i32( + sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) + ) < ((stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE)) + ( + running_row_max, + running_row_anchor, + carried_dq1024, + carried_old_scale, + running_part_sum, + ) = _runtime_t336_producer_tile_v23( sAtomicRunningRowMax, sP, sPSF, - sScoreSlotRowMax, - sRowScaleRing, qk_full_mbar, pair_overlap_sync_mbars, leader_pair_overlap_sync_mbars, leader_pair_overlap_sync_mbars_dsmem, p4_smem_slot0_producer, - p4_smem_slot0_mbar_ptr, - tmem_score_slot0_consumer, - tmem_score_slot0_mbar_ptr, pv_full_mbar, corr_o_acc_tmem, warp_idx, @@ -7080,57 +7308,80 @@ def _run_mla_decode_body( running_row_anchor, carried_dq1024, carried_old_scale, + running_part_sum, stream_li_idx, stream_li_idx > ctm.Int32(0), - tail_needs_mask, - True, + producer_needs_mask, + False, use_mixed_imlp, + folded_query_heads > 0 or (packed_query_heads > 0 and query_len_per_seq > 1), ) - ) - if is_compute_warp: - final_direct_p_li_idx = stream_li_total - ctm.Int32(1) - final_direct_p_slot_idx = final_direct_p_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - final_direct_p_phase = ( - final_direct_p_li_idx - // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - final_direct_p_ready_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_DIRECT_P_READY_MBAR_OFFSET + final_direct_p_slot_idx - ) - while not prims.mbarrier_try_wait_parity( - final_direct_p_ready_mbar, final_direct_p_phase, time_limit=10000000 - ): - pass - cute.arch.mbarrier_arrive( - p4_smem_slot0_mbar_ptr + p4_smem_slot0_producer.current_handle().index, - cutlass.Int32(0), - ) - p4_smem_slot0_producer.advance() - cute.arch.mbarrier_arrive( - tmem_score_slot0_mbar_ptr - + SMEM_P4_TMEM_SCORE_PIPELINE_STAGES - + tmem_score_slot0_consumer.current_handle().index, - cutlass.Int32(0), + stream_li_idx = stream_li_idx + ctm.Int32(1) + stream_li_remaining = stream_li_remaining - ctm.Int32(1) + final_tile_k_end = (stream_li_idx + ctm.Int32(1)) * ctm.Int32(KV_TILE) + tail_needs_mask = valid_k_for_l < final_tile_k_end + if cutlass.const_expr(folded_query_heads > 0): + tail_needs_mask = folded_first_valid_k < final_tile_k_end + if cutlass.const_expr(packed_query_heads > 0 and query_len_per_seq > 1): + tail_needs_mask = ( + _load_smem_i32( + sRuntimeScalePair.data_ptr() + ctm.Int32(SMEM_P4_RUNTIME_MASK_BOUND_WORD) + ) + < final_tile_k_end ) - tmem_score_slot0_consumer.advance() + ( + running_row_max, + running_row_anchor, + carried_dq1024, + carried_old_scale, + running_part_sum, + ) = _runtime_t336_producer_tile_v23( + sAtomicRunningRowMax, + sP, + sPSF, + qk_full_mbar, + pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars, + leader_pair_overlap_sync_mbars_dsmem, + p4_smem_slot0_producer, + pv_full_mbar, + corr_o_acc_tmem, + warp_idx, + tidx, + p4_tmem_base_col_id, + p4_tmem_base_row_id, + producer_valid_k, + softmax_scale_log2, + pv_psf_rescale, + running_row_max, + running_row_anchor, + carried_dq1024, + carried_old_scale, + running_part_sum, + stream_li_idx, + stream_li_idx > ctm.Int32(0), + tail_needs_mask, + True, + use_mixed_imlp, + ) + # Publish each warp's denominator slice: part = score half * 2 + column band. + part_local_warp_idx = warp_idx & ctm.Int32(3) + part_idx = (warp_idx >> ctm.Int32(2)) * ctm.Int32(2) + (part_local_warp_idx >> ctm.Int32(1)) + part_row = (part_local_warp_idx & ctm.Int32(1)) * ctm.Int32(32) + (tidx & ctm.Int32(31)) + sScoreSlotRowMax[part_idx * ctm.Int32(SMEM_P4_BMM1_M) + part_row] = running_part_sum + prims.barrier( + barrier_id=SMEM_P4_DUAL_HALF_BAR_ID, + number_of_threads=SMEM_P4_DUAL_HALF_BAR_THREADS, + ) if is_correction_warp: - last_rowsum_li_idx = stream_li_total - ctm.Int32(1) - last_rowsum_qk_slot_idx = last_rowsum_li_idx % ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - last_rowsum_empty_mbar = pair_overlap_sync_mbars.subview( - SMEM_P4_PAIR_HALF1_ROWSUM_READY_MBAR_OFFSET + last_rowsum_qk_slot_idx - ) - last_rowsum_empty_phase = ( - last_rowsum_li_idx - // ctm.Int32(SMEM_P4_QK_PIPELINE_SLOTS) - % ctm.Int32(SMEM_P4_MBAR_PARITY_PHASES) - ) - while not prims.mbarrier_try_wait_parity( - last_rowsum_empty_mbar, last_rowsum_empty_phase, time_limit=10000000 - ): - pass final_row_state_local_row = corr_row_band * ctm.Int32(32) + (tidx & ctm.Int32(31)) - final_anchor_row_sum = sFinalAnchorRowSum[final_row_state_local_row] + final_anchor_row_sum = ( + sScoreSlotRowMax[final_row_state_local_row] + + sScoreSlotRowMax[ctm.Int32(2 * SMEM_P4_BMM1_M) + final_row_state_local_row] + ) + ( + sScoreSlotRowMax[ctm.Int32(SMEM_P4_BMM1_M) + final_row_state_local_row] + + sScoreSlotRowMax[ctm.Int32(3 * SMEM_P4_BMM1_M) + final_row_state_local_row] + ) # Packed rows may be empty while their tile is not (Helix tokens # whose KV lives on another rank); judge emptiness per row. row_valid_k = valid_k_for_l From 358b1cc40667cfe7ceacc95916e97ba0bf35c853 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Thu, 22 Oct 2026 00:44:47 -0700 Subject: [PATCH 08/16] [None][perf] FP4 MLA packed decode: simulate the waves to pick the task length Signed-off-by: Lanyu Liao Signed-off-by: Mingyang Hao --- .../fp4_mla/fp4_mla_cutedsl_mufu16.py | 18 +- .../backends/fp4_mla/fp4_mla_packed_reduce.py | 338 +- .../attention_backend/fp4_mla/__init__.py | 5135 +++++++++++++++++ .../_torch/attention/test_fp4_mla_packed.py | 135 +- 4 files changed, 5551 insertions(+), 75 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py index 301ef2e2148a..d76a685232fa 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py @@ -6271,10 +6271,10 @@ def _run_mla_decode_body( page_batch = ctm.Int32(0) valid_k_for_l = ctm.Int32(0) q_batch = bidz - # Packed rows: grid.z enumerates (split, request) tasks, request fastest. - # Rows of this M tile are [query token, head] pairs of one request; the - # split owns KV tiles [split_tile_begin, split_tile_begin + split_tiles) - # and every bound below is relative to split_base_k. + # Packed rows: grid.z is the request when num_splits == 1, else a slot of the plan's + # tables, in grid order (packed_plan_kernel). Rows of this M tile are [query token, + # head] pairs of one request; the split owns KV tiles [split_tile_begin, + # split_tile_begin + split_tiles) and every bound below is relative to split_base_k. split_tile_begin = ctm.Int32(0) split_base_k = ctm.Int32(0) first_valid_k = ctm.Int32(0) @@ -6282,6 +6282,8 @@ def _run_mla_decode_body( packed_request_valid_k = ctm.Int32(0) packed_single_task = ctm.Boolean(False) c_batch_idx = bidz + # Partials and row stats go to begins[request] + split, where the reduction reads them. + store_task_idx = bidz if cutlass.const_expr(packed_query_heads > 0): # Balanced plan (one task per request when num_splits == 1, else the # plan kernel's tables): the request and KV split this task owns. @@ -6296,6 +6298,8 @@ def _run_mla_decode_body( task_request = ctm.Int32(mPlan_i[plan_tables + bidz]) task_split = ctm.Int32(mPlan_i[plan_tables + task_capacity + bidz]) task_count = ctm.Int32(mPlan_i[plan_tables + task_capacity + task_capacity + bidz]) + store_task_idx = ctm.Int32(mPlan_i[ctm.Int32(2) + task_request]) + task_split + store_task_idx = cute.arch.make_warp_uniform(store_task_idx) tiles_per_split = cute.arch.make_warp_uniform(tiles_per_split) page_batch = cute.arch.make_warp_uniform(task_request) split_idx = cute.arch.make_warp_uniform(task_split) @@ -7453,7 +7457,7 @@ def _run_mla_decode_body( tidx, bidx, bidy, - bidz, + store_task_idx, m, n, final_row_max=final_row_max, @@ -7479,7 +7483,7 @@ def _run_mla_decode_body( tidx, bidx, bidy, - bidz, + store_task_idx, m, n, final_row_max=final_row_max, @@ -7504,7 +7508,7 @@ def _run_mla_decode_body( tidx, bidx, bidy, - bidz, + store_task_idx, m, n, final_row_max=final_row_max, diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_packed_reduce.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_packed_reduce.py index 7c72afcb3bd6..3f4b9b158b3b 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_packed_reduce.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_packed_reduce.py @@ -11,6 +11,7 @@ scales and KV lengths are read on the device, so CUDA graph replays follow them. """ +import heapq import math from collections.abc import Callable @@ -32,17 +33,30 @@ ) _PACKED_REDUCE_COMPILE_CACHE: dict[tuple[int, bool, torch.dtype], Callable] = {} -_PACKED_PLAN_COMPILE_CACHE: dict[tuple[int, bool, bool], Callable] = {} +_PACKED_PLAN_COMPILE_CACHE: dict[tuple[int, bool, bool, bool], Callable] = {} _REDUCE_THREADS = 128 -_PLAN_THREADS = 128 +# One warp per candidate task length. +_PLAN_THREADS = 1024 # Requests one plan CTA stages in shared memory (4 KiB of tiles + begins). PLAN_MAX_REQUESTS = 1024 +# The plan simulates the decode grid for this many task lengths, one warp each: +# the rule's (warp 0) and the longest request cut into 1..PLAN_SIM_CANDIDATES - 1 pieces. +PLAN_SIM_CANDIDATES = _PLAN_THREADS // 32 +# Waves of more slots (four per lane) or batches of more requests keep the rule's length. +PLAN_SIM_MAX_SLOTS = 128 +PLAN_SIM_MAX_REQUESTS = 64 +# Candidates of more tasks are skipped: the simulation leaves out the reduction, whose +# cost grows with the partials. +PLAN_SIM_MAX_TASKS = 256 +# Fixed cost of one task (Q load, pipeline fill, epilogue) in KV tiles, and the extra cost of +# a task whose request is split (FP32 partial and row stats, merged by the reduction). +PLAN_TASK_COST_TILES = 12 +PLAN_SPLIT_COST_TILES = 24 _VALUE_DIM = 512 # Tasks whose partial loads a warp issues back-to-back before any is consumed # with a full 512-column band (16 vector loads per lane; narrower bands take # proportionally more tasks per chunk). The reduction is latency-bound, so -# this sets its memory-level parallelism; 32 loads in flight measured slower -# (register pressure), as did issuing all loads ahead of the weight math. +# this sets its memory-level parallelism. _TASKS_PER_CHUNK = 4 # Persistent reducer grid: up to this many CTAs per SM, one warp per output row. _REDUCE_CTAS_PER_SM = 4 @@ -358,6 +372,104 @@ def packed_split_reduce_host( ) +@cute.jit +def _plan_simulated_makespan( + sTiles: cute.Tensor, + num_requests: cutlass.Int32, + tiles_per_task: cutlass.Int32, + wave_tasks: cutlass.Int32, + lane: cutlass.Int32, +) -> cutlass.Int32: + """Makespan (KV tiles) of the decode grid for one task length, simulated by one warp. + + The grid runs the tasks in the plan's table order (``plan_task_order``): every + request's tasks but its last, all ``tiles_per_task`` long, then the last tasks, + longest first. Each task takes the slot of ``wave_tasks`` that frees first and + holds it for its KV tiles plus PLAN_TASK_COST_TILES, plus PLAN_SPLIT_COST_TILES + when its request has several tasks. Slot j lives in lane j % 32, register + j // 32; request r's last task in lane r % 32, bit r // 32. + """ + # Every task but a request's last belongs to a split request. + task_cost = tiles_per_task + cutlass.Int32(PLAN_TASK_COST_TILES + PLAN_SPLIT_COST_TILES) + full_tasks = cutlass.Int32(0) + for request_idx in cutlass.range(lane, num_requests, 32, unroll=1): + full_tasks = ( + full_tasks + + (sTiles[request_idx] + tiles_per_task - cutlass.Int32(1)) // tiles_per_task + - cutlass.Int32(1) + ) + full_tasks = cute.arch.warp_redux_sync(full_tasks, "add") + # Equal tasks fill the slots round-robin. + full_rounds = full_tasks // wave_tasks + full_extra = full_tasks - full_rounds * wave_tasks + idle = cutlass.Int32(1 << 30) + slot0 = idle + slot1 = idle + slot2 = idle + slot3 = idle + if lane < wave_tasks: + slot0 = full_rounds * task_cost + if lane < full_extra: + slot0 = slot0 + task_cost + if lane + cutlass.Int32(32) < wave_tasks: + slot1 = full_rounds * task_cost + if lane + cutlass.Int32(32) < full_extra: + slot1 = slot1 + task_cost + if lane + cutlass.Int32(64) < wave_tasks: + slot2 = full_rounds * task_cost + if lane + cutlass.Int32(64) < full_extra: + slot2 = slot2 + task_cost + if lane + cutlass.Int32(96) < wave_tasks: + slot3 = full_rounds * task_cost + if lane + cutlass.Int32(96) < full_extra: + slot3 = slot3 + task_cost + makespan = full_rounds * task_cost + if full_extra > cutlass.Int32(0): + makespan = makespan + task_cost + taken = cutlass.Int32(0) + for _ in cutlass.range(0, num_requests, 1, unroll=1): + # The longest last task left (ties: lower request), as tail * 1024 + (1023 - request). + local_key = cutlass.Int32(-1) + for request_idx in cutlass.range(lane, num_requests, 32, unroll=1): + taken_bit = (taken >> (request_idx // cutlass.Int32(32))) & cutlass.Int32(1) + if taken_bit == cutlass.Int32(0): + request_tiles = sTiles[request_idx] + request_tasks = ( + request_tiles + tiles_per_task - cutlass.Int32(1) + ) // tiles_per_task + tail = request_tiles - (request_tasks - cutlass.Int32(1)) * tiles_per_task + local_key = cutlass.max( + local_key, + tail * cutlass.Int32(PLAN_MAX_REQUESTS) + + cutlass.Int32(PLAN_MAX_REQUESTS - 1) + - request_idx, + ) + key = cute.arch.warp_redux_sync(local_key, "max") + request = cutlass.Int32(PLAN_MAX_REQUESTS - 1) - key % cutlass.Int32(PLAN_MAX_REQUESTS) + if lane == request % cutlass.Int32(32): + taken = taken | (cutlass.Int32(1) << (request // cutlass.Int32(32))) + tail_cost = key // cutlass.Int32(PLAN_MAX_REQUESTS) + cutlass.Int32(PLAN_TASK_COST_TILES) + if sTiles[request] > tiles_per_task: + tail_cost = tail_cost + cutlass.Int32(PLAN_SPLIT_COST_TILES) + local_free = cutlass.min(cutlass.min(slot0, slot1), cutlass.min(slot2, slot3)) + start = cute.arch.warp_redux_sync(local_free, "min") + owner_key = cutlass.Int32(32) + if local_free == start: + owner_key = lane + owner = cute.arch.warp_redux_sync(owner_key, "min") + if lane == owner: + if slot0 == start: + slot0 = start + tail_cost + elif slot1 == start: + slot1 = start + tail_cost + elif slot2 == start: + slot2 = start + tail_cost + else: + slot3 = start + tail_cost + makespan = cutlass.max(makespan, start + tail_cost) + return makespan + + @cute.kernel def packed_plan_kernel( mValidK_l: cute.Tensor, @@ -374,6 +486,7 @@ def packed_plan_kernel( query_len_per_seq: cutlass.Constexpr, use_helix_kv_bounds: cutlass.Constexpr, write_generation_lengths: cutlass.Constexpr, + wave_plan: cutlass.Constexpr, ) -> None: """One CTA: the balanced task plan both decode and reduction kernels read. @@ -385,9 +498,11 @@ def packed_plan_kernel( ``plan = [tiles_per_task, used_tasks, begins[0..num_requests], request[task_capacity], split[task_capacity], count[task_capacity]]``. - Request ``r`` owns ``ceil(tiles_r / tiles_per_task)`` consecutive tasks; - slots past ``used_tasks`` map to the last request with an out-of-range - split and are never scheduled (the decode kernel exits them). + Request ``r`` owns ``ceil(tiles_r / tiles_per_task)`` tasks, whose partials + sit at ``begins[r] + split``; the tables list them in grid order (every + request's tasks but its last, then the last ones longest first, see + ``plan_task_order``). Slots past ``used_tasks`` map to the last request + with an out-of-range split and are never scheduled (the decode kernel exits them). Task length: a request is only split when it is at least 1.5x longer than the batch's work spread over one wave, ``ceil(total / wave_tasks)``, i.e. @@ -404,8 +519,15 @@ def packed_plan_kernel( ``ceil(total / (capacity - requests))`` and the per-request share ``ceil(longest / splits)`` that fits is used. Everything is clamped to ``max_tiles_per_task``, the KV the compiled runtime profile covers (the - host sizes it so the clamp never needs more than the capacity). The host - mirror is ``plan_packed_tasks``. + host sizes it so the clamp never needs more than the capacity). + + That length is only a candidate: one warp each, the plan + simulates the grid (``_plan_simulated_makespan``) for it and for the + longest request cut into 1..PLAN_SIM_CANDIDATES - 1 equal pieces, and + takes the length that finishes first (ties keep the longer tasks). + Without ``wave_plan`` (``TRTLLM_FP4_MLA_PACKED_WAVE_PLAN=0``) the plan keeps + the rule's length and lists every request's tasks in request order. + The host mirror is ``plan_packed_tasks``. """ tidx, _, _ = cute.arch.thread_idx() sTiles = cute.make_tensor( @@ -416,6 +538,11 @@ def packed_plan_kernel( cute.make_layout(PLAN_MAX_REQUESTS + 1), ) sScalars = cute.make_tensor(cute.arch.alloc_smem(cutlass.Int32, 2), cute.make_layout(2)) + # Per candidate warp: simulated makespan, then task length. + sCandidates = cute.make_tensor( + cute.arch.alloc_smem(cutlass.Int32, 2 * PLAN_SIM_CANDIDATES), + cute.make_layout(2 * PLAN_SIM_CANDIDATES), + ) for request_idx in cutlass.range(tidx, num_requests, _PLAN_THREADS, unroll=1): request_valid_k = cutlass.Int32(mValidK_l[request_idx]) if cutlass.const_expr(write_generation_lengths): @@ -500,36 +627,108 @@ def packed_plan_kernel( tiles_per_task = cutlass.min(tiles_per_task, longest) tiles_per_task = cutlass.min(tiles_per_task, max_tiles_per_task) if tidx == cutlass.Int32(0): - task_begin = cutlass.Int32(0) - for request_idx in cutlass.range(0, num_requests, 1, unroll=1): - sBegins[request_idx] = task_begin - mPlan_i[cutlass.Int32(2) + request_idx] = task_begin - task_begin = ( - task_begin - + (sTiles[request_idx] + tiles_per_task - cutlass.Int32(1)) // tiles_per_task - ) - sBegins[num_requests] = task_begin - mPlan_i[cutlass.Int32(2) + num_requests] = task_begin - mPlan_i[0] = tiles_per_task - mPlan_i[1] = task_begin sScalars[0] = tiles_per_task + sScalars[1] = longest cute.arch.sync_threads() - tiles_per_task = sScalars[0] - tables = cutlass.Int32(3) + num_requests - for task_idx in cutlass.range(tidx, task_capacity, _PLAN_THREADS, unroll=1): - request = num_requests - cutlass.Int32(1) - split_idx = cutlass.Int32(1 << 15) - request_tasks = cutlass.Int32(1 << 15) + if cutlass.const_expr(wave_plan): + warp_idx = tidx // cutlass.Int32(32) + lane = tidx % cutlass.Int32(32) + fits = (wave_tasks <= cutlass.Int32(PLAN_SIM_MAX_SLOTS)) & ( + num_requests <= cutlass.Int32(PLAN_SIM_MAX_REQUESTS) + ) + simulate = (task_capacity > num_requests) & fits + if simulate: + longest = sScalars[1] + candidate = sScalars[0] + if warp_idx > cutlass.Int32(0): + candidate = cutlass.min( + cutlass.max( + (longest + warp_idx - cutlass.Int32(1)) // warp_idx, + cutlass.Int32(PACKED_MIN_TILES_PER_TASK), + ), + cutlass.min(longest, max_tiles_per_task), + ) + candidate_tasks = cutlass.Int32(0) + for request_idx in cutlass.range(lane, num_requests, 32, unroll=1): + candidate_tasks = ( + candidate_tasks + + (sTiles[request_idx] + candidate - cutlass.Int32(1)) // candidate + ) + candidate_tasks = cute.arch.warp_redux_sync(candidate_tasks, "add") + candidate_makespan = cutlass.Int32(-1) + if candidate_tasks <= cutlass.min(task_capacity, cutlass.Int32(PLAN_SIM_MAX_TASKS)): + candidate_makespan = _plan_simulated_makespan( + sTiles, num_requests, candidate, wave_tasks, lane + ) + if lane == cutlass.Int32(0): + sCandidates[warp_idx] = candidate_makespan + sCandidates[PLAN_SIM_CANDIDATES + warp_idx] = candidate + cute.arch.sync_threads() + if simulate & (tidx == cutlass.Int32(0)): + best_makespan = cutlass.Int32(-1) + best_tiles = sScalars[0] + for candidate_idx in cutlass.range(0, PLAN_SIM_CANDIDATES, 1, unroll=1): + candidate_makespan = sCandidates[candidate_idx] + candidate = sCandidates[PLAN_SIM_CANDIDATES + candidate_idx] + if candidate_makespan >= cutlass.Int32(0): + better = ( + (best_makespan < cutlass.Int32(0)) + | (candidate_makespan < best_makespan) + | ((candidate_makespan == best_makespan) & (candidate > best_tiles)) + ) + if better: + best_makespan = candidate_makespan + best_tiles = candidate + sScalars[0] = best_tiles + cute.arch.sync_threads() + if tidx == cutlass.Int32(0): + tiles_per_task = sScalars[0] + task_begin = cutlass.Int32(0) for request_idx in cutlass.range(0, num_requests, 1, unroll=1): - begin = sBegins[request_idx] - end = sBegins[request_idx + cutlass.Int32(1)] - if task_idx >= begin and task_idx < end: - request = request_idx - split_idx = task_idx - begin - request_tasks = end - begin - mPlan_i[tables + task_idx] = request - mPlan_i[tables + task_capacity + task_idx] = split_idx - mPlan_i[tables + task_capacity + task_capacity + task_idx] = request_tasks + sBegins[request_idx] = task_begin + mPlan_i[cutlass.Int32(2) + request_idx] = task_begin + task_begin = ( + task_begin + + (sTiles[request_idx] + tiles_per_task - cutlass.Int32(1)) // tiles_per_task + ) + sBegins[num_requests] = task_begin + mPlan_i[cutlass.Int32(2) + num_requests] = task_begin + mPlan_i[0] = tiles_per_task + mPlan_i[1] = task_begin + cute.arch.sync_threads() + task_tiles = sScalars[0] + tables = cutlass.Int32(3) + num_requests + used_tasks = sBegins[num_requests] + # Grid order with wave_plan (plan_task_order): non-last tasks in request order from + # begins[r] - r, then last tasks longest first. Partials stay at begins[r] + split. + full_tasks = used_tasks - num_requests + for request_idx in cutlass.range(tidx, num_requests, _PLAN_THREADS, unroll=1): + begin = sBegins[request_idx] + request_tasks = sBegins[request_idx + cutlass.Int32(1)] - begin + first_task = begin + last_task = begin + request_tasks - cutlass.Int32(1) + if cutlass.const_expr(wave_plan): + request_tiles = sTiles[request_idx] + tail = request_tiles - (request_tasks - cutlass.Int32(1)) * task_tiles + rank = cutlass.Int32(0) + for other_idx in cutlass.range(0, num_requests, 1, unroll=1): + other_tasks = sBegins[other_idx + cutlass.Int32(1)] - sBegins[other_idx] + other_tail = sTiles[other_idx] - (other_tasks - cutlass.Int32(1)) * task_tiles + if (other_tail > tail) | ((other_tail == tail) & (other_idx < request_idx)): + rank = rank + cutlass.Int32(1) + first_task = begin - request_idx + last_task = full_tasks + rank + for split_idx in cutlass.range(0, request_tasks, 1, unroll=1): + task_idx = last_task + if split_idx + cutlass.Int32(1) < request_tasks: + task_idx = first_task + split_idx + mPlan_i[tables + task_idx] = request_idx + mPlan_i[tables + task_capacity + task_idx] = split_idx + mPlan_i[tables + task_capacity + task_capacity + task_idx] = request_tasks + for task_idx in cutlass.range(used_tasks + tidx, task_capacity, _PLAN_THREADS, unroll=1): + mPlan_i[tables + task_idx] = num_requests - cutlass.Int32(1) + mPlan_i[tables + task_capacity + task_idx] = cutlass.Int32(1 << 15) + mPlan_i[tables + task_capacity + task_capacity + task_idx] = cutlass.Int32(1 << 15) @cute.jit @@ -549,6 +748,7 @@ def packed_plan_host( query_len_per_seq: cutlass.Constexpr, use_helix_kv_bounds: cutlass.Constexpr, write_generation_lengths: cutlass.Constexpr, + wave_plan: cutlass.Constexpr, ) -> None: request_layout = cute.make_layout((num_requests,), stride=(1,)) mValidK_l = cute.make_tensor(kv_lens_ptr, request_layout) @@ -577,6 +777,7 @@ def packed_plan_host( query_len_per_seq, use_helix_kv_bounds, write_generation_lengths, + wave_plan, ).launch(grid=(1, 1, 1), block=(_PLAN_THREADS, 1, 1), stream=stream) @@ -596,6 +797,7 @@ def launch_packed_plan( query_len: int, helix_kv_bounds: torch.Tensor | None = None, generation_lengths: tuple[torch.Tensor, torch.Tensor, torch.Tensor, int] | None = None, + wave_plan: bool = True, ) -> None: """Write the balanced task plan for ``num_requests * num_splits`` slots into ``plan``. @@ -606,6 +808,8 @@ def launch_packed_plan( generation_len)`` the same launch first writes the per-request generation lengths ``kv_lens - prompt_lens + generation_len`` and ``generation_len`` (``kv_lens`` then being the raw lengths) and plans on the corrected ones. + ``wave_plan=False`` keeps the rule's task length and request order (no + wave simulation, no tail reordering). """ task_capacity = num_requests * num_splits if num_requests <= 0 or num_splits <= 0 or wave_tasks <= 0 or max_tiles_per_task <= 0: @@ -660,7 +864,7 @@ def launch_packed_plan( make_ptr(cutlass.Int32, plan.data_ptr(), gmem, assumed_align=4), ) stream = _current_cu_stream() - cache_key = (query_len, use_helix, write_lengths) + cache_key = (query_len, use_helix, write_lengths, wave_plan) plan_fn = _PACKED_PLAN_COMPILE_CACHE.get(cache_key) if plan_fn is None: plan_fn = _compile_cutedsl( @@ -671,6 +875,7 @@ def launch_packed_plan( query_len, use_helix, write_lengths, + wave_plan, ) _PACKED_PLAN_COMPILE_CACHE[cache_key] = plan_fn plan_fn( @@ -719,13 +924,16 @@ def plan_packed_tasks( max_tiles_per_task: int = 1 << 30, min_tiles_per_task: int = PACKED_MIN_TILES_PER_TASK, one_wave: bool = False, + wave_plan: bool = True, ) -> tuple[int, list[int]]: """Host mirror of the kernels' balanced task plan: (tiles per task, task begins). ``request_tiles`` are the per-request KV tile counts (>= 1), ``task_capacity`` the grid.z slots (``num_requests * num_splits``) and ``wave_tasks`` the tasks that fill one wave of clusters; ``one_wave`` (Helix) never takes the two-wave - target. Returns the KV tiles each task may own and the + target. With ``wave_plan`` the length competes + with the longest request cut into 1..PLAN_SIM_CANDIDATES - 1 pieces by + ``plan_simulated_makespan``. Returns the KV tiles each task may own and the ``len(request_tiles) + 1`` task prefix, whose last element is the number of tasks actually used (<= ``task_capacity``). See ``_packed_tiles_per_task``. """ @@ -761,12 +969,66 @@ def used(tiles: int) -> int: tiles_per_task = even tiles_per_task = min(tiles_per_task, longest) tiles_per_task = min(tiles_per_task, max_tiles_per_task) + if ( + wave_plan + and task_capacity > num_requests + and wave_tasks <= PLAN_SIM_MAX_SLOTS + and num_requests <= PLAN_SIM_MAX_REQUESTS + ): + candidates = [tiles_per_task] + [ + min(max(math.ceil(longest / pieces), min_tiles_per_task), longest, max_tiles_per_task) + for pieces in range(1, PLAN_SIM_CANDIDATES) + ] + best = None + for candidate in candidates: + if used(candidate) > min(task_capacity, PLAN_SIM_MAX_TASKS): + continue + makespan = plan_simulated_makespan(request_tiles, candidate, wave_tasks) + if best is None or (makespan, -candidate) < best: + best = (makespan, -candidate) + if best is not None: + tiles_per_task = -best[1] begins = [0] for tiles in request_tiles: begins.append(begins[-1] + math.ceil(tiles / tiles_per_task)) return tiles_per_task, begins +def plan_task_order( + request_tiles: list[int], tiles_per_task: int, wave_plan: bool = True +) -> list[tuple[int, int]]: + """Grid order of the plan's tasks as (request, split). + + Every request's tasks but its last in request order, then the last tasks + longest first (ties: lower request); without ``wave_plan`` all in request order. + """ + tasks = [math.ceil(tiles / tiles_per_task) for tiles in request_tiles] + if not wave_plan: + return [(r, split) for r, n in enumerate(tasks) for split in range(n)] + order = [(r, split) for r, n in enumerate(tasks) for split in range(n - 1)] + tails = [ + (tiles - (n - 1) * tiles_per_task, r) + for r, (tiles, n) in enumerate(zip(request_tiles, tasks)) + ] + order += [(r, tasks[r] - 1) for _, r in sorted(tails, key=lambda tail: (-tail[0], tail[1]))] + return order + + +def plan_simulated_makespan(request_tiles: list[int], tiles_per_task: int, wave_tasks: int) -> int: + """Host mirror of ``_plan_simulated_makespan``: grid makespan in KV tiles for one length.""" + slots = [0] * wave_tasks + makespan = 0 + for request, split in plan_task_order(request_tiles, tiles_per_task): + start = heapq.heappop(slots) + tiles = min(tiles_per_task, request_tiles[request] - split * tiles_per_task) + end = start + tiles + PLAN_TASK_COST_TILES + if request_tiles[request] > tiles_per_task: + end += PLAN_SPLIT_COST_TILES + makespan = max(makespan, end) + heapq.heappush(slots, end) + return makespan + + def request_kv_tiles( kv_lens: list[int], query_len: int, helix_kv_bounds: list[int] | None = None ) -> list[int]: diff --git a/tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py new file mode 100644 index 000000000000..17208d66d0f2 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py @@ -0,0 +1,5135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared MLA FP4 KV-cache helpers. + +The high-precision (HP) BF16 KV pool is a standalone circular buffer used +alongside the paged FP4 KV pool when MLA models run with NVFP4 KV cache. +The pool retains one 16-token quantization tile plus speculative-rewind slack +per sequence. + +Used by the TRTLLM attention backend FP4 MLA FMHA path. +""" + +import importlib.util +import os +from typing import Any, Literal, NamedTuple, Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._utils import get_sm_version, prefer_pinned +from tensorrt_llm.bindings import DataType + +from .fp4_mla_kernels import ( + _fp4_mla_chunked_cache_gather_kernel, + _fp4_mla_context_cache_update_kernel, + _fp4_mla_generation_fused_qk_rope_cache_update_kernel, + _fp4_mla_rebuild_v_scale_from_k_scale_kernel, +) + +HP_BLOCK_SIZE: int = 16 +FP4_BLOCK_SIZE: int = 16 +FP4_MLA_TOKENS_PER_BLOCK: int = 128 +FP4_MLA_SCALE_ROW_GROUP: int = 128 +FP4_MLA_SCALE_COL_GROUP: int = 4 +FP4_MLA_P_GLOBAL_SCALE: float = 448.0 * 6.0 +FP4_MLA_Q_STATIC_AMAX: float = 400.0 +FP4_MLA_KV_STATIC_AMAX: float = 30.0 +FP4_MLA_Q_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_Q_STATIC_AMAX +FP4_MLA_KV_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_KV_STATIC_AMAX +# Max finite e4m3 magnitude for FP4 MLA block-scale clamping. +FP4_MLA_E4M3_MAX: float = 448.0 +FP4_MLA_Q_RESIDUAL_DIM: int = 64 +FP4_MLA_K_RESIDUAL_DIM: int = FP4_MLA_Q_RESIDUAL_DIM +FP4_MLA_Q_PREFIX_DIM: int = 512 +FP4_MLA_Q_PREFIX_BLOCK_DIM: int = 256 +FP4_MLA_Q1_PREFIX_BLOCK_DIM: int = 512 +FP4_MLA_Q_LOGICAL_DIM: int = FP4_MLA_Q_PREFIX_DIM + 2 * FP4_MLA_Q_RESIDUAL_DIM +FP4_MLA_Q_PACKED_DIM: int = FP4_MLA_Q_LOGICAL_DIM // 2 +FP4_MLA_Q_SF_GROUPS: int = FP4_MLA_Q_LOGICAL_DIM // FP4_BLOCK_SIZE +FP4_MLA_ATTENTION_BACKEND_ENV = "TRTLLM_FP4_MLA_ATTENTION_BACKEND" +# Retained for compatibility; mufu16 is now the only non-fused CuTeDSL variant. +FP4_MLA_CUTEDSL_MUFU16_ENV = "TRTLLM_FP4_MLA_CUTEDSL_MUFU16" +FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV = "TRTLLM_FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE" +_FP4_MLA_CUTEDSL_BACKEND = "cutedsl" +_FP4_MLA_K_RESIDUAL_BACKENDS = ("triton", _FP4_MLA_CUTEDSL_BACKEND) +# Architectures on which the K-residual Triton path can run the fused Q-quant +# cache update. SM107 is not listed because it selects CuTeDSL instead, which +# carries its own check at launch; SM103 (Blackwell-Ultra, GB300) is the part +# this was first exercised on. +_FP4_MLA_TRITON_Q_QUANT_FUSION_SMS = (103,) +_FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH = 4 +_FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH = 16 +_FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH = 32 +_FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD = 256 +_FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD = 512 +_FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD = 640 +_FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD = 768 +_HPUpdatePhase = Literal["context", "generation"] +_FP4_MLA_TRITON_PRELOAD_KEYS = "_fp4_mla_triton_preload_keys" +_FP4_MLA_PAGE_TABLE_TILE_SIZE = 128 +_FP4_MLA_MAX_GRID_Z = 65_535 + + +# Environment helpers + + +def _env_enabled_default(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None or value == "": + return default + return value.lower() in ( + "1", + "true", + "yes", + "on", + ) + + +def _fp4_mla_cutedsl_fused_v_transpose_enabled() -> bool: + return _env_enabled_default(FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, False) + + +FP4_MLA_PACKED_Q_ENV = "TRTLLM_FP4_MLA_PACKED_Q" +# Packed decode: each request owns a 128-row-padded block of [query token, head] +# Q rows, the kernel splits KV across grid.z tasks and reads per-row causal +# bounds, so no Q padding/copies, no Q replication, and Helix stays split. +_FP4_MLA_PACKED_MAX_M = 1024 +_FP4_MLA_PACKED_MAX_SPLITS = 16 +_FP4_MLA_PACKED_MIN_TILES_PER_SPLIT = 2 +_FP4_MLA_SPLIT_KV_ENV = "TRTLLM_FP4_MLA_SPLIT_KV" +# Rollback switch (default on): 0 makes the packed plan keep the rule's task length +# and run the tasks in request order instead of simulating the waves. +FP4_MLA_PACKED_WAVE_PLAN_ENV = "TRTLLM_FP4_MLA_PACKED_WAVE_PLAN" + + +def _packed_wave_tasks(device: torch.device, m_tiles: int) -> int: + """Tasks that fill one wave of 2-CTA clusters.""" + clusters_per_wave = max(1, _get_sm_count(device) // 2) + return max(1, clusters_per_wave // max(1, m_tiles)) + + +def _fp4_mla_packed_q_enabled() -> bool: + """Whether decode uses the packed per-request Q layout (mufu16 CuTe DSL only).""" + return ( + _env_enabled_default(FP4_MLA_PACKED_Q_ENV, True) + and _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND + and not _fp4_mla_cutedsl_fused_v_transpose_enabled() + ) + + +def _fp4_mla_packed_rows_per_seq(num_heads: int, query_len: int) -> int: + """Rows one request owns in the packed Q layout, or 0 when packing is off.""" + if not _fp4_mla_packed_q_enabled() or num_heads <= 0 or query_len <= 0: + return 0 + rows = _ceil_div(num_heads * query_len, FP4_MLA_SCALE_ROW_GROUP) * FP4_MLA_SCALE_ROW_GROUP + return rows if rows <= _FP4_MLA_PACKED_MAX_M else 0 + + +def _select_packed_split_count( + num_requests: int, + m_tiles: int, + device: torch.device, + *, + max_kv_len: int, + is_cuda_graph: bool, +) -> int: + """Task slots per request (grid.z = requests * splits): enough for two + waves of 2-CTA clusters plus one task per request, and never fewer than + two so a long request in a large batch can still be split (the kernel + shares the slots by KV length and leaves them unused when the batch is + even). + + The count depends only on the request count and tile count, so a CUDA + graph bucket always replays the same geometry. Eager calls know the real + longest KV and keep at least _FP4_MLA_PACKED_MIN_TILES_PER_SPLIT KV tiles + per split. ``TRTLLM_FP4_MLA_SPLIT_KV`` overrides the count for testing. + """ + override = os.getenv(_FP4_MLA_SPLIT_KV_ENV, "auto") + if override != "auto": + return max(1, min(int(override), _FP4_MLA_PACKED_MAX_SPLITS)) + # Slots for two waves of tasks (the plan's finest target) plus one rounding + # task per request, so the target share fits whenever the batch is skewed. + wave_tasks = _packed_wave_tasks(device, m_tiles) + splits = _ceil_div(2 * wave_tasks, max(1, num_requests)) + 1 + splits = max(2, min(_FP4_MLA_PACKED_MAX_SPLITS, splits)) + if not is_cuda_graph: + kv_tiles = _ceil_div(max(max_kv_len, 1), 256) + splits = max(1, min(splits, _ceil_div(kv_tiles, _FP4_MLA_PACKED_MIN_TILES_PER_SPLIT))) + return splits + + +class _PackedPlanGeometry(NamedTuple): + """Per-step launch geometry of the packed split-KV decode.""" + + num_requests: int + query_len: int + splits: int + wave_tasks: int + profile_tiles: int + task_capacity: int + split_kv_len: int + helix_kv_bounds: Optional[torch.Tensor] + plan: Optional[torch.Tensor] + key: Optional[tuple] + + +def _fp4_mla_packed_plan_geometry( + metadata: Any, + *, + num_heads: int, + query_len: int, + num_requests: int, + device: torch.device, + helix_kv_bounds: Optional[torch.Tensor], +) -> Optional[_PackedPlanGeometry]: + """Slots, task-length cap and plan buffer of this step's packed decode. + + None when the packed path is off. The geometry depends only on the batch + shape and the page-table capacity, so the generation cache update (the + first FP4 MLA kernel of a step) and the attention decode derive the same + one, and whichever runs first writes the step's plan fused with the + generation lengths (``_fp4_mla_uniform_generation_lengths``). ``key`` is + None when every request keeps a single task and no plan is needed. + """ + rows_per_seq = _fp4_mla_packed_rows_per_seq(num_heads, query_len) + if rows_per_seq == 0 or num_requests <= 0: + return None + if not getattr(metadata, "_fp4_mla_device_page_table", False): + return None + # The fixed-stride generation page-table capacity (see _max_generation_pages). + max_pages = int(getattr(metadata, "fp4_mla_page_table_stride", 0)) + if max_pages <= 0: + return None + max_kv_len = max_pages * metadata.page_size + is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) + m_tiles = rows_per_seq // FP4_MLA_SCALE_ROW_GROUP + splits = _select_packed_split_count( + num_requests, + m_tiles, + device, + max_kv_len=max_kv_len, + is_cuda_graph=is_cuda_graph, + ) + # Tasks that fill one wave of 2-CTA clusters: the plan's task-length reference. + wave_tasks = _packed_wave_tasks(device, m_tiles) + # The runtime KV profile only has to cover one task: the kernel clamps the + # plan's task length to it, and ceil(longest / (splits - 1)) tiles per task + # always fits the splits * requests slots. A batch that fills the machine + # with one task per request keeps the full profile so even batches never + # have to split. + kv_tiles = _ceil_div(max(max_kv_len, 1), 256) + if splits == 1 or wave_tasks <= 2 * num_requests: + task_tiles = kv_tiles + else: + task_tiles = min(kv_tiles, max(_ceil_div(kv_tiles, splits - 1), 2)) + split_kv_len = max(task_tiles * 256, 256) + # The kernels clamp the plan to the compiled profile (a bucket >= split_kv_len). + module = _fp4_mla_cutedsl_kernel_module() + profile_tiles = module._select_runtime_kv_profile(split_kv_len) // 256 + if splits == 1: + return _PackedPlanGeometry( + num_requests, + query_len, + 1, + wave_tasks, + profile_tiles, + num_requests, + split_kv_len, + helix_kv_bounds, + None, + None, + ) + from .fp4_mla_packed_reduce import packed_plan_length + + tasks = splits * num_requests + # Graph buckets hold exactly their tasks; eager metadata grows on demand. + task_capacity = tasks if is_cuda_graph else max(tasks, 32) + plan = _ensure_workspace_tensor( + metadata, + "_fp4_mla_packed_plan", + (packed_plan_length(num_requests, task_capacity),), + dtype=torch.int32, + device=device, + ) + kv_lens = getattr(metadata, "fp4_mla_generation_kv_lens", None) + key = ( + num_requests, + splits, + wave_tasks, + profile_tiles, + query_len, + task_capacity, + kv_lens.data_ptr() if isinstance(kv_lens, torch.Tensor) else 0, + helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else 0, + plan.data_ptr(), + ) + return _PackedPlanGeometry( + num_requests, + query_len, + splits, + wave_tasks, + profile_tiles, + task_capacity, + split_kv_len, + helix_kv_bounds, + plan, + key, + ) + + +def _fp4_mla_helix_plan_bounds( + metadata: Any, num_tokens: int, *, use_helix: bool +) -> Optional[torch.Tensor]: + """Per-token Helix KV bounds the decode plans on, or None (validated at decode).""" + if not use_helix or not bool(getattr(metadata, "_helix_spec_tokens_valid", False)): + return None + bounds = getattr(metadata, "helix_kv_bounds", None) + if ( + not isinstance(bounds, torch.Tensor) + or bounds.dtype != torch.int32 + or bounds.ndim != 1 + or bounds.numel() < num_tokens + or not bounds.is_contiguous() + ): + return None + return bounds[:num_tokens] + + +def _launch_fp4_mla_packed_plan( + geometry: _PackedPlanGeometry, + kv_lens: torch.Tensor, + metadata: Any, + *, + generation_lengths: Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = None, +) -> None: + """Write the step's packed split plan, fused with the generation lengths when given.""" + from .fp4_mla_packed_reduce import launch_packed_plan + + launch_packed_plan( + kv_lens, + geometry.plan, + num_requests=geometry.num_requests, + num_splits=geometry.splits, + wave_tasks=geometry.wave_tasks, + max_tiles_per_task=geometry.profile_tiles, + query_len=geometry.query_len, + helix_kv_bounds=geometry.helix_kv_bounds, + generation_lengths=generation_lengths, + wave_plan=_env_enabled_default(FP4_MLA_PACKED_WAVE_PLAN_ENV, True), + ) + metadata._fp4_mla_packed_plan_key = geometry.key + + +def _env_int(name: str) -> Optional[int]: + value = os.environ.get(name) + if value is None or value == "": + return None + return int(value) + + +def _fp4_mla_attention_backend() -> str: + backend = os.getenv(FP4_MLA_ATTENTION_BACKEND_ENV) + if backend: + return backend.lower() + return _FP4_MLA_CUTEDSL_BACKEND if get_sm_version() == 107 else "triton" + + +def _q_quant_fusion_supported() -> bool: + """Whether the selected backend can fuse Q quantization on this GPU. + + The backend selector already encodes the architecture: SM107 gets CuTeDSL, + everything else gets the Triton K-residual path, and the CuTeDSL launch + re-checks SM107 for itself. Gating this predicate on SM107 as well ruled + out the Triton path on architectures that select it, which is not a + no-op: ``uses_fp4_mla_attention`` carries no SM check, so on GB300 + (SM103) it stayed true while this returned false, and the caller treats + that disagreement as a hard error ("FP4 MLA fused Q quantization + eligibility changed before launch") rather than a fallback. + """ + backend = _fp4_mla_attention_backend() + if backend == _FP4_MLA_CUTEDSL_BACKEND: + return get_sm_version() == 107 + return get_sm_version() in _FP4_MLA_TRITON_Q_QUANT_FUSION_SMS + + +def _cutedsl_backend_available() -> bool: + try: + return all( + importlib.util.find_spec(module) is not None + for module in ("ctm", "cutlass", "cuda.bindings.driver") + ) + except ModuleNotFoundError: + return False + + +def _fp4_mla_cutedsl_kernel_module() -> Any: + if _fp4_mla_cutedsl_fused_v_transpose_enabled(): + from . import fp4_mla_cutedsl_mufu16_fused_v_transpose + + return fp4_mla_cutedsl_mufu16_fused_v_transpose + + from . import fp4_mla_cutedsl_mufu16 + + return fp4_mla_cutedsl_mufu16 + + +def _ceil_div(lhs: int, rhs: int) -> int: + return (lhs + rhs - 1) // rhs + + +def _fp4_mla_q1_kv_blocks_per_program(num_gen: int, v_head_dim: int) -> int: + """Select the Q1 KV work per program without changing the launch boundary.""" + large_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH * FP4_BLOCK_SIZE + if num_gen >= _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD and v_head_dim % large_batch_block_dim == 0: + return _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH + medium_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH * FP4_BLOCK_SIZE + if ( + num_gen >= _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD + and v_head_dim % medium_batch_block_dim == 0 + ): + return _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH + small_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH * FP4_BLOCK_SIZE + if v_head_dim % small_batch_block_dim == 0: + return _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH + return 1 + + +def _fp4_mla_q1_prefix_blocks_per_program( + num_gen: int, + q1_kv_blocks_per_program: int, +) -> int: + """Select rolled Q-prefix work only when the batch keeps it efficient.""" + max_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM + if ( + num_gen >= _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD + and q1_kv_blocks_per_program == _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH + ): + if num_gen >= _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD: + return min(4, max_prefix_blocks) + return min(2, max_prefix_blocks) + return 1 + + +def _fp4_mla_q1_preload_variants( + max_num_sequences: int, + v_head_dim: int, +) -> tuple[tuple[int, int], ...]: + """Return every Q1 tuning variant reachable by the configured batch limit.""" + batch_sizes = [1] + for threshold in ( + _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD, + _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD, + _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD, + _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD, + ): + if threshold <= max_num_sequences: + batch_sizes.append(threshold) + + variants = [] + for batch_size in batch_sizes: + kv_blocks = _fp4_mla_q1_kv_blocks_per_program(batch_size, v_head_dim) + prefix_blocks = _fp4_mla_q1_prefix_blocks_per_program( + batch_size, + kv_blocks, + ) + variant = (kv_blocks, prefix_blocks) + if variant not in variants: + variants.append(variant) + return tuple(variants) + + +def _fp4_mla_triton_preload_key_set(metadata: Any) -> set[tuple[object, ...]]: + """Return the engine-scoped set of Triton variants loaded during warmup.""" + owner = getattr(metadata, "kv_cache_manager", None) + if owner is None: + owner = metadata + preload_keys = getattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, None) + if preload_keys is None: + preload_keys = set() + setattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, preload_keys) + return preload_keys + + +@triton.jit +def _fp4_mla_store_sequence_append_metadata( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + sequence_idx, + num_tokens, + PREFIX_BLOCK: tl.constexpr, + TOKEN_BLOCK: tl.constexpr, +): + append_len = tl.load(append_lens_ptr + sequence_idx) + prefix_offsets = tl.arange(0, PREFIX_BLOCK) + token_start = append_len - append_len + for prefix_start in tl.range(0, sequence_idx, PREFIX_BLOCK): + preceding_sequences = prefix_start + prefix_offsets + preceding_lens = tl.load( + append_lens_ptr + preceding_sequences, + mask=preceding_sequences < sequence_idx, + other=0, + ) + token_start += tl.sum(preceding_lens) + + cached_len = tl.load(kv_lens_ptr + sequence_idx) - append_len + + token_offsets = tl.arange(0, TOKEN_BLOCK) + for block_start in tl.range(0, append_len, TOKEN_BLOCK): + local_offsets = block_start + token_offsets + token_mask = (local_offsets < append_len) & (token_start + local_offsets < num_tokens) + output_offsets = token_start + local_offsets + tl.store( + batch_indices_ptr + output_offsets, + sequence_idx, + mask=token_mask, + ) + tl.store( + positions_ptr + output_offsets, + cached_len + local_offsets, + mask=token_mask, + ) + + +@triton.jit( + do_not_specialize=[ + "num_tokens", + "num_contexts", + "num_generation_sequences", + ], + do_not_specialize_on_alignment=[ + "num_tokens", + "num_contexts", + "num_generation_sequences", + ], +) +def _fp4_mla_append_metadata_kernel( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + num_tokens, + num_contexts, + num_generation_sequences, + ONE_TOKEN_GENERATION: tl.constexpr, + PREFIX_BLOCK: tl.constexpr, + TOKEN_BLOCK: tl.constexpr, + GENERATION_BLOCK: tl.constexpr, +): + program_idx = tl.program_id(0) + if ONE_TOKEN_GENERATION: + if program_idx < num_contexts: + _fp4_mla_store_sequence_append_metadata( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + program_idx, + num_tokens, + PREFIX_BLOCK, + TOKEN_BLOCK, + ) + else: + generation_offsets = (program_idx - num_contexts) * GENERATION_BLOCK + tl.arange( + 0, GENERATION_BLOCK + ) + generation_mask = generation_offsets < num_generation_sequences + sequence_indices = num_contexts + generation_offsets + output_offsets = num_tokens - num_generation_sequences + generation_offsets + generation_mask = generation_mask & (output_offsets < num_tokens) + generation_positions = ( + tl.load( + kv_lens_ptr + sequence_indices, + mask=generation_mask, + other=1, + ) + - 1 + ) + tl.store( + batch_indices_ptr + output_offsets, + sequence_indices, + mask=generation_mask, + ) + tl.store( + positions_ptr + output_offsets, + generation_positions, + mask=generation_mask, + ) + else: + _fp4_mla_store_sequence_append_metadata( + append_lens_ptr, + kv_lens_ptr, + batch_indices_ptr, + positions_ptr, + program_idx, + num_tokens, + PREFIX_BLOCK, + TOKEN_BLOCK, + ) + + +def populate_fp4_mla_append_metadata( + append_lens: torch.Tensor, + kv_lens: torch.Tensor, + batch_indices: torch.Tensor, + positions: torch.Tensor, + *, + num_tokens: int, + num_sequences: int, + num_contexts: int, + num_context_tokens: int, +) -> None: + """Populate FP4 MLA token-to-sequence metadata in one Triton launch. + + Mixed batches vectorize their one-token generation rows. Multi-token MTP + and fallback shapes use the generic per-sequence path in the same kernel. + """ + if num_sequences <= 0 or num_tokens <= 0: + return + if not 0 <= num_contexts <= num_sequences: + raise ValueError( + f"FP4 MLA num_contexts must be in [0, {num_sequences}], got {num_contexts}." + ) + if not 0 <= num_context_tokens <= num_tokens: + raise ValueError( + f"FP4 MLA num_context_tokens must be in [0, {num_tokens}], got {num_context_tokens}." + ) + + tensors = ( + append_lens, + kv_lens, + batch_indices, + positions, + ) + if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): + raise ValueError("FP4 MLA append metadata tensors must be contiguous and one-dimensional.") + if any(tensor.dtype != torch.int32 for tensor in tensors): + raise TypeError("FP4 MLA append metadata tensors must use int32.") + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("FP4 MLA append metadata tensors must be CUDA tensors.") + if any(tensor.device != append_lens.device for tensor in tensors[1:]): + raise ValueError("FP4 MLA append metadata tensors must be on the same device.") + sequence_tensors = (append_lens, kv_lens) + if any(tensor.numel() < num_sequences for tensor in sequence_tensors): + raise ValueError( + f"FP4 MLA sequence metadata tensors need at least {num_sequences} entries." + ) + token_tensors = (batch_indices, positions) + if any(tensor.numel() < num_tokens for tensor in token_tensors): + raise ValueError(f"FP4 MLA token metadata tensors need at least {num_tokens} entries.") + + num_generation_sequences = num_sequences - num_contexts + # Each scheduled generation sequence appends at least one token. Equality + # therefore identifies the common mixed batch with one token per decode + # row without reading the device append lengths back on the host. + one_token_generation = num_tokens == num_context_tokens + num_generation_sequences + generation_block = 128 + grid = num_sequences + if one_token_generation: + grid = num_contexts + triton.cdiv(num_generation_sequences, generation_block) + + _fp4_mla_append_metadata_kernel[(grid,)]( + append_lens, + kv_lens, + batch_indices, + positions, + num_tokens, + num_contexts, + num_generation_sequences, + ONE_TOKEN_GENERATION=one_token_generation, + PREFIX_BLOCK=128, + TOKEN_BLOCK=256, + GENERATION_BLOCK=generation_block, + num_warps=4, + ) + + +@triton.jit( + do_not_specialize=["num_gen", "generation_len"], + do_not_specialize_on_alignment=["num_gen", "generation_len"], +) +def _fp4_mla_generation_lengths_kernel( + kv_lens_ptr, + prompt_lens_ptr, + corrected_kv_lens_ptr, + generation_lens_ptr, + num_gen, + generation_len, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < num_gen + kv_lens = tl.load(kv_lens_ptr + offsets, mask=mask, other=0) + prompt_lens = tl.load(prompt_lens_ptr + offsets, mask=mask, other=0) + tl.store( + corrected_kv_lens_ptr + offsets, + kv_lens - prompt_lens + generation_len, + mask=mask, + ) + tl.store(generation_lens_ptr + offsets, generation_len, mask=mask) + + +def populate_fp4_mla_generation_lengths( + kv_lens: torch.Tensor, + prompt_lens: torch.Tensor, + corrected_kv_lens: torch.Tensor, + generation_lens: torch.Tensor, + *, + num_gen_tokens: int, + num_gen: int, +) -> None: + """Populate reusable FP4 MLA generation lengths in one Triton launch.""" + if num_gen <= 0 or num_gen_tokens % num_gen != 0: + raise ValueError( + "FP4 MLA generation lengths require a positive sequence count and " + f"uniform token count, got {num_gen_tokens} tokens for {num_gen} sequences." + ) + tensors = (kv_lens, prompt_lens, corrected_kv_lens, generation_lens) + if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): + raise ValueError( + "FP4 MLA generation length tensors must be contiguous and one-dimensional." + ) + if any(tensor.dtype != torch.int32 for tensor in tensors): + raise TypeError("FP4 MLA generation length tensors must use int32.") + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("FP4 MLA generation length tensors must be CUDA tensors.") + if any(tensor.device != kv_lens.device for tensor in tensors[1:]): + raise ValueError("FP4 MLA generation length tensors must be on the same device.") + if any(tensor.numel() < num_gen for tensor in tensors): + raise ValueError(f"FP4 MLA generation length tensors need at least {num_gen} entries.") + + block = 128 + _fp4_mla_generation_lengths_kernel[(triton.cdiv(num_gen, block),)]( + kv_lens, + prompt_lens, + corrected_kv_lens, + generation_lens, + num_gen, + num_gen_tokens // num_gen, + BLOCK=block, + num_warps=4, + ) + + +@triton.jit( + do_not_specialize=["num_tokens", "tokens_per_seq"], + do_not_specialize_on_alignment=["num_tokens", "tokens_per_seq"], +) +def _helix_spec_buffers_kernel( + positions_ptr, + local_slots_ptr, + kv_bounds_ptr, + kv_lens_ptr, + num_tokens, + tokens_per_seq, + tokens_per_block, + cp_size, + cp_rank, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < num_tokens + position = tl.load(positions_ptr + offsets, mask=mask, other=0) + # Rank-local prefix length of the first g global tokens: ledger page b lives + # on CP rank b % cp_size, so count the full ledgers plus this rank's clamped + # share of the partial one (mirrors KVCacheManagerV2._helix_local_len). + ledger = tokens_per_block * cp_size + rank_base = cp_rank * tokens_per_block + full_before = position // ledger + local_before = full_before * tokens_per_block + tl.minimum( + tl.maximum(position - full_before * ledger - rank_base, 0), tokens_per_block + ) + after = position + 1 + full_after = after // ledger + bound = full_after * tokens_per_block + tl.minimum( + tl.maximum(after - full_after * ledger - rank_base, 0), tokens_per_block + ) + owned = (position // tokens_per_block) % cp_size == cp_rank + tl.store(local_slots_ptr + offsets, tl.where(owned, local_before, -1), mask=mask) + tl.store(kv_bounds_ptr + offsets, bound, mask=mask) + # The last token of each verify group sets its sequence's rank-local KV length. + last_of_seq = offsets % tokens_per_seq == tokens_per_seq - 1 + tl.store(kv_lens_ptr + offsets // tokens_per_seq, bound, mask=mask & last_of_seq) + + +def populate_helix_spec_buffers( + positions: torch.Tensor, + local_slots: torch.Tensor, + kv_bounds: torch.Tensor, + kv_lens: torch.Tensor, + *, + num_tokens: int, + tokens_per_seq: int, + tokens_per_block: int, + cp_size: int, + cp_rank: int, +) -> None: + """Derive the Helix speculative per-token buffers in one launch. + + For generation token ``t`` at global position ``positions[t]``: + ``local_slots[t]`` is its rank-local KV write slot (-1 when another CP rank + owns the page), ``kv_bounds[t]`` the rank-local KV entries it attends to + (the local length of ``position + 1``) and ``kv_lens[s]`` the bound of + sequence ``s``'s last token. Static shapes, safe under CUDA graph capture. + """ + if num_tokens <= 0: + return + if tokens_per_seq <= 0 or num_tokens % tokens_per_seq != 0: + raise ValueError( + f"Helix verify groups must be uniform: {num_tokens} tokens for groups of " + f"{tokens_per_seq}." + ) + if tokens_per_block <= 0 or cp_size <= 0 or not 0 <= cp_rank < cp_size: + raise ValueError("Helix spec buffers need a positive page size and a valid CP rank.") + num_seqs = num_tokens // tokens_per_seq + tensors = (positions, local_slots, kv_bounds, kv_lens) + if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): + raise ValueError("Helix spec buffers must be contiguous and one-dimensional.") + if any(tensor.dtype != torch.int32 for tensor in tensors): + raise TypeError("Helix spec buffers must use int32.") + if any(not tensor.is_cuda or tensor.device != positions.device for tensor in tensors): + raise ValueError("Helix spec buffers must be CUDA tensors on one device.") + if any(tensor.numel() < num_tokens for tensor in tensors[:3]) or kv_lens.numel() < num_seqs: + raise ValueError("Helix spec buffers are too small for the generation tokens.") + + block = 256 + _helix_spec_buffers_kernel[(triton.cdiv(num_tokens, block),)]( + positions, + local_slots, + kv_bounds, + kv_lens, + num_tokens, + tokens_per_seq, + tokens_per_block, + cp_size, + cp_rank, + BLOCK=block, + num_warps=4, + ) + + +def _fp4_mla_page_table_spec(kv_cache_manager: Any) -> Any: + get_spec = getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None) + if not callable(get_spec): + raise RuntimeError("FP4 MLA requires V2 cache-layout page metadata.") + spec = get_spec() + for field_name in ( + "cache_pool_id", + "cache_page_index_scale", + "hp_pool_id", + "hp_page_index_scale", + ): + value = getattr(spec, field_name, None) + if not isinstance(value, int) or value < 0: + raise ValueError( + f"FP4 MLA page-table spec requires non-negative {field_name}, got {value}." + ) + if spec.cache_page_index_scale <= 0 or spec.hp_page_index_scale <= 0: + raise ValueError("FP4 MLA page-index scales must be positive.") + return spec + + +# Mixed batches frequently vary by one sequence. Keep per-forward dimensions +# out of Triton's specialization key and tile page rows at one fixed width so +# those shape changes cannot trigger JIT compilation on the critical path. +@triton.jit( + do_not_specialize=["num_sequences", "num_contexts", "max_pages"], + do_not_specialize_on_alignment=["num_sequences", "num_contexts", "max_pages"], +) +def _fp4_mla_materialize_page_table_kernel( + page_ids_ptr, + paged_kv_indptr_ptr, + paged_kv_indptr_decode_ptr, + block_offsets_ptr, + kv_lens_ptr, + generation_kv_lens_ptr, + block_offsets_stride, + num_sequences, + num_contexts, + max_pages, + PAGE_SIZE: tl.constexpr, + PAGE_INDEX_SCALE: tl.constexpr, + PAGE_TILE_SIZE: tl.constexpr, +): + sequence_idx = tl.program_id(0) + page_tile_idx = tl.program_id(1) + page_offsets = page_tile_idx * PAGE_TILE_SIZE + tl.arange(0, PAGE_TILE_SIZE) + generation_idx = sequence_idx - num_contexts + is_generation = generation_idx >= 0 + context_kv_len = tl.load(kv_lens_ptr + sequence_idx) + generation_kv_len = tl.load( + generation_kv_lens_ptr + generation_idx, + mask=is_generation, + other=0, + ) + kv_len = tl.maximum(tl.where(is_generation, generation_kv_len, context_kv_len), 0) + num_active_pages = tl.minimum( + (kv_len + PAGE_SIZE - 1) // PAGE_SIZE, + max_pages, + ) + active_page_mask = page_offsets < num_active_pages + encoded_page_offsets = tl.load( + block_offsets_ptr + sequence_idx * block_offsets_stride + page_offsets, + mask=active_page_mask, + other=-1, + ) + decoded_page_ids = tl.where( + encoded_page_offsets >= 0, + encoded_page_offsets // PAGE_INDEX_SCALE, + encoded_page_offsets, + ) + page_ids = tl.where(active_page_mask, decoded_page_ids, 0) + table_offset = sequence_idx * max_pages + page_offsets + # Fixed-stride indptrs expose the whole row. Initialize inactive slots so + # masked or prefetched page-table reads cannot observe stale page IDs. + tl.store( + page_ids_ptr + table_offset, + page_ids, + mask=page_offsets < max_pages, + ) + + first_lane = page_offsets == 0 + sequence_start = sequence_idx * max_pages + tl.store( + paged_kv_indptr_ptr + sequence_idx + page_offsets, + sequence_start, + mask=first_lane, + ) + tl.store( + paged_kv_indptr_decode_ptr + generation_idx + page_offsets, + generation_idx * max_pages, + mask=first_lane & is_generation, + ) + final_sequence = sequence_idx == num_sequences - 1 + table_end = num_sequences * max_pages + num_generation_sequences = num_sequences - num_contexts + tl.store( + paged_kv_indptr_ptr + num_sequences + page_offsets, + table_end, + mask=first_lane & final_sequence, + ) + tl.store( + paged_kv_indptr_decode_ptr + num_generation_sequences + page_offsets, + num_generation_sequences * max_pages, + mask=first_lane & final_sequence, + ) + + +def configure_fp4_mla_device_page_table( + metadata: Any, + kv_lens: Optional[torch.Tensor] = None, +) -> bool: + """Configure the fixed-stride, device-materialized page table. + + Eager context and generation batches receive the full block-offset + table on the GPU. The materialization kernel decodes V2 page indices and + refreshes rows from the final device KV lengths before cache update. + """ + metadata._fp4_mla_device_page_table = False + metadata._fp4_mla_device_page_table_valid = False + metadata.fp4_mla_page_table_stride = 0 + metadata.fp4_mla_context_repack_max_touched_pages = 1 + + kv_cache_manager = getattr(metadata, "kv_cache_manager", None) + num_contexts = int(getattr(metadata, "num_contexts", 0)) + num_sequences = int(getattr(metadata, "num_seqs", 0)) + num_generation_sequences = num_sequences - num_contexts + num_tokens = int(getattr(metadata, "num_tokens", 0)) + num_context_tokens = int(getattr(metadata, "num_ctx_tokens", 0)) + num_generation_tokens = num_tokens - num_context_tokens + block_offsets = getattr(metadata, "kv_cache_block_offsets", None) + page_ids = getattr(metadata, "_paged_kv_indices", None) + paged_kv_indptr = getattr(metadata, "_paged_kv_indptr", None) + paged_kv_indptr_decode = getattr(metadata, "paged_kv_indptr_decode", None) + max_page_capacity = int(getattr(kv_cache_manager, "max_blocks_per_seq", 0) or 0) + page_spec = _fp4_mla_page_table_spec(kv_cache_manager) + page_index_scale = int(page_spec.cache_page_index_scale) + + tensors = (block_offsets, page_ids, paged_kv_indptr, paged_kv_indptr_decode) + is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) + generation_only = num_contexts == 0 + eager_context = not is_cuda_graph and num_contexts > 0 + has_valid_generation = num_generation_sequences == 0 or ( + num_generation_tokens >= num_generation_sequences + and num_generation_tokens % num_generation_sequences == 0 + ) + # NVFP4 exposes one data pool plus its paired block-scale pool. The + # materializer reads encoded data offsets from pool 0. + supported = ( + (generation_only or eager_context) + and kv_cache_manager is not None + and has_valid_generation + and int(getattr(metadata, "beam_width", 1)) == 1 + and not bool(getattr(metadata, "is_spec_dec_tree", False)) + and not bool(getattr(metadata, "locality_domain_enabled", False)) + and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK + and max_page_capacity > 0 + and page_index_scale > 0 + and all(isinstance(tensor, torch.Tensor) for tensor in tensors) + and all(tensor.dtype == torch.int32 for tensor in tensors) + and all(tensor.is_cuda for tensor in tensors) + ) + if not supported: + return False + + hp_page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) + max_pool_id = max(page_spec.cache_pool_id, page_spec.hp_pool_id) + if ( + not isinstance(hp_page_ids, torch.Tensor) + or hp_page_ids.dtype != torch.int32 + or not hp_page_ids.is_cuda + or block_offsets.shape[0] <= max_pool_id + ): + return False + metadata._fp4_mla_cache_pool_id = int(page_spec.cache_pool_id) + metadata._fp4_mla_cache_page_index_scale = int(page_spec.cache_page_index_scale) + metadata._fp4_mla_hp_pool_id = int(page_spec.hp_pool_id) + metadata._fp4_mla_hp_page_index_scale = int(page_spec.hp_page_index_scale) + + max_pages = max_page_capacity + host_kv_lens_available = ( + isinstance(kv_lens, torch.Tensor) + and kv_lens.device.type == "cpu" + and kv_lens.ndim == 1 + and kv_lens.numel() >= num_sequences + ) + if eager_context and num_generation_sequences > 0 and not host_kv_lens_available: + return False + if not is_cuda_graph and host_kv_lens_available: + generation_tokens_per_sequence = ( + num_generation_tokens // num_generation_sequences if num_generation_sequences > 0 else 0 + ) + # Eager execution can narrow the fixed row stride to the current + # batch. CUDA Graph metadata retains full configured capacity so a + # replay never changes tensor addresses or launch dimensions. + max_kv_len = int(kv_lens[:num_sequences].max().item()) + max( + 0, + generation_tokens_per_sequence - 1, + ) + max_pages = min( + max_page_capacity, + max(1, _ceil_div(max_kv_len, FP4_MLA_TOKENS_PER_BLOCK)), + ) + if num_contexts > 0: + max_context_len = int(kv_lens[:num_contexts].max().item()) + max_context_pages = _ceil_div( + max_context_len, + FP4_MLA_TOKENS_PER_BLOCK, + ) + metadata.fp4_mla_context_repack_max_touched_pages = min( + max_pages, + triton.next_power_of_2(max(1, max_context_pages)), + ) + + assert isinstance(block_offsets, torch.Tensor) + assert isinstance(page_ids, torch.Tensor) + assert isinstance(paged_kv_indptr, torch.Tensor) + assert isinstance(paged_kv_indptr_decode, torch.Tensor) + required_page_ids = num_sequences * max_pages + buffers_cover_table = ( + block_offsets.ndim == 4 + and block_offsets.shape[0] >= 1 + and block_offsets.shape[1] >= num_sequences + and block_offsets.shape[2] >= 1 + and block_offsets.shape[3] >= max_pages + and page_ids.ndim == 1 + and page_ids.numel() >= required_page_ids + and paged_kv_indptr.ndim == 1 + and paged_kv_indptr.numel() >= num_sequences + 1 + and paged_kv_indptr_decode.ndim == 1 + and paged_kv_indptr_decode.numel() >= num_generation_sequences + 1 + ) + if not buffers_cover_table: + return False + if metadata._fp4_mla_hp_page_indices.numel() < required_page_ids: + return False + + metadata._fp4_mla_device_page_table = True + metadata.fp4_mla_page_table_stride = max_pages + metadata.num_blocks = None + metadata.num_context_blocks = num_contexts * max_pages + metadata.num_generation_blocks = num_generation_sequences * max_pages + return True + + +def materialize_fp4_mla_device_page_table( + metadata: Any, + kv_lens: torch.Tensor, + generation_kv_lens: Optional[torch.Tensor] = None, +) -> None: + """Refresh the fixed-stride context and generation page table once per forward.""" + if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): + raise RuntimeError("FP4 MLA requires fixed-stride device page metadata.") + if bool(getattr(metadata, "_fp4_mla_device_page_table_valid", False)): + return + + num_contexts = int(metadata.num_contexts) + num_sequences = int(metadata.num_seqs) + num_generation_sequences = num_sequences - num_contexts + max_pages = int(metadata.fp4_mla_page_table_stride) + if num_sequences <= 0 or max_pages <= 0: + raise RuntimeError( + "FP4 MLA device page metadata requires positive sequence and page capacities." + ) + if ( + kv_lens.dtype != torch.int32 + or not kv_lens.is_cuda + or kv_lens.ndim != 1 + or kv_lens.stride(0) != 1 + or kv_lens.numel() < num_sequences + ): + raise ValueError( + "FP4 MLA device page metadata requires a contiguous CUDA int32 " + f"KV-length tensor with at least {num_sequences} entries." + ) + + if generation_kv_lens is None: + generation_kv_lens = kv_lens[num_contexts:num_sequences] + if ( + generation_kv_lens.dtype != torch.int32 + or not generation_kv_lens.is_cuda + or generation_kv_lens.ndim != 1 + or generation_kv_lens.stride(0) != 1 + or generation_kv_lens.numel() < num_generation_sequences + ): + raise ValueError( + "FP4 MLA device page metadata requires a contiguous CUDA int32 " + "generation KV-length tensor with at least " + f"{num_generation_sequences} entries." + ) + + cache_pool_id = int(getattr(metadata, "_fp4_mla_cache_pool_id", 0)) + block_offsets = metadata.kv_cache_block_offsets[ + cache_pool_id, + :num_sequences, + 0, + :max_pages, + ] + page_ids = metadata._paged_kv_indices[: num_sequences * max_pages] + page_index_scale = int(metadata._fp4_mla_cache_page_index_scale) + if page_index_scale <= 0: + raise RuntimeError("FP4 MLA device page metadata requires a positive page-index scale.") + grid = ( + num_sequences, + triton.cdiv(max_pages, _FP4_MLA_PAGE_TABLE_TILE_SIZE), + ) + _fp4_mla_materialize_page_table_kernel[grid]( + page_ids, + metadata._paged_kv_indptr, + metadata.paged_kv_indptr_decode, + block_offsets, + kv_lens, + generation_kv_lens, + block_offsets.stride(0), + num_sequences, + num_contexts, + max_pages, + PAGE_SIZE=metadata.page_size, + PAGE_INDEX_SCALE=page_index_scale, + PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, + num_warps=4, + ) + hp_page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) + if not isinstance(hp_page_ids, torch.Tensor): + raise RuntimeError("FP4 MLA requires an HP page-table output tensor.") + hp_pool_id = int(metadata._fp4_mla_hp_pool_id) + hp_page_index_scale = int(metadata._fp4_mla_hp_page_index_scale) + hp_block_offsets = metadata.kv_cache_block_offsets[ + hp_pool_id, + :num_sequences, + 0, + :max_pages, + ] + _fp4_mla_materialize_page_table_kernel[grid]( + hp_page_ids, + metadata._paged_kv_indptr, + metadata.paged_kv_indptr_decode, + hp_block_offsets, + kv_lens, + generation_kv_lens, + hp_block_offsets.stride(0), + num_sequences, + num_contexts, + max_pages, + PAGE_SIZE=metadata.page_size, + PAGE_INDEX_SCALE=hp_page_index_scale, + PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, + num_warps=4, + ) + metadata._fp4_mla_device_page_table_valid = True + + +@triton.jit +def _cutedsl_swizzled_sf_offset(row_idx, col_idx, sf_cols: tl.constexpr): + padded_cols = ((sf_cols + 3) // 4) * 4 + return ( + col_idx % 4 + + (col_idx // 4) * (4 * 128) + + (row_idx % 32) * 16 + + ((row_idx % 128) // 32) * 4 + + (row_idx // 128) * (128 * padded_cols) + ) + + +@triton.jit +def _cutedsl_pad_q_and_sf_kernel( + q_padded_ptr, + q_ptr, + q_sf_padded_ptr, + q_sf_ptr, + num_heads, + output_heads: tl.constexpr, + packed_dim: tl.constexpr, + block_bytes: tl.constexpr, + sf_cols: tl.constexpr, + sf_cols_per_byte_block: tl.constexpr, +): + query_idx = tl.program_id(0) + byte_block = tl.program_id(1) + head_offsets = tl.arange(0, output_heads) + byte_offsets = byte_block * block_bytes + tl.arange(0, block_bytes) + head_mask = head_offsets < num_heads + byte_mask = byte_offsets < packed_dim + source_rows = query_idx * num_heads + head_offsets + destination_rows = query_idx * output_heads + head_offsets + values = tl.load( + q_ptr + source_rows[:, None] * packed_dim + byte_offsets[None, :], + mask=head_mask[:, None] & byte_mask[None, :], + other=0, + ) + tl.store( + q_padded_ptr + destination_rows[:, None] * packed_dim + byte_offsets[None, :], + values, + mask=byte_mask[None, :], + ) + sf_col_offsets = byte_block * sf_cols_per_byte_block + tl.arange(0, sf_cols_per_byte_block) + sf_col_mask = sf_col_offsets < sf_cols + source_offsets = _cutedsl_swizzled_sf_offset( + source_rows[:, None], sf_col_offsets[None, :], sf_cols + ) + destination_offsets = _cutedsl_swizzled_sf_offset( + destination_rows[:, None], sf_col_offsets[None, :], sf_cols + ) + sf_values = tl.load( + q_sf_ptr + source_offsets, + mask=head_mask[:, None] & sf_col_mask[None, :], + other=1.0, + ) + tl.store( + q_sf_padded_ptr + destination_offsets, + sf_values, + mask=sf_col_mask[None, :], + ) + + +_SM_COUNT_CACHE: dict[int, int] = {} + + +def _get_sm_count(device: torch.device) -> int: + """Return the SM (multiprocessor) count for ``device``, cached per index.""" + index = device.index if device.index is not None else torch.cuda.current_device() + count = _SM_COUNT_CACHE.get(index) + if count is None: + count = torch.cuda.get_device_properties(index).multi_processor_count + _SM_COUNT_CACHE[index] = count + return count + + +def _validate_fp4_mla_context_rope( + latent_cache: torch.Tensor, + rotary_cos_sin: torch.Tensor, + v_head_dim: int, +) -> int: + head_dim = latent_cache.shape[-1] + rope_dim = head_dim - v_head_dim + if rope_dim <= 0 or rope_dim % 2 != 0: + raise ValueError( + "FP4 MLA fused context K-RoPE requires a positive even RoPE dimension, " + f"got head_dim={head_dim}, v_head_dim={v_head_dim}." + ) + if rotary_cos_sin.device != latent_cache.device: + raise ValueError("FP4 MLA context latent cache and RoPE table must use the same device.") + if rotary_cos_sin.dtype != torch.float32: + raise TypeError( + f"FP4 MLA fused context K-RoPE requires a float32 table, got {rotary_cos_sin.dtype}." + ) + if not rotary_cos_sin.is_contiguous(): + raise ValueError("FP4 MLA fused context K-RoPE requires a contiguous RoPE table.") + table_row_size = rope_dim * 2 + if rotary_cos_sin.numel() < table_row_size or rotary_cos_sin.numel() % table_row_size != 0: + raise ValueError( + "FP4 MLA context RoPE table size must be a positive multiple of " + f"{table_row_size}, got {rotary_cos_sin.numel()}." + ) + return rope_dim + + +def _host_int_list_during_forward(value: Any, start: int, end: int) -> Optional[list[int]]: + if torch.cuda.is_current_stream_capturing(): + return None + return _host_int_list(value, start, end) + + +# FP4 MLA scale-layout helpers + + +def get_fp4_mla_v_scale_pool_size(v_head_dim: int, page_size: int) -> int: + """Return elements per page for the swizzled FP4 MLA V-scale pool. + + The PV matmul treats V as a RHS matrix shaped ``[v_head_dim, kv_tokens]``. + NVFP4 block scales therefore group along the token/K axis, not along the + latent dimension as the K-view cache does. The physical layout matches the + Triton block-scaled matmul scale layout: + ``[ceil(v_head_dim / 128), ceil(page_size / 16 / 4), 32, 16]``. + """ + return _get_fp4_mla_swizzled_scale_size(v_head_dim, page_size) + + +def _get_fp4_mla_swizzled_scale_size(rows: int, cols: int) -> int: + scale_cols = _ceil_div(cols, FP4_BLOCK_SIZE) + row_groups = _ceil_div(rows, FP4_MLA_SCALE_ROW_GROUP) + col_groups = _ceil_div(scale_cols, FP4_MLA_SCALE_COL_GROUP) + return row_groups * col_groups * 32 * 16 + + +def can_fuse_fp4_mla_q_quant( + metadata: Any, + q: torch.Tensor, + q_pe: torch.Tensor, + latent_cache: torch.Tensor, +) -> bool: + """Return whether generation can quantize Q in the fused cache update.""" + num_gen = metadata.num_seqs - metadata.num_contexts + return bool( + _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS + and _q_quant_fusion_supported() + and num_gen > 0 + and getattr(metadata, "kv_cache_manager", None) is not None + and q.shape[0] > 0 + and q.shape[0] % num_gen == 0 + and q.is_cuda + and q_pe.is_cuda + and latent_cache.is_cuda + and q.device == q_pe.device == latent_cache.device + and q.dtype == torch.bfloat16 + and q.is_contiguous() + and q.ndim == 3 + and 0 < q.shape[1] <= 128 + and q.shape[2] == FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM + and q_pe.dtype == torch.bfloat16 + and tuple(q_pe.shape) == (q.shape[0], q.shape[1], FP4_MLA_Q_RESIDUAL_DIM) + and latent_cache.dtype == torch.bfloat16 + and tuple(latent_cache.shape) == (q.shape[0], FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM) + and metadata.page_size == FP4_MLA_TOKENS_PER_BLOCK + ) + + +def _prepare_fp4_mla_q_buffers( + metadata: Any, + num_queries: int, + num_heads: int, + device: torch.device, + rows_per_seq: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Return manager-owned, fixed-capacity packed-Q staging buffers. + + With ``rows_per_seq`` the buffers hold ``max_num_sequences`` request blocks + of that many rows (zero-filled so padded rows stay finite); otherwise one + row per (token, head) up to the token capacity. + """ + if num_queries <= 0: + raise ValueError(f"FP4 MLA Q buffer preparation needs queries, got {num_queries}.") + owner = getattr(metadata, "kv_cache_manager", None) + if owner is None: + raise RuntimeError("Fused FP4 MLA Q quantization requires a KV cache manager.") + if num_heads <= 0 or num_heads > 128: + raise ValueError(f"FP4 MLA Q buffers require 1-128 local heads, got {num_heads}.") + buffers = getattr(owner, "_fp4_mla_q_buffers", None) + if buffers is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Cannot create FP4 MLA Q buffers while capturing a CUDA graph. " + "Run a warmup forward first." + ) + buffers = {} + setattr(owner, "_fp4_mla_q_buffers", buffers) + + max_num_tokens = int(getattr(metadata, "max_num_tokens", num_queries) or num_queries) + max_num_sequences = int( + getattr(metadata, "max_num_sequences", None) + or getattr(metadata, "max_num_requests", num_queries) + or num_queries + ) + max_query_width = 1 + int(getattr(metadata, "max_total_draft_tokens", None) or 0) + capacity = min( + max_num_tokens, + max_num_sequences * max_query_width, + _FP4_MLA_MAX_GRID_Z, + ) + if num_queries > capacity: + raise ValueError( + f"FP4 MLA active queries exceed the configured Q capacity: {num_queries} > {capacity}." + ) + + device_index = device.index if device.index is not None else torch.cuda.current_device() + canonical_device = torch.device("cuda", device_index) + if rows_per_seq > 0: + q_rows = max_num_sequences * rows_per_seq + q_key = f"q_{num_heads}_p{rows_per_seq}" + q_sf_key = f"q_sf_{num_heads}_p{rows_per_seq}" + else: + q_rows = capacity * num_heads + q_key = f"q_{num_heads}" + q_sf_key = f"q_sf_{num_heads}" + expected_q_shape = (q_rows, FP4_MLA_Q_PACKED_DIM) + expected_q_sf_shape = (_get_fp4_mla_swizzled_scale_size(q_rows, FP4_MLA_Q_LOGICAL_DIM),) + q_storage = buffers.get(q_key) + q_sf_storage = buffers.get(q_sf_key) + if q_storage is None and q_sf_storage is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Cannot allocate FP4 MLA Q buffers while capturing a CUDA graph. " + "Run a warmup forward first." + ) + allocate = torch.zeros if rows_per_seq > 0 else torch.empty + q_storage = allocate(expected_q_shape, dtype=torch.uint8, device=canonical_device) + q_sf_storage = allocate( + expected_q_sf_shape, + dtype=torch.float8_e4m3fn, + device=canonical_device, + ) + buffers[q_key] = q_storage + buffers[q_sf_key] = q_sf_storage + if ( + q_storage is None + or q_storage.dtype != torch.uint8 + or q_storage.device != canonical_device + or tuple(q_storage.shape) != expected_q_shape + or not q_storage.is_contiguous() + or q_sf_storage is None + or q_sf_storage.dtype != torch.float8_e4m3fn + or q_sf_storage.device != canonical_device + or tuple(q_sf_storage.shape) != expected_q_sf_shape + or not q_sf_storage.is_contiguous() + ): + raise RuntimeError("FP4 MLA packed-Q buffers do not match the configured capacity.") + return q_storage, q_sf_storage, capacity + + +def _get_fp4_mla_context_start_positions(metadata: Any, num_contexts: int) -> torch.Tensor: + kv_cache_params = getattr(metadata, "kv_cache_params", None) + cached_token_lens = getattr(kv_cache_params, "num_cached_tokens_per_seq", None) + if cached_token_lens is not None: + return torch.as_tensor(cached_token_lens[:num_contexts], dtype=torch.int64, device="cpu") + + return ( + ( + metadata.kv_lens_cuda_runtime[:num_contexts] + - metadata.prompt_lens_cuda_runtime[:num_contexts] + ) + .detach() + .cpu() + ) + + +def _validate_fp4_mla_context_start_alignment( + metadata: Any, + num_contexts: int, + *, + alignment: int = HP_BLOCK_SIZE, +) -> None: + context_start_positions = _get_fp4_mla_context_start_positions(metadata, num_contexts) + bad_start = (context_start_positions < 0) | ((context_start_positions % alignment) != 0) + if bool(torch.any(bad_start).item()): + starts = context_start_positions.detach().cpu().tolist() + raise ValueError( + "FP4 MLA shared-tile context update requires every context " + f"start position to be {alignment}-token aligned, got " + f"start positions {starts}." + ) + + +def get_fp4_mla_v_scale_pool_shape( + num_layers: int, + num_pages: int, + v_head_dim: int, + page_size: int, +) -> tuple[int, int, int, int, int, int]: + """Return the logical swizzled V-scale view shape. + + The leading dimensions are ``[layer, physical_page]``. The remaining + dimensions are the preshuffled ``[N // 128, K // 16 // 4, 32, 16]`` shape + consumed by Triton block-scaled matmul for the V/PV RHS operand. + """ + token_scale_cols = _ceil_div(page_size, FP4_BLOCK_SIZE) + return ( + num_layers, + num_pages, + _ceil_div(v_head_dim, FP4_MLA_SCALE_ROW_GROUP), + _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP), + 32, + 16, + ) + + +def get_fp4_mla_v_scale_pool_view( + metadata: Any, + *, + v_head_dim: int, +) -> torch.Tensor: + """View the auxiliary MLA V-scale pool in Triton's block-scaled layout.""" + pool = getattr(metadata, "fp4_mla_v_scale_pool", None) + if pool is None: + raise RuntimeError("FP4 MLA V scale pool is not allocated.") + + elems_per_page = get_fp4_mla_v_scale_pool_size(v_head_dim, metadata.page_size) + if pool.shape[-1] < elems_per_page: + raise RuntimeError( + f"FP4 MLA V scale pool page stride is too small: got " + f"{pool.shape[-1]}, need {elems_per_page}." + ) + + token_scale_cols = _ceil_div(metadata.page_size, FP4_BLOCK_SIZE) + col_groups = _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP) + shape = get_fp4_mla_v_scale_pool_shape( + pool.shape[0], pool.shape[1], v_head_dim, metadata.page_size + ) + strides = ( + pool.stride(0), + pool.stride(1), + col_groups * 32 * 16, + 32 * 16, + 16, + 1, + ) + return torch.as_strided(pool, size=shape, stride=strides) + + +def _rebuild_fp4_mla_v_scales_from_k_scales( + sf_cache: torch.Tensor, + v_scale_pool: torch.Tensor, + page_ids: torch.Tensor, + page_valid_tokens: torch.Tensor, + *, + local_layer: int, + v_head_dim: int, + page_size: int, +) -> None: + """Bit-exactly rebuild imported MLA V scales from transferred K scales.""" + if page_ids.numel() == 0: + return + if page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + "FP4 MLA imported V-scale rebuild requires " + f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." + ) + for name, tensor in ( + ("sf_cache", sf_cache), + ("v_scale_pool", v_scale_pool), + ("page_ids", page_ids), + ("page_valid_tokens", page_valid_tokens), + ): + if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor.") + if page_ids.dtype != torch.int32 or page_valid_tokens.dtype != torch.int32: + raise TypeError("FP4 MLA imported page IDs and valid-token counts must use int32.") + if page_ids.ndim != 1 or page_valid_tokens.ndim != 1: + raise ValueError("FP4 MLA imported page metadata must be one-dimensional.") + if page_ids.numel() != page_valid_tokens.numel(): + raise ValueError("FP4 MLA imported page IDs and valid-token counts must have equal length.") + if not page_ids.is_contiguous() or not page_valid_tokens.is_contiguous(): + raise ValueError("FP4 MLA imported page metadata must be contiguous.") + if not (sf_cache.device == v_scale_pool.device == page_ids.device == page_valid_tokens.device): + raise ValueError("FP4 MLA imported cache tensors must be on the same device.") + + k_sf_bytes = sf_cache.view(torch.uint8) + v_sf_bytes = v_scale_pool.view(torch.uint8) + if k_sf_bytes.ndim < 2 or v_sf_bytes.ndim < 3: + raise ValueError( + "FP4 MLA imported scale pools require per-page K storage and " + "per-layer/per-page V storage." + ) + num_layers = int(v_sf_bytes.shape[0]) + num_pages = int(v_sf_bytes.shape[1]) + if not 0 <= local_layer < num_layers: + raise IndexError( + f"local_layer={local_layer} is outside the V-scale pool with {num_layers} layers." + ) + if int(k_sf_bytes.shape[0]) != num_pages: + raise ValueError( + "FP4 MLA K/V scale pools disagree on their physical page count: " + f"{int(k_sf_bytes.shape[0])} != {num_pages}." + ) + sf_per_token = int(k_sf_bytes.shape[-1]) + required_sf_per_token = _ceil_div(v_head_dim, FP4_BLOCK_SIZE) + if sf_per_token < required_sf_per_token: + raise ValueError( + "FP4 MLA K-scale storage is too narrow for the compressed V head: " + f"{sf_per_token} < {required_sf_per_token}." + ) + required_v_page_elems = get_fp4_mla_v_scale_pool_size(v_head_dim, page_size) + if int(v_sf_bytes.shape[-1]) < required_v_page_elems: + raise ValueError( + "FP4 MLA V-scale page storage is too small for import rebuild: " + f"{int(v_sf_bytes.shape[-1])} < {required_v_page_elems}." + ) + + token_groups = page_size // HP_BLOCK_SIZE + _fp4_mla_rebuild_v_scale_from_k_scale_kernel[ + (page_ids.numel(), triton.cdiv(v_head_dim, FP4_BLOCK_SIZE)) + ]( + k_sf_bytes, + v_sf_bytes, + page_ids, + page_valid_tokens, + page_ids.numel(), + num_pages, + num_layers, + local_layer, + page_size, + k_sf_bytes.stride(0), + v_sf_bytes.stride(0), + v_sf_bytes.stride(1), + V_HEAD_D=v_head_dim, + HP_BLOCK=HP_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=token_groups, + BLOCK_TOKEN_GROUPS=triton.next_power_of_2(token_groups), + num_warps=4, + ) + + +def _stage_fp4_mla_import_page_metadata( + prompt_block_ids: list[int], + *, + prompt_len: int, + page_size: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + """Stage imported page metadata without synchronizing the CUDA stream.""" + if device.type != "cuda": + raise ValueError("FP4 MLA disaggregated import requires a CUDA device.") + num_prompt_pages = len(prompt_block_ids) + pin_memory = prefer_pinned() + page_ids_host = torch.tensor( + prompt_block_ids, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + page_valid_tokens_host = torch.full( + (num_prompt_pages,), + page_size, + dtype=torch.int32, + device="cpu", + pin_memory=pin_memory, + ) + page_valid_tokens_host[-1] = prompt_len - (num_prompt_pages - 1) * page_size + # Constructing a CUDA tensor directly from a Python list synchronizes the + # current stream. During overlap scheduling that stream already waits for + # the previous forward, exposing the import rebuild as an inter-step gap. + # Explicit non-blocking copies keep the CPU free to enqueue every rebuild + # and the next forward while preserving their existing stream order. + page_ids = torch.empty_like(page_ids_host, device=device) + page_valid_tokens = torch.empty_like(page_valid_tokens_host, device=device) + page_ids.copy_(page_ids_host, non_blocking=True) + page_valid_tokens.copy_(page_valid_tokens_host, non_blocking=True) + return page_ids, page_valid_tokens + + +def rebuild_fp4_mla_disagg_imported_cache( + kv_cache_manager: Any, + request_id: int, + prompt_len: int, +) -> bool: + """Rebuild GEN-local FP4 MLA sidecars after a disaggregated KV import. + + The disaggregated payload carries the native V2 K, K-scale, and BF16 HP + roles. V scales and CuTeDSL's V-packed layout are deterministic + process-local views, so rebuilding them here avoids transfer bandwidth and + guarantees they are ready before a first decode step that may execute + through a pre-captured CUDA graph. + """ + if ( + kv_cache_manager is None + or getattr(kv_cache_manager, "dtype", None) != DataType.NVFP4 + or getattr(kv_cache_manager, "kv_factor", None) != 1 + or getattr(kv_cache_manager, "mla_v_scale_head_dim", None) is None + or not callable(getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None)) + ): + return False + if not isinstance(prompt_len, int) or prompt_len < 0: + raise ValueError( + f"FP4 MLA disaggregated import needs a nonnegative prompt_len, got {prompt_len}." + ) + # Helix assigns whole pages round-robin, so a rank may own no prompt pages. + # There are no process-local V sidecars to rebuild on that rank. + if prompt_len == 0: + return True + + page_size = int(kv_cache_manager.tokens_per_block) + if page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + "FP4 MLA disaggregated import requires " + f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." + ) + pp_layers = list(getattr(kv_cache_manager, "pp_layers", ())) + num_local_layers = int(getattr(kv_cache_manager, "num_local_layers", len(pp_layers))) + if len(pp_layers) != num_local_layers: + raise RuntimeError( + "FP4 MLA disaggregated import cannot map local to global layers: " + f"{len(pp_layers)} PP layers for {num_local_layers} local layers." + ) + fp4_local_layers = list( + getattr(kv_cache_manager, "_fp4_mla_compact_to_local", range(num_local_layers)) + ) + if not fp4_local_layers: + raise RuntimeError("FP4 MLA disaggregated import found no local MLA layers.") + if any(local_layer < 0 or local_layer >= num_local_layers for local_layer in fp4_local_layers): + raise RuntimeError( + "FP4 MLA disaggregated import has invalid compact-to-local layer mapping: " + f"{fp4_local_layers}." + ) + + num_prompt_pages = _ceil_div(prompt_len, page_size) + first_attention_layer = pp_layers[fp4_local_layers[0]] + block_ids_per_seq = kv_cache_manager.get_batch_cache_indices( + [int(request_id)], layer_idx=first_attention_layer + ) + if len(block_ids_per_seq) != 1 or len(block_ids_per_seq[0]) < num_prompt_pages: + available = len(block_ids_per_seq[0]) if block_ids_per_seq else 0 + raise RuntimeError( + "FP4 MLA disaggregated import is missing prompt pages for request " + f"{request_id}: need {num_prompt_pages}, have {available}." + ) + prompt_block_ids = [int(block_id) for block_id in block_ids_per_seq[0][:num_prompt_pages]] + + v_scale_pool = kv_cache_manager.get_mla_v_scale_pool() + if not isinstance(v_scale_pool, torch.Tensor): + raise RuntimeError("FP4 MLA disaggregated import requires the manager V-scale pool.") + page_ids, page_valid_tokens = _stage_fp4_mla_import_page_metadata( + prompt_block_ids, + prompt_len=prompt_len, + page_size=page_size, + device=v_scale_pool.device, + ) + + v_scale_head_dim = int(kv_cache_manager.mla_v_scale_head_dim) + cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND + for compact_layer, local_layer in enumerate(fp4_local_layers): + layer_idx = pp_layers[local_layer] + kv_cache, sf_cache = kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) + _rebuild_fp4_mla_v_scales_from_k_scales( + sf_cache, + v_scale_pool, + page_ids, + page_valid_tokens, + local_layer=compact_layer, + v_head_dim=v_scale_head_dim, + page_size=page_size, + ) + if cutedsl_backend and not _fp4_mla_cutedsl_fused_v_transpose_enabled(): + v_head_dim = getattr(kv_cache_manager, "mla_v_head_dim", None) + if v_head_dim is None: + raise RuntimeError( + "CuTeDSL FP4 MLA disaggregated import requires a persistent V head dimension." + ) + v_packed = kv_cache_manager.get_mla_v_packed_pool(compact_layer) + if not isinstance(v_packed, torch.Tensor): + raise RuntimeError( + "CuTeDSL FP4 MLA disaggregated import requires the persistent V-packed pool." + ) + _repack_cutedsl_v_packed_cache( + v_packed, + kv_cache, + page_ids, + v_head_dim=int(v_head_dim), + page_size=page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + ) + return True + + +# Python launch helpers + + +def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: + global_scale = getattr(metadata, "_fp4_mla_kv_global_scale", None) + if ( + not isinstance(global_scale, torch.Tensor) + or global_scale.device != device + or global_scale.dtype != torch.float32 + or global_scale.numel() != 1 + ): + raise RuntimeError("FP4 MLA requires a preallocated FP32 KV global-scale tensor.") + return global_scale + + +def _get_fp4_mla_q_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: + global_scale = getattr(metadata, "_fp4_mla_q_global_scale", None) + if ( + not isinstance(global_scale, torch.Tensor) + or global_scale.device != device + or global_scale.dtype != torch.float32 + or global_scale.numel() != 1 + ): + raise RuntimeError("FP4 MLA requires a preallocated FP32 Q global-scale tensor.") + return global_scale + + +def _get_fp4_mla_kv_cache_tensors( + metadata: Any, layer_idx: int +) -> tuple[torch.Tensor, torch.Tensor]: + return metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) + + +def _get_fp4_mla_hp_pool_layout( + metadata: Any, + pool: torch.Tensor, +) -> tuple[int, int]: + """Return the manager-owned HP ring size and per-token head dimension.""" + manager = getattr(metadata, "kv_cache_manager", None) + if manager is None or not hasattr(manager, "_fp4_mla_hp_pool_size"): + raise ValueError("FP4 MLA requires a V2 manager-owned HP ring.") + hp_pool_size = int(manager._fp4_mla_hp_pool_size) + if ( + hp_pool_size < HP_BLOCK_SIZE + or pool.ndim != 4 + or pool.shape[2] < 1 + or pool.shape[-1] % hp_pool_size != 0 + ): + raise ValueError( + "FP4 MLA high-precision KV pool does not match its configured " + f"ring: shape={tuple(pool.shape)}, ring_size={hp_pool_size}." + ) + return hp_pool_size, pool.shape[-1] // hp_pool_size + + +def _validate_fp4_mla_hp_generation_width( + hp_pool_size: int, + generation_len: int, +) -> None: + """Ensure one target plus rewindable drafts fit without clobbering the live tail.""" + max_rewind_len = hp_pool_size - HP_BLOCK_SIZE + if generation_len <= 0 or generation_len - 1 > max_rewind_len: + raise RuntimeError( + "FP4 MLA generation exceeds the HP ring's rewind slack: " + f"generation={generation_len}, max_rewind={max_rewind_len}." + ) + + +def _validate_fp4_mla_kv_storage_shape( + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + *, + head_dim: int, + backend: str, +) -> int: + """Validate the backend-specific physical KV and scale strides.""" + residual_dim = FP4_MLA_K_RESIDUAL_DIM if backend in _FP4_MLA_K_RESIDUAL_BACKENDS else 0 + expected_storage_head_dim = head_dim + residual_dim + storage_head_dim = kv_cache.shape[-1] * 2 + if storage_head_dim != expected_storage_head_dim: + raise RuntimeError( + "FP4 MLA KV cache storage head dimension does not match the selected backend: " + f"got {storage_head_dim}, expected {expected_storage_head_dim}. Recreate the engine " + f"after setting {FP4_MLA_ATTENTION_BACKEND_ENV}." + ) + + expected_scale_columns = expected_storage_head_dim // FP4_BLOCK_SIZE + if sf_cache.shape[-1] != expected_scale_columns: + raise RuntimeError( + "FP4 MLA KV cache scale storage does not match the contiguous data layout: " + f"got {sf_cache.shape[-1]} columns, expected {expected_scale_columns}." + ) + return storage_head_dim + + +def _scatter_fp4_mla_kv_cache_2d_context( + metadata: Any, + latent_cache: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_sf: Optional[torch.Tensor], + global_scale: torch.Tensor, + rotary_cos_sin: Optional[torch.Tensor], + q_context: Optional[torch.Tensor], + q_nope_head_dim: Optional[int], + *, + token_offset: int, + local_layer: int, + v_head_dim: int, + head_dim: int, + num_tokens: int, + num_dim_blocks: int, + sf_per_token: int, + sf_per_page: int, + v_packed_base: Optional[torch.Tensor] = None, + v_page_offset: int = 0, +) -> bool: + num_contexts = metadata.num_contexts + if num_contexts > 0: + prompt_lens_cpu = metadata.prompt_lens_cpu_runtime[:num_contexts] + ctx_token_count = int(prompt_lens_cpu.sum().item()) + if num_tokens != ctx_token_count: + raise RuntimeError( + f"FP4 MLA 2D context scatter needs {ctx_token_count} context tokens, got " + f"{num_tokens}." + ) + _validate_fp4_mla_context_start_alignment(metadata, num_contexts, alignment=FP4_BLOCK_SIZE) + + apply_k_rope = rotary_cos_sin is not None + rope_dim = ( + _validate_fp4_mla_context_rope(latent_cache, rotary_cos_sin, v_head_dim) + if rotary_cos_sin is not None + else 0 + ) + rotary_cos_sin_ptr = rotary_cos_sin if rotary_cos_sin is not None else latent_cache + + apply_q_rope = q_context is not None + block_q_heads = 16 + if apply_q_rope: + if rotary_cos_sin is None or q_nope_head_dim is None: + raise ValueError("FP4 MLA fused context Q-RoPE requires a rotary table and Q layout.") + q_head_dim = q_nope_head_dim + rope_dim + if ( + q_context.dtype != torch.bfloat16 + or q_context.device != latent_cache.device + or q_context.ndim != 2 + or q_context.shape[0] != num_tokens + or q_context.shape[1] <= 0 + or q_nope_head_dim <= 0 + or q_context.shape[1] % q_head_dim != 0 + or not q_context.is_contiguous() + ): + raise ValueError( + "FP4 MLA fused context Q-RoPE requires a contiguous same-device BF16 " + f"tensor shaped [tokens, heads * ({q_nope_head_dim} + {rope_dim})]." + ) + num_q_heads = q_context.shape[1] // q_head_dim + q_context_view = q_context.view(num_tokens, num_q_heads, q_head_dim) + q_head_blocks = triton.cdiv(num_q_heads, block_q_heads) + else: + num_q_heads = 0 + q_context_view = latent_cache + q_head_blocks = 0 + + hp_pool = getattr(metadata, "high_precision_kv_pool", None) + if not isinstance(hp_pool, torch.Tensor): + raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") + if hp_pool.device != latent_cache.device: + raise ValueError("FP4 MLA latent cache and high-precision pool must share a device.") + if hp_pool.dtype != torch.bfloat16: + raise TypeError(f"FP4 MLA high-precision KV pool must use BF16, got {hp_pool.dtype}.") + hp_pool_size, pool_head_dim = _get_fp4_mla_hp_pool_layout(metadata, hp_pool) + if pool_head_dim < head_dim: + raise RuntimeError( + f"FP4 MLA HP pool head dimension is too small: got " + f"{pool_head_dim}, need at least {head_dim}." + ) + if local_layer < 0 or local_layer >= hp_pool.shape[1]: + raise ValueError( + f"FP4 MLA local layer {local_layer} is outside the HP pool's {hp_pool.shape[1]} layers." + ) + if hp_pool.stride(-1) != 1: + raise ValueError("FP4 MLA high-precision KV pool must be contiguous in head_dim.") + hp_page_ids = metadata._fp4_mla_hp_page_indices + if not isinstance(hp_page_ids, torch.Tensor): + raise RuntimeError("FP4 MLA context cache update requires HP page metadata.") + store_hp_tail = num_contexts > 0 + num_hp_pages = hp_pool.shape[0] + pool_s0 = hp_pool.stride(0) + pool_s1 = hp_pool.stride(1) + + write_v_packed = v_packed_base is not None + v_packed_output = v_packed_base if write_v_packed else kv_cache + v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 + v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 + + _fp4_mla_context_cache_update_kernel[ + ( + num_tokens, + num_dim_blocks + q_head_blocks, + ) + ]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + latent_cache, + q_context_view, + global_scale, + rotary_cos_sin_ptr, + hp_pool, + hp_page_ids, + metadata.batch_indices, + metadata.positions, + metadata.paged_kv_indices, + metadata.paged_kv_indptr, + metadata.paged_kv_indices.shape[0], + metadata.paged_kv_indptr.shape[0], + metadata.batch_indices.shape[0], + v_sf.shape[1] if v_sf is not None else kv_cache.shape[0], + v_sf.shape[0] if v_sf is not None else hp_pool.shape[1], + num_contexts, + num_hp_pages, + token_offset, + num_tokens, + local_layer, + v_page_offset if write_v_packed else 0, + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + latent_cache.stride(0), + latent_cache.stride(1), + q_context_view.stride(0), + q_context_view.stride(1) if apply_q_rope else 0, + q_context_view.stride(2) if apply_q_rope else 0, + v_sf.stride(0) if v_sf is not None else 0, + v_sf.stride(1) if v_sf is not None else 0, + v_packed_s0, + v_packed_s1, + pool_s0, + pool_s1, + HEAD_D=head_dim, + V_HEAD_D=v_head_dim, + HP_BLOCK=FP4_BLOCK_SIZE, + HP_POOL_SIZE=hp_pool_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=sf_per_page, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + STORE_K_RESIDUAL=(_fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS), + ROPE_DIM=rope_dim, + APPLY_K_ROPE=apply_k_rope, + APPLY_Q_ROPE=apply_q_rope, + NUM_DIM_BLOCKS=num_dim_blocks, + NUM_Q_HEADS=num_q_heads, + Q_NOPE_DIM=q_nope_head_dim if q_nope_head_dim is not None else 0, + BLOCK_Q_HEADS=block_q_heads, + POOL_HEAD_D=pool_head_dim, + STORE_HP_TAIL=store_hp_tail, + WRITE_V_PACKED=write_v_packed, + STORE_V_SCALE=v_sf is not None, + ) + return store_hp_tail + + +def _fp4_mla_generation_num_blocks_device(metadata: Any) -> torch.Tensor: + """Device-side scalar view holding the generation page-table capacity. + + ``paged_kv_indptr_decode[num_gen]`` is the fixed-stride generation-table + endpoint. Device kernels combine it with the live KV lengths, so inactive + slots are never consumed. + """ + num_gen = metadata.num_seqs - metadata.num_contexts + return metadata.paged_kv_indptr_decode[num_gen : num_gen + 1] + + +def _fp4_mla_uniform_generation_lengths( + metadata: Any, + num_gen_tokens: int, + num_gen: int, + plan_geometry: Optional[_PackedPlanGeometry] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return preallocated CUDA generation KV and append lengths. + + With ``plan_geometry`` the launch that writes the lengths also writes the + step's packed split plan (one CuTe DSL kernel instead of two launches), and + a stale plan is rewritten even when the lengths are already current. + """ + if num_gen <= 0 or num_gen_tokens % num_gen != 0: + raise RuntimeError("FP4 MLA generation requires a non-empty uniform request batch.") + if plan_geometry is not None and plan_geometry.num_requests != num_gen: + raise RuntimeError("FP4 MLA packed plan geometry does not match the generation batch.") + + num_contexts = metadata.num_contexts + num_seqs = metadata.num_seqs + kv_lens_gen = metadata.kv_lens_cuda_runtime[num_contexts:num_seqs] + prompt_lens_gen = metadata.prompt_lens_cuda_runtime[num_contexts:num_seqs] + corrected_kv_lens = getattr(metadata, "fp4_mla_generation_kv_lens", None) + generation_lens = getattr(metadata, "fp4_mla_generation_append_lens", None) + tensors = ( + kv_lens_gen, + prompt_lens_gen, + corrected_kv_lens, + generation_lens, + ) + if ( + not all(isinstance(tensor, torch.Tensor) for tensor in tensors) + or corrected_kv_lens.numel() < num_gen + or generation_lens.numel() < num_gen + or not all(tensor.is_cuda for tensor in tensors) + ): + raise RuntimeError("FP4 MLA generation lengths require preallocated CUDA buffers.") + + record_for_capture = bool( + getattr(metadata, "is_cuda_graph", False) + and torch.cuda.is_current_stream_capturing() + and not getattr( + metadata, + "_fp4_mla_generation_lengths_capture_recorded", + False, + ) + ) + precomputed = ( + not record_for_capture + and metadata.fp4_mla_generation_lengths_num_tokens == num_gen_tokens + and metadata.fp4_mla_generation_lengths_num_seqs == num_gen + and metadata.fp4_mla_generation_lengths_num_contexts == num_contexts + ) + plan_pending = ( + plan_geometry is not None + and plan_geometry.key is not None + and getattr(metadata, "_fp4_mla_packed_plan_key", None) != plan_geometry.key + ) + if not precomputed: + if plan_pending: + _launch_fp4_mla_packed_plan( + plan_geometry, + kv_lens_gen, + metadata, + generation_lengths=( + prompt_lens_gen, + corrected_kv_lens[:num_gen], + generation_lens[:num_gen], + num_gen_tokens // num_gen, + ), + ) + else: + populate_fp4_mla_generation_lengths( + kv_lens_gen, + prompt_lens_gen, + corrected_kv_lens[:num_gen], + generation_lens[:num_gen], + num_gen_tokens=num_gen_tokens, + num_gen=num_gen, + ) + metadata.fp4_mla_generation_lengths_num_tokens = num_gen_tokens + metadata.fp4_mla_generation_lengths_num_seqs = num_gen + metadata.fp4_mla_generation_lengths_num_contexts = num_contexts + if record_for_capture: + metadata._fp4_mla_generation_lengths_capture_recorded = True + elif plan_pending: + _launch_fp4_mla_packed_plan(plan_geometry, corrected_kv_lens[:num_gen], metadata) + return corrected_kv_lens[:num_gen], generation_lens[:num_gen] + + +def _materialize_fp4_mla_device_page_table_for_forward( + metadata: Any, + generation_kv_lens: Optional[torch.Tensor] = None, +) -> None: + """Materialize all fixed-stride rows from final per-forward device lengths.""" + if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): + raise RuntimeError("FP4 MLA cache update requires fixed-stride device page metadata.") + if bool(getattr(metadata, "_fp4_mla_device_page_table_valid", False)): + return + + num_contexts = int(metadata.num_contexts) + num_sequences = int(metadata.num_seqs) + num_generation_sequences = num_sequences - num_contexts + if generation_kv_lens is None: + if num_generation_sequences > 0: + num_generation_tokens = int(metadata.num_tokens) - int(metadata.num_ctx_tokens) + generation_kv_lens, _ = _fp4_mla_uniform_generation_lengths( + metadata, + num_generation_tokens, + num_generation_sequences, + ) + else: + generation_kv_lens = metadata.kv_lens_cuda_runtime[num_contexts:num_sequences] + materialize_fp4_mla_device_page_table( + metadata, + metadata.kv_lens_cuda_runtime[:num_sequences], + generation_kv_lens, + ) + + +def _scatter_fp4_mla_kv_cache_2d_generation( + metadata: Any, + latent_cache: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + *, + token_offset: int, + local_layer: int, + v_head_dim: int, + head_dim: int, + num_tokens: int, + num_dim_blocks: int, + sf_per_token: int, + sf_per_page: int, + rotary_cos_sin: torch.Tensor, + q_pe: torch.Tensor, + q_rope_out: torch.Tensor, + q_quant_input: torch.Tensor, + q_fp4_out: torch.Tensor, + q_sf_out: torch.Tensor, + v_packed_base: Optional[torch.Tensor], + v_page_offset: int, + helix_position_offsets: Optional[torch.Tensor], + helix_is_inactive_rank: Optional[torch.Tensor], +) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + num_contexts = metadata.num_contexts + num_seqs = metadata.num_seqs + num_gen = num_seqs - num_contexts + if num_gen <= 0: + return + if num_tokens < num_gen: + raise RuntimeError( + f"FP4 MLA 2D generation scatter needs at least {num_gen} generation " + f"tokens, got {num_tokens}." + ) + if num_tokens % num_gen != 0: + raise NotImplementedError( + "FP4 MLA no-dequant generation scatter requires a uniform linear MTP " + f"generation length, got {num_tokens} tokens for {num_gen} sequences." + ) + # The prompt_lens/kv_lens runtime aliases can lag at the decode anchor + # (seq_lens == 1) under CUDA graph / one-engine MTP while each generation + # sequence really appends num_tokens // num_gen tokens this step. Recover the + # true per-sequence lengths for the no-dequant kernel below (a no-op when the + # aliases already match). As the first FP4 MLA kernel of the step, the same + # launch also writes the packed split plan the decode layers reuse. + plan_geometry = None + if q_pe.ndim == 3 and q_pe.shape[1] > 0: + plan_geometry = _fp4_mla_packed_plan_geometry( + metadata, + num_heads=q_pe.shape[1], + query_len=num_tokens // num_gen, + num_requests=num_gen, + device=latent_cache.device, + helix_kv_bounds=_fp4_mla_helix_plan_bounds( + metadata, + num_tokens, + use_helix=helix_position_offsets is not None or helix_is_inactive_rank is not None, + ), + ) + kv_lens_gen, gen_lens_gen = _fp4_mla_uniform_generation_lengths( + metadata, num_tokens, num_gen, plan_geometry=plan_geometry + ) + _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens_gen) + + pool = getattr(metadata, "high_precision_kv_pool", None) + if pool is None: + raise RuntimeError("FP4 MLA 2D generation scatter requires the HP KV pool.") + try: + hp_pool_size, hp_head_dim = _get_fp4_mla_hp_pool_layout(metadata, pool) + except ValueError as error: + raise RuntimeError(str(error)) from error + if hp_head_dim < head_dim: + raise RuntimeError( + f"FP4 MLA 2D generation scatter needs at least {head_dim} HP channels, got " + f"{hp_head_dim}." + ) + hp_page_ids = _fp4_mla_generation_hp_page_ids(metadata, num_gen) + if not isinstance(hp_page_ids, torch.Tensor): + raise RuntimeError("FP4 MLA generation requires HP page metadata.") + num_hp_pages = pool.shape[0] + + max_gen_len = num_tokens // num_gen + use_helix = helix_position_offsets is not None or helix_is_inactive_rank is not None + use_helix_local_slots = use_helix and bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + if use_helix: + if helix_position_offsets is None or helix_is_inactive_rank is None: + raise RuntimeError( + "FP4 MLA Helix requires both position-offset and inactive-rank metadata." + ) + if max_gen_len != 1 and not use_helix_local_slots: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + if ( + helix_position_offsets.dtype != torch.int32 + or helix_position_offsets.device != latent_cache.device + or helix_position_offsets.ndim != 1 + or helix_position_offsets.numel() < num_tokens + or not helix_position_offsets.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix position offsets must be a contiguous same-device " + "int32 tensor covering every generation token." + ) + if ( + helix_is_inactive_rank.dtype != torch.bool + or helix_is_inactive_rank.device != latent_cache.device + or helix_is_inactive_rank.ndim != 1 + or helix_is_inactive_rank.numel() < num_gen + or not helix_is_inactive_rank.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix inactive-rank metadata must be a contiguous " + "same-device bool tensor covering every generation sequence." + ) + if use_helix_local_slots: + helix_local_slots = getattr(metadata, "helix_local_slots", None) + if ( + not isinstance(helix_local_slots, torch.Tensor) + or helix_local_slots.dtype != torch.int32 + or helix_local_slots.device != latent_cache.device + or helix_local_slots.ndim != 1 + or helix_local_slots.numel() < num_tokens + or not helix_local_slots.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix local slots must be a contiguous " + "same-device int32 tensor covering every generation token." + ) + else: + helix_local_slots = helix_position_offsets + else: + helix_position_offsets = kv_lens_gen + helix_local_slots = kv_lens_gen + helix_is_inactive_rank = gen_lens_gen + _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) + max_rewind_len = hp_pool_size - HP_BLOCK_SIZE + page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) + rope_dim = head_dim - v_head_dim + block_q_heads = 32 + q1_kv_blocks_per_program = 1 + # Grouped Q1 kernels reuse K's packed codes for V. Keep one dimension + # block per program until their warp-specialized V quantizer is split out. + if max_gen_len == 1: + q1_kv_blocks_per_program = _fp4_mla_q1_kv_blocks_per_program(num_gen, v_head_dim) + q_prefix_block_dim = ( + FP4_MLA_Q1_PREFIX_BLOCK_DIM if max_gen_len == 1 else FP4_MLA_Q_PREFIX_BLOCK_DIM + ) + q_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // q_prefix_block_dim + q_prefix_blocks_per_program = _fp4_mla_q1_prefix_blocks_per_program( + num_gen, + q1_kv_blocks_per_program, + ) + q_work_blocks = q_prefix_blocks // q_prefix_blocks_per_program + 1 + if latent_cache.dtype != torch.bfloat16: + raise TypeError( + "Fused FP4 MLA Q/K RoPE and cache storage requires BF16 latent KV, " + f"got {latent_cache.dtype}." + ) + if rope_dim <= 0 or rope_dim % 2 != 0: + raise ValueError( + "Fused FP4 MLA K RoPE requires a positive even K tail, " + f"got head_dim={head_dim} v_head_dim={v_head_dim}." + ) + if ( + rotary_cos_sin is None + or rotary_cos_sin.device != latent_cache.device + or rotary_cos_sin.dtype != torch.float32 + or rotary_cos_sin.numel() % (rope_dim * 2) != 0 + ): + raise ValueError( + "Fused FP4 MLA K RoPE requires a same-device FP32 rotary " + f"table with rows of {rope_dim * 2} values." + ) + if ( + q_pe is None + or q_rope_out is None + or q_pe.dtype != torch.bfloat16 + or q_rope_out.dtype != torch.bfloat16 + or q_pe.device != latent_cache.device + or q_rope_out.device != latent_cache.device + or q_pe.ndim != 3 + or q_rope_out.shape != q_pe.shape + or q_pe.shape[0] != num_tokens + or q_pe.shape[1] <= 0 + or q_pe.shape[2] != rope_dim + ): + raise ValueError( + "Fused FP4 MLA Q RoPE requires same-device BF16 q_pe and " + f"q_rope_out tensors shaped [tokens, heads, {rope_dim}]." + ) + num_q_heads = q_pe.shape[1] + q_head_blocks = _ceil_div(num_q_heads, block_q_heads) + q_rows_per_seq = _fp4_mla_packed_rows_per_seq(num_q_heads, max_gen_len) + if q_rows_per_seq and q_fp4_out.shape[0] < num_gen * q_rows_per_seq: + raise RuntimeError( + "FP4 MLA packed Q buffer is smaller than the generation batch: " + f"{q_fp4_out.shape[0]} rows < {num_gen} x {q_rows_per_seq}." + ) + max_gen_tiles = _ceil_div(max_gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) + rotary_table = rotary_cos_sin + q_global_scale = _get_fp4_mla_q_global_scale(metadata, latent_cache.device) + q_pe_input = q_pe + q_rope_output = q_rope_out + q_full_input = q_quant_input + q_fp4_output = q_fp4_out + q_sf_output = q_sf_out + write_v_packed = v_packed_base is not None + v_packed_output = v_packed_base if write_v_packed else kv_cache + v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 + v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 + kv_work_blocks = ( + v_head_dim // FP4_BLOCK_SIZE // q1_kv_blocks_per_program + 1 + if max_gen_len == 1 + else num_dim_blocks + ) + launch_grid = ( + num_gen, + max( + kv_work_blocks, + max_gen_len * q_head_blocks * q_work_blocks, + ), + ) + store_k_residual = _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS + + def launch_generation_update( + grid: tuple[int, ...], + *, + page_ids_len: int, + indptr_len: int, + max_gen_tiles_variant: int, + q_prefix_block_dim_variant: int, + q_prefix_blocks_variant: int, + q_prefix_blocks_per_program_variant: int, + q1_kv_blocks_per_program_variant: int, + ) -> None: + q_work_blocks_variant = q_prefix_blocks_variant // q_prefix_blocks_per_program_variant + 1 + _fp4_mla_generation_fused_qk_rope_cache_update_kernel[grid]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + pool, + latent_cache, + global_scale, + q_global_scale, + rotary_table, + q_pe_input, + q_rope_output, + q_full_input, + q_fp4_output, + q_sf_output, + kv_lens_gen, + gen_lens_gen, + helix_position_offsets, + helix_local_slots, + helix_is_inactive_rank, + page_ids, + hp_page_ids, + metadata.paged_kv_indptr_decode, + page_ids_len, + hp_page_ids.numel(), + indptr_len, + v_sf.shape[1], + num_hp_pages, + v_sf.shape[0], + local_layer, + v_page_offset if write_v_packed else 0, + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + pool.stride(0), + pool.stride(1), + v_sf.stride(0), + v_sf.stride(1), + v_packed_s0, + v_packed_s1, + q_pe_input.stride(0), + q_pe_input.stride(1) if q_pe_input.ndim > 1 else 0, + q_pe_input.stride(2) if q_pe_input.ndim > 2 else 0, + q_rope_output.stride(0), + q_rope_output.stride(1) if q_rope_output.ndim > 1 else 0, + q_rope_output.stride(2) if q_rope_output.ndim > 2 else 0, + HEAD_D=hp_head_dim, + V_HEAD_D=v_head_dim, + HP_BLOCK=FP4_BLOCK_SIZE, + HP_POOL_SIZE=hp_pool_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=sf_per_page, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + STORE_K_RESIDUAL=store_k_residual, + FUSE_ROPE_CACHE_STORE=True, + USE_HELIX=use_helix, + USE_HELIX_LOCAL_SLOTS=use_helix_local_slots, + WRITE_V_PACKED=write_v_packed, + MAX_GEN_TILES=max_gen_tiles_variant, + ROPE_DIM=rope_dim, + ROPE_PAIR_BLOCK=triton.next_power_of_2(rope_dim // 2), + NUM_DIM_BLOCKS=num_dim_blocks, + NUM_Q_HEADS=num_q_heads, + Q_HEAD_BLOCKS=max(q_head_blocks, 1), + BLOCK_Q_HEADS=block_q_heads, + Q_PREFIX_D=FP4_MLA_Q_PREFIX_DIM, + Q_PREFIX_BLOCK_D=q_prefix_block_dim_variant, + Q_PREFIX_BLOCKS=q_prefix_blocks_variant, + Q_PREFIX_BLOCKS_PER_PROGRAM=q_prefix_blocks_per_program_variant, + Q_WORK_BLOCKS=q_work_blocks_variant, + Q_SF_COLS=FP4_MLA_Q_SF_GROUPS, + WRITE_Q=True, + Q1_KV_BLOCKS_PER_PROGRAM=q1_kv_blocks_per_program_variant, + Q_ROWS_PER_SEQ=q_rows_per_seq, + maxnreg=56, + ) + + # Triton compiles and loads a CUDA module on first launch. Use runtime-zero + # work here so every reachable static tuning variant is resident before + # warmup hands the engine to serving. + if getattr(metadata, "is_warmup", False) and not torch.cuda.is_current_stream_capturing(): + configured_generation_len = int(getattr(metadata, "max_total_draft_tokens", 0) or 0) + 1 + max_preload_generation_len = max(max_gen_len, configured_generation_len) + if max_preload_generation_len - 1 > max_rewind_len: + raise NotImplementedError( + "FP4 MLA finite Triton preload exceeds the HP ring's rewind slack: " + f"max_rewind={max_rewind_len}, generation=" + f"{max_preload_generation_len}." + ) + max_num_sequences = int(getattr(metadata, "max_num_sequences", num_seqs) or num_seqs) + q1_variants = _fp4_mla_q1_preload_variants( + max_num_sequences, + v_head_dim, + ) + multi_token_tiles = tuple( + sorted( + { + _ceil_div(gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) + for gen_len in range(2, max_preload_generation_len + 1) + } + ) + ) + preload_key = ( + "generation-cache-update", + str(latent_cache.device), + hp_pool_size, + hp_head_dim, + v_head_dim, + num_q_heads, + q_rows_per_seq, + metadata.page_size, + write_v_packed, + store_k_residual, + use_helix, + use_helix_local_slots, + tuple(q1_variants), + multi_token_tiles, + tuple(kv_cache.stride()), + tuple(sf_cache.stride()), + tuple(pool.stride()), + tuple(v_sf.stride()), + tuple(v_packed_output.stride()), + tuple(q_pe_input.stride()), + tuple(q_rope_output.stride()), + str(kv_cache.dtype), + str(latent_cache.dtype), + str(q_pe_input.dtype), + str(q_fp4_output.dtype), + str(q_sf_output.dtype), + ) + preload_keys = _fp4_mla_triton_preload_key_set(metadata) + if preload_key not in preload_keys: + context_page_ids = getattr(metadata, "_paged_kv_indices", None) + if not isinstance(context_page_ids, torch.Tensor): + context_page_ids = page_ids + context_indptr = getattr(metadata, "_paged_kv_indptr", None) + if not isinstance(context_indptr, torch.Tensor): + context_indptr = metadata.paged_kv_indptr_decode + context_batch_indices = getattr(metadata, "batch_indices", None) + if not isinstance(context_batch_indices, torch.Tensor): + context_batch_indices = hp_page_ids + context_positions = getattr(metadata, "positions", None) + if not isinstance(context_positions, torch.Tensor): + context_positions = hp_page_ids + _fp4_mla_context_cache_update_kernel[(1, num_dim_blocks)]( + kv_cache, + sf_cache, + v_sf, + v_packed_output, + latent_cache, + latent_cache, + global_scale, + rotary_table, + pool, + hp_page_ids, + context_batch_indices, + context_positions, + context_page_ids, + context_indptr, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + latent_cache.stride(0), + latent_cache.stride(1), + latent_cache.stride(0), + 0, + 0, + v_sf.stride(0), + v_sf.stride(1), + v_packed_s0, + v_packed_s1, + pool.stride(0), + pool.stride(1), + HEAD_D=head_dim, + V_HEAD_D=v_head_dim, + HP_BLOCK=FP4_BLOCK_SIZE, + HP_POOL_SIZE=hp_pool_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=sf_per_token, + SF_PER_PAGE=sf_per_page, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + STORE_K_RESIDUAL=store_k_residual, + ROPE_DIM=rope_dim, + APPLY_K_ROPE=True, + APPLY_Q_ROPE=False, + NUM_DIM_BLOCKS=num_dim_blocks, + NUM_Q_HEADS=0, + Q_NOPE_DIM=0, + BLOCK_Q_HEADS=16, + POOL_HEAD_D=hp_head_dim, + STORE_HP_TAIL=True, + WRITE_V_PACKED=write_v_packed, + STORE_V_SCALE=True, + ) + q1_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM + for q1_kv_blocks, q1_prefix_blocks_per_program in q1_variants: + launch_generation_update( + (1, 1), + page_ids_len=0, + indptr_len=0, + max_gen_tiles_variant=1, + q_prefix_block_dim_variant=FP4_MLA_Q1_PREFIX_BLOCK_DIM, + q_prefix_blocks_variant=q1_prefix_blocks, + q_prefix_blocks_per_program_variant=q1_prefix_blocks_per_program, + q1_kv_blocks_per_program_variant=q1_kv_blocks, + ) + multi_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q_PREFIX_BLOCK_DIM + for multi_token_tile in multi_token_tiles: + launch_generation_update( + (1, 1), + page_ids_len=0, + indptr_len=0, + max_gen_tiles_variant=multi_token_tile, + q_prefix_block_dim_variant=FP4_MLA_Q_PREFIX_BLOCK_DIM, + q_prefix_blocks_variant=multi_prefix_blocks, + q_prefix_blocks_per_program_variant=1, + q1_kv_blocks_per_program_variant=1, + ) + torch.cuda.synchronize(latent_cache.device) + preload_keys.add(preload_key) + + launch_generation_update( + launch_grid, + page_ids_len=page_ids.shape[0], + indptr_len=metadata.paged_kv_indptr_decode.shape[0], + max_gen_tiles_variant=max(max_gen_tiles, 1), + q_prefix_block_dim_variant=q_prefix_block_dim, + q_prefix_blocks_variant=q_prefix_blocks, + q_prefix_blocks_per_program_variant=q_prefix_blocks_per_program, + q1_kv_blocks_per_program_variant=q1_kv_blocks_per_program, + ) + return kv_lens_gen, gen_lens_gen, page_ids + + +# Public cache update and decode entry points + + +def load_fp4_mla_chunked_kv_cache( + metadata: Any, + layer_idx: int, + *, + num_ctx_cached_tokens: int, + cu_chunked_seq_len: torch.Tensor, + chunked_global_offset: torch.Tensor, + chunked_max_seq_len: int, + out_dtype: torch.dtype, + kv_lora_rank: int, + qk_rope_head_dim: int, + output_format: Literal["dense", "mxfp8"] = "dense", +) -> tuple[torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor]: + """Gather one cached-prefix partition from dense FP4 MLA V2 storage. + + ``mxfp8`` preserves BF16 reconstruction, returning ((E4M3 [M,512], + UE8M0 [pad_up(M,128)*16]), E4M3 [M,64]). Scales use R128c4 K32 + storage with K128 quantization; RoPE output uses unit FMHA scale. + """ + if output_format not in ("dense", "mxfp8"): + raise ValueError(f"Unsupported FP4 MLA gather output format: {output_format}.") + quantize_mxfp8 = output_format == "mxfp8" + if quantize_mxfp8 and (out_dtype != torch.bfloat16 or kv_lora_rank != 512): + raise ValueError("MXFP8 gather requires BF16 reconstruction and kv_lora_rank=512.") + if out_dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError(f"FP4 MLA chunk gather does not support output dtype {out_dtype}.") + if num_ctx_cached_tokens < 0: + raise ValueError("FP4 MLA chunk gather token count must be non-negative.") + if chunked_max_seq_len < 0: + raise ValueError("FP4 MLA chunk gather max sequence length must be non-negative.") + if num_ctx_cached_tokens > 0 and chunked_max_seq_len == 0: + raise ValueError("A non-empty FP4 MLA chunk gather requires a positive max length.") + if ( + kv_lora_rank <= 0 + or kv_lora_rank % FP4_BLOCK_SIZE != 0 + or qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM + ): + raise ValueError( + "FP4 MLA chunk gather requires a positive kv_lora_rank and " + f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got " + f"{kv_lora_rank} and {qk_rope_head_dim}." + ) + + num_contexts = int(metadata.num_contexts) + if num_ctx_cached_tokens > 0 and num_contexts <= 0: + raise ValueError("A non-empty FP4 MLA chunk gather requires context requests.") + tensors = (cu_chunked_seq_len, chunked_global_offset) + if any(not isinstance(tensor, torch.Tensor) or not tensor.is_cuda for tensor in tensors): + raise ValueError("FP4 MLA chunk gather metadata must use CUDA tensors.") + if ( + cu_chunked_seq_len.dtype != torch.int64 + or cu_chunked_seq_len.ndim != 1 + or cu_chunked_seq_len.numel() < num_contexts + 1 + or not cu_chunked_seq_len.is_contiguous() + ): + raise ValueError( + "FP4 MLA chunk gather requires a contiguous int64 cumulative-length tensor." + ) + if ( + chunked_global_offset.dtype != torch.int64 + or chunked_global_offset.ndim != 1 + or chunked_global_offset.numel() < num_contexts + or not chunked_global_offset.is_contiguous() + ): + raise ValueError("FP4 MLA chunk gather requires contiguous int64 global offsets.") + if cu_chunked_seq_len.device != chunked_global_offset.device: + raise ValueError("FP4 MLA chunk gather metadata tensors must share one device.") + + storage_dtype = torch.float8_e4m3fn if quantize_mxfp8 else out_dtype + compressed_kv = torch.empty( + (num_ctx_cached_tokens, kv_lora_rank), + dtype=storage_dtype, + device=cu_chunked_seq_len.device, + ) + k_pe = torch.empty( + (num_ctx_cached_tokens, qk_rope_head_dim), + dtype=storage_dtype, + device=cu_chunked_seq_len.device, + ) + if quantize_mxfp8: + compressed_sf = torch.empty( + (triton.cdiv(num_ctx_cached_tokens, 128) * 128 * (kv_lora_rank // 32),), + dtype=torch.uint8, + device=cu_chunked_seq_len.device, + ) + if num_ctx_cached_tokens == 0: + return ((compressed_kv, compressed_sf), k_pe) if quantize_mxfp8 else (compressed_kv, k_pe) + + if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): + raise RuntimeError("FP4 MLA chunk gather requires fixed-stride device page metadata.") + _materialize_fp4_mla_device_page_table_for_forward(metadata) + page_table_stride = int(metadata.fp4_mla_page_table_stride) + page_ids = metadata._paged_kv_indices + if ( + page_table_stride <= 0 + or not isinstance(page_ids, torch.Tensor) + or page_ids.dtype != torch.int32 + or not page_ids.is_cuda + or page_ids.device != cu_chunked_seq_len.device + or page_ids.numel() < num_contexts * page_table_stride + ): + raise RuntimeError("FP4 MLA chunk gather received invalid device page metadata.") + + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + head_dim = kv_lora_rank + qk_rope_head_dim + storage_head_dim = _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=_fp4_mla_attention_backend(), + ) + if storage_head_dim != head_dim + FP4_MLA_K_RESIDUAL_DIM: + raise RuntimeError( + "FP4 MLA chunk gather requires K residual storage for BF16 reconstruction." + ) + sf_cache = sf_cache.view(torch.float8_e4m3fn) + global_scale = _get_fp4_mla_global_scale(metadata, kv_cache.device) + if quantize_mxfp8: + from tensorrt_llm._torch.cute_dsl_kernels.rubin.fp4_mla_context_gather import gather_mxfp8 + + gather_mxfp8( + compressed_kv, + compressed_sf, + k_pe, + kv_cache, + sf_cache, + page_ids, + cu_chunked_seq_len, + chunked_global_offset, + global_scale, + chunked_max_seq_len, + metadata.page_size, + page_table_stride=page_table_stride, + num_contexts=num_contexts, + ) + return (compressed_kv, compressed_sf), k_pe + token_block = 16 + grid = ( + triton.cdiv(chunked_max_seq_len, token_block), + num_contexts, + triton.cdiv(head_dim, FP4_BLOCK_SIZE), + ) + _fp4_mla_chunked_cache_gather_kernel[grid]( + compressed_kv, + k_pe, + kv_cache, + sf_cache, + page_ids, + cu_chunked_seq_len, + chunked_global_offset, + global_scale, + chunked_max_seq_len, + page_table_stride, + kv_cache.shape[0], + metadata.page_size, + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + compressed_kv.stride(0), + k_pe.stride(0), + KV_LORA_RANK=kv_lora_rank, + QK_ROPE_HEAD_DIM=qk_rope_head_dim, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_TOKEN=storage_head_dim // FP4_BLOCK_SIZE, + TOKEN_BLOCK=token_block, + num_warps=4, + ) + return compressed_kv, k_pe + + +def scatter_fp4_mla_kv_cache( + metadata: Any, + latent_cache: torch.Tensor, + layer_idx: int, + *, + token_offset: int, + phase: _HPUpdatePhase, + local_layer: int, + v_head_dim: int, + rotary_cos_sin: Optional[torch.Tensor] = None, + q_pe: Optional[torch.Tensor] = None, + q_rope_out: Optional[torch.Tensor] = None, + q_quant_input: Optional[torch.Tensor] = None, + helix_position_offsets: Optional[torch.Tensor] = None, + helix_is_inactive_rank: Optional[torch.Tensor] = None, + q_context: Optional[torch.Tensor] = None, + q_nope_head_dim: Optional[int] = None, +) -> bool: + """Quantize MLA latent tokens and scatter them into the paged FP4 cache. + + Contract: this helper scatters exactly ``latent_cache.shape[0]`` tokens, + reading index metadata at ``batch_indices[token_offset : token_offset + N]`` + and ``positions[token_offset : token_offset + N]``. Callers must pass a + latent_cache pre-sliced to the current phase (context or generation) so + that ``shape[0]`` matches the number of index entries they intend to + consume. ``MLA.forward_impl`` (tensorrt_llm/_torch/modules/attention.py) + slices ``latent_cache[:num_ctx_tokens]`` for context and + ``latent_cache[num_ctx_tokens:]`` for generation before dispatching. + + Callers must pass ``phase``, ``local_layer``, and ``v_head_dim``. Context + scatter writes the final FP4 tile representation directly. Dimensions + below ``v_head_dim`` share one 16-token by 16-dim FP4 tile between K + and V, with the scale written into K's token-major and V's dim-major + layouts. Tail K-only dimensions use K's per-token 1D scales. For + exclusively owned CuTeDSL pages, context scatter also writes the + persistent packed-V sidecar. + Context scatter can rotate Q in place and rotate the K tail directly from + the unassembled latent tensor. When context Q is supplied for chunked + prefill, it also writes rotated K back for current-chunk attention. + Generation scatter rewrites each touched 16-token tile by reading + old tokens from the HP pool and new tokens from ``latent_cache``. The + static-scale generation + specialization can also rotate Q and new K tails while updating the HP pool. + The context kernel also stores the final incomplete tile in the BF16 HP + pool. When ``q_quant_input`` is supplied, the generation kernel also emits + backend-ready residual FP4 Q. The return value reports whether the current + phase updated the HP pool. + """ + if phase == "generation": + metadata._fp4_mla_prequantized_q = None + metadata._fp4_mla_prequantized_q_sf = None + metadata._fp4_mla_q_batch_capacity = None + if latent_cache.numel() == 0: + raise ValueError("FP4 MLA cache scatter requires at least one latent token.") + if q_context is not None and not latent_cache.is_contiguous(): + raise ValueError( + "FP4 MLA fused context Q/K RoPE requires contiguous latent_cache " + "storage for in-place current-K update." + ) + + latent_cache = latent_cache.reshape(latent_cache.shape[0], -1).contiguous() + num_tokens = latent_cache.shape[0] + head_dim = latent_cache.shape[-1] + if head_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA KV head_dim must be divisible by {FP4_BLOCK_SIZE}, got {head_dim}." + ) + indices_len = metadata.batch_indices.shape[0] + positions_len = metadata.positions.shape[0] + if token_offset + num_tokens > indices_len or token_offset + num_tokens > positions_len: + raise RuntimeError( + f"FP4 MLA scatter would read batch_indices[{token_offset}:" + f"{token_offset + num_tokens}] / positions[{token_offset}:" + f"{token_offset + num_tokens}], but only {indices_len} / " + f"{positions_len} entries are available. This indicates " + "latent_cache was not pre-sliced to the current phase's token " + "range (see MLA.forward_impl)." + ) + + _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) + + backend = _fp4_mla_attention_backend() + global_scale = _get_fp4_mla_global_scale(metadata, latent_cache.device) + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + storage_head_dim = _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=backend, + ) + sf_per_token = storage_head_dim // FP4_BLOCK_SIZE + + if phase not in ("context", "generation"): + raise ValueError("FP4 MLA scatter requires phase='context' or 'generation'.") + context_only_cache = getattr(metadata.kv_cache_manager, "fp4_mla_context_only_cache", False) + if context_only_cache and phase != "context": + raise RuntimeError("FP4 MLA context-only cache cannot execute generation.") + if not context_only_cache and getattr(metadata, "fp4_mla_v_scale_pool", None) is None: + raise RuntimeError("FP4 MLA scatter requires the auxiliary V scale pool.") + if metadata.page_size % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA scatter requires page_size divisible by " + f"{FP4_BLOCK_SIZE}, got {metadata.page_size}." + ) + if v_head_dim > head_dim: + raise ValueError(f"FP4 MLA v_head_dim={v_head_dim} cannot exceed head_dim={head_dim}.") + if head_dim - v_head_dim != FP4_MLA_K_RESIDUAL_DIM: + raise ValueError( + "FP4 MLA K residual quantization requires the K-only tail to match " + f"the {FP4_MLA_K_RESIDUAL_DIM}-channel residual, got " + f"head_dim={head_dim} v_head_dim={v_head_dim}." + ) + if v_head_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA v_head_dim must be divisible by {FP4_BLOCK_SIZE}, got {v_head_dim}." + ) + + sf_cache = sf_cache.view(torch.float8_e4m3fn) + v_sf = ( + None + if context_only_cache + else get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=v_head_dim) + ) + num_dim_blocks = triton.cdiv(head_dim, FP4_BLOCK_SIZE) + sf_per_page = metadata.page_size // FP4_BLOCK_SIZE + + generation_state = None + generation_inputs = (rotary_cos_sin, q_pe, q_rope_out, q_quant_input) + q_fp4_out = None + q_sf_out = None + if phase == "context": + if any(arg is not None for arg in (q_pe, q_rope_out, q_quant_input)): + raise ValueError("FP4 MLA context cache update does not accept generation Q tensors.") + if (q_context is None) != (q_nope_head_dim is None): + raise ValueError("FP4 MLA context Q and q_nope_head_dim must be provided together.") + hp_pool_updated = False + else: + if q_context is not None or q_nope_head_dim is not None: + raise ValueError("FP4 MLA generation cache update does not accept context Q tensors.") + if (helix_position_offsets is None) != (helix_is_inactive_rank is None): + raise ValueError( + "FP4 MLA Helix position-offset and inactive-rank metadata " + "must be provided together." + ) + if not all(arg is not None for arg in generation_inputs): + raise ValueError( + "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " + "and q_quant_input for fused RoPE, cache update, and Q quantization." + ) + if not can_fuse_fp4_mla_q_quant(metadata, q_quant_input, q_pe, latent_cache): + raise ValueError( + "Fused FP4 MLA Q quantization received an unsupported shape, dtype, " + "scale mode, or backend." + ) + num_gen_for_q = metadata.num_seqs - metadata.num_contexts + q_rows_per_seq = 0 + if num_gen_for_q > 0 and num_tokens % num_gen_for_q == 0: + q_rows_per_seq = _fp4_mla_packed_rows_per_seq( + q_quant_input.shape[1], num_tokens // num_gen_for_q + ) + q_fp4_out, q_sf_out, q_batch_capacity = _prepare_fp4_mla_q_buffers( + metadata, + num_tokens, + q_quant_input.shape[1], + q_quant_input.device, + rows_per_seq=q_rows_per_seq, + ) + metadata._fp4_mla_prequantized_q = q_fp4_out + metadata._fp4_mla_prequantized_q_sf = q_sf_out + metadata._fp4_mla_q_batch_capacity = q_batch_capacity + metadata._fp4_mla_prequantized_q_rows_per_seq = q_rows_per_seq + hp_pool_updated = True + cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND + fused_v_transpose = cutedsl_backend and _fp4_mla_cutedsl_fused_v_transpose_enabled() + kv_cache_manager = getattr(metadata, "kv_cache_manager", None) + persistent_v_packed = None + v_packed_base = None + v_page_offset = 0 + direct_v_packed_write = False + if cutedsl_backend and not fused_v_transpose and not context_only_cache: + persistent_v_packed = _get_cutedsl_persistent_v_packed_cache( + metadata, + local_layer, + kv_cache, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + ) + # Reused/imported pages may not carry this process-local sidecar. Write + # packed V directly only when the cache pages are exclusively owned. + # The fused generation kernel handles uniform linear-MTP batches and + # updates every 16-token tile touched by the verification window. + num_gen = metadata.num_seqs - metadata.num_contexts + block_reuse = getattr(kv_cache_manager, "enable_block_reuse", True) + direct_context_v_packed_write = ( + phase == "context" + and metadata.num_contexts > 0 + and num_tokens > 0 + and block_reuse is False + ) + direct_generation_v_packed_write = ( + phase == "generation" + and num_gen > 0 + and num_tokens >= num_gen + and num_tokens % num_gen == 0 + and block_reuse is False + ) + direct_v_packed_write = direct_context_v_packed_write or direct_generation_v_packed_write + if direct_v_packed_write: + v_packed_base = _get_fp4_mla_v_packed_pool_base(metadata) + get_v_page_offset = getattr(kv_cache_manager, "get_mla_v_packed_page_offset", None) + v_page_offset = ( + int(get_v_page_offset(local_layer)) + if callable(get_v_page_offset) + else local_layer * kv_cache.shape[0] + ) + expected_row_width = metadata.page_size // 2 + required_base_rows = (v_page_offset + kv_cache.shape[0]) * v_head_dim + if ( + not isinstance(v_packed_base, torch.Tensor) + or v_packed_base.dtype != torch.uint8 + or v_packed_base.device != kv_cache.device + or v_packed_base.ndim != 2 + or v_packed_base.shape[0] < required_base_rows + or v_packed_base.shape[1] != expected_row_width + or not v_packed_base.is_contiguous() + ): + raise RuntimeError( + "FP4 MLA direct V-packed cache update requires the " + "stable full-pool base to be a contiguous uint8 tensor " + f"with at least {required_base_rows} rows and " + f"{expected_row_width} columns on {kv_cache.device}." + ) + expected_layer_ptr = v_packed_base.data_ptr() + ( + v_page_offset * v_head_dim * expected_row_width + ) + if expected_layer_ptr != persistent_v_packed.data_ptr(): + raise RuntimeError( + "FP4 MLA V-packed layer view does not match its stable " + "full-pool base and page offset." + ) + v_pack_num_valid = None + if phase == "context": + _materialize_fp4_mla_device_page_table_for_forward(metadata) + hp_pool_updated = _scatter_fp4_mla_kv_cache_2d_context( + metadata, + latent_cache, + kv_cache, + sf_cache, + v_sf, + global_scale, + rotary_cos_sin, + q_context, + q_nope_head_dim, + token_offset=token_offset, + local_layer=local_layer, + v_head_dim=v_head_dim, + head_dim=head_dim, + num_tokens=num_tokens, + num_dim_blocks=num_dim_blocks, + sf_per_token=sf_per_token, + sf_per_page=sf_per_page, + v_packed_base=v_packed_base, + v_page_offset=v_page_offset, + ) + v_pack_page_ids = metadata.paged_kv_indices + else: + generation_state = _scatter_fp4_mla_kv_cache_2d_generation( + metadata, + latent_cache, + kv_cache, + sf_cache, + v_sf, + global_scale, + token_offset=token_offset, + local_layer=local_layer, + v_head_dim=v_head_dim, + head_dim=head_dim, + num_tokens=num_tokens, + num_dim_blocks=num_dim_blocks, + sf_per_token=sf_per_token, + sf_per_page=sf_per_page, + rotary_cos_sin=rotary_cos_sin, + q_pe=q_pe, + q_rope_out=q_rope_out, + q_quant_input=q_quant_input, + q_fp4_out=q_fp4_out, + q_sf_out=q_sf_out, + v_packed_base=v_packed_base, + v_page_offset=v_page_offset, + helix_position_offsets=helix_position_offsets, + helix_is_inactive_rank=helix_is_inactive_rank, + ) + v_pack_page_ids = _fp4_mla_generation_page_ids( + metadata, metadata.num_seqs - metadata.num_contexts + ) + if getattr(metadata, "is_cuda_graph", False): + # Frozen launch grids cannot follow the per-replay page count; + # the repack kernels stride over this device-side count instead. + v_pack_num_valid = _fp4_mla_generation_num_blocks_device(metadata) + cutedsl_repack_page_indptr = None + cutedsl_repack_kv_lens = None + cutedsl_repack_generation_lens = None + cutedsl_repack_max_touched_pages = 1 + if ( + cutedsl_backend + and not fused_v_transpose + and not direct_v_packed_write + and not context_only_cache + ): + if phase == "context": + num_contexts = metadata.num_contexts + cutedsl_v_pack_page_ids = metadata.paged_kv_indices[: metadata.num_context_blocks] + cutedsl_repack_page_indptr = metadata.paged_kv_indptr[: num_contexts + 1] + cutedsl_repack_kv_lens = metadata.kv_lens_cuda_runtime[:num_contexts] + cutedsl_repack_generation_lens = metadata.prompt_lens_cuda_runtime[:num_contexts] + cutedsl_repack_max_touched_pages = int( + metadata.fp4_mla_context_repack_max_touched_pages + ) + elif generation_state is not None: + kv_lens_gen, gen_lens_gen, generation_page_ids = generation_state + num_gen = kv_lens_gen.numel() + cutedsl_v_pack_page_ids = generation_page_ids + cutedsl_repack_page_indptr = metadata.paged_kv_indptr_decode + cutedsl_repack_kv_lens = kv_lens_gen + cutedsl_repack_generation_lens = gen_lens_gen + cutedsl_repack_max_touched_pages = _ceil_div( + num_tokens // num_gen + metadata.page_size - 1, + metadata.page_size, + ) + else: + cutedsl_v_pack_page_ids = v_pack_page_ids + _repack_cutedsl_v_packed_cache( + persistent_v_packed, + kv_cache, + cutedsl_v_pack_page_ids, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + block_v=FP4_MLA_SCALE_ROW_GROUP, + page_indptr=cutedsl_repack_page_indptr, + kv_lens=cutedsl_repack_kv_lens, + generation_lens=cutedsl_repack_generation_lens, + max_touched_pages=cutedsl_repack_max_touched_pages, + ) + if context_only_cache: + return hp_pool_updated + _maybe_update_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + v_pack_page_ids, + num_queries=num_tokens, + v_head_dim=v_head_dim, + page_size=metadata.page_size, + local_layer=local_layer, + v_sf=v_sf[local_layer], + num_valid_pages=v_pack_num_valid, + ) + return hp_pool_updated + + +def _validate_fp4_mla_cache_shape(page_size: int, head_dim: int) -> None: + if page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + f"FP4 MLA KV cache requires tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK} " + f"for swizzled block scales, got {page_size}." + ) + + sf_per_token = head_dim // FP4_BLOCK_SIZE + if head_dim % FP4_BLOCK_SIZE != 0 or sf_per_token % 4 != 0: + raise ValueError( + f"FP4 MLA KV head_dim must produce a scale column count divisible by 4; " + f"got head_dim={head_dim}, scale_columns={sf_per_token}." + ) + + +def _validate_fp4_mla_attention_q_shape(head_dim: int, q_residual_dim: int) -> None: + if q_residual_dim % FP4_BLOCK_SIZE != 0: + raise ValueError( + f"FP4 MLA Q residual_dim must be divisible by {FP4_BLOCK_SIZE}, got {q_residual_dim}." + ) + if q_residual_dim <= 0 or q_residual_dim > head_dim: + raise ValueError( + f"FP4 MLA Q residual_dim must be in (0, head_dim], got " + f"residual_dim={q_residual_dim}, head_dim={head_dim}." + ) + + q_head_dim = head_dim + q_residual_dim + q_sf_per_token = q_head_dim // FP4_BLOCK_SIZE + if q_head_dim % FP4_BLOCK_SIZE != 0 or q_sf_per_token % FP4_MLA_SCALE_COL_GROUP != 0: + raise ValueError( + f"FP4 MLA residual Q must produce a scale column count divisible " + f"by {FP4_MLA_SCALE_COL_GROUP}; got q_head_dim={q_head_dim}, " + f"scale_columns={q_sf_per_token}." + ) + + +def _ensure_workspace_tensor( + metadata: Any, + attr_name: str, + shape: tuple[int, ...], + *, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + tensor = getattr(metadata, attr_name, None) + needs_alloc = ( + tensor is None + or tensor.dtype != dtype + or tensor.device != device + or len(tensor.shape) != len(shape) + or any(tensor.shape[idx] < dim for idx, dim in enumerate(shape)) + ) + if needs_alloc: + if torch.cuda.is_current_stream_capturing(): + raise ValueError( + f"Cannot allocate {attr_name} while capturing a CUDA graph. " + "Run a warmup prepare/forward first." + ) + tensor = torch.empty(shape, dtype=dtype, device=device) + setattr(metadata, attr_name, tensor) + + slices = tuple(slice(0, dim) for dim in shape) + return tensor[slices] + + +def _shared_v_pack_storage_enabled() -> bool: + return os.getenv("TRTLLM_FP4_MLA_SHARE_V_PACK_STORAGE", "1").lower() not in ( + "0", + "false", + "no", + "off", + ) + + +def _select_triton_block_v(num_queries: int, *, prefer_prepacked_v: bool = False) -> int: + env_block_v = _env_int("TRTLLM_FP4_MLA_BLOCK_V") + if env_block_v is not None: + return env_block_v + if prefer_prepacked_v: + return 128 + return 32 if num_queries <= 32 else 128 + + +def _v_packed_shape( + kv_cache: torch.Tensor, + v_head_dim: int, + page_size: int, + block_v: int, +) -> tuple[int, int]: + return (kv_cache.shape[0] * _ceil_div(v_head_dim, block_v) * block_v, page_size // 2) + + +def _get_fp4_mla_v_packed_pool(metadata: Any, local_layer: int) -> Optional[torch.Tensor]: + return metadata.kv_cache_manager.get_mla_v_packed_pool(local_layer) + + +def _get_fp4_mla_v_packed_pool_base(metadata: Any) -> Optional[torch.Tensor]: + return metadata.kv_cache_manager.get_mla_v_packed_pool_base() + + +def _get_fp4_mla_v_scale_pool_base(metadata: Any) -> Optional[torch.Tensor]: + return metadata.kv_cache_manager.get_mla_v_scale_pool_base() + + +def _get_cutedsl_persistent_v_packed_cache( + metadata: Any, + local_layer: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + block_v: int, +) -> torch.Tensor: + v_packed = _get_fp4_mla_v_packed_pool(metadata, local_layer) + if v_packed is None: + raise RuntimeError( + "CuTeDSL FP4 MLA requires the manager-owned persistent V-packed " + "pool; the scratch full-repack fallback has been removed." + ) + expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) + if ( + not isinstance(v_packed, torch.Tensor) + or v_packed.dtype != torch.uint8 + or v_packed.device != kv_cache.device + or tuple(v_packed.shape) != expected_shape + or not v_packed.is_contiguous() + ): + raise RuntimeError( + "FP4 MLA persistent V-packed pool must be a contiguous uint8 tensor " + f"with shape {expected_shape} on {kv_cache.device}; got " + f"{type(v_packed).__name__}, " + f"shape={getattr(v_packed, 'shape', None)}, " + f"dtype={getattr(v_packed, 'dtype', None)}, " + f"device={getattr(v_packed, 'device', None)}." + ) + return v_packed + + +def _repack_cutedsl_v_packed_cache( + v_packed: torch.Tensor, + kv_cache: torch.Tensor, + page_ids: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + block_v: int, + page_indptr: Optional[torch.Tensor] = None, + kv_lens: Optional[torch.Tensor] = None, + generation_lens: Optional[torch.Tensor] = None, + max_touched_pages: int = 1, +) -> None: + if page_ids.numel() == 0: + return + from .fp4_mla_cutedsl_v_repack import fp4_mla_repack_v_cache + + fp4_mla_repack_v_cache( + v_packed, + kv_cache, + page_ids, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + page_indptr=page_indptr, + kv_lens=kv_lens, + generation_lens=generation_lens, + max_touched_pages=max_touched_pages, + ) + + +def _v_packed_cache_tag( + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> tuple[Any, ...]: + v_sf_tag = ( + None + if v_sf is None + else ( + int(v_sf.data_ptr()), + str(v_sf.device), + str(v_sf.dtype), + tuple(int(dim) for dim in v_sf.shape), + tuple(int(stride) for stride in v_sf.stride()), + ) + ) + page_ids_tag = ( + None + if page_ids is None + else ( + int(page_ids.data_ptr()), + str(page_ids.device), + str(page_ids.dtype), + tuple(int(dim) for dim in page_ids.shape), + tuple(int(stride) for stride in page_ids.stride()), + ) + ) + return ( + int(layer_idx), + None if local_layer is None else int(local_layer), + int(kv_cache.data_ptr()), + str(kv_cache.device), + str(kv_cache.dtype), + tuple(int(dim) for dim in kv_cache.shape), + tuple(int(stride) for stride in kv_cache.stride()), + int(v_head_dim), + int(page_size), + int(block_v), + v_sf_tag, + page_ids_tag, + ) + + +def _triton_prepack_v_enabled() -> bool: + if _fp4_mla_attention_backend() != "triton": + return False + default = _env_enabled_default("TRTLLM_FP4_MLA_PREPACK_V", True) + return _env_enabled_default("TRTLLM_FP4_MLA_TRITON_PREPACK_V", default) + + +def _triton_can_prepack_v(v_head_dim: int, page_size: int, block_v: int) -> bool: + return ( + _triton_prepack_v_enabled() + and hasattr(tl, "make_tensor_descriptor") + and block_v in (32, 128) + and v_head_dim % block_v == 0 + and page_size == FP4_MLA_TOKENS_PER_BLOCK + ) + + +def _triton_v_packed_attr(layer_idx: int) -> str: + if _shared_v_pack_storage_enabled(): + return "_fp4_mla_triton_attention_v_packed_buf" + return f"_fp4_mla_triton_attention_v_packed_buf_l{layer_idx}" + + +def _triton_v_packed_valid_attr(layer_idx: int) -> str: + return f"_fp4_mla_triton_attention_v_packed_valid_l{layer_idx}" + + +def _triton_shared_v_packed_valid_attr() -> str: + return "_fp4_mla_triton_attention_v_packed_valid_tag" + + +def _triton_v_packed_cache_tag( + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> tuple[Any, ...]: + return ( + "triton", + _v_packed_cache_tag( + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ), + ) + + +def _set_triton_v_packed_cache_valid( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> None: + valid_attr = ( + _triton_shared_v_packed_valid_attr() + if _shared_v_pack_storage_enabled() + else _triton_v_packed_valid_attr(layer_idx) + ) + setattr( + metadata, + valid_attr, + _triton_v_packed_cache_tag( + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ), + ) + + +def _is_triton_v_packed_cache_valid( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> bool: + valid_attr = ( + _triton_shared_v_packed_valid_attr() + if _shared_v_pack_storage_enabled() + else _triton_v_packed_valid_attr(layer_idx) + ) + return getattr(metadata, valid_attr, None) == _triton_v_packed_cache_tag( + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ) + + +def _get_triton_v_packed_cache( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + page_ids: Optional[torch.Tensor] = None, + block_v: int = 128, +) -> Optional[torch.Tensor]: + if not _triton_can_prepack_v(v_head_dim, page_size, block_v): + return None + if not _is_triton_v_packed_cache_valid( + metadata, + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ): + return None + v_packed = getattr(metadata, _triton_v_packed_attr(layer_idx), None) + expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) + if ( + v_packed is None + or v_packed.dtype != torch.uint8 + or v_packed.device != kv_cache.device + or len(v_packed.shape) != 2 + or v_packed.shape[0] < expected_shape[0] + or v_packed.shape[1] < expected_shape[1] + ): + return None + return v_packed[: expected_shape[0], : expected_shape[1]] + + +def _update_triton_v_packed_cache( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + page_ids: torch.Tensor, + *, + v_head_dim: int, + page_size: int, + block_v: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + num_valid_pages: Optional[torch.Tensor] = None, +) -> Optional[torch.Tensor]: + if not _triton_can_prepack_v(v_head_dim, page_size, block_v): + return None + if page_ids.numel() == 0: + return None + from .fp4_mla_triton import fp4_mla_repack_v_cache + + def _tma_alloc(size: int, alignment: int, stream): + return torch.empty(size, device=kv_cache.device, dtype=torch.int8) + + triton.set_allocator(_tma_alloc) + attr_name = _triton_v_packed_attr(layer_idx) + v_packed = _ensure_workspace_tensor( + metadata, + attr_name, + _v_packed_shape(kv_cache, v_head_dim, page_size, block_v), + dtype=torch.uint8, + device=kv_cache.device, + ) + fp4_mla_repack_v_cache( + v_packed, + kv_cache, + page_ids, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + num_valid_pages=num_valid_pages, + ) + _set_triton_v_packed_cache_valid( + metadata, + layer_idx, + kv_cache, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=page_ids, + ) + return v_packed + + +def _maybe_update_triton_v_packed_cache( + metadata: Any, + layer_idx: int, + kv_cache: torch.Tensor, + page_ids: torch.Tensor, + *, + num_queries: int, + v_head_dim: int, + page_size: int, + local_layer: Optional[int] = None, + v_sf: Optional[torch.Tensor] = None, + num_valid_pages: Optional[torch.Tensor] = None, +) -> None: + block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) + _update_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + page_ids, + v_head_dim=v_head_dim, + page_size=page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + num_valid_pages=num_valid_pages, + ) + + +def _max_generation_pages(metadata: Any) -> int: + num_gen = metadata.num_seqs - metadata.num_contexts + if num_gen <= 0: + return 0 + if not getattr(metadata, "_fp4_mla_device_page_table", False): + raise RuntimeError("FP4 MLA generation requires fixed-stride device page metadata.") + max_pages = int(metadata.fp4_mla_page_table_stride) + if max_pages <= 0: + raise RuntimeError("FP4 MLA device page-table stride must be positive.") + return max_pages + + +def _fp4_mla_generation_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: + """Return the fixed-stride generation page-table view.""" + expected_num_gen = metadata.num_seqs - metadata.num_contexts + if num_gen_seqs != expected_num_gen: + raise RuntimeError( + "FP4 MLA generation sequence count does not match metadata: " + f"{num_gen_seqs} != {expected_num_gen}." + ) + max_pages = _max_generation_pages(metadata) + page_ids = getattr(metadata, "_paged_kv_indices", None) + start = metadata.num_contexts * max_pages + end = start + num_gen_seqs * max_pages + if ( + not isinstance(page_ids, torch.Tensor) + or page_ids.ndim != 1 + or page_ids.dtype != torch.int32 + or not page_ids.is_contiguous() + or page_ids.numel() < end + ): + raise RuntimeError("FP4 MLA fixed-stride generation page-table backing is invalid.") + return page_ids[start:end] + + +def _fp4_mla_generation_hp_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: + """Return generation rows from the fixed-stride V2 HP page table.""" + expected_num_gen = metadata.num_seqs - metadata.num_contexts + if num_gen_seqs != expected_num_gen: + raise RuntimeError( + "FP4 MLA generation sequence count does not match metadata: " + f"{num_gen_seqs} != {expected_num_gen}." + ) + max_pages = _max_generation_pages(metadata) + page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) + start = metadata.num_contexts * max_pages + end = start + num_gen_seqs * max_pages + if ( + not isinstance(page_ids, torch.Tensor) + or page_ids.ndim != 1 + or page_ids.dtype != torch.int32 + or not page_ids.is_contiguous() + or page_ids.numel() < end + ): + raise RuntimeError("FP4 MLA fixed-stride generation HP page-table backing is invalid.") + return page_ids[start:end] + + +def _host_int_list(value: Any, start: int, end: int) -> Optional[list[int]]: + if value is None: + return None + if isinstance(value, torch.Tensor): + if value.is_cuda: + return None + return [int(item) for item in value[start:end].tolist()] + try: + return [int(item) for item in value[start:end]] + except (TypeError, ValueError): + return None + + +def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> bool: + if getattr(metadata, "is_cuda_graph", False): + return False + + start = metadata.num_contexts + end = metadata.num_seqs + block_counts = _host_int_list(getattr(metadata, "num_blocks", None), start, end) + if block_counts is not None and ( + not block_counts or min(block_counts) != max_pages or max(block_counts) != max_pages + ): + return False + + kv_lens_cuda = getattr(metadata, "kv_lens_cuda_runtime", None) + if isinstance(kv_lens_cuda, torch.Tensor): + cache_key = ( + start, + end, + max_pages, + page_size, + tuple(block_counts) if block_counts is not None else None, + kv_lens_cuda.data_ptr(), + ) + cache = getattr(metadata, "_fp4_mla_full_pages_cache", None) + if cache is not None and cache[0] == cache_key: + return bool(cache[1]) + kv_lens = [int(item) for item in kv_lens_cuda[start:end].detach().cpu().tolist()] + result = bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size + setattr(metadata, "_fp4_mla_full_pages_cache", (cache_key, result)) + return result + + kv_cache_params = getattr(metadata, "kv_cache_params", None) + cached_token_lens = _host_int_list( + getattr(kv_cache_params, "num_cached_tokens_per_seq", None), + start, + end, + ) + seq_lens_kv = _host_int_list(getattr(metadata, "seq_lens_kv", None), start, end) + if cached_token_lens is not None and seq_lens_kv is not None: + if len(cached_token_lens) != len(seq_lens_kv): + return False + kv_lens = [ + cached_len + seq_len for cached_len, seq_len in zip(cached_token_lens, seq_lens_kv) + ] + elif kv_cache_params is None: + kv_lens = _host_int_list(getattr(metadata, "prompt_lens_cpu_runtime", None), start, end) + else: + return False + + return bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size + + +def _get_linear_mtp_query_len_per_seq( + metadata: Any, + *, + num_queries: int, + num_gen_seqs: int, +) -> int: + """Return the uniform generation query length required by linear MTP. + + Derives the length from the real query-token count (``num_queries``, taken + from the q shape) and the generation sequence count, which are reliable in + every representation. The host ``prompt_lens``/``seq_lens`` mirror can lag at + the decode anchor (== 1) under CUDA graph / one-engine MTP, so it is only + consulted to produce a precise diagnostic when the counts do not divide + evenly (a genuinely non-uniform batch, which the no-dequant path does not + support). + """ + if num_gen_seqs <= 0: + return 1 + + if num_queries % num_gen_seqs == 0: + return num_queries // num_gen_seqs + + start = metadata.num_contexts + end = metadata.num_seqs + query_lens = _host_int_list_during_forward( + getattr(metadata, "prompt_lens_cpu_runtime", None), start, end + ) + if query_lens is None: + query_lens = _host_int_list_during_forward(getattr(metadata, "seq_lens", None), start, end) + raise NotImplementedError( + "FP4 MLA no-dequant attention requires a uniform linear MTP generation " + f"query length; got {num_queries} query tokens for {num_gen_seqs} " + f"sequences (per-sequence lengths {query_lens})." + ) + + +def _run_triton_attention_decode( + *, + metadata: Any, + layer_idx: int, + local_layer: int, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_sf: torch.Tensor, + global_scale: torch.Tensor, + src_page_ids: torch.Tensor, + kv_lens: torch.Tensor, + p_fp4: torch.Tensor, + p_sf: torch.Tensor, + max_scores: torch.Tensor, + denom: torch.Tensor, + output: torch.Tensor, + num_queries: int, + num_heads: int, + head_dim: int, + kv_lora_rank: int, + q_residual_dim: int, + query_len_per_seq: int, + max_pages: int, + sm_scale: float, + q_global_scale: torch.Tensor, +) -> None: + """Dispatch the ``triton`` FP4 MLA decode pipeline. + + Mirrors the four-stage layout used by ``fp4_mla_cutile.py`` + (page-stats with packed P -> reduce-stats -> prob-scale -> PV) but + routes through the self-contained kernels in + ``fp4_mla_triton.py``. Threads through the constexpr assume flags, + TMA descriptors, occupancy/num-warps launch meta, and pipelined PV loop. + """ + from .fp4_mla_triton import ( + _fp4_mla_attention_group_reduce_stats_kernel as _attn_group_reduce_stats_kernel, + ) + from .fp4_mla_triton import _fp4_mla_attention_page_stats_kernel as _attn_page_stats_kernel + from .fp4_mla_triton import _fp4_mla_attention_prob_scale_kernel as _attn_prob_scale_kernel + from .fp4_mla_triton import _fp4_mla_attention_pv_kernel as _attn_pv_kernel + from .fp4_mla_triton import ( + _fp4_mla_attention_pv_prepacked_v_kernel as _attn_pv_prepacked_v_kernel, + ) + from .fp4_mla_triton import _fp4_mla_attention_pv_reduce_kernel as _attn_pv_reduce_kernel + from .fp4_mla_triton import _fp4_mla_attention_reduce_stats_kernel as _attn_reduce_stats_kernel + + block_h = 128 + block_t = metadata.page_size + # Adaptive BLOCK_V: the fallback PV path uses a finer V split at small batch + # on B200 (~148 SMs). PV grid = num_queries * num_head_blocks(1) * + # (kv_lora_rank / BLOCK_V). We want >= ~2*num_SMs programs so that >1 CTA + # lands per SM and hides the L1TEX scoreboard stalls. Empirically (sweep): + # bs<=32 -> BLOCK_V=32; bs>=64 -> BLOCK_V=128. + # (BLOCK_V=16 is rejected by the V TMA descriptor min-stride requirement.) + # With prepacked V, BLOCK_V=128 avoids reloading the same P tile four times + # and matches the cutile prepacked-V tile shape. + block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) + q_storage_head_dim = head_dim + q_residual_dim + # The virtual GEMM tail evaluates QK + Q_r K + Q K_r in one reduction. + # Q and Q_r still occupy the 640-channel interleaved physical Q buffer; + # the final Q term reuses Q's main tail groups while K_r comes from the + # contiguous 64-channel tail of the primary paged KV cache. + q_head_dim = head_dim + q_residual_dim + FP4_MLA_K_RESIDUAL_DIM + # BLOCK_K = 512 aligns the K-window with the 512-channel non-residual prefix. + block_k = 512 + full_block_end = (q_head_dim // block_k) * block_k + tail_k = q_head_dim - full_block_end + tail_block_k = 1 << (tail_k - 1).bit_length() if tail_k > 0 else block_k + q_sf_per_token = q_storage_head_dim // FP4_BLOCK_SIZE + k_sf_per_token = (head_dim + FP4_MLA_K_RESIDUAL_DIM) // FP4_BLOCK_SIZE + sf_per_page = metadata.page_size // FP4_BLOCK_SIZE + num_head_blocks = triton.cdiv(num_heads, block_h) + + assume_full_heads = num_heads % block_h == 0 + assume_full_v = kv_lora_rank % block_v == 0 + # Match the cutile path: only mark pages "full" when we can prove every + # generation sequence has the same number of cached tokens AND + # query_len_per_seq == 1 (so the kv_len adjustment is a no-op). + assume_full_pages = ( + _infer_assume_full_pages(metadata, max_pages, metadata.page_size) and query_len_per_seq == 1 + ) + # Leave validity checks on. Matches cutile's default and is correctness- + # safe. The perfect-shape PV fast path (tl.ext.make_view + load_view_tko) + # remains gated off — when measured on the TileIR backend (ENABLE_TILE=1) + # it was net-slower on the bench, so the cost of enabling it isn't worth + # the win on the FP4 MLA shapes we care about. + assume_valid_pages = False + num_gen_seqs = num_queries // query_len_per_seq + if ( + not assume_valid_pages + and assume_full_pages + and src_page_ids.numel() == num_gen_seqs * max_pages + ): + assume_valid_pages = True + # cutile checks only `make_tensor_descriptor`; on the nvt backend the + # presence of TMA descriptors implies `tl.ext.make_view` is available too. + use_tma_data_load = hasattr(triton.language, "make_tensor_descriptor") + + # Install the device-side scratch allocator on every call. Triton stores + # the allocator in a ContextVar (triton.runtime._allocation), so a single + # process-wide install is not visible from worker threads / asyncio tasks + # that run with a different Context — the kernel launch would then hit the + # default NullAllocator and raise. Matches the cutile path. + if use_tma_data_load: + + def _tma_alloc(size: int, alignment: int, stream): + return torch.empty(size, device=q_fp4.device, dtype=torch.int8) + + triton.set_allocator(_tma_alloc) + + # cutile-equivalent launch meta. occupancy=2 lets two CTAs land per SM + # which improves wave-tail efficiency at the bs=32 hot point. + # NOTE: num_stages=2 (instead of the Triton 3.6 default of 3) sidesteps + # the TritonGPUAutomaticWarpSpecialization + NVWSInsertTmemAref pass that + # ICEs on the page_stats kernel under Triton 3.6.0 / sm_100. + launch_meta = {"occupancy": 2} + # The matmul kernels (page-stats QK and PV) are register-limited: at the + # Triton default of num_warps=4 the [BLOCK_H, BLOCK_T] epilogue spills the + # register file down to ~2 CTAs/SM (12.5% occupancy), so there are too few + # warps to hide the QK/PV load latency (ncu: ~0.3 eligible warps/scheduler). + # Spreading the tile epilogue over num_warps=8 halves the per-thread + # register need and roughly doubles resident warps. Matches the cutile + # ("nvt") backend, which launches page-stats at num_warps=8. Both are + # overridable for tuning. + sm_count = _get_sm_count(q_fp4.device) + # page-stats num_warps: the full-pages fast path (uniform q_len==1 decode) + # benefits from num_warps=8 (more warps hide the QK load latency); the + # masked path (q_len>1 / ragged lengths) carries extra per-thread state and + # measured markedly faster at num_warps=4 (e.g. bs256 q_len4: 131->95ms). + page_stats_num_warps = _env_int("TRTLLM_FP4_MLA_PAGE_STATS_NUM_WARPS") + if page_stats_num_warps is None: + page_stats_num_warps = 8 if assume_full_pages else 4 + page_stats_launch_meta = {"occupancy": 2, "num_warps": page_stats_num_warps} + # PV benefits from num_warps=8 across shapes measured. + pv_num_warps = _env_int("TRTLLM_FP4_MLA_PV_NUM_WARPS") or 8 + pv_launch_meta = {"occupancy": 2, "num_warps": pv_num_warps} + # PV loop pipelining. With TMA loads, num_stages>=2 lets the next page's + # loads overlap with the current MMA via mbarrier. The PV report shows + # long_scoreboard=4.5 cycles avg on V loads at PV_LOOP_STAGES=2; bumping the + # depth pays off when the grid is small enough that occupancy can absorb + # the extra in-flight tile state — i.e. medium batch / large max_pages. + # Larger pipelines hurt at small batch (more live state, fewer dim blocks). + if num_queries <= 16 or max_pages <= 4: + pv_loop_stages = 2 + else: + pv_loop_stages = 3 + + # Page-stats kernel: per (query, head_block, page) program, does QK, + # softmax stats, and packs probs into FP4 with the per-page local-max + # scaling trick. The page-max correction is applied later by + # prob_scale_kernel via p_sf in-place rescaling. + page_stats_shape = (num_queries, max_pages, num_heads) + page_max = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_page_max_buf", + page_stats_shape, + dtype=torch.float32, + device=q_fp4.device, + ) + page_sum = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_page_sum_buf", + page_stats_shape, + dtype=torch.float32, + device=q_fp4.device, + ) + + pack_prob_in_page_stats = True + _attn_page_stats_kernel[(num_queries, num_head_blocks, max_pages)]( + page_max, + page_sum, + p_fp4, + p_sf, + q_fp4, + q_sf, + kv_cache, + sf_cache, + global_scale, + q_global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + q_fp4.stride(0), + q_fp4.stride(1), + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + sf_cache.stride(0), + page_max.stride(0), + page_max.stride(1), + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + q_fp4.shape[0], + sm_scale, + NUM_HEADS=num_heads, + Q_HEAD_D=q_head_dim, + Q_STORAGE_HEAD_D=q_storage_head_dim, + K_HEAD_D=head_dim, + Q_RESIDUAL_D=q_residual_dim, + K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + Q_SF_PER_TOKEN=q_sf_per_token, + K_SF_PER_TOKEN=k_sf_per_token, + SF_PER_PAGE=sf_per_page, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + BLOCK_H=block_h, + BLOCK_T=block_t, + BLOCK_K=block_k, + FULL_BLOCK_END=full_block_end, + TAIL_BLOCK_K=tail_block_k, + USE_TMA_DATA_LOAD=use_tma_data_load, + PACK_PROBS=pack_prob_in_page_stats, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_VALID_PAGES=assume_valid_pages, + **page_stats_launch_meta, + ) + # Two-level softmax-stats reduction. The single-level reduce launched only + # (num_queries * num_head_blocks) CTAs, each serially walking all max_pages + # twice -- at small batch that handful of CTAs left the GPU almost idle and + # the reduce cost more than the QK matmul. Level 1 parallelizes the page + # reduction across a page-group axis (online-softmax partials, pipelined); + # level 2 reuses the existing reduce kernel to fold the few groups into the + # global (max, denom). When the (query, head) grid already fills the GPU the + # group count collapses to 1 and this degenerates to the original reduce. + seqhead_ctas = num_queries * num_head_blocks + # Aim for ~3 waves of level-1 CTAs so page loads have enough memory-level + # parallelism to hide latency, while keeping the group count small enough + # that the level-2 combine loop stays short. + target_l1_ctas = 3 * sm_count + num_reduce_groups = _ceil_div(target_l1_ctas, max(seqhead_ctas, 1)) + num_reduce_groups = max(1, min(num_reduce_groups, max_pages, 64)) + # The grouped (two-level) reduce needs an auxiliary workspace, and + # _ensure_workspace_tensor can only (re)allocate it outside CUDA graph + # capture. If a warmup forward did not already size that workspace (e.g. the + # warmup batch took the single-level path), fall back to the single-level + # reduce during capture so we never allocate mid-capture. The single-level + # reduce is numerically identical (it just launches fewer CTAs). + if num_reduce_groups > 1 and torch.cuda.is_current_stream_capturing(): + gmax = getattr(metadata, "_fp4_mla_attention_group_max_buf", None) + gsum = getattr(metadata, "_fp4_mla_attention_group_sum_buf", None) + groups_ready = ( + gmax is not None + and gsum is not None + and gmax.shape[0] >= num_queries + and gmax.shape[1] >= num_reduce_groups + and gmax.shape[2] >= num_heads + and gsum.shape[0] >= num_queries + and gsum.shape[1] >= num_reduce_groups + and gsum.shape[2] >= num_heads + ) + if not groups_ready: + num_reduce_groups = 1 + if num_reduce_groups <= 1: + _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( + max_scores, + denom, + page_max, + page_sum, + max_pages, + max_scores.stride(0), + page_max.stride(0), + page_max.stride(1), + NUM_HEADS=num_heads, + MAX_PAGES=max_pages, + BLOCK_H=block_h, + **launch_meta, + ) + else: + group_pages = _ceil_div(max_pages, num_reduce_groups) + num_reduce_groups = _ceil_div(max_pages, group_pages) + group_max = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_group_max_buf", + (num_queries, num_reduce_groups, num_heads), + dtype=torch.float32, + device=q_fp4.device, + ) + group_sum = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_group_sum_buf", + (num_queries, num_reduce_groups, num_heads), + dtype=torch.float32, + device=q_fp4.device, + ) + _attn_group_reduce_stats_kernel[(num_queries, num_head_blocks, num_reduce_groups)]( + group_max, + group_sum, + page_max, + page_sum, + max_pages, + group_max.stride(0), + group_max.stride(1), + page_max.stride(0), + page_max.stride(1), + NUM_HEADS=num_heads, + GROUP_PAGES=group_pages, + BLOCK_H=block_h, + PIPELINE_STAGES=min(group_pages, 4), + **launch_meta, + ) + _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( + max_scores, + denom, + group_max, + group_sum, + num_reduce_groups, + max_scores.stride(0), + group_max.stride(0), + group_max.stride(1), + NUM_HEADS=num_heads, + MAX_PAGES=num_reduce_groups, + BLOCK_H=block_h, + **launch_meta, + ) + _attn_prob_scale_kernel[(num_queries, num_head_blocks, max_pages)]( + p_sf, + max_scores, + denom, + page_max, + metadata.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + max_scores.stride(0), + page_max.stride(0), + page_max.stride(1), + NUM_HEADS=num_heads, + PAGE_SIZE=metadata.page_size, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + BLOCK_H=block_h, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_VALID_PAGES=assume_valid_pages, + **launch_meta, + ) + num_dim_blocks = triton.cdiv(kv_lora_rank, block_v) + v_packed = _get_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + v_head_dim=kv_lora_rank, + page_size=metadata.page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + page_ids=src_page_ids, + ) + if ( + v_packed is None + and _triton_can_prepack_v(kv_lora_rank, metadata.page_size, block_v) + and not torch.cuda.is_current_stream_capturing() + ): + v_packed = _update_triton_v_packed_cache( + metadata, + layer_idx, + kv_cache, + src_page_ids, + v_head_dim=kv_lora_rank, + page_size=metadata.page_size, + block_v=block_v, + local_layer=local_layer, + v_sf=v_sf, + ) + use_triton_v_packed_cache = v_packed is not None + + # PV page split: partition the page range across additional programs and + # reduce in a follow-up kernel. ncu showed PV at waves/SM=0.49 for bs=32 — + # PV is L1-bandwidth bound, so raising in-flight CTAs is the lever. + # BLOCK_V is bounded below by the 16-byte TMA descriptor min-stride. + # PV page split: ncu shows that with the current shape (bs=32, max_pages=256) + # the PV kernel is L1-cache-throughput bound (long_scoreboard=4.5 cycles + # avg, L1 global LD hit-rate <40%). Increasing the program count via page + # splitting reduced waves/SM idle time but did NOT improve wall-time at + # current shapes — the per-CTA L1 thrash is the limit. Gate the split off + # by default; re-enable only for very small grids where occupancy is the + # bottleneck rather than per-CTA L1 pressure. + page_split = 1 + base_grid = num_queries * num_head_blocks * num_dim_blocks + if max_pages >= 16 and base_grid < 148: + for p in (8, 4, 2): + if max_pages % p == 0 and max_pages // p >= 16 and base_grid * p <= 148 * 4: + page_split = p + break + # The page-split PV path needs a partial-output workspace, which + # _ensure_workspace_tensor can only (re)allocate outside CUDA graph capture. + # Fall back to the unsplit PV (numerically identical) during capture unless a + # warmup forward already sized that workspace, so capture never allocates. + if page_split > 1 and torch.cuda.is_current_stream_capturing(): + pbuf = getattr(metadata, "_fp4_mla_attention_pv_partial_buf", None) + partial_ready = ( + pbuf is not None + and pbuf.shape[0] >= num_queries + and pbuf.shape[1] >= page_split + and pbuf.shape[2] >= num_heads + and pbuf.shape[3] >= kv_lora_rank + ) + if not partial_ready: + page_split = 1 + if page_split > 1: + pages_per_split = max_pages // page_split + partial_out = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_pv_partial_buf", + (num_queries, page_split, num_heads, kv_lora_rank), + dtype=torch.float32, + device=q_fp4.device, + ) + if use_triton_v_packed_cache: + _attn_pv_prepacked_v_kernel[ + (num_queries, num_head_blocks, num_dim_blocks * page_split) + ]( + output, + p_fp4, + p_sf, + v_packed, + v_sf, + global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + PAGE_SPLIT=page_split, + PAGES_PER_SPLIT=pages_per_split, + PARTIAL_OUT=True, + partial_out_ptr=partial_out, + partial_s0=partial_out.stride(0), + partial_s1=partial_out.stride(1), + partial_s2=partial_out.stride(2), + partial_s3=partial_out.stride(3), + **pv_launch_meta, + ) + else: + _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks * page_split)]( + output, + p_fp4, + p_sf, + kv_cache, + kv_cache, + v_sf, + global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, + USE_PREPACKED_V=False, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + PAGE_SPLIT=page_split, + PAGES_PER_SPLIT=pages_per_split, + PARTIAL_OUT=True, + partial_out_ptr=partial_out, + partial_s0=partial_out.stride(0), + partial_s1=partial_out.stride(1), + partial_s2=partial_out.stride(2), + partial_s3=partial_out.stride(3), + **pv_launch_meta, + ) + _attn_pv_reduce_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( + output, + partial_out, + global_scale, + output.stride(0), + output.stride(1), + output.stride(2), + partial_out.stride(0), + partial_out.stride(1), + partial_out.stride(2), + partial_out.stride(3), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SPLIT=page_split, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_V=assume_full_v, + **launch_meta, + ) + else: + if use_triton_v_packed_cache: + _attn_pv_prepacked_v_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( + output, + p_fp4, + p_sf, + v_packed, + v_sf, + global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + **pv_launch_meta, + ) + else: + _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( + output, + p_fp4, + p_sf, + kv_cache, + kv_cache, + v_sf, + global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode, + kv_lens, + src_page_ids.shape[0], + kv_cache.shape[0], + output.stride(0), + output.stride(1), + output.stride(2), + output.shape[0] * output.shape[1], + p_fp4.stride(0), + p_fp4.stride(1), + p_fp4.shape[0], + kv_cache.stride(0), + kv_cache.stride(2), + kv_cache.stride(4), + v_sf.stride(0), + NUM_HEADS=num_heads, + V_HEAD_D=kv_lora_rank, + PAGE_SIZE=metadata.page_size, + FP4_BLOCK=FP4_BLOCK_SIZE, + SF_PER_PAGE=sf_per_page, + QUERY_LEN_PER_SEQ=query_len_per_seq, + MAX_PAGES=max_pages, + P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, + BLOCK_H=block_h, + BLOCK_V=block_v, + USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, + USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, + USE_PREPACKED_V=False, + PV_LOOP_STAGES=pv_loop_stages, + ASSUME_FULL_HEADS=assume_full_heads, + ASSUME_FULL_PAGES=assume_full_pages, + ASSUME_FULL_V=assume_full_v, + ASSUME_VALID_PAGES=assume_valid_pages, + **pv_launch_meta, + ) + + +def _run_cutedsl_packed_decode( + module: Any, + metadata: Any, + q_fp4: torch.Tensor, + q_sf: torch.Tensor, + kv_cache: torch.Tensor, + sf_cache: torch.Tensor, + v_packed: torch.Tensor, + v_sf: torch.Tensor, + kv_global_scale: torch.Tensor, + q_global_scale: torch.Tensor, + src_page_ids: torch.Tensor, + page_indptr: torch.Tensor, + kv_lens: torch.Tensor, + output: torch.Tensor, + *, + num_heads: int, + query_len: int, + num_requests: int, + rows_per_seq: int, + sm_scale: float, + v_page_offset: int, + helix_kv_bounds: Optional[torch.Tensor], + softmax_stats_tensor: Optional[torch.Tensor], + plan_geometry: Optional[_PackedPlanGeometry] = None, +) -> None: + """Packed CuTe DSL decode: one kernel over (split, request) tasks. + + A single split writes the output (and Helix stats) directly; several + splits write FP32 partials that one reduction kernel folds into the + output and stats. Nothing pads, replicates, or copies Q or the output. + """ + valid_m, physical_m = module.packed_query_rows(num_heads, query_len) + if physical_m != rows_per_seq: + raise RuntimeError( + f"FP4 MLA packed Q rows {rows_per_seq} do not match the kernel geometry {physical_m}." + ) + if v_packed is None: + raise RuntimeError("FP4 MLA packed decode requires the persistent packed-V cache.") + if not output.is_contiguous() or output.numel() != num_requests * valid_m * 512: + raise RuntimeError( + "FP4 MLA packed decode requires a contiguous [tokens, heads, 512] output." + ) + device = output.device + if plan_geometry is None: + plan_geometry = _fp4_mla_packed_plan_geometry( + metadata, + num_heads=num_heads, + query_len=query_len, + num_requests=num_requests, + device=device, + helix_kv_bounds=helix_kv_bounds, + ) + if plan_geometry is None: + raise RuntimeError("FP4 MLA packed decode requires generation cache pages.") + splits = plan_geometry.splits + output_rows = output.view(num_requests * valid_m, 512) + stats_rows = None + if softmax_stats_tensor is not None: + stats_rows = softmax_stats_tensor.view(num_requests * valid_m, 2) + common = dict( + max_kv_len=plan_geometry.split_kv_len, + sm_scale=sm_scale, + num_heads=num_heads, + query_len=query_len, + num_requests=num_requests, + num_splits=splits, + q_global_scale=q_global_scale, + v_page_offset=v_page_offset, + helix_kv_bounds=helix_kv_bounds, + ) + if splits == 1: + module.run_trtllm_fp4_mla_decode_packed( + q_fp4, + q_sf, + kv_cache, + sf_cache, + v_packed, + v_sf, + kv_global_scale, + src_page_ids, + page_indptr, + kv_lens, + output_rows, + softmax_stats=stats_rows, + **common, + ) + return + from .fp4_mla_packed_reduce import reduce_packed_splits + + tasks = splits * num_requests + task_capacity = plan_geometry.task_capacity + plan = plan_geometry.plan + # The first FP4 MLA kernel of the step wrote the plan fused with the + # generation lengths; a stale key means this call plans on other bounds. + if getattr(metadata, "_fp4_mla_packed_plan_key", None) != plan_geometry.key: + _launch_fp4_mla_packed_plan(plan_geometry, kv_lens, metadata) + partial = _ensure_workspace_tensor( + metadata, + "_fp4_mla_packed_split_partial", + (task_capacity, physical_m, 512), + dtype=torch.float32, + device=device, + )[:tasks] + maxima = _ensure_workspace_tensor( + metadata, + "_fp4_mla_packed_split_maxima", + (task_capacity, physical_m), + dtype=torch.float32, + device=device, + )[:tasks] + sums = _ensure_workspace_tensor( + metadata, + "_fp4_mla_packed_split_sums", + (task_capacity, physical_m), + dtype=torch.float32, + device=device, + )[:tasks] + module.run_trtllm_fp4_mla_decode_packed( + q_fp4, + q_sf, + kv_cache, + sf_cache, + v_packed, + v_sf, + kv_global_scale, + src_page_ids, + page_indptr, + kv_lens, + output_rows, + split_kv_buffers=(partial, maxima, sums), + plan=plan, + **common, + ) + reduce_packed_splits( + partial, + maxima, + sums, + output_rows, + q_global_scale, + kv_global_scale, + num_requests=num_requests, + num_splits=splits, + valid_m=valid_m, + physical_m=physical_m, + sm_scale=sm_scale, + plan=plan, + output_scale=1.0 / module.FP4_MLA_P_GLOBAL_SCALE, + stats=stats_rows, + sm_count=_get_sm_count(device), + ) + + +def run_fp4_mla_attention_decode( + metadata: Any, + layer_idx: int, + local_layer: int, + q: torch.Tensor, + output: torch.Tensor, + *, + sm_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, + prequantized_q: torch.Tensor, + prequantized_q_sf: torch.Tensor, + q_batch_capacity: int, + softmax_stats_tensor: Optional[torch.Tensor] = None, +) -> None: + """Run MLA decode with FP4 QK and FP4 PV tensor-core matmuls. + + Q is supplied in its assembled ``[latent, RoPE]`` layout and quantized to + FP4 directly. QK reads ``[KV-nope, K-RoPE, K-RoPE-residual]`` contiguously + from the primary cache with swizzled block scales. Softmax probabilities + are quantized to FP4 per page, and PV repacks V nibbles from the shared KV + cache while reading the auxiliary V-view scale pool. No BF16 dequantized + KV workspace is materialized on this path. Callers must supply the packed Q + and scales produced by the fused generation cache update. + """ + head_dim = kv_lora_rank + qk_rope_head_dim + if qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM: + raise ValueError( + "FP4 MLA K residual attention requires " + f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got {qk_rope_head_dim}." + ) + _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) + if metadata.page_size != FP4_MLA_TOKENS_PER_BLOCK: + raise ValueError( + f"FP4 MLA attention decode requires page_size={FP4_MLA_TOKENS_PER_BLOCK}, " + f"got {metadata.page_size}." + ) + + if q.ndim != 3 or q.shape[-1] != head_dim: + raise ValueError( + "FP4 MLA attention Q must have shape " + f"[tokens, heads, {head_dim}], got {tuple(q.shape)}." + ) + if not q.is_contiguous(): + raise ValueError("FP4 MLA attention Q must be contiguous.") + + num_queries = q.shape[0] + if num_queries == 0: + raise ValueError("FP4 MLA attention decode requires at least one query token.") + num_gen_seqs = metadata.num_seqs - metadata.num_contexts + query_len_per_seq = _get_linear_mtp_query_len_per_seq( + metadata, + num_queries=num_queries, + num_gen_seqs=num_gen_seqs, + ) + + num_heads = q.shape[1] + if output.shape[:2] != (num_queries, num_heads): + raise ValueError("FP4 MLA attention output batch dimensions do not match.") + + backend = _fp4_mla_attention_backend() + helix_spec_tokens_valid = bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + helix_kv_bounds = None + if softmax_stats_tensor is not None: + if backend != _FP4_MLA_CUTEDSL_BACKEND: + raise NotImplementedError( + "FP4 MLA Helix softmax stats require the cutedsl attention backend." + ) + if query_len_per_seq != 1 and not helix_spec_tokens_valid: + raise NotImplementedError( + "FP4 MLA multi-token Helix requires speculative per-token metadata." + ) + expected_stats_shape = (num_queries, num_heads, 2) + if ( + softmax_stats_tensor.shape != expected_stats_shape + or softmax_stats_tensor.dtype != torch.float32 + or softmax_stats_tensor.device != q.device + or not softmax_stats_tensor.is_contiguous() + ): + raise ValueError( + "FP4 MLA Helix requires contiguous same-device float32 softmax " + f"stats with shape {expected_stats_shape}." + ) + if helix_spec_tokens_valid: + helix_kv_bounds = getattr(metadata, "helix_kv_bounds", None) + if ( + not isinstance(helix_kv_bounds, torch.Tensor) + or helix_kv_bounds.dtype != torch.int32 + or helix_kv_bounds.device != q.device + or helix_kv_bounds.ndim != 1 + or helix_kv_bounds.numel() < num_queries + or not helix_kv_bounds.is_contiguous() + ): + raise ValueError( + "FP4 MLA speculative Helix KV bounds must be a contiguous " + "same-device int32 tensor covering every query token." + ) + helix_kv_bounds = helix_kv_bounds[:num_queries] + if getattr(metadata, "fp4_mla_v_scale_pool", None) is None: + raise RuntimeError( + "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." + ) + + global_scale = _get_fp4_mla_global_scale(metadata, q.device) + q_residual_dim = FP4_MLA_Q_RESIDUAL_DIM + _validate_fp4_mla_attention_q_shape(head_dim, q_residual_dim) + + if prequantized_q is None or prequantized_q_sf is None or q_batch_capacity is None: + raise RuntimeError( + "FP4 MLA decode requires Q prequantized by the fused generation cache update." + ) + + capacity = int(q_batch_capacity) + # The fused generation update records the Q layout it produced; metadata + # that never went through it (contract tests) carries the legacy layout. + recorded_rows_per_seq = getattr(metadata, "_fp4_mla_prequantized_q_rows_per_seq", None) + packed_rows_per_seq = int(recorded_rows_per_seq or 0) + if recorded_rows_per_seq is not None and packed_rows_per_seq != _fp4_mla_packed_rows_per_seq( + num_heads, query_len_per_seq + ): + raise RuntimeError( + "FP4 MLA prequantized Q layout does not match the packed decode " + f"configuration (rows per request {packed_rows_per_seq})." + ) + if packed_rows_per_seq: + active_q_rows = num_gen_seqs * packed_rows_per_seq + else: + active_q_rows = num_queries * num_heads + active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( + active_q_rows, + FP4_MLA_Q_LOGICAL_DIM, + ) + if packed_rows_per_seq: + q_contract_ok = ( + prequantized_q.ndim == 2 + and prequantized_q.shape[0] >= active_q_rows + and prequantized_q.shape[1] == FP4_MLA_Q_PACKED_DIM + and prequantized_q_sf.numel() >= active_q_sf_bytes + ) + else: + expected_q_shape = (capacity * num_heads, FP4_MLA_Q_PACKED_DIM) + expected_q_sf_shape = ( + _get_fp4_mla_swizzled_scale_size( + capacity * num_heads, + FP4_MLA_Q_LOGICAL_DIM, + ), + ) + q_contract_ok = ( + tuple(prequantized_q.shape) == expected_q_shape + and tuple(prequantized_q_sf.shape) == expected_q_sf_shape + ) + if ( + q.dtype != torch.bfloat16 + or capacity <= 0 + or num_queries > capacity + or not q_contract_ok + or prequantized_q.dtype != torch.uint8 + or not prequantized_q.is_contiguous() + or prequantized_q_sf.dtype != torch.float8_e4m3fn + or not prequantized_q_sf.is_contiguous() + ): + raise ValueError("FP4 MLA prequantized Q does not satisfy the fused-Q contract.") + q_fp4 = prequantized_q[:active_q_rows] + q_sf = prequantized_q_sf[:active_q_sf_bytes] + q_global_scale = _get_fp4_mla_q_global_scale(metadata, q.device) + + kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) + _validate_fp4_mla_kv_storage_shape( + kv_cache, + sf_cache, + head_dim=head_dim, + backend=backend, + ) + sf_cache = sf_cache.view(torch.float8_e4m3fn) + + v_sf_pool = metadata.fp4_mla_v_scale_pool + v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=kv_lora_rank)[local_layer].view( + torch.float8_e4m3fn + ) + # The kv_lens runtime alias can lag at the decode anchor (seq_lens == 1) under + # CUDA graph / one-engine MTP; recover the true total per sequence so the + # per-query causal masking sees the full 1 + draft_len window (no-op when the + # alias already matches). + plan_geometry = _fp4_mla_packed_plan_geometry( + metadata, + num_heads=num_heads, + query_len=query_len_per_seq, + num_requests=num_gen_seqs, + device=q.device, + helix_kv_bounds=helix_kv_bounds, + ) + kv_lens, _ = _fp4_mla_uniform_generation_lengths( + metadata, num_queries, num_gen_seqs, plan_geometry=plan_geometry + ) + _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens) + src_page_ids = _fp4_mla_generation_page_ids(metadata, num_gen_seqs) + max_pages = _max_generation_pages(metadata) + if max_pages == 0: + raise RuntimeError("FP4 MLA attention decode requires generation cache pages.") + if backend == _FP4_MLA_CUTEDSL_BACKEND: + if get_sm_version() != 107: + raise RuntimeError( + "FP4 MLA cutedsl attention backend requires Rubin SM107; " + f"current architecture is SM{get_sm_version()}." + ) + if not _cutedsl_backend_available(): + raise RuntimeError( + "FP4 MLA cutedsl attention backend requires the Rubin CTM and " + "CuTeDSL runtime packages." + ) + if not 0 < num_heads <= 128 or kv_lora_rank != 512: + raise ValueError( + "FP4 MLA cutedsl attention requires 1-128 local heads and " + f"kv_lora_rank=512, got num_heads={num_heads}, " + f"kv_lora_rank={kv_lora_rank}." + ) + + cutedsl_kernel = _fp4_mla_cutedsl_kernel_module() + QK_LOGICAL_DIM = cutedsl_kernel.QK_LOGICAL_DIM + QK_SF_GROUPS = cutedsl_kernel.QK_SF_GROUPS + SMEM_P4_V_N_PER_CTA = cutedsl_kernel.SMEM_P4_V_N_PER_CTA + run_trtllm_fp4_mla_decode_page_native_from_raw = ( + cutedsl_kernel.run_trtllm_fp4_mla_decode_page_native_from_raw + ) + + # Keep general cases on the original path. This specialization uses + # the prepacked-V mufu16 core and linear six-query MTP only. + fold_queries = ( + not packed_rows_per_seq + and num_heads == 6 + and query_len_per_seq == 6 + and num_gen_seqs == 1 + and helix_kv_bounds is None + and softmax_stats_tensor is None + and getattr(cutedsl_kernel, "SUPPORTS_FOLDED_QUERIES", False) + and not _fp4_mla_cutedsl_fused_v_transpose_enabled() + and getattr(metadata, "_fp4_mla_fold_q", True) + and os.environ.get("TRTLLM_FP4_MLA_FOLD_Q", "1") != "0" + ) + if fold_queries: + from .fp4_mla_splitkv import select_split_count + + fold_queries = select_split_count(max_pages * metadata.page_size) > 1 + physical_heads = 128 + kernel_q = q_fp4 + kernel_q_sf = q_sf + if not packed_rows_per_seq and num_heads < physical_heads and not fold_queries: + kernel_q_storage, kernel_q_sf_storage, q_batch_capacity = _prepare_fp4_mla_q_buffers( + metadata, + num_queries, + physical_heads, + q.device, + ) + active_q_rows = num_queries * physical_heads + active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( + active_q_rows, + QK_LOGICAL_DIM, + ) + kernel_q = kernel_q_storage[:active_q_rows] + kernel_q_sf = kernel_q_sf_storage[:active_q_sf_bytes] + _cutedsl_pad_q_and_sf_kernel[(num_queries, _ceil_div(QK_LOGICAL_DIM // 2, 64))]( + kernel_q, + q_fp4, + kernel_q_sf, + q_sf, + num_heads, + output_heads=physical_heads, + packed_dim=QK_LOGICAL_DIM // 2, + block_bytes=64, + sf_cols=QK_SF_GROUPS, + sf_cols_per_byte_block=8, + ) + + fused_v_transpose = _fp4_mla_cutedsl_fused_v_transpose_enabled() + if fused_v_transpose: + # The fusion kernel reads V from the canonical KV cache and uses + # the current layer V scales directly. Keep a None placeholder so + # the mufu16 and fused-V launchers share one Python call site. + core_v_packed = None + core_v_sf = v_sf + v_page_offset = 0 + else: + v_packed = _get_cutedsl_persistent_v_packed_cache( + metadata, + local_layer, + kv_cache, + v_head_dim=kv_lora_rank, + page_size=metadata.page_size, + block_v=SMEM_P4_V_N_PER_CTA, + ) + core_v_packed = _get_fp4_mla_v_packed_pool_base(metadata) + if core_v_packed is None: + raise RuntimeError("Persistent FP4 MLA V packing requires a stable full-pool base.") + get_v_page_offset = getattr( + metadata.kv_cache_manager, "get_mla_v_packed_page_offset", None + ) + v_page_offset = ( + int(get_v_page_offset(local_layer)) + if callable(get_v_page_offset) + else local_layer * kv_cache.shape[0] + ) + page_bytes = kv_lora_rank * (metadata.page_size // 2) + expected_layer_ptr = core_v_packed.data_ptr() + v_page_offset * page_bytes + if expected_layer_ptr != v_packed.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-packed layer view does not match its " + "full-pool base and page offset." + ) + + v_sf_pool_base = _get_fp4_mla_v_scale_pool_base(metadata) + if v_sf_pool_base is None: + v_sf_pool_base = v_sf_pool.flatten(0, 1).view(torch.uint8) + if v_sf_pool_base.data_ptr() != v_sf_pool.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-scale pool must flatten without a copy." + ) + if ( + not isinstance(v_sf_pool_base, torch.Tensor) + or v_sf_pool_base.dtype != torch.uint8 + or v_sf_pool_base.device != v_sf.device + or v_sf_pool_base.ndim != 2 + or v_sf_pool_base.shape[1] != v_sf_pool.shape[-1] + or not v_sf_pool_base.is_contiguous() + ): + raise RuntimeError( + "Persistent FP4 MLA V-scale pool base must be a contiguous " + "two-dimensional uint8 tensor with the configured page stride." + ) + core_v_sf = v_sf_pool_base.view(torch.float8_e4m3fn) + get_v_sf_page_offset = getattr( + metadata.kv_cache_manager, "get_mla_v_scale_page_offset", None + ) + v_sf_page_offset = ( + int(get_v_sf_page_offset(local_layer)) + if callable(get_v_sf_page_offset) + else local_layer * kv_cache.shape[0] + ) + if v_sf_page_offset != v_page_offset: + raise RuntimeError( + "Persistent FP4 MLA V-packed and V-scale pools require " + "matching encoded layer offsets." + ) + expected_v_sf_ptr = ( + core_v_sf.data_ptr() + + v_sf_page_offset * v_sf_pool.stride(1) * v_sf_pool.element_size() + ) + if expected_v_sf_ptr != v_sf.data_ptr(): + raise RuntimeError( + "Persistent FP4 MLA V-scale layer view does not match its " + "full-pool base and page offset." + ) + + if packed_rows_per_seq: + _run_cutedsl_packed_decode( + cutedsl_kernel, + metadata, + q_fp4, + q_sf, + kv_cache, + sf_cache, + core_v_packed, + core_v_sf, + global_scale, + q_global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], + kv_lens, + output, + num_heads=num_heads, + query_len=query_len_per_seq, + num_requests=num_gen_seqs, + rows_per_seq=packed_rows_per_seq, + sm_scale=float(sm_scale), + v_page_offset=v_page_offset, + helix_kv_bounds=helix_kv_bounds, + softmax_stats_tensor=softmax_stats_tensor, + plan_geometry=plan_geometry, + ) + return + + from .fp4_mla_splitkv import run_fp4_mla_splitkv, select_split_count + + if fold_queries: + from .fp4_mla_foldq import run_fp4_mla_foldq + + run_fp4_mla_foldq( + cutedsl_kernel, + metadata, + q_fp4, + q_sf, + kv_cache, + sf_cache, + core_v_packed, + core_v_sf, + global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], + kv_lens, + output, + splits=select_split_count(max_pages * metadata.page_size), + max_kv_len=max_pages * metadata.page_size, + sm_scale=float(sm_scale), + q_global_scale=q_global_scale, + query_len_per_seq=query_len_per_seq, + v_page_offset=v_page_offset, + ) + return + + # Helix also needs per-query bounds and softmax stats from the native path. + splits = ( + select_split_count(max_pages * metadata.page_size) + if helix_kv_bounds is None and softmax_stats_tensor is None + else 1 + ) + direct_split_output = splits > 1 and getattr( + metadata, "_fp4_mla_splitkv_direct_output", True + ) + kernel_output = output + kernel_softmax_stats = None + if softmax_stats_tensor is not None: + kernel_softmax_stats = _ensure_workspace_tensor( + metadata, + "_fp4_mla_cutedsl_softmax_stats_buf", + (2, num_queries, physical_heads), + dtype=torch.float32, + device=output.device, + ) + if num_heads < physical_heads and not direct_split_output: + kernel_output = _ensure_workspace_tensor( + metadata, + "_fp4_mla_cutedsl_output_buf", + (q_batch_capacity, physical_heads, kv_lora_rank), + dtype=output.dtype, + device=output.device, + )[:num_queries] + + if splits > 1: + run_fp4_mla_splitkv( + _fp4_mla_cutedsl_kernel_module(), + metadata, + kernel_q, + kernel_q_sf, + kv_cache, + sf_cache, + core_v_packed, + core_v_sf, + global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], + kv_lens, + kernel_output, + splits=splits, + max_kv_len=max_pages * metadata.page_size, + sm_scale=float(sm_scale), + q_global_scale=q_global_scale, + query_len_per_seq=query_len_per_seq, + v_page_offset=v_page_offset, + q_batch_capacity=q_batch_capacity, + ) + else: + run_trtllm_fp4_mla_decode_page_native_from_raw( + kernel_q, + kernel_q_sf, + kv_cache, + sf_cache, + core_v_packed, + core_v_sf, + global_scale, + src_page_ids, + metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], + kv_lens, + kernel_output, + max_kv_len=max_pages * metadata.page_size, + sm_scale=float(sm_scale), + num_heads=physical_heads, + q_global_scale=q_global_scale, + page_size=metadata.page_size, + query_len_per_seq=query_len_per_seq, + v_page_offset=v_page_offset, + q_batch_capacity=q_batch_capacity, + partition_runtime_valid_k=bool(getattr(metadata, "is_cuda_graph", False)), + softmax_row_max=(None if kernel_softmax_stats is None else kernel_softmax_stats[0]), + softmax_row_sum=(None if kernel_softmax_stats is None else kernel_softmax_stats[1]), + helix_kv_bounds=helix_kv_bounds, + ) + if kernel_output is not output: + output.copy_(kernel_output[:, :num_heads]) + if kernel_softmax_stats is not None: + softmax_stats_tensor[..., 0].copy_(kernel_softmax_stats[0, :, :num_heads]) + softmax_stats_tensor[..., 1].copy_(kernel_softmax_stats[1, :, :num_heads]) + return + + total_p_rows = num_queries * max_pages * num_heads + p_fp4 = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_p_buf", + (max(total_p_rows, 1), metadata.page_size // 2), + dtype=torch.uint8, + device=q.device, + )[:total_p_rows] + p_sf = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_p_sf_buf", + (max(_get_fp4_mla_swizzled_scale_size(total_p_rows, metadata.page_size), 1),), + dtype=torch.float8_e4m3fn, + device=q.device, + ) + stats_shape = (num_queries, num_heads) + max_scores = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_max_buf", + stats_shape, + dtype=torch.float32, + device=q.device, + ) + denom = _ensure_workspace_tensor( + metadata, + "_fp4_mla_attention_denom_buf", + stats_shape, + dtype=torch.float32, + device=q.device, + ) + + if backend != "triton": + raise ValueError( + f"Unsupported FP4 MLA attention backend '{backend}'. " + f"Set {FP4_MLA_ATTENTION_BACKEND_ENV} to 'triton' or " + f"'{_FP4_MLA_CUTEDSL_BACKEND}'." + ) + + # Self-contained public-Triton path: TMA-loaded QK + fused page-stats pack, + # reduce-stats, prob-scale, and PV with an optional prepacked V cache. + _run_triton_attention_decode( + metadata=metadata, + layer_idx=layer_idx, + local_layer=local_layer, + q_fp4=q_fp4, + q_sf=q_sf.contiguous().view(-1), + kv_cache=kv_cache, + sf_cache=sf_cache, + v_sf=v_sf, + global_scale=global_scale, + src_page_ids=src_page_ids, + kv_lens=kv_lens, + p_fp4=p_fp4, + p_sf=p_sf, + max_scores=max_scores, + denom=denom, + output=output, + num_queries=num_queries, + num_heads=num_heads, + head_dim=head_dim, + kv_lora_rank=kv_lora_rank, + q_residual_dim=q_residual_dim, + query_len_per_seq=query_len_per_seq, + max_pages=max_pages, + sm_scale=float(sm_scale), + q_global_scale=q_global_scale, + ) diff --git a/tests/unittest/_torch/attention/test_fp4_mla_packed.py b/tests/unittest/_torch/attention/test_fp4_mla_packed.py index 753f8cf0981c..e03e3cc1db91 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla_packed.py +++ b/tests/unittest/_torch/attention/test_fp4_mla_packed.py @@ -69,19 +69,19 @@ def _plan_mod(): [ # Even batch under one wave (helix16 B4, 6 m-tiles -> 17 wave tasks): one wave of 32-tile tasks. ([128, 128, 128, 128], 24, 17, 32, [0, 4, 8, 12, 16]), - # One long request among short ones takes most of the slots (two-wave target, fit by the fallback share). - ([1, 79, 3, 176], 24, 17, 13, [0, 1, 8, 9, 23]), - # Full machine, longest 1.5x above the wave share: two-wave target, the short requests also halve. - ([128] * 15 + [200], 48, 17, 67, list(range(0, 32, 2)) + [33]), + # One long request among short ones: 16-tile tasks; the one task past the wave is a short tail. + ([1, 79, 3, 176], 24, 17, 16, [0, 1, 6, 7, 18]), + # Full machine, longest 1.5x above the wave share: 100-tile tasks, the tails in a second wave. + ([128] * 15 + [200], 48, 17, 100, list(range(0, 34, 2))), # Full machine, longest within 1.5x of the wave share: no split at all. ([128] * 15 + [130], 48, 17, 130, list(range(17))), - # Strong skew in a full machine: the long request gets twelve tasks. - ([128] * 15 + [1024], 48, 17, 86, list(range(0, 32, 2)) + [42]), - # Strong skew in a small batch: two-wave target, fit by the fallback share. - ([128, 128, 128, 1024], 24, 17, 71, [0, 2, 4, 6, 21]), - # A single request fills one wave with 16 tasks, or honours the minimum. + # Strong skew in a full machine: eleven 94-tile tasks for the long request, short tails behind. + ([128] * 15 + [1024], 48, 17, 94, list(range(0, 32, 2)) + [41]), + # Strong skew in a small batch: 94-tile tasks, 11 for the long request, fill one wave. + ([128, 128, 128, 1024], 24, 17, 94, [0, 2, 4, 6, 17]), + # A single request fills one wave with 16 tasks; a short one is not worth splitting. ([128], 16, 17, 8, [0, 16]), - ([8], 16, 17, 2, [0, 4]), + ([8], 16, 17, 8, [0, 1]), # No spare slots: one task per request, however long. ([5, 900], 2, 17, 900, list(range(3))), ], @@ -102,17 +102,63 @@ def test_plan_packed_tasks_balances_by_kv_length( assert (end - begin - 1) * tiles_per_task < request_tiles <= (end - begin) * tiles_per_task -def test_plan_packed_tasks_one_wave_for_helix() -> None: +def test_plan_packed_tasks_simulates_helix_ranks() -> None: plan = _plan_mod().plan_packed_tasks - # helix16 bs16 (6 m-tiles, 17 wave tasks), per-rank KV skewed 3.5x: the two-wave - # target cuts 12-tile pieces (42 tasks); a Helix rank keeps one wave of 28-tile tasks. + # helix16 bs16 (6 m-tiles, 17 wave tasks), per-rank KV skewed 3.5x: the Helix rule targets one + # wave of 28-tile tasks (21 tasks, a second wave); the simulation keeps the 16 requests whole. tiles = [16] * 12 + [40, 48, 56, 56] - assert plan(tiles, 48, 17) == (12, list(range(0, 26, 2)) + [28, 32, 37, 42]) - assert plan(tiles, 48, 17, one_wave=True) == (28, list(range(13)) + [14, 16, 18, 20]) + assert plan(tiles, 48, 17, one_wave=True, wave_plan=False) == ( + 28, + list(range(13)) + [14, 16, 18, 20], + ) + assert plan(tiles, 48, 17, one_wave=True) == (56, list(range(17))) # Even Helix batches are unchanged. assert plan([128] * 4, 24, 17, one_wave=True) == plan([128] * 4, 24, 17) +def test_plan_packed_tasks_simulates_the_waves() -> None: + reduce_mod = _plan_mod() + plan = reduce_mod.plan_packed_tasks + # Skewed request lengths exercise the task-rounding overhead across waves. + tiles = [1848, 1250, 977, 743, 586, 430, 313, 196, 106] + assert plan(tiles, 9 * 9, 35) == (206, [0, 9, 16, 21, 25, 28, 31, 33, 34, 35]) + # Nine equal requests: 36 one-wave tasks spill one task; three 239-tile pieces each fit one wave. + assert plan([715] * 9, 9 * 9, 35) == (239, list(range(0, 28, 3))) + simulate = reduce_mod.plan_simulated_makespan + cost = reduce_mod.PLAN_TASK_COST_TILES + reduce_mod.PLAN_SPLIT_COST_TILES + # 36 equal 179-tile tasks on 35 slots: the last one waits for a whole task. + assert simulate([716] * 9, 179, 35) == 2 * (179 + cost) + # The 36th task (a 178-tile tail) starts when the first tail frees its slot. + assert simulate([715] * 9, 179, 35) == (178 + cost) + (178 + cost) + # 2-tile tails free their slots early, so a spilled tail does not extend the makespan. + assert simulate([181] * 18, 179, 35) == 179 + cost + # Requests that keep one task pay no split cost. + assert simulate([100, 300], 300, 35) == 300 + reduce_mod.PLAN_TASK_COST_TILES + + +def test_plan_task_order_runs_full_tasks_then_longest_tails() -> None: + order = _plan_mod().plan_task_order + # Tails: request 0 -> 4 tiles, request 1 -> 10 (a full task), request 2 -> 3, request 3 -> 4. + full = [(0, 0), (0, 1), (1, 0), (1, 1), (3, 0)] + tails = [(1, 2), (0, 2), (3, 1), (2, 0)] + assert order([24, 30, 3, 14], 10) == full + tails + # One task per request: longest first, ties in request order. + assert order([5, 9, 5, 2], 10) == [(1, 0), (0, 0), (2, 0), (3, 0)] + + +def test_plan_without_wave_plan_keeps_the_rule_and_request_order() -> None: + reduce_mod = _plan_mod() + plan = reduce_mod.plan_packed_tasks + # TRTLLM_FP4_MLA_PACKED_WAVE_PLAN=0: the rule's length, as before the wave simulation. + assert plan([1, 79, 3, 176], 24, 17, wave_plan=False) == (13, [0, 1, 8, 9, 23]) + assert plan([128] * 15 + [1024], 48, 17, wave_plan=False) == ( + 86, + list(range(0, 32, 2)) + [42], + ) + order = reduce_mod.plan_task_order([24, 30, 3, 14], 10, wave_plan=False) + assert order == [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (3, 0), (3, 1)] + + def test_plan_packed_tasks_honours_the_runtime_profile_cap() -> None: plan = _plan_mod().plan_packed_tasks # An even full machine would not split, but the profile only covers 512 tiles. @@ -140,22 +186,33 @@ def test_plan_packed_tasks_random_batches_fit_their_budget() -> None: ) +# Skewed request lengths for planner coverage. +_DEP16_RANK_KV_LENS = [473000, 320000, 250000, 190000, 150000, 110000, 80000, 50000, 27000] + + @pytest.mark.skipif( torch.cuda.get_device_capability() != (10, 7) if torch.cuda.is_available() else True, reason="the CuTe DSL plan kernel is compiled for Rubin SM107", ) @pytest.mark.parametrize( - "kv_lens,splits,wave_tasks,query_len,helix", + "kv_lens,splits,wave_tasks,query_len,helix,wave_plan", [ - ([32768] * 15 + [262144], 4, 17, 8, False), - ([131, 20000, 600, 45000], 10, 17, 8, False), - ([131, 20000, 600, 45000], 10, 17, 8, True), - ([300, 70000, 2200], 16, 106, 1, False), - ([5, 5, 5], 2, 106, 4, False), + ([32768] * 15 + [262144], 4, 17, 8, False, True), + ([131, 20000, 600, 45000], 10, 17, 8, False, True), + ([131, 20000, 600, 45000], 10, 17, 8, True, True), + ([300, 70000, 2200], 16, 106, 1, False, True), + ([5, 5, 5], 2, 106, 4, False, True), + # DEP16 rank shapes (3 m-tiles, 35 wave tasks): the simulated task length. + (_DEP16_RANK_KV_LENS, 9, 35, 4, False, True), + ([183000] * 9, 9, 35, 4, False, True), + ([103000] * 16, 6, 35, 4, False, True), + # TRTLLM_FP4_MLA_PACKED_WAVE_PLAN=0, off and on Helix: the rule's length, request order. + (_DEP16_RANK_KV_LENS, 9, 35, 4, False, False), + ([131, 20000, 600, 45000], 10, 17, 8, True, False), ], ) def test_packed_plan_kernel_matches_host_mirror( - kv_lens: list[int], splits: int, wave_tasks: int, query_len: int, helix: bool + kv_lens: list[int], splits: int, wave_tasks: int, query_len: int, helix: bool, wave_plan: bool ) -> None: reduce_mod = _plan_mod() requests = len(kv_lens) @@ -174,7 +231,7 @@ def test_packed_plan_kernel_matches_host_mirror( tiles = reduce_mod.request_kv_tiles(kv_lens, query_len, bounds_list) cap = max(tiles) expected_tiles, expected_begins = reduce_mod.plan_packed_tasks( - tiles, tasks, wave_tasks, cap, one_wave=helix + tiles, tasks, wave_tasks, cap, one_wave=helix, wave_plan=wave_plan ) plan = torch.full( (reduce_mod.packed_plan_length(requests, tasks),), -7, dtype=torch.int32, device="cuda" @@ -188,6 +245,7 @@ def test_packed_plan_kernel_matches_host_mirror( max_tiles_per_task=cap, query_len=query_len, helix_kv_bounds=bounds, + wave_plan=wave_plan, ) torch.cuda.synchronize() plan = plan.tolist() @@ -200,13 +258,15 @@ def test_packed_plan_kernel_matches_host_mirror( plan[tables + tasks : tables + 2 * tasks], plan[tables + 2 * tasks : tables + 3 * tasks], ) - for request, (begin, end) in enumerate(zip(expected_begins, expected_begins[1:])): - for task in range(begin, end): - assert (request_of[task], split_of[task], count_of[task]) == ( - request, - task - begin, - end - begin, - ) + # Tasks are listed in grid order; each request still owns begins[r + 1] - begins[r] of them. + order = reduce_mod.plan_task_order(tiles, expected_tiles, wave_plan) + assert len(order) == expected_begins[-1] + for task, (request, split) in enumerate(order): + assert (request_of[task], split_of[task], count_of[task]) == ( + request, + split, + expected_begins[request + 1] - expected_begins[request], + ) # Unused slots point past every request's tiles and are never scheduled. for task in range(expected_begins[-1], tasks): assert request_of[task] == requests - 1 and split_of[task] * expected_tiles >= tiles[-1] @@ -322,6 +382,21 @@ def test_uniform_generation_lengths_fuse_the_packed_plan(monkeypatch) -> None: uniform(metadata, 12, 3, plan_geometry=geometry) +def test_packed_plan_wave_plan_env_rolls_back(monkeypatch) -> None: + calls = [] + monkeypatch.setattr( + _plan_mod(), "launch_packed_plan", lambda *args, **kwargs: calls.append(kwargs["wave_plan"]) + ) + geometry = fp4_mla_backend._PackedPlanGeometry(2, 4, 3, 17, 64, 6, 1024, None, None, ("key",)) + kv_lens = torch.zeros(2, dtype=torch.int32, device="cuda") + metadata = SimpleNamespace(_fp4_mla_packed_plan_key=None) + monkeypatch.delenv(fp4_mla_backend.FP4_MLA_PACKED_WAVE_PLAN_ENV, raising=False) + fp4_mla_backend._launch_fp4_mla_packed_plan(geometry, kv_lens, metadata) + monkeypatch.setenv(fp4_mla_backend.FP4_MLA_PACKED_WAVE_PLAN_ENV, "0") + fp4_mla_backend._launch_fp4_mla_packed_plan(geometry, kv_lens, metadata) + assert calls == [True, False] + + def test_packed_plan_geometry_is_shared_by_cache_update_and_decode(monkeypatch) -> None: monkeypatch.delenv("TRTLLM_FP4_MLA_SPLIT_KV", raising=False) monkeypatch.setenv(fp4_mla_config.FP4_MLA_PACKED_Q_ENV, "1") From f189f034030c1b04d6d36467071bcdf98a0882a7 Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Tue, 27 Oct 2026 23:15:59 -0700 Subject: [PATCH 09/16] [None][fix] Adapt imported runtime and kernel changes to current main Keep the source implementations in the preceding commits and collect their integration changes here. Relocate the attention, MoE, executor, and KV-cache changes to the current module layout. Preserve upstream interfaces, cache lifecycle behavior, transceiver synchronization, and byte accounting. Combine shared speculative-decoding buffers and drafter mappings with the unified draft cache. Preserve layer-offset decoding and keep unified caches on the preparation path that records history positions and checks capacity. Update bindings, imports, tests, and formatting for the current interfaces. Signed-off-by: Mingyang Hao --- .../kv_cache_manager_v2/config.h | 6 +- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 6 +- .../kv_cache_manager_v2/lifeCycleRegistry.h | 38 +- .../kv_cache_manager_v2/storage/config.cpp | 23 +- cpp/tensorrt_llm/kernels/helixAllToAll.cu | 6 +- .../batch_manager/kvCacheManagerV2.cpp | 20 +- .../unit_tests/batch_manager/CMakeLists.txt | 85 +- .../kvCacheManagerV2ConfigTest.cpp | 148 +- .../backends}/dspark.py | 0 .../backends/fp4_mla/cache_manager.py | 6 +- .../attention/backends/fp4_mla/config.py | 1 + .../attention/backends/fp4_mla/packed_plan.py | 3 + .../sparse/deepseek_v4/cache_manager.py | 19 +- .../backends/sparse/dsa/cache_manager.py | 6 +- .../kernels/blackwell/dspark/attention.py | 7 +- .../blackwell/dspark/attention_kernel.py | 222 +- .../kernels/blackwell/dspark_rmsnorm_rope.py | 247 +- .../attention_backend/fp4_mla/__init__.py | 5135 ---------------- .../sparse/minimax_m3/cache_manager.py | 473 -- .../custom_ops/dspark_attention_custom_op.py | 268 +- .../dspark_rmsnorm_rope_custom_op.py | 356 +- .../disaggregation/native/bounce/impl.py | 2 +- .../native/mixers/attention/peer.py | 23 +- .../_torch/disaggregation/native/peer.py | 69 +- .../_torch/disaggregation/native/rank_info.py | 6 + .../_torch/disaggregation/native/transfer.py | 3766 +++++++----- .../disaggregation/resource/cache_reuse.py | 8 +- .../disaggregation/resource/kv_extractor.py | 28 +- .../_torch/disaggregation/resource/page.py | 4 +- .../_torch/disaggregation/transceiver.py | 239 +- tensorrt_llm/_torch/models/modeling_dflash.py | 53 +- tensorrt_llm/_torch/models/modeling_dspark.py | 1853 +++++- .../_torch/models/modeling_kimi_linear.py | 16 +- .../fused_moe/fused_moe_cute_dsl_b12x.py | 289 - .../fused_moe/fused_moe_cute_dsl_fc12.py | 264 - .../moe/fused_moe/cutedsl/fused_fc12.py | 6 + .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 2 + tensorrt_llm/_torch/pyexecutor/_util.py | 113 +- .../engine/runners/decoder/runner.py | 30 + .../kv_cache/kv_cache_manager_v2.py | 810 ++- .../kv_cache/mamba_cache_manager/common.py | 22 + .../mamba_cache_manager_v2.py | 23 +- .../_torch/pyexecutor/kv_cache_manager_v2.py | 5326 ----------------- .../_torch/pyexecutor/mamba_cache_manager.py | 4723 --------------- .../_torch/pyexecutor/model_engine.py | 899 +-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 4 + .../_torch/pyexecutor/sampler/sampler.py | 17 +- .../pyexecutor/standalone_draft_cache.py | 2 +- tensorrt_llm/_torch/speculative/dflash.py | 878 ++- tensorrt_llm/_torch/speculative/dspark.py | 611 +- tensorrt_llm/_torch/speculative/interface.py | 17 +- tensorrt_llm/_torch/speculative/utils.py | 2 - .../runtime/kv_cache_manager_v2/__init__.py | 16 +- .../runtime/kv_cache_manager_v2/__init__.pyi | 13 +- .../_torch/attention/test_fp4_mla_packed.py | 14 +- .../kv_cache/test_kv_cache_estimation.py | 44 + .../executor/test_kv_cache_estimation.py | 1010 ---- .../hw_agnostic/test_dflash_worker.py | 79 + .../hw_agnostic/test_dspark_attention.py | 680 ++- .../hw_agnostic/test_dspark_cuda_graph.py | 231 +- .../hw_agnostic/test_dspark_worker.py | 362 +- .../test_dspark_cute_dsl_attention.py | 367 +- 62 files changed, 8657 insertions(+), 21339 deletions(-) rename tensorrt_llm/_torch/{attention_backend => attention/backends}/dspark.py (100%) delete mode 100644 tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py delete mode 100644 tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py delete mode 100644 tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py delete mode 100644 tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py delete mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py delete mode 100644 tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py delete mode 100644 tests/unittest/_torch/executor/test_kv_cache_estimation.py diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index e071b04e7b6a..983073986b11 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -204,10 +204,12 @@ struct AttentionLayerConfig // nullopt or 0 = no sink tokens. std::optional numSinkTokens; - // Layers in different ownership domains must not share lifecycle or storage pools. + // Private working pages never constrain prefix matching or enter the tree. + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + CacheDomain cacheDomain = CacheDomain::kTarget; - std::optional windowSize() const noexcept + [[nodiscard]] std::optional windowSize() const noexcept { return slidingWindowSize; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index 891210d92a9c..1d798cccd7a7 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -39,7 +39,11 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) return SsmLifeCycle{}; } else - return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, cfg.cacheDomain); + { + bool const isSparse = !cfg.buffers.empty() && cfg.buffers.front().isSparse; + return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, + cfg.reusePolicy, cfg.cacheDomain); + } }, layer); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index a5ae27b78d89..6dbecf343e68 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -41,8 +41,25 @@ struct AttnLifeCycle { std::optional windowSize; // nullopt = no sliding window int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) + bool isSparse = false; + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; CacheDomain cacheDomain = CacheDomain::kTarget; + bool isPrivate() const noexcept + { + return reusePolicy == AttentionReusePolicy::PRIVATE; + } + + bool requiresPrefixCoverage() const noexcept + { + return reusePolicy == AttentionReusePolicy::REQUIRED; + } + + bool allowsCheckpoint() const noexcept + { + return reusePolicy == AttentionReusePolicy::OPTIONAL; + } + HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const { int numBlocks = divUp(historyLength, tokensPerBlock); @@ -59,16 +76,29 @@ struct AttnLifeCycle bool operator==(AttnLifeCycle const& o) const noexcept { - return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && cacheDomain == o.cacheDomain; + return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse + && reusePolicy == o.reusePolicy && cacheDomain == o.cacheDomain; } bool operator<(AttnLifeCycle const& o) const noexcept { - return std::tie(windowSize, numSinkBlocks, cacheDomain) - < std::tie(o.windowSize, o.numSinkBlocks, o.cacheDomain); + if (windowSize != o.windowSize) + { + return windowSize < o.windowSize; + } + if (numSinkBlocks != o.numSinkBlocks) + { + return numSinkBlocks < o.numSinkBlocks; + } + if (isSparse != o.isSparse) + { + return isSparse < o.isSparse; + } + return std::tie(reusePolicy, cacheDomain) < std::tie(o.reusePolicy, o.cacheDomain); } static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock, + bool isSparse = false, AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED, CacheDomain cacheDomain = CacheDomain::kTarget) { TLLM_CHECK_DEBUG(tokensPerBlock > 0); @@ -76,7 +106,7 @@ struct AttnLifeCycle TLLM_CHECK_DEBUG(!numSinkTokens.has_value() || *numSinkTokens >= 0); TLLM_CHECK_DEBUG((!numSinkTokens.has_value() || *numSinkTokens == 0) || ws.has_value()); int sinkBlocks = divUp(numSinkTokens.value_or(0), tokensPerBlock); - return AttnLifeCycle{ws, sinkBlocks, cacheDomain}; + return AttnLifeCycle{ws, sinkBlocks, isSparse, reusePolicy, cacheDomain}; } }; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index c19e63250ffd..d2bbad0920c9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -206,14 +206,31 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) slotGroups.push_back(std::move(var)); } - // Only matching ownership domains and slot layouts may share pools. - std::map>, std::vector> poolGroupsBySizes; + // Distinct sparsity and reuse policies need separate GPU pools and eviction queues. + struct PoolGroupKey + { + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + std::vector slotSizes; + bool isSparse = false; + CacheDomain cacheDomain = CacheDomain::kTarget; + + bool operator<(PoolGroupKey const& other) const + { + return std::tie(reusePolicy, slotSizes, isSparse, cacheDomain) + < std::tie(other.reusePolicy, other.slotSizes, other.isSparse, other.cacheDomain); + } + }; + + std::map> poolGroups; for (auto& sg : slotGroups) { auto const sizes = sg.slotSizeList(); auto const* attn = std::get_if(®istry.getLifeCycle(sg.lifeCycleId)); + bool const isSparse = attn != nullptr && attn->isSparse; + auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; CacheDomain const domain = attn ? attn->cacheDomain : CacheDomain::kTarget; - poolGroupsBySizes[{domain, sizes.raw()}].push_back(std::move(sg)); + poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .isSparse = isSparse, .cacheDomain = domain}] + .push_back(std::move(sg)); } StorageConfig out; diff --git a/cpp/tensorrt_llm/kernels/helixAllToAll.cu b/cpp/tensorrt_llm/kernels/helixAllToAll.cu index 9f9ed7fac6fd..acd8140aa02b 100644 --- a/cpp/tensorrt_llm/kernels/helixAllToAll.cu +++ b/cpp/tensorrt_llm/kernels/helixAllToAll.cu @@ -542,9 +542,9 @@ __global__ void helixAllToAllKernel(HelixAllToAllParams params) LL128Proto::protoUnpack( shmem, tail, singlePacked128ByteCount, fifoEntry128ByteIndexBase, loaded128ByteCount, laneId); - // protoUnpack restores the u64 under each line's flag with ordinary - // stores, and the bulk store below reads shared memory through the - // async proxy: without this fence it can still read the flag value. + // protoUnpack restores the words displaced by the LL128 flags through the generic proxy. Make + // those writes visible to the async proxy before cp.async.bulk reads the same buffer for S2G; + // otherwise the store can copy a flag (the FIFO step) in place of a restored data word. tensorrt_llm::common::fence_view_async_shared(); __syncwarp(); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 42a585fab383..3ac482b6f064 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1411,12 +1411,20 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("OPTIONAL", kv::AttentionReusePolicy::OPTIONAL); nb::class_(m, "AttnLifeCycle") - .def(nb::init, int, kv::CacheDomain>(), nb::arg("window_size").none(), - nb::arg("num_sink_blocks"), nb::arg("cache_domain") = kv::CacheDomain::kTarget) + .def(nb::init, int, bool, kv::AttentionReusePolicy, kv::CacheDomain>(), + nb::arg("window_size").none(), nb::arg("num_sink_blocks"), nb::arg("is_sparse") = false, + nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, + nb::arg("cache_domain") = kv::CacheDomain::kTarget) + // Sink tokens round up to whole blocks. Bound rather than repeated in Python so the + // connector's view of a life cycle is built by the same code as the allocator's. .def_static("make", &kv::AttnLifeCycle::make, nb::arg("window_size").none(), nb::arg("num_sink_tokens").none(), - nb::arg("tokens_per_block"), nb::arg("cache_domain") = kv::CacheDomain::kTarget) + nb::arg("tokens_per_block"), nb::arg("is_sparse") = false, + nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, + nb::arg("cache_domain") = kv::CacheDomain::kTarget) .def_prop_ro("window_size", [](kv::AttnLifeCycle const& self) { return self.windowSize; }) .def_ro("num_sink_blocks", &kv::AttnLifeCycle::numSinkBlocks) + .def_ro("is_sparse", &kv::AttnLifeCycle::isSparse) + .def_ro("reuse_policy", &kv::AttnLifeCycle::reusePolicy) .def_ro("cache_domain", &kv::AttnLifeCycle::cacheDomain) .def("get_stale_range", &kv::AttnLifeCycle::getStaleRange, nb::arg("history_length"), nb::arg("tokens_per_block")) @@ -1626,13 +1634,15 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, - kv::CacheDomain>(), + kv::AttentionReusePolicy, kv::CacheDomain>(), nb::arg("layer_id"), nb::arg("buffers"), nb::arg("sliding_window_size") = std::nullopt, - nb::arg("num_sink_tokens") = std::nullopt, nb::arg("cache_domain") = kv::CacheDomain::kTarget) + nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, + nb::arg("cache_domain") = kv::CacheDomain::kTarget) .def_rw("layer_id", &kv::AttentionLayerConfig::layerId) .def_rw("buffers", &kv::AttentionLayerConfig::buffers) .def_rw("sliding_window_size", &kv::AttentionLayerConfig::slidingWindowSize) .def_rw("num_sink_tokens", &kv::AttentionLayerConfig::numSinkTokens) + .def_rw("reuse_policy", &kv::AttentionLayerConfig::reusePolicy) .def_rw("cache_domain", &kv::AttentionLayerConfig::cacheDomain) .def_prop_ro("window_size", &kv::AttentionLayerConfig::windowSize) DEF_COPY(kv::AttentionLayerConfig); diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 30b6ea318e39..0f2879b1dd77 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -290,18 +290,73 @@ add_tllm_gtest( add_tllm_gtest( kvCacheManagerV2DigestPoolTest - PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager - ${PROJECT_SOURCE_DIR}/tensorrt_llm/common/sha256) -add_gtest(kvCacheUtilsTest kvCacheUtilsTest.cpp) -add_gtest(llmRequestTest llmRequestTest.cpp) -add_gtest(microBatchSchedulerTest microBatchSchedulerTest.cpp) -add_gtest(peftCacheManagerTest peftCacheManagerTest.cpp) -add_gtest(staticThreadPoolTest staticThreadPoolTest.cpp) -add_gtest(rnnCacheFormatterTest rnnCacheFormatterTest.cpp) -add_gtest(agentTreeTest agentTreeTest.cpp) -add_gtest(truncateBlocksTest truncateBlocksTest.cpp) - -add_gtest(kvCacheManagerV2ConfigTest kvCacheManagerV2ConfigTest.cpp) -target_include_directories( - kvCacheManagerV2ConfigTest - PRIVATE ${PROJECT_SOURCE_DIR}/tensorrt_llm/batch_manager) + OWNER + batch-manager + SOURCES + kvCacheManagerV2DigestPoolTest.cpp + LINK_LIBRARIES + tllm::kv_cache + $ + LABELS + "unit;batch-manager;host") + +add_tllm_gtest( + kvCacheUtilsTest + OWNER + batch-manager + SOURCES + kvCacheUtilsTest.cpp + LINK_LIBRARIES + tllm::kv_cache_legacy + LABELS + "unit;batch-manager;gpu") + +add_tllm_gtest( + llmRequestTest + OWNER + batch-manager + SOURCES + llmRequestTest.cpp + LINK_LIBRARIES + tllm::kv_cache_legacy + LABELS + "unit;batch-manager;gpu") + +add_tllm_gtest( + microBatchSchedulerTest + OWNER + batch-manager + SOURCES + microBatchSchedulerTest.cpp + LINK_LIBRARIES + tllm::batch_scheduler + LABELS + "unit;batch-manager;gpu") + +add_tllm_gtest(peftCacheManagerTest SOURCES peftCacheManagerTest.cpp FULL_STACK + LABELS "unit;full-stack;gpu") + +add_tllm_gtest(rnnCacheFormatterTest SOURCES rnnCacheFormatterTest.cpp + FULL_STACK LABELS "unit;full-stack;host") + +add_tllm_gtest( + agentTreeTest + OWNER + batch-manager + SOURCES + agentTreeTest.cpp + LINK_LIBRARIES + tllm::batch_scheduler + LABELS + "unit;batch-manager;host") + +add_tllm_gtest( + truncateBlocksTest + OWNER + batch-manager + SOURCES + truncateBlocksTest.cpp + LINK_LIBRARIES + tllm::kv_cache_legacy + LABELS + "unit;batch-manager;gpu") diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index ef9b0d7ab3df..bb43f5c618e8 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp @@ -15,25 +15,165 @@ * limitations under the License. */ -#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h" -#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.h" +#include "kvCacheManagerV2TestUtils.h" +#include "kv_cache_manager_v2/lifeCycleRegistry.h" +#include "kv_cache_manager_v2/storage/config.h" #include +#include +#include +#include + namespace { + using namespace tensorrt_llm::batch_manager::kv_cache_manager_v2; +using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeConfig; +using tensorrt_llm::batch_manager::kv_cache_manager_v2::test::makeTieredConfig; + +AttentionLayerConfig makeAttentionLayer(LayerId layerId, bool isSparse, size_t bufferSize = 4096) +{ + return AttentionLayerConfig{ + .layerId = layerId, .buffers = {BufferConfig{.role = "key", .size = bufferSize, .isSparse = isSparse}}}; +} + +TEST(KvCacheManagerV2ConfigTest, ExistingBufferConfigurationDefaultsToDense) +{ + BufferConfig const buffer{"key", 4096, 2}; + EXPECT_FALSE(buffer.isSparse); + EXPECT_EQ(buffer.tokensPerBlockOverride, 2); + EXPECT_NO_THROW(makeConfig().validate()); +} + +TEST(KvCacheManagerV2ConfigTest, SparseBuffersRequireHostMemoryAtLevelOne) +{ + auto config = makeConfig(); + config.layers = {makeAttentionLayer(0, true)}; + EXPECT_THROW(config.validate(), std::invalid_argument); + + config.cacheTiers.emplace_back(DiskCacheTierConfig{4 << 20, "/tmp"}); + EXPECT_THROW(config.validate(), std::invalid_argument); + + config.cacheTiers[1] = GpuCacheTierConfig{4 << 20}; + config.cacheTiers.emplace_back(HostCacheTierConfig{4 << 20}); + EXPECT_THROW(config.validate(), std::invalid_argument); + + config.cacheTiers[1] = HostCacheTierConfig{4 << 20}; + EXPECT_NO_THROW(config.validate()); +} + +TEST(KvCacheManagerV2ConfigTest, MixedBuffersInOneLayerAreRejected) +{ + auto config = makeTieredConfig(); + for (bool const firstIsSparse : {false, true}) + { + auto layer = makeAttentionLayer(0, firstIsSparse); + layer.buffers.push_back(BufferConfig{.role = "value", .size = 4096, .isSparse = !firstIsSparse}); + config.layers = {layer}; + EXPECT_THROW(config.validate(), std::invalid_argument); + EXPECT_THROW(createStorageConfig(config), std::invalid_argument); + } +} + +TEST(KvCacheManagerV2ConfigTest, SparseSsmBuffersAreRejected) +{ + auto config = makeTieredConfig(); + config.commitMinSnapshot = true; + config.layers + = {SsmLayerConfig{.layerId = 0, .buffers = {BufferConfig{.role = "state", .size = 4096, .isSparse = true}}}}; + EXPECT_THROW(config.validate(), std::invalid_argument); + EXPECT_THROW(createStorageConfig(config), std::invalid_argument); +} + +TEST(KvCacheManagerV2ConfigTest, SparseAndDenseBuffersHaveSeparateLifecyclesAndPools) +{ + auto config = makeTieredConfig(); + config.layers = {makeAttentionLayer(0, false), makeAttentionLayer(1, true), makeAttentionLayer(2, false), + makeAttentionLayer(3, true)}; + ASSERT_NO_THROW(config.validate()); + + LifeCycleRegistry const registry(config); + auto const storage = createStorageConfig(config); + auto const attributes = storage.bufferAttributes(); + auto const grouping = storage.lifeCycleGrouping(); + auto const& dense = attributes.at(BufferId{0, "key"}); + auto const& sparse = attributes.at(BufferId{1, "key"}); + + EXPECT_EQ(registry.size(), LifeCycleId{2}); + EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{2}); + EXPECT_NE(dense.lifeCycleId, sparse.lifeCycleId); + EXPECT_NE(grouping[dense.lifeCycleId], grouping[sparse.lifeCycleId]); + EXPECT_FALSE(std::get(registry[dense.lifeCycleId]).isSparse); + EXPECT_TRUE(std::get(registry[sparse.lifeCycleId]).isSparse); + EXPECT_FALSE(registry[dense.lifeCycleId] == registry[sparse.lifeCycleId]); + + auto const& otherDense = attributes.at(BufferId{2, "key"}); + auto const& otherSparse = attributes.at(BufferId{3, "key"}); + EXPECT_EQ(dense.lifeCycleId, otherDense.lifeCycleId); + EXPECT_EQ(sparse.lifeCycleId, otherSparse.lifeCycleId); + EXPECT_EQ(sparse.poolIndex, otherSparse.poolIndex); + EXPECT_EQ(otherSparse.offset, sparse.offset + sparse.size); + EXPECT_EQ(otherDense.offset, dense.offset + dense.size); +} + +TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsWithinEachSparsity) +{ + auto config = makeTieredConfig(); + config.layers.clear(); + for (LayerId layerId = 0; layerId < 4; ++layerId) + { + auto layer = makeAttentionLayer(layerId, layerId >= 2); + layer.slidingWindowSize = layerId % 2 == 0 ? 128 : 256; + config.layers.push_back(layer); + } + ASSERT_NO_THROW(config.validate()); + + auto const storage = createStorageConfig(config); + auto const attributes = storage.bufferAttributes(); + auto const grouping = storage.lifeCycleGrouping(); + auto poolGroup = [&](LayerId layerId) { return grouping[attributes.at(BufferId{layerId, "key"}).lifeCycleId]; }; + EXPECT_EQ(storage.numLifeCycles(), LifeCycleId{4}); + EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{2}); + EXPECT_EQ(poolGroup(0), poolGroup(1)); + EXPECT_EQ(poolGroup(2), poolGroup(3)); + EXPECT_NE(poolGroup(0), poolGroup(2)); +} + +TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceByExpandedSize) +{ + auto config = makeTieredConfig(); + auto firstLayer = makeAttentionLayer(0, true, 2048); + firstLayer.buffers.front().tokensPerBlockOverride = 2; + config.layers = {firstLayer, makeAttentionLayer(1, true, 4096)}; + ASSERT_NO_THROW(config.validate()); + + auto const storage = createStorageConfig(config); + auto const attributes = storage.bufferAttributes(); + auto const& first = attributes.at(BufferId{0, "key"}); + auto const& second = attributes.at(BufferId{1, "key"}); + EXPECT_EQ(storage.numLifeCycles(), LifeCycleId{1}); + EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{1}); + EXPECT_EQ(first.lifeCycleId, second.lifeCycleId); + EXPECT_EQ(first.poolIndex, second.poolIndex); + EXPECT_EQ(first.size, second.size); + EXPECT_EQ(second.offset, first.offset + first.size); + EXPECT_EQ(first.expansion, 2); + EXPECT_EQ(second.expansion, 1); +} TEST(KvCacheManagerV2ConfigTest, AttentionLifecycleFactoryPreservesCacheDomain) { auto const target = AttnLifeCycle::make(128, 5, 4); - auto const draft = AttnLifeCycle::make(128, 5, 4, CacheDomain::kStandaloneDraft); + auto const draft + = AttnLifeCycle::make(128, 5, 4, false, AttentionReusePolicy::REQUIRED, CacheDomain::kStandaloneDraft); EXPECT_EQ(target.numSinkBlocks, 2); EXPECT_EQ(target.cacheDomain, CacheDomain::kTarget); EXPECT_EQ(draft.cacheDomain, CacheDomain::kStandaloneDraft); EXPECT_FALSE(target == draft); EXPECT_TRUE(target < draft || draft < target); - EXPECT_TRUE(draft == AttnLifeCycle::make(128, 8, 4, CacheDomain::kStandaloneDraft)); + EXPECT_TRUE( + draft == AttnLifeCycle::make(128, 8, 4, false, AttentionReusePolicy::REQUIRED, CacheDomain::kStandaloneDraft)); } TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsOnlyWithinEachCacheDomain) diff --git a/tensorrt_llm/_torch/attention_backend/dspark.py b/tensorrt_llm/_torch/attention/backends/dspark.py similarity index 100% rename from tensorrt_llm/_torch/attention_backend/dspark.py rename to tensorrt_llm/_torch/attention/backends/dspark.py diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index e19dae50c301..76d9d5806be0 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -326,8 +326,8 @@ def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: local_layer = self.layer_offsets[layer_idx] if self._is_standalone_draft_layer(local_layer): return self.dtype_per_layer[local_layer] - if local_layer in self._fp8_mla_local_layer_indices(): - return DataType.FP8 + if local_layer in self._bf16_mla_local_layer_indices(): + return DataType.BF16 return self.dtype def get_kv_cache_num_blocks(self, layer_idx: int) -> int: @@ -1073,7 +1073,7 @@ def get_cache_size_per_token( * num_bf16_layers ) fixed_bytes = hp_bytes * max_batch_size * mapping.pp_size - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _estimate_draft_cache_size_components, _get_draft_layer_sizes, _get_generation_kv_capacity, diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py index d3d89d640f97..b2f13e20d0ad 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/config.py @@ -137,6 +137,7 @@ def _fp4_mla_cutedsl_fused_v_transpose_enabled() -> bool: _FP4_MLA_PACKED_MAX_SPLITS = 16 _FP4_MLA_PACKED_MIN_TILES_PER_SPLIT = 2 _FP4_MLA_SPLIT_KV_ENV = "TRTLLM_FP4_MLA_SPLIT_KV" +FP4_MLA_PACKED_WAVE_PLAN_ENV = "TRTLLM_FP4_MLA_PACKED_WAVE_PLAN" def _fp4_mla_packed_q_enabled() -> bool: diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py index adf3b723cfcb..1a0725ea46c9 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py @@ -11,8 +11,10 @@ _FP4_MLA_PACKED_MAX_SPLITS, _FP4_MLA_PACKED_MIN_TILES_PER_SPLIT, _FP4_MLA_SPLIT_KV_ENV, + FP4_MLA_PACKED_WAVE_PLAN_ENV, FP4_MLA_SCALE_ROW_GROUP, _ceil_div, + _env_enabled_default, _fp4_mla_cutedsl_kernel_module, _fp4_mla_packed_rows_per_seq, ) @@ -227,5 +229,6 @@ def _launch_fp4_mla_packed_plan( query_len=geometry.query_len, helix_kv_bounds=geometry.helix_kv_bounds, generation_lengths=generation_lengths, + wave_plan=_env_enabled_default(FP4_MLA_PACKED_WAVE_PLAN_ENV, True), ) metadata.fp4_mla_state.packed_plan_key = geometry.key diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index 492e965e80c4..784b02992fcd 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -21,13 +21,19 @@ import torch from tensorrt_llm._torch.pyexecutor import llm_request -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + _GUARD_PAGE_REQUEST_ID, + _RESERVED_REQUEST_IDS, + CUDA_GRAPH_DUMMY_REQUEST_ID, GPU_LEVEL, KVCacheManagerV2, _estimate_draft_cache_size_components, + _fill_kv_pages, _get_draft_layer_sizes, _get_generation_kv_capacity, ) +from tensorrt_llm._torch.pyexecutor.resource_manager import scan_kv_buffer_for_invalid_values +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import ( TensorWrapper, convert_to_torch_tensor, @@ -950,7 +956,16 @@ def _get_cache_cost_components(self) -> tuple[int, int, int, int]: """Return context/generation bytes per token, then fixed bytes per request.""" compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] has_fp8_kv_cache = self.dtype == DataType.FP8 - non_sliding_attn_size_per_token = self.get_cache_bytes_per_token() + non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + has_nvfp4_compress=self._use_nvfp4_compress, + nvfp4_residual_dim=self.nvfp4_residual_dim, + ) ( context_swa_size_per_token, context_swa_size_per_request, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py index 4dc3a34a7c8e..9a1c8820758a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py @@ -688,8 +688,8 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) - cache_bytes += self.mla_kv_cache_residual_dim // 16 if ( data_role == Role.ALL - and local_layer_idx < len(self.indexer_k_cache_local_layer_mask) - and self.indexer_k_cache_local_layer_mask[local_layer_idx] + and not self._is_standalone_draft_layer(local_layer_idx) + and self._indexer_k_cache_layer_mask[self.pp_layers[local_layer_idx]] ): cache_bytes += self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY) return cache_bytes @@ -727,7 +727,7 @@ def get_cache_size_per_token( model_config, mapping, num_layers ) cache_bytes += num_attention_layers * (residual_dim // 2 + math.ceil(residual_dim / 16)) - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( _estimate_draft_cache_size_components, _get_draft_layer_sizes, _get_generation_kv_capacity, diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py index 15902ef1cb4e..6b29cb9c2871 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py @@ -5,7 +5,6 @@ from typing import Optional, Tuple, Type -import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cuda.bindings import driver as cuda @@ -273,6 +272,12 @@ def __init__( raise ValueError("DSpark history page size must be 16, 32, 64, 128, or 256") if inverse_rope_dim not in (0, 64): raise ValueError(f"DSpark inverse_rope_dim must be 0 or 64, got {inverse_rope_dim}") + if output_groups is not None and ( + output_groups <= 0 or self.num_heads % output_groups != 0 + ): + raise ValueError( + f"DSpark output groups must divide {self.num_heads}, got {output_groups}" + ) super().__init__( acc_dtype, diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py index 4e0edae9bbf3..1f38601547bd 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py @@ -37,9 +37,10 @@ import cutlass.pipeline as pipeline import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass._mlir.dialects import llvm from cutlass.cute.arch import Arch from cutlass.cute.nvgpu import OperandMajorMode, tcgen05 -from cutlass.cutlass_dsl import BaseDSL +from cutlass.cutlass_dsl import BaseDSL, Int32, T, dsl_user_op from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait try: @@ -68,6 +69,73 @@ LOG2_E = 1.4426950408889634074 +@dsl_user_op +def _store_global_b64(gmem_ptr, value_lo: Int32, value_hi: Int32, *, loc=None, ip=None) -> None: + gmem_addr = gmem_ptr.toint(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ + gmem_addr.ir_value(loc=loc, ip=ip), + Int32(value_lo).ir_value(loc=loc, ip=ip), + Int32(value_hi).ir_value(loc=loc, ip=ip), + ], + "{\n .reg .b64 packed;\n mov.b64 packed, {$1, $2};\n st.global.u64 [$0], packed;\n}", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _load_shared_b32(smem_ptr, *, loc=None, ip=None) -> Int32: + smem_addr = smem_ptr.toint(loc=loc, ip=ip) + return Int32( + llvm.inline_asm( + T.i32(), + [smem_addr.ir_value(loc=loc, ip=ip)], + "ld.shared.b32 $0, [$1];", + "=r,r", + has_side_effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _store_shared_b128( + smem_ptr, + value0: Int32, + value1: Int32, + value2: Int32, + value3: Int32, + *, + loc=None, + ip=None, +) -> None: + smem_addr = smem_ptr.toint(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ + smem_addr.ir_value(loc=loc, ip=ip), + Int32(value0).ir_value(loc=loc, ip=ip), + Int32(value1).ir_value(loc=loc, ip=ip), + Int32(value2).ir_value(loc=loc, ip=ip), + Int32(value3).ir_value(loc=loc, ip=ip), + ], + "st.shared.v4.b32 [$0], {$1, $2, $3, $4};", + "r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + class DSparkAttentionKernel: # The DSpark accumulator/correction layouts occupy a fixed 512-column # TMEM footprint on every supported architecture. @@ -222,6 +290,7 @@ def __call__( c_latent_draft: cute.Tensor, page_table_win: cute.Tensor, o: cute.Tensor, + o_scale: Optional[cute.Tensor], cache_seqs: Optional[cute.Tensor], window_valid_lens: Optional[cute.Tensor], capacities: cute.Tensor, @@ -244,7 +313,7 @@ def __call__( self.q_dtype = q_latent.element_type self.k_dtype = c_latent_draft.element_type self.v_dtype = c_latent_draft.element_type - self.o_dtype = o.element_type + self.o_dtype = cutlass.BFloat16 # check type consistency if cutlass.const_expr( @@ -260,8 +329,18 @@ def __call__( raise ValueError("c_latent_draft must have leading dimension 1") if cutlass.const_expr(c_latent_win.stride[1] != 1): raise ValueError("c_latent_win must have leading dimension 1") - if cutlass.const_expr(o.stride[1] != 1): - raise ValueError("o must have leading dimension 1") + if cutlass.const_expr(self.output_groups is None): + if cutlass.const_expr(o.element_type is not cutlass.BFloat16 or o.stride[1] != 1): + raise ValueError("o must be leading-dimension-1 BF16") + else: + if cutlass.const_expr( + o.element_type is not cutlass.Float8E4M3FN + or o_scale is None + or o_scale.element_type is not cutlass.Float32 + or o.stride[2] != 1 + or o_scale.stride[2] != 1 + ): + raise ValueError("quantized o must be contiguous FP8 with contiguous FP32 scales") c_latent_draft_transpose_layout = cute.select(c_latent_draft.layout, mode=[1, 0, 2]) c_latent_draft_transpose = cute.make_tensor( @@ -453,7 +532,7 @@ def __call__( self.tma_copy_kc_bytes = kc_copy_size tile_sched_params, grid = self._compute_grid( - o, self.cluster_shape_mnk, self.max_active_clusters + q_latent, self.cluster_shape_mnk, self.max_active_clusters ) @cute.struct @@ -478,6 +557,17 @@ class AttentionKernelSharedStorage: epilogue_smem_exchange: cute.struct.MemRange[ self.acc_dtype, self.num_compute_warps * self.threads_per_warp ] + fp8_epilogue_scratch: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Int32, + ( + self.num_compute_warps * self.threads_per_warp * 32 + if self.output_groups is not None + else 1 + ), + ], + 128, + ] smem_q_latent: cute.struct.Align[ cute.struct.MemRange[self.q_dtype, cute.cosize(q_latent_smem_layout_staged)], 1024, @@ -516,6 +606,7 @@ class AttentionKernelSharedStorage: tma_tensor_c_latent_transpose_draft, page_table_win, o, + o_scale, cache_seqs, window_valid_lens, capacities, @@ -592,6 +683,7 @@ def kernel( mCLT_draft: cute.Tensor, mPT_win: cute.Tensor, mO: Optional[cute.Tensor], + mOScale: Optional[cute.Tensor], cache_seqs: cute.Tensor, window_valid_lens: cute.Tensor, capacities: cute.Tensor, @@ -715,6 +807,14 @@ def kernel( epilogue_smem_exchange = storage.epilogue_smem_exchange.get_tensor( cute.make_layout(self.num_compute_warps * self.threads_per_warp) ) + fp8_epilogue_scratch = None + if cutlass.const_expr(self.output_groups is not None): + fp8_epilogue_scratch = storage.fp8_epilogue_scratch.get_tensor( + cute.make_layout( + (self.num_compute_warps, self.threads_per_warp, 32), + stride=(self.threads_per_warp * 32, 32, 1), + ) + ) # # Cluster wait before tensor memory alloc @@ -937,7 +1037,7 @@ def kernel( compute_common_params = SimpleNamespace( blk_coord=blk_coord, smem_exchange=softmax_smem_exchange, - mO=mO, + mO=mQL if cutlass.const_expr(self.output_groups is not None) else mO, K=self.get_cache_seq_len(cache_seqs, blk_coord[2]), window_valid_len=self.get_window_valid_len( window_valid_lens, capacities, cache_seqs, blk_coord[2] @@ -992,7 +1092,7 @@ def kernel( compute_common_params = SimpleNamespace( blk_coord=blk_coord, smem_exchange=epilogue_smem_exchange, - mO=mO, + mO=mQL if cutlass.const_expr(self.output_groups is not None) else mO, K=self.get_cache_seq_len(cache_seqs, blk_coord[2]), L=mCL_draft.shape[1], H=mQL.shape[0], @@ -1006,6 +1106,9 @@ def kernel( output_scale=output_scale, softmax_scale_log2=softmax_scale_log2, mFreqs=inverse_rope_freqs, + mOQuant=mO, + mOScale=mOScale, + sOQuant=fp8_epilogue_scratch, ) p_cor_consumer_state, mma_o_consumer_state = self.correction( compute_common_params, @@ -1970,7 +2073,10 @@ def softmax( # reduction for row_max row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) - elif cutlass.const_expr(arch >= Arch.sm_103 and arch <= Arch.sm_103f): + elif cutlass.const_expr( + (arch >= Arch.sm_103 and arch <= Arch.sm_103f) + or (arch >= Arch.sm_107 and arch <= Arch.sm_107f) + ): tmem_load_red_atom = cute.make_copy_atom( tcgen05.copy.LdRed32x32bOp( tcgen05.copy.Repetition(64), redOp=tcgen05.TmemLoadRedOp.MAX @@ -2368,12 +2474,102 @@ def epilogue( tR2G_rO_src[i] = (re * cos + im * sin).to(self.o_dtype) tR2G_rO_src[i + 1] = (im * cos - re * sin).to(self.o_dtype) + if cutlass.const_expr(self.output_groups is not None and cute.size(tR2G_rO_src) != 128): + raise ValueError("DSpark FP8 output requires one 128-value block per thread") + if cute.elem_less(tTR_cO[0][0], common_params.H): - cute.autovec_copy( - tR2G_rO_src, - tR2G_rO_dst, - l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, - ) + if cutlass.const_expr(self.output_groups is None): + cute.autovec_copy( + tR2G_rO_src, + tR2G_rO_dst, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.NO_ALLOCATE, + ) + else: + # Each correction thread owns one complete 1x128 quant + # block. Use eight independent chains so the amax is not + # serialized behind 128 dependent fmax instructions. + partial_absmax = cute.make_rmem_tensor((8,), cutlass.Float32) + for i in cutlass.range_constexpr(8): + partial_absmax[i] = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(0, cute.size(tR2G_rO_src), 8): + for lane in cutlass.range_constexpr(8): + value = cutlass.Float32(tR2G_rO_src[i + lane]) + partial_absmax[lane] = cute.arch.fmax( + partial_absmax[lane], + cute.arch.fmax(value, -value), + ) + for i in cutlass.range_constexpr(4): + partial_absmax[i] = cute.arch.fmax(partial_absmax[i], partial_absmax[i + 4]) + for i in cutlass.range_constexpr(2): + partial_absmax[i] = cute.arch.fmax(partial_absmax[i], partial_absmax[i + 2]) + block_absmax = cute.arch.fmax(partial_absmax[0], partial_absmax[1]) + block_absmax = cute.arch.fmax(block_absmax, cutlass.Float32(1e-10)) + quant_scale = cutlass.Float32(448.0) / block_absmax + + head = tTR_cO[0][0] + group_size = self.num_heads // self.output_groups + group = head // group_size + row = common_params.blk_coord[2] * self.seq_len_q + common_params.blk_coord[1] + scale_block = ((head % group_size) * self.latent_dim + tTR_cO[0][1]) // 128 + epilogue_params.mOScale[group, scale_block, row] = ( + cutlass.Float32(1.0) / quant_scale + ) + quantized = cute.make_rmem_tensor( + (cute.size(tR2G_rO_src),), cutlass.Float8E4M3FN + ) + for i in cutlass.range( + cute.size(tR2G_rO_src), vectorize=True, unroll_full=True + ): + quantized[i] = (cutlass.Float32(tR2G_rO_src[i]) * quant_scale).to( + cutlass.Float8E4M3FN + ) + packed_output = cute.recast_tensor(quantized, cutlass.Int32) + lane_idx = tidx % self.threads_per_warp + epilogue_warp_idx = tidx // self.threads_per_warp + scratch_row = cute.domain_offset( + (epilogue_warp_idx, lane_idx, 0), epilogue_params.sOQuant + ).iterator + for i in cutlass.range_constexpr(0, 32, 4): + _store_shared_b128( + scratch_row + i, + packed_output[i], + packed_output[i + 1], + packed_output[i + 2], + packed_output[i + 3], + ) + + # Four adjacent lanes cooperatively drain one scratch row, + # producing fully coalesced 32-byte global sectors. + cute.arch.sync_warp() + row_group_idx = lane_idx // 4 + row_lane_idx = lane_idx % 4 + warp_head_base = head - lane_idx + for row_batch_idx in cutlass.range_constexpr(4): + output_row = row_batch_idx * 8 + row_group_idx + output_head = warp_head_base + output_row + output_group = output_head // group_size + output_group_k = (output_head % group_size) * self.latent_dim + tTR_cO[0][1] + source_row = cute.domain_offset( + (epilogue_warp_idx, output_row, 0), + epilogue_params.sOQuant, + ).iterator + for segment_idx in cutlass.range_constexpr(4): + word_offset = segment_idx * 8 + row_lane_idx * 2 + value_lo = _load_shared_b32(source_row + word_offset) + value_hi = _load_shared_b32(source_row + word_offset + 1) + _store_global_b64( + cute.domain_offset( + ( + output_group, + row, + output_group_k + segment_idx * 32 + row_lane_idx * 8, + ), + epilogue_params.mOQuant, + ).iterator, + value_lo, + value_hi, + ) + cute.arch.sync_warp() cute.arch.fence_view_async_tmem_load() common_params.mma_o_pipeline.consumer_release(mma_o_consumer_state) diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py index b35be0314cc1..2ce278e5d0dd 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py @@ -171,6 +171,174 @@ def kernel( output[row, out_imag] = (imag * cos + real * sin).to(output.element_type) +class DSparkRMSNormRoPECacheWriteKernel(DSparkRMSNormRoPEKernel): + """Apply RMSNorm/RoPE and scatter each row into the rolling KV cache.""" + + def __init__(self, hidden_dim: int, rope_dim: int, eps: float, window_size: int): + super().__init__(hidden_dim, rope_dim, 1, eps, True, True, False) + self.window_size = window_size + + @cute.jit + def __call__( + self, + x: cute.Tensor, + weight: cute.Tensor, + freqs: cute.Tensor, + kv_cache: cute.Tensor, + slots: cute.Tensor, + start_pos: cute.Tensor, + slots_i32: cute.Tensor, + cache_seqs: cute.Tensor, + stream: cuda.CUstream, + ): + self.kernel( + x, + weight, + freqs, + kv_cache, + slots, + start_pos, + slots_i32, + cache_seqs, + ).launch( + grid=[x.shape[0], 1, 1], + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x: cute.Tensor, + weight: cute.Tensor, + freqs: cute.Tensor, + kv_cache: cute.Tensor, + slots: cute.Tensor, + start_pos: cute.Tensor, + slots_i32: cute.Tensor, + cache_seqs: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + row, _, _ = cute.arch.block_idx() + slot = cutlass.Int32(slots[row]) + cache_seq = cutlass.Int32(start_pos[row]) + cache_pos = cache_seq % self.window_size + + if tidx == 0: + slots_i32[row] = slot + cache_seqs[row] = cache_seq + + sum_sq = cutlass.Float32(0.0) + for item in cutlass.range_constexpr(self.elements_per_thread): + dim = tidx + item * self.num_threads + value = cutlass.Float32(x[row, dim]) + sum_sq += value * value + sum_sq = cute.arch.warp_reduction_sum(sum_sq) + inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) + + for item in cutlass.range_constexpr(self.nope_elements_per_thread): + dim = tidx + item * self.num_threads + value = cutlass.Float32(x[row, dim]) * inverse_rms + value *= cutlass.Float32(weight[dim]) + kv_cache[slot, cache_pos, dim] = value.to(kv_cache.element_type) + + for item in cutlass.range_constexpr(self.pairs_per_thread): + pair = tidx + item * self.num_threads + real_dim = self.nope_dim + pair * 2 + imag_dim = real_dim + 1 + real = cutlass.Float32(x[row, real_dim]) * inverse_rms + imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms + real *= cutlass.Float32(weight[real_dim]) + imag *= cutlass.Float32(weight[imag_dim]) + cos = cutlass.Float32(freqs[row, pair, 0]) + sin = cutlass.Float32(freqs[row, pair, 1]) + kv_cache[slot, cache_pos, real_dim] = (real * cos - imag * sin).to( + kv_cache.element_type + ) + kv_cache[slot, cache_pos, imag_dim] = (imag * cos + real * sin).to( + kv_cache.element_type + ) + + +class DSparkRMSNormRoPEDraftBlockKernel(DSparkRMSNormRoPEKernel): + """Apply RMSNorm/RoPE and materialize a fixed-size zero-padded draft block.""" + + def __init__( + self, + hidden_dim: int, + rope_dim: int, + eps: float, + block_size: int, + storage_size: int, + ): + super().__init__(hidden_dim, rope_dim, 1, eps, True, True, False) + self.block_size = block_size + self.storage_size = storage_size + + @cute.jit + def __call__( + self, + x: cute.Tensor, + weight: cute.Tensor, + freqs: cute.Tensor, + output: cute.Tensor, + stream: cuda.CUstream, + ): + self.kernel(x, weight, freqs, output).launch( + grid=[output.shape[0], self.storage_size, 1], + block=[self.num_threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x: cute.Tensor, + weight: cute.Tensor, + freqs: cute.Tensor, + output: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + batch, page_row, _ = cute.arch.block_idx() + + if page_row < self.block_size: + row = batch * self.block_size + page_row + sum_sq = cutlass.Float32(0.0) + for item in cutlass.range_constexpr(self.elements_per_thread): + dim = tidx + item * self.num_threads + value = cutlass.Float32(x[row, dim]) + sum_sq += value * value + sum_sq = cute.arch.warp_reduction_sum(sum_sq) + inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) + + for item in cutlass.range_constexpr(self.nope_elements_per_thread): + dim = tidx + item * self.num_threads + value = cutlass.Float32(x[row, dim]) * inverse_rms + value *= cutlass.Float32(weight[dim]) + output[batch, page_row, dim] = value.to(output.element_type) + + for item in cutlass.range_constexpr(self.pairs_per_thread): + pair = tidx + item * self.num_threads + real_dim = self.nope_dim + pair * 2 + imag_dim = real_dim + 1 + real = cutlass.Float32(x[row, real_dim]) * inverse_rms + imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms + real *= cutlass.Float32(weight[real_dim]) + imag *= cutlass.Float32(weight[imag_dim]) + cos = cutlass.Float32(freqs[row, pair, 0]) + sin = cutlass.Float32(freqs[row, pair, 1]) + output[batch, page_row, real_dim] = (real * cos - imag * sin).to( + output.element_type + ) + output[batch, page_row, imag_dim] = (imag * cos + real * sin).to( + output.element_type + ) + else: + for item in cutlass.range_constexpr(self.elements_per_thread): + dim = tidx + item * self.num_threads + output[batch, page_row, dim] = cutlass.BFloat16(0.0) + + class DSparkRMSNormFP8QuantKernel: """Apply RMSNorm and emit packed 1x128 MXFP8 activations.""" @@ -442,82 +610,3 @@ def kernel( kv_cache[cutlass.Int64(page), cache_pos, imag_dim] = (imag * cos + real * sin).to( kv_cache.element_type ) - - -class DSparkRMSNormRoPEDraftBlockKernel(DSparkRMSNormRoPEKernel): - """Apply RMSNorm/RoPE and materialize a zero-padded physical KV page.""" - - def __init__( - self, - hidden_dim: int, - rope_dim: int, - eps: float, - block_size: int, - storage_size: int, - ): - super().__init__(hidden_dim, rope_dim, 1, eps, True, True, False) - self.block_size = block_size - self.storage_size = storage_size - - @cute.jit - def __call__( - self, - x: cute.Tensor, - weight: cute.Tensor, - freqs: cute.Tensor, - output: cute.Tensor, - stream: cuda.CUstream, - ): - self.kernel(x, weight, freqs, output).launch( - grid=[output.shape[0], self.storage_size, 1], - block=[self.num_threads, 1, 1], - stream=stream, - ) - - @cute.kernel - def kernel( - self, - x: cute.Tensor, - weight: cute.Tensor, - freqs: cute.Tensor, - output: cute.Tensor, - ): - tidx, _, _ = cute.arch.thread_idx() - batch, page_row, _ = cute.arch.block_idx() - - if page_row < self.block_size: - row = batch * self.block_size + page_row - sum_sq = cutlass.Float32(0.0) - for item in cutlass.range_constexpr(self.elements_per_thread): - dim = tidx + item * self.num_threads - value = cutlass.Float32(x[row, dim]) - sum_sq += value * value - sum_sq = cute.arch.warp_reduction_sum(sum_sq) - inverse_rms = cute.math.rsqrt(sum_sq / self.hidden_dim + self.eps) - - for item in cutlass.range_constexpr(self.nope_elements_per_thread): - dim = tidx + item * self.num_threads - value = cutlass.Float32(x[row, dim]) * inverse_rms - value *= cutlass.Float32(weight[dim]) - output[batch, page_row, dim] = value.to(output.element_type) - - for item in cutlass.range_constexpr(self.pairs_per_thread): - pair = tidx + item * self.num_threads - real_dim = self.nope_dim + pair * 2 - imag_dim = real_dim + 1 - real = cutlass.Float32(x[row, real_dim]) * inverse_rms - imag = cutlass.Float32(x[row, imag_dim]) * inverse_rms - real *= cutlass.Float32(weight[real_dim]) - imag *= cutlass.Float32(weight[imag_dim]) - cos = cutlass.Float32(freqs[row, pair, 0]) - sin = cutlass.Float32(freqs[row, pair, 1]) - output[batch, page_row, real_dim] = (real * cos - imag * sin).to( - output.element_type - ) - output[batch, page_row, imag_dim] = (imag * cos + real * sin).to( - output.element_type - ) - else: - for item in cutlass.range_constexpr(self.elements_per_thread): - dim = tidx + item * self.num_threads - output[batch, page_row, dim] = cutlass.BFloat16(0.0) diff --git a/tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py b/tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py deleted file mode 100644 index 17208d66d0f2..000000000000 --- a/tensorrt_llm/_torch/attention_backend/fp4_mla/__init__.py +++ /dev/null @@ -1,5135 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Shared MLA FP4 KV-cache helpers. - -The high-precision (HP) BF16 KV pool is a standalone circular buffer used -alongside the paged FP4 KV pool when MLA models run with NVFP4 KV cache. -The pool retains one 16-token quantization tile plus speculative-rewind slack -per sequence. - -Used by the TRTLLM attention backend FP4 MLA FMHA path. -""" - -import importlib.util -import os -from typing import Any, Literal, NamedTuple, Optional - -import torch -import triton -import triton.language as tl - -from tensorrt_llm._utils import get_sm_version, prefer_pinned -from tensorrt_llm.bindings import DataType - -from .fp4_mla_kernels import ( - _fp4_mla_chunked_cache_gather_kernel, - _fp4_mla_context_cache_update_kernel, - _fp4_mla_generation_fused_qk_rope_cache_update_kernel, - _fp4_mla_rebuild_v_scale_from_k_scale_kernel, -) - -HP_BLOCK_SIZE: int = 16 -FP4_BLOCK_SIZE: int = 16 -FP4_MLA_TOKENS_PER_BLOCK: int = 128 -FP4_MLA_SCALE_ROW_GROUP: int = 128 -FP4_MLA_SCALE_COL_GROUP: int = 4 -FP4_MLA_P_GLOBAL_SCALE: float = 448.0 * 6.0 -FP4_MLA_Q_STATIC_AMAX: float = 400.0 -FP4_MLA_KV_STATIC_AMAX: float = 30.0 -FP4_MLA_Q_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_Q_STATIC_AMAX -FP4_MLA_KV_GLOBAL_SCALE: float = FP4_MLA_P_GLOBAL_SCALE / FP4_MLA_KV_STATIC_AMAX -# Max finite e4m3 magnitude for FP4 MLA block-scale clamping. -FP4_MLA_E4M3_MAX: float = 448.0 -FP4_MLA_Q_RESIDUAL_DIM: int = 64 -FP4_MLA_K_RESIDUAL_DIM: int = FP4_MLA_Q_RESIDUAL_DIM -FP4_MLA_Q_PREFIX_DIM: int = 512 -FP4_MLA_Q_PREFIX_BLOCK_DIM: int = 256 -FP4_MLA_Q1_PREFIX_BLOCK_DIM: int = 512 -FP4_MLA_Q_LOGICAL_DIM: int = FP4_MLA_Q_PREFIX_DIM + 2 * FP4_MLA_Q_RESIDUAL_DIM -FP4_MLA_Q_PACKED_DIM: int = FP4_MLA_Q_LOGICAL_DIM // 2 -FP4_MLA_Q_SF_GROUPS: int = FP4_MLA_Q_LOGICAL_DIM // FP4_BLOCK_SIZE -FP4_MLA_ATTENTION_BACKEND_ENV = "TRTLLM_FP4_MLA_ATTENTION_BACKEND" -# Retained for compatibility; mufu16 is now the only non-fused CuTeDSL variant. -FP4_MLA_CUTEDSL_MUFU16_ENV = "TRTLLM_FP4_MLA_CUTEDSL_MUFU16" -FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV = "TRTLLM_FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE" -_FP4_MLA_CUTEDSL_BACKEND = "cutedsl" -_FP4_MLA_K_RESIDUAL_BACKENDS = ("triton", _FP4_MLA_CUTEDSL_BACKEND) -# Architectures on which the K-residual Triton path can run the fused Q-quant -# cache update. SM107 is not listed because it selects CuTeDSL instead, which -# carries its own check at launch; SM103 (Blackwell-Ultra, GB300) is the part -# this was first exercised on. -_FP4_MLA_TRITON_Q_QUANT_FUSION_SMS = (103,) -_FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH = 4 -_FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH = 16 -_FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH = 32 -_FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD = 256 -_FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD = 512 -_FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD = 640 -_FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD = 768 -_HPUpdatePhase = Literal["context", "generation"] -_FP4_MLA_TRITON_PRELOAD_KEYS = "_fp4_mla_triton_preload_keys" -_FP4_MLA_PAGE_TABLE_TILE_SIZE = 128 -_FP4_MLA_MAX_GRID_Z = 65_535 - - -# Environment helpers - - -def _env_enabled_default(name: str, default: bool) -> bool: - value = os.getenv(name) - if value is None or value == "": - return default - return value.lower() in ( - "1", - "true", - "yes", - "on", - ) - - -def _fp4_mla_cutedsl_fused_v_transpose_enabled() -> bool: - return _env_enabled_default(FP4_MLA_CUTEDSL_FUSED_V_TRANSPOSE_ENV, False) - - -FP4_MLA_PACKED_Q_ENV = "TRTLLM_FP4_MLA_PACKED_Q" -# Packed decode: each request owns a 128-row-padded block of [query token, head] -# Q rows, the kernel splits KV across grid.z tasks and reads per-row causal -# bounds, so no Q padding/copies, no Q replication, and Helix stays split. -_FP4_MLA_PACKED_MAX_M = 1024 -_FP4_MLA_PACKED_MAX_SPLITS = 16 -_FP4_MLA_PACKED_MIN_TILES_PER_SPLIT = 2 -_FP4_MLA_SPLIT_KV_ENV = "TRTLLM_FP4_MLA_SPLIT_KV" -# Rollback switch (default on): 0 makes the packed plan keep the rule's task length -# and run the tasks in request order instead of simulating the waves. -FP4_MLA_PACKED_WAVE_PLAN_ENV = "TRTLLM_FP4_MLA_PACKED_WAVE_PLAN" - - -def _packed_wave_tasks(device: torch.device, m_tiles: int) -> int: - """Tasks that fill one wave of 2-CTA clusters.""" - clusters_per_wave = max(1, _get_sm_count(device) // 2) - return max(1, clusters_per_wave // max(1, m_tiles)) - - -def _fp4_mla_packed_q_enabled() -> bool: - """Whether decode uses the packed per-request Q layout (mufu16 CuTe DSL only).""" - return ( - _env_enabled_default(FP4_MLA_PACKED_Q_ENV, True) - and _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND - and not _fp4_mla_cutedsl_fused_v_transpose_enabled() - ) - - -def _fp4_mla_packed_rows_per_seq(num_heads: int, query_len: int) -> int: - """Rows one request owns in the packed Q layout, or 0 when packing is off.""" - if not _fp4_mla_packed_q_enabled() or num_heads <= 0 or query_len <= 0: - return 0 - rows = _ceil_div(num_heads * query_len, FP4_MLA_SCALE_ROW_GROUP) * FP4_MLA_SCALE_ROW_GROUP - return rows if rows <= _FP4_MLA_PACKED_MAX_M else 0 - - -def _select_packed_split_count( - num_requests: int, - m_tiles: int, - device: torch.device, - *, - max_kv_len: int, - is_cuda_graph: bool, -) -> int: - """Task slots per request (grid.z = requests * splits): enough for two - waves of 2-CTA clusters plus one task per request, and never fewer than - two so a long request in a large batch can still be split (the kernel - shares the slots by KV length and leaves them unused when the batch is - even). - - The count depends only on the request count and tile count, so a CUDA - graph bucket always replays the same geometry. Eager calls know the real - longest KV and keep at least _FP4_MLA_PACKED_MIN_TILES_PER_SPLIT KV tiles - per split. ``TRTLLM_FP4_MLA_SPLIT_KV`` overrides the count for testing. - """ - override = os.getenv(_FP4_MLA_SPLIT_KV_ENV, "auto") - if override != "auto": - return max(1, min(int(override), _FP4_MLA_PACKED_MAX_SPLITS)) - # Slots for two waves of tasks (the plan's finest target) plus one rounding - # task per request, so the target share fits whenever the batch is skewed. - wave_tasks = _packed_wave_tasks(device, m_tiles) - splits = _ceil_div(2 * wave_tasks, max(1, num_requests)) + 1 - splits = max(2, min(_FP4_MLA_PACKED_MAX_SPLITS, splits)) - if not is_cuda_graph: - kv_tiles = _ceil_div(max(max_kv_len, 1), 256) - splits = max(1, min(splits, _ceil_div(kv_tiles, _FP4_MLA_PACKED_MIN_TILES_PER_SPLIT))) - return splits - - -class _PackedPlanGeometry(NamedTuple): - """Per-step launch geometry of the packed split-KV decode.""" - - num_requests: int - query_len: int - splits: int - wave_tasks: int - profile_tiles: int - task_capacity: int - split_kv_len: int - helix_kv_bounds: Optional[torch.Tensor] - plan: Optional[torch.Tensor] - key: Optional[tuple] - - -def _fp4_mla_packed_plan_geometry( - metadata: Any, - *, - num_heads: int, - query_len: int, - num_requests: int, - device: torch.device, - helix_kv_bounds: Optional[torch.Tensor], -) -> Optional[_PackedPlanGeometry]: - """Slots, task-length cap and plan buffer of this step's packed decode. - - None when the packed path is off. The geometry depends only on the batch - shape and the page-table capacity, so the generation cache update (the - first FP4 MLA kernel of a step) and the attention decode derive the same - one, and whichever runs first writes the step's plan fused with the - generation lengths (``_fp4_mla_uniform_generation_lengths``). ``key`` is - None when every request keeps a single task and no plan is needed. - """ - rows_per_seq = _fp4_mla_packed_rows_per_seq(num_heads, query_len) - if rows_per_seq == 0 or num_requests <= 0: - return None - if not getattr(metadata, "_fp4_mla_device_page_table", False): - return None - # The fixed-stride generation page-table capacity (see _max_generation_pages). - max_pages = int(getattr(metadata, "fp4_mla_page_table_stride", 0)) - if max_pages <= 0: - return None - max_kv_len = max_pages * metadata.page_size - is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) - m_tiles = rows_per_seq // FP4_MLA_SCALE_ROW_GROUP - splits = _select_packed_split_count( - num_requests, - m_tiles, - device, - max_kv_len=max_kv_len, - is_cuda_graph=is_cuda_graph, - ) - # Tasks that fill one wave of 2-CTA clusters: the plan's task-length reference. - wave_tasks = _packed_wave_tasks(device, m_tiles) - # The runtime KV profile only has to cover one task: the kernel clamps the - # plan's task length to it, and ceil(longest / (splits - 1)) tiles per task - # always fits the splits * requests slots. A batch that fills the machine - # with one task per request keeps the full profile so even batches never - # have to split. - kv_tiles = _ceil_div(max(max_kv_len, 1), 256) - if splits == 1 or wave_tasks <= 2 * num_requests: - task_tiles = kv_tiles - else: - task_tiles = min(kv_tiles, max(_ceil_div(kv_tiles, splits - 1), 2)) - split_kv_len = max(task_tiles * 256, 256) - # The kernels clamp the plan to the compiled profile (a bucket >= split_kv_len). - module = _fp4_mla_cutedsl_kernel_module() - profile_tiles = module._select_runtime_kv_profile(split_kv_len) // 256 - if splits == 1: - return _PackedPlanGeometry( - num_requests, - query_len, - 1, - wave_tasks, - profile_tiles, - num_requests, - split_kv_len, - helix_kv_bounds, - None, - None, - ) - from .fp4_mla_packed_reduce import packed_plan_length - - tasks = splits * num_requests - # Graph buckets hold exactly their tasks; eager metadata grows on demand. - task_capacity = tasks if is_cuda_graph else max(tasks, 32) - plan = _ensure_workspace_tensor( - metadata, - "_fp4_mla_packed_plan", - (packed_plan_length(num_requests, task_capacity),), - dtype=torch.int32, - device=device, - ) - kv_lens = getattr(metadata, "fp4_mla_generation_kv_lens", None) - key = ( - num_requests, - splits, - wave_tasks, - profile_tiles, - query_len, - task_capacity, - kv_lens.data_ptr() if isinstance(kv_lens, torch.Tensor) else 0, - helix_kv_bounds.data_ptr() if helix_kv_bounds is not None else 0, - plan.data_ptr(), - ) - return _PackedPlanGeometry( - num_requests, - query_len, - splits, - wave_tasks, - profile_tiles, - task_capacity, - split_kv_len, - helix_kv_bounds, - plan, - key, - ) - - -def _fp4_mla_helix_plan_bounds( - metadata: Any, num_tokens: int, *, use_helix: bool -) -> Optional[torch.Tensor]: - """Per-token Helix KV bounds the decode plans on, or None (validated at decode).""" - if not use_helix or not bool(getattr(metadata, "_helix_spec_tokens_valid", False)): - return None - bounds = getattr(metadata, "helix_kv_bounds", None) - if ( - not isinstance(bounds, torch.Tensor) - or bounds.dtype != torch.int32 - or bounds.ndim != 1 - or bounds.numel() < num_tokens - or not bounds.is_contiguous() - ): - return None - return bounds[:num_tokens] - - -def _launch_fp4_mla_packed_plan( - geometry: _PackedPlanGeometry, - kv_lens: torch.Tensor, - metadata: Any, - *, - generation_lengths: Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = None, -) -> None: - """Write the step's packed split plan, fused with the generation lengths when given.""" - from .fp4_mla_packed_reduce import launch_packed_plan - - launch_packed_plan( - kv_lens, - geometry.plan, - num_requests=geometry.num_requests, - num_splits=geometry.splits, - wave_tasks=geometry.wave_tasks, - max_tiles_per_task=geometry.profile_tiles, - query_len=geometry.query_len, - helix_kv_bounds=geometry.helix_kv_bounds, - generation_lengths=generation_lengths, - wave_plan=_env_enabled_default(FP4_MLA_PACKED_WAVE_PLAN_ENV, True), - ) - metadata._fp4_mla_packed_plan_key = geometry.key - - -def _env_int(name: str) -> Optional[int]: - value = os.environ.get(name) - if value is None or value == "": - return None - return int(value) - - -def _fp4_mla_attention_backend() -> str: - backend = os.getenv(FP4_MLA_ATTENTION_BACKEND_ENV) - if backend: - return backend.lower() - return _FP4_MLA_CUTEDSL_BACKEND if get_sm_version() == 107 else "triton" - - -def _q_quant_fusion_supported() -> bool: - """Whether the selected backend can fuse Q quantization on this GPU. - - The backend selector already encodes the architecture: SM107 gets CuTeDSL, - everything else gets the Triton K-residual path, and the CuTeDSL launch - re-checks SM107 for itself. Gating this predicate on SM107 as well ruled - out the Triton path on architectures that select it, which is not a - no-op: ``uses_fp4_mla_attention`` carries no SM check, so on GB300 - (SM103) it stayed true while this returned false, and the caller treats - that disagreement as a hard error ("FP4 MLA fused Q quantization - eligibility changed before launch") rather than a fallback. - """ - backend = _fp4_mla_attention_backend() - if backend == _FP4_MLA_CUTEDSL_BACKEND: - return get_sm_version() == 107 - return get_sm_version() in _FP4_MLA_TRITON_Q_QUANT_FUSION_SMS - - -def _cutedsl_backend_available() -> bool: - try: - return all( - importlib.util.find_spec(module) is not None - for module in ("ctm", "cutlass", "cuda.bindings.driver") - ) - except ModuleNotFoundError: - return False - - -def _fp4_mla_cutedsl_kernel_module() -> Any: - if _fp4_mla_cutedsl_fused_v_transpose_enabled(): - from . import fp4_mla_cutedsl_mufu16_fused_v_transpose - - return fp4_mla_cutedsl_mufu16_fused_v_transpose - - from . import fp4_mla_cutedsl_mufu16 - - return fp4_mla_cutedsl_mufu16 - - -def _ceil_div(lhs: int, rhs: int) -> int: - return (lhs + rhs - 1) // rhs - - -def _fp4_mla_q1_kv_blocks_per_program(num_gen: int, v_head_dim: int) -> int: - """Select the Q1 KV work per program without changing the launch boundary.""" - large_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH * FP4_BLOCK_SIZE - if num_gen >= _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD and v_head_dim % large_batch_block_dim == 0: - return _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH - medium_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH * FP4_BLOCK_SIZE - if ( - num_gen >= _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD - and v_head_dim % medium_batch_block_dim == 0 - ): - return _FP4_MLA_Q1_KV_BLOCKS_MEDIUM_BATCH - small_batch_block_dim = _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH * FP4_BLOCK_SIZE - if v_head_dim % small_batch_block_dim == 0: - return _FP4_MLA_Q1_KV_BLOCKS_SMALL_BATCH - return 1 - - -def _fp4_mla_q1_prefix_blocks_per_program( - num_gen: int, - q1_kv_blocks_per_program: int, -) -> int: - """Select rolled Q-prefix work only when the batch keeps it efficient.""" - max_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM - if ( - num_gen >= _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD - and q1_kv_blocks_per_program == _FP4_MLA_Q1_KV_BLOCKS_LARGE_BATCH - ): - if num_gen >= _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD: - return min(4, max_prefix_blocks) - return min(2, max_prefix_blocks) - return 1 - - -def _fp4_mla_q1_preload_variants( - max_num_sequences: int, - v_head_dim: int, -) -> tuple[tuple[int, int], ...]: - """Return every Q1 tuning variant reachable by the configured batch limit.""" - batch_sizes = [1] - for threshold in ( - _FP4_MLA_Q1_KV_MEDIUM_BATCH_THRESHOLD, - _FP4_MLA_Q1_KV_LARGE_BATCH_THRESHOLD, - _FP4_MLA_Q1_PREFIX_PAIR_BATCH_THRESHOLD, - _FP4_MLA_Q1_PREFIX_GROUP4_BATCH_THRESHOLD, - ): - if threshold <= max_num_sequences: - batch_sizes.append(threshold) - - variants = [] - for batch_size in batch_sizes: - kv_blocks = _fp4_mla_q1_kv_blocks_per_program(batch_size, v_head_dim) - prefix_blocks = _fp4_mla_q1_prefix_blocks_per_program( - batch_size, - kv_blocks, - ) - variant = (kv_blocks, prefix_blocks) - if variant not in variants: - variants.append(variant) - return tuple(variants) - - -def _fp4_mla_triton_preload_key_set(metadata: Any) -> set[tuple[object, ...]]: - """Return the engine-scoped set of Triton variants loaded during warmup.""" - owner = getattr(metadata, "kv_cache_manager", None) - if owner is None: - owner = metadata - preload_keys = getattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, None) - if preload_keys is None: - preload_keys = set() - setattr(owner, _FP4_MLA_TRITON_PRELOAD_KEYS, preload_keys) - return preload_keys - - -@triton.jit -def _fp4_mla_store_sequence_append_metadata( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - sequence_idx, - num_tokens, - PREFIX_BLOCK: tl.constexpr, - TOKEN_BLOCK: tl.constexpr, -): - append_len = tl.load(append_lens_ptr + sequence_idx) - prefix_offsets = tl.arange(0, PREFIX_BLOCK) - token_start = append_len - append_len - for prefix_start in tl.range(0, sequence_idx, PREFIX_BLOCK): - preceding_sequences = prefix_start + prefix_offsets - preceding_lens = tl.load( - append_lens_ptr + preceding_sequences, - mask=preceding_sequences < sequence_idx, - other=0, - ) - token_start += tl.sum(preceding_lens) - - cached_len = tl.load(kv_lens_ptr + sequence_idx) - append_len - - token_offsets = tl.arange(0, TOKEN_BLOCK) - for block_start in tl.range(0, append_len, TOKEN_BLOCK): - local_offsets = block_start + token_offsets - token_mask = (local_offsets < append_len) & (token_start + local_offsets < num_tokens) - output_offsets = token_start + local_offsets - tl.store( - batch_indices_ptr + output_offsets, - sequence_idx, - mask=token_mask, - ) - tl.store( - positions_ptr + output_offsets, - cached_len + local_offsets, - mask=token_mask, - ) - - -@triton.jit( - do_not_specialize=[ - "num_tokens", - "num_contexts", - "num_generation_sequences", - ], - do_not_specialize_on_alignment=[ - "num_tokens", - "num_contexts", - "num_generation_sequences", - ], -) -def _fp4_mla_append_metadata_kernel( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - num_tokens, - num_contexts, - num_generation_sequences, - ONE_TOKEN_GENERATION: tl.constexpr, - PREFIX_BLOCK: tl.constexpr, - TOKEN_BLOCK: tl.constexpr, - GENERATION_BLOCK: tl.constexpr, -): - program_idx = tl.program_id(0) - if ONE_TOKEN_GENERATION: - if program_idx < num_contexts: - _fp4_mla_store_sequence_append_metadata( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - program_idx, - num_tokens, - PREFIX_BLOCK, - TOKEN_BLOCK, - ) - else: - generation_offsets = (program_idx - num_contexts) * GENERATION_BLOCK + tl.arange( - 0, GENERATION_BLOCK - ) - generation_mask = generation_offsets < num_generation_sequences - sequence_indices = num_contexts + generation_offsets - output_offsets = num_tokens - num_generation_sequences + generation_offsets - generation_mask = generation_mask & (output_offsets < num_tokens) - generation_positions = ( - tl.load( - kv_lens_ptr + sequence_indices, - mask=generation_mask, - other=1, - ) - - 1 - ) - tl.store( - batch_indices_ptr + output_offsets, - sequence_indices, - mask=generation_mask, - ) - tl.store( - positions_ptr + output_offsets, - generation_positions, - mask=generation_mask, - ) - else: - _fp4_mla_store_sequence_append_metadata( - append_lens_ptr, - kv_lens_ptr, - batch_indices_ptr, - positions_ptr, - program_idx, - num_tokens, - PREFIX_BLOCK, - TOKEN_BLOCK, - ) - - -def populate_fp4_mla_append_metadata( - append_lens: torch.Tensor, - kv_lens: torch.Tensor, - batch_indices: torch.Tensor, - positions: torch.Tensor, - *, - num_tokens: int, - num_sequences: int, - num_contexts: int, - num_context_tokens: int, -) -> None: - """Populate FP4 MLA token-to-sequence metadata in one Triton launch. - - Mixed batches vectorize their one-token generation rows. Multi-token MTP - and fallback shapes use the generic per-sequence path in the same kernel. - """ - if num_sequences <= 0 or num_tokens <= 0: - return - if not 0 <= num_contexts <= num_sequences: - raise ValueError( - f"FP4 MLA num_contexts must be in [0, {num_sequences}], got {num_contexts}." - ) - if not 0 <= num_context_tokens <= num_tokens: - raise ValueError( - f"FP4 MLA num_context_tokens must be in [0, {num_tokens}], got {num_context_tokens}." - ) - - tensors = ( - append_lens, - kv_lens, - batch_indices, - positions, - ) - if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): - raise ValueError("FP4 MLA append metadata tensors must be contiguous and one-dimensional.") - if any(tensor.dtype != torch.int32 for tensor in tensors): - raise TypeError("FP4 MLA append metadata tensors must use int32.") - if any(not tensor.is_cuda for tensor in tensors): - raise ValueError("FP4 MLA append metadata tensors must be CUDA tensors.") - if any(tensor.device != append_lens.device for tensor in tensors[1:]): - raise ValueError("FP4 MLA append metadata tensors must be on the same device.") - sequence_tensors = (append_lens, kv_lens) - if any(tensor.numel() < num_sequences for tensor in sequence_tensors): - raise ValueError( - f"FP4 MLA sequence metadata tensors need at least {num_sequences} entries." - ) - token_tensors = (batch_indices, positions) - if any(tensor.numel() < num_tokens for tensor in token_tensors): - raise ValueError(f"FP4 MLA token metadata tensors need at least {num_tokens} entries.") - - num_generation_sequences = num_sequences - num_contexts - # Each scheduled generation sequence appends at least one token. Equality - # therefore identifies the common mixed batch with one token per decode - # row without reading the device append lengths back on the host. - one_token_generation = num_tokens == num_context_tokens + num_generation_sequences - generation_block = 128 - grid = num_sequences - if one_token_generation: - grid = num_contexts + triton.cdiv(num_generation_sequences, generation_block) - - _fp4_mla_append_metadata_kernel[(grid,)]( - append_lens, - kv_lens, - batch_indices, - positions, - num_tokens, - num_contexts, - num_generation_sequences, - ONE_TOKEN_GENERATION=one_token_generation, - PREFIX_BLOCK=128, - TOKEN_BLOCK=256, - GENERATION_BLOCK=generation_block, - num_warps=4, - ) - - -@triton.jit( - do_not_specialize=["num_gen", "generation_len"], - do_not_specialize_on_alignment=["num_gen", "generation_len"], -) -def _fp4_mla_generation_lengths_kernel( - kv_lens_ptr, - prompt_lens_ptr, - corrected_kv_lens_ptr, - generation_lens_ptr, - num_gen, - generation_len, - BLOCK: tl.constexpr, -): - offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < num_gen - kv_lens = tl.load(kv_lens_ptr + offsets, mask=mask, other=0) - prompt_lens = tl.load(prompt_lens_ptr + offsets, mask=mask, other=0) - tl.store( - corrected_kv_lens_ptr + offsets, - kv_lens - prompt_lens + generation_len, - mask=mask, - ) - tl.store(generation_lens_ptr + offsets, generation_len, mask=mask) - - -def populate_fp4_mla_generation_lengths( - kv_lens: torch.Tensor, - prompt_lens: torch.Tensor, - corrected_kv_lens: torch.Tensor, - generation_lens: torch.Tensor, - *, - num_gen_tokens: int, - num_gen: int, -) -> None: - """Populate reusable FP4 MLA generation lengths in one Triton launch.""" - if num_gen <= 0 or num_gen_tokens % num_gen != 0: - raise ValueError( - "FP4 MLA generation lengths require a positive sequence count and " - f"uniform token count, got {num_gen_tokens} tokens for {num_gen} sequences." - ) - tensors = (kv_lens, prompt_lens, corrected_kv_lens, generation_lens) - if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): - raise ValueError( - "FP4 MLA generation length tensors must be contiguous and one-dimensional." - ) - if any(tensor.dtype != torch.int32 for tensor in tensors): - raise TypeError("FP4 MLA generation length tensors must use int32.") - if any(not tensor.is_cuda for tensor in tensors): - raise ValueError("FP4 MLA generation length tensors must be CUDA tensors.") - if any(tensor.device != kv_lens.device for tensor in tensors[1:]): - raise ValueError("FP4 MLA generation length tensors must be on the same device.") - if any(tensor.numel() < num_gen for tensor in tensors): - raise ValueError(f"FP4 MLA generation length tensors need at least {num_gen} entries.") - - block = 128 - _fp4_mla_generation_lengths_kernel[(triton.cdiv(num_gen, block),)]( - kv_lens, - prompt_lens, - corrected_kv_lens, - generation_lens, - num_gen, - num_gen_tokens // num_gen, - BLOCK=block, - num_warps=4, - ) - - -@triton.jit( - do_not_specialize=["num_tokens", "tokens_per_seq"], - do_not_specialize_on_alignment=["num_tokens", "tokens_per_seq"], -) -def _helix_spec_buffers_kernel( - positions_ptr, - local_slots_ptr, - kv_bounds_ptr, - kv_lens_ptr, - num_tokens, - tokens_per_seq, - tokens_per_block, - cp_size, - cp_rank, - BLOCK: tl.constexpr, -): - offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) - mask = offsets < num_tokens - position = tl.load(positions_ptr + offsets, mask=mask, other=0) - # Rank-local prefix length of the first g global tokens: ledger page b lives - # on CP rank b % cp_size, so count the full ledgers plus this rank's clamped - # share of the partial one (mirrors KVCacheManagerV2._helix_local_len). - ledger = tokens_per_block * cp_size - rank_base = cp_rank * tokens_per_block - full_before = position // ledger - local_before = full_before * tokens_per_block + tl.minimum( - tl.maximum(position - full_before * ledger - rank_base, 0), tokens_per_block - ) - after = position + 1 - full_after = after // ledger - bound = full_after * tokens_per_block + tl.minimum( - tl.maximum(after - full_after * ledger - rank_base, 0), tokens_per_block - ) - owned = (position // tokens_per_block) % cp_size == cp_rank - tl.store(local_slots_ptr + offsets, tl.where(owned, local_before, -1), mask=mask) - tl.store(kv_bounds_ptr + offsets, bound, mask=mask) - # The last token of each verify group sets its sequence's rank-local KV length. - last_of_seq = offsets % tokens_per_seq == tokens_per_seq - 1 - tl.store(kv_lens_ptr + offsets // tokens_per_seq, bound, mask=mask & last_of_seq) - - -def populate_helix_spec_buffers( - positions: torch.Tensor, - local_slots: torch.Tensor, - kv_bounds: torch.Tensor, - kv_lens: torch.Tensor, - *, - num_tokens: int, - tokens_per_seq: int, - tokens_per_block: int, - cp_size: int, - cp_rank: int, -) -> None: - """Derive the Helix speculative per-token buffers in one launch. - - For generation token ``t`` at global position ``positions[t]``: - ``local_slots[t]`` is its rank-local KV write slot (-1 when another CP rank - owns the page), ``kv_bounds[t]`` the rank-local KV entries it attends to - (the local length of ``position + 1``) and ``kv_lens[s]`` the bound of - sequence ``s``'s last token. Static shapes, safe under CUDA graph capture. - """ - if num_tokens <= 0: - return - if tokens_per_seq <= 0 or num_tokens % tokens_per_seq != 0: - raise ValueError( - f"Helix verify groups must be uniform: {num_tokens} tokens for groups of " - f"{tokens_per_seq}." - ) - if tokens_per_block <= 0 or cp_size <= 0 or not 0 <= cp_rank < cp_size: - raise ValueError("Helix spec buffers need a positive page size and a valid CP rank.") - num_seqs = num_tokens // tokens_per_seq - tensors = (positions, local_slots, kv_bounds, kv_lens) - if any(tensor.ndim != 1 or tensor.stride(0) != 1 for tensor in tensors): - raise ValueError("Helix spec buffers must be contiguous and one-dimensional.") - if any(tensor.dtype != torch.int32 for tensor in tensors): - raise TypeError("Helix spec buffers must use int32.") - if any(not tensor.is_cuda or tensor.device != positions.device for tensor in tensors): - raise ValueError("Helix spec buffers must be CUDA tensors on one device.") - if any(tensor.numel() < num_tokens for tensor in tensors[:3]) or kv_lens.numel() < num_seqs: - raise ValueError("Helix spec buffers are too small for the generation tokens.") - - block = 256 - _helix_spec_buffers_kernel[(triton.cdiv(num_tokens, block),)]( - positions, - local_slots, - kv_bounds, - kv_lens, - num_tokens, - tokens_per_seq, - tokens_per_block, - cp_size, - cp_rank, - BLOCK=block, - num_warps=4, - ) - - -def _fp4_mla_page_table_spec(kv_cache_manager: Any) -> Any: - get_spec = getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None) - if not callable(get_spec): - raise RuntimeError("FP4 MLA requires V2 cache-layout page metadata.") - spec = get_spec() - for field_name in ( - "cache_pool_id", - "cache_page_index_scale", - "hp_pool_id", - "hp_page_index_scale", - ): - value = getattr(spec, field_name, None) - if not isinstance(value, int) or value < 0: - raise ValueError( - f"FP4 MLA page-table spec requires non-negative {field_name}, got {value}." - ) - if spec.cache_page_index_scale <= 0 or spec.hp_page_index_scale <= 0: - raise ValueError("FP4 MLA page-index scales must be positive.") - return spec - - -# Mixed batches frequently vary by one sequence. Keep per-forward dimensions -# out of Triton's specialization key and tile page rows at one fixed width so -# those shape changes cannot trigger JIT compilation on the critical path. -@triton.jit( - do_not_specialize=["num_sequences", "num_contexts", "max_pages"], - do_not_specialize_on_alignment=["num_sequences", "num_contexts", "max_pages"], -) -def _fp4_mla_materialize_page_table_kernel( - page_ids_ptr, - paged_kv_indptr_ptr, - paged_kv_indptr_decode_ptr, - block_offsets_ptr, - kv_lens_ptr, - generation_kv_lens_ptr, - block_offsets_stride, - num_sequences, - num_contexts, - max_pages, - PAGE_SIZE: tl.constexpr, - PAGE_INDEX_SCALE: tl.constexpr, - PAGE_TILE_SIZE: tl.constexpr, -): - sequence_idx = tl.program_id(0) - page_tile_idx = tl.program_id(1) - page_offsets = page_tile_idx * PAGE_TILE_SIZE + tl.arange(0, PAGE_TILE_SIZE) - generation_idx = sequence_idx - num_contexts - is_generation = generation_idx >= 0 - context_kv_len = tl.load(kv_lens_ptr + sequence_idx) - generation_kv_len = tl.load( - generation_kv_lens_ptr + generation_idx, - mask=is_generation, - other=0, - ) - kv_len = tl.maximum(tl.where(is_generation, generation_kv_len, context_kv_len), 0) - num_active_pages = tl.minimum( - (kv_len + PAGE_SIZE - 1) // PAGE_SIZE, - max_pages, - ) - active_page_mask = page_offsets < num_active_pages - encoded_page_offsets = tl.load( - block_offsets_ptr + sequence_idx * block_offsets_stride + page_offsets, - mask=active_page_mask, - other=-1, - ) - decoded_page_ids = tl.where( - encoded_page_offsets >= 0, - encoded_page_offsets // PAGE_INDEX_SCALE, - encoded_page_offsets, - ) - page_ids = tl.where(active_page_mask, decoded_page_ids, 0) - table_offset = sequence_idx * max_pages + page_offsets - # Fixed-stride indptrs expose the whole row. Initialize inactive slots so - # masked or prefetched page-table reads cannot observe stale page IDs. - tl.store( - page_ids_ptr + table_offset, - page_ids, - mask=page_offsets < max_pages, - ) - - first_lane = page_offsets == 0 - sequence_start = sequence_idx * max_pages - tl.store( - paged_kv_indptr_ptr + sequence_idx + page_offsets, - sequence_start, - mask=first_lane, - ) - tl.store( - paged_kv_indptr_decode_ptr + generation_idx + page_offsets, - generation_idx * max_pages, - mask=first_lane & is_generation, - ) - final_sequence = sequence_idx == num_sequences - 1 - table_end = num_sequences * max_pages - num_generation_sequences = num_sequences - num_contexts - tl.store( - paged_kv_indptr_ptr + num_sequences + page_offsets, - table_end, - mask=first_lane & final_sequence, - ) - tl.store( - paged_kv_indptr_decode_ptr + num_generation_sequences + page_offsets, - num_generation_sequences * max_pages, - mask=first_lane & final_sequence, - ) - - -def configure_fp4_mla_device_page_table( - metadata: Any, - kv_lens: Optional[torch.Tensor] = None, -) -> bool: - """Configure the fixed-stride, device-materialized page table. - - Eager context and generation batches receive the full block-offset - table on the GPU. The materialization kernel decodes V2 page indices and - refreshes rows from the final device KV lengths before cache update. - """ - metadata._fp4_mla_device_page_table = False - metadata._fp4_mla_device_page_table_valid = False - metadata.fp4_mla_page_table_stride = 0 - metadata.fp4_mla_context_repack_max_touched_pages = 1 - - kv_cache_manager = getattr(metadata, "kv_cache_manager", None) - num_contexts = int(getattr(metadata, "num_contexts", 0)) - num_sequences = int(getattr(metadata, "num_seqs", 0)) - num_generation_sequences = num_sequences - num_contexts - num_tokens = int(getattr(metadata, "num_tokens", 0)) - num_context_tokens = int(getattr(metadata, "num_ctx_tokens", 0)) - num_generation_tokens = num_tokens - num_context_tokens - block_offsets = getattr(metadata, "kv_cache_block_offsets", None) - page_ids = getattr(metadata, "_paged_kv_indices", None) - paged_kv_indptr = getattr(metadata, "_paged_kv_indptr", None) - paged_kv_indptr_decode = getattr(metadata, "paged_kv_indptr_decode", None) - max_page_capacity = int(getattr(kv_cache_manager, "max_blocks_per_seq", 0) or 0) - page_spec = _fp4_mla_page_table_spec(kv_cache_manager) - page_index_scale = int(page_spec.cache_page_index_scale) - - tensors = (block_offsets, page_ids, paged_kv_indptr, paged_kv_indptr_decode) - is_cuda_graph = bool(getattr(metadata, "is_cuda_graph", False)) - generation_only = num_contexts == 0 - eager_context = not is_cuda_graph and num_contexts > 0 - has_valid_generation = num_generation_sequences == 0 or ( - num_generation_tokens >= num_generation_sequences - and num_generation_tokens % num_generation_sequences == 0 - ) - # NVFP4 exposes one data pool plus its paired block-scale pool. The - # materializer reads encoded data offsets from pool 0. - supported = ( - (generation_only or eager_context) - and kv_cache_manager is not None - and has_valid_generation - and int(getattr(metadata, "beam_width", 1)) == 1 - and not bool(getattr(metadata, "is_spec_dec_tree", False)) - and not bool(getattr(metadata, "locality_domain_enabled", False)) - and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK - and max_page_capacity > 0 - and page_index_scale > 0 - and all(isinstance(tensor, torch.Tensor) for tensor in tensors) - and all(tensor.dtype == torch.int32 for tensor in tensors) - and all(tensor.is_cuda for tensor in tensors) - ) - if not supported: - return False - - hp_page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) - max_pool_id = max(page_spec.cache_pool_id, page_spec.hp_pool_id) - if ( - not isinstance(hp_page_ids, torch.Tensor) - or hp_page_ids.dtype != torch.int32 - or not hp_page_ids.is_cuda - or block_offsets.shape[0] <= max_pool_id - ): - return False - metadata._fp4_mla_cache_pool_id = int(page_spec.cache_pool_id) - metadata._fp4_mla_cache_page_index_scale = int(page_spec.cache_page_index_scale) - metadata._fp4_mla_hp_pool_id = int(page_spec.hp_pool_id) - metadata._fp4_mla_hp_page_index_scale = int(page_spec.hp_page_index_scale) - - max_pages = max_page_capacity - host_kv_lens_available = ( - isinstance(kv_lens, torch.Tensor) - and kv_lens.device.type == "cpu" - and kv_lens.ndim == 1 - and kv_lens.numel() >= num_sequences - ) - if eager_context and num_generation_sequences > 0 and not host_kv_lens_available: - return False - if not is_cuda_graph and host_kv_lens_available: - generation_tokens_per_sequence = ( - num_generation_tokens // num_generation_sequences if num_generation_sequences > 0 else 0 - ) - # Eager execution can narrow the fixed row stride to the current - # batch. CUDA Graph metadata retains full configured capacity so a - # replay never changes tensor addresses or launch dimensions. - max_kv_len = int(kv_lens[:num_sequences].max().item()) + max( - 0, - generation_tokens_per_sequence - 1, - ) - max_pages = min( - max_page_capacity, - max(1, _ceil_div(max_kv_len, FP4_MLA_TOKENS_PER_BLOCK)), - ) - if num_contexts > 0: - max_context_len = int(kv_lens[:num_contexts].max().item()) - max_context_pages = _ceil_div( - max_context_len, - FP4_MLA_TOKENS_PER_BLOCK, - ) - metadata.fp4_mla_context_repack_max_touched_pages = min( - max_pages, - triton.next_power_of_2(max(1, max_context_pages)), - ) - - assert isinstance(block_offsets, torch.Tensor) - assert isinstance(page_ids, torch.Tensor) - assert isinstance(paged_kv_indptr, torch.Tensor) - assert isinstance(paged_kv_indptr_decode, torch.Tensor) - required_page_ids = num_sequences * max_pages - buffers_cover_table = ( - block_offsets.ndim == 4 - and block_offsets.shape[0] >= 1 - and block_offsets.shape[1] >= num_sequences - and block_offsets.shape[2] >= 1 - and block_offsets.shape[3] >= max_pages - and page_ids.ndim == 1 - and page_ids.numel() >= required_page_ids - and paged_kv_indptr.ndim == 1 - and paged_kv_indptr.numel() >= num_sequences + 1 - and paged_kv_indptr_decode.ndim == 1 - and paged_kv_indptr_decode.numel() >= num_generation_sequences + 1 - ) - if not buffers_cover_table: - return False - if metadata._fp4_mla_hp_page_indices.numel() < required_page_ids: - return False - - metadata._fp4_mla_device_page_table = True - metadata.fp4_mla_page_table_stride = max_pages - metadata.num_blocks = None - metadata.num_context_blocks = num_contexts * max_pages - metadata.num_generation_blocks = num_generation_sequences * max_pages - return True - - -def materialize_fp4_mla_device_page_table( - metadata: Any, - kv_lens: torch.Tensor, - generation_kv_lens: Optional[torch.Tensor] = None, -) -> None: - """Refresh the fixed-stride context and generation page table once per forward.""" - if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): - raise RuntimeError("FP4 MLA requires fixed-stride device page metadata.") - if bool(getattr(metadata, "_fp4_mla_device_page_table_valid", False)): - return - - num_contexts = int(metadata.num_contexts) - num_sequences = int(metadata.num_seqs) - num_generation_sequences = num_sequences - num_contexts - max_pages = int(metadata.fp4_mla_page_table_stride) - if num_sequences <= 0 or max_pages <= 0: - raise RuntimeError( - "FP4 MLA device page metadata requires positive sequence and page capacities." - ) - if ( - kv_lens.dtype != torch.int32 - or not kv_lens.is_cuda - or kv_lens.ndim != 1 - or kv_lens.stride(0) != 1 - or kv_lens.numel() < num_sequences - ): - raise ValueError( - "FP4 MLA device page metadata requires a contiguous CUDA int32 " - f"KV-length tensor with at least {num_sequences} entries." - ) - - if generation_kv_lens is None: - generation_kv_lens = kv_lens[num_contexts:num_sequences] - if ( - generation_kv_lens.dtype != torch.int32 - or not generation_kv_lens.is_cuda - or generation_kv_lens.ndim != 1 - or generation_kv_lens.stride(0) != 1 - or generation_kv_lens.numel() < num_generation_sequences - ): - raise ValueError( - "FP4 MLA device page metadata requires a contiguous CUDA int32 " - "generation KV-length tensor with at least " - f"{num_generation_sequences} entries." - ) - - cache_pool_id = int(getattr(metadata, "_fp4_mla_cache_pool_id", 0)) - block_offsets = metadata.kv_cache_block_offsets[ - cache_pool_id, - :num_sequences, - 0, - :max_pages, - ] - page_ids = metadata._paged_kv_indices[: num_sequences * max_pages] - page_index_scale = int(metadata._fp4_mla_cache_page_index_scale) - if page_index_scale <= 0: - raise RuntimeError("FP4 MLA device page metadata requires a positive page-index scale.") - grid = ( - num_sequences, - triton.cdiv(max_pages, _FP4_MLA_PAGE_TABLE_TILE_SIZE), - ) - _fp4_mla_materialize_page_table_kernel[grid]( - page_ids, - metadata._paged_kv_indptr, - metadata.paged_kv_indptr_decode, - block_offsets, - kv_lens, - generation_kv_lens, - block_offsets.stride(0), - num_sequences, - num_contexts, - max_pages, - PAGE_SIZE=metadata.page_size, - PAGE_INDEX_SCALE=page_index_scale, - PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, - num_warps=4, - ) - hp_page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) - if not isinstance(hp_page_ids, torch.Tensor): - raise RuntimeError("FP4 MLA requires an HP page-table output tensor.") - hp_pool_id = int(metadata._fp4_mla_hp_pool_id) - hp_page_index_scale = int(metadata._fp4_mla_hp_page_index_scale) - hp_block_offsets = metadata.kv_cache_block_offsets[ - hp_pool_id, - :num_sequences, - 0, - :max_pages, - ] - _fp4_mla_materialize_page_table_kernel[grid]( - hp_page_ids, - metadata._paged_kv_indptr, - metadata.paged_kv_indptr_decode, - hp_block_offsets, - kv_lens, - generation_kv_lens, - hp_block_offsets.stride(0), - num_sequences, - num_contexts, - max_pages, - PAGE_SIZE=metadata.page_size, - PAGE_INDEX_SCALE=hp_page_index_scale, - PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, - num_warps=4, - ) - metadata._fp4_mla_device_page_table_valid = True - - -@triton.jit -def _cutedsl_swizzled_sf_offset(row_idx, col_idx, sf_cols: tl.constexpr): - padded_cols = ((sf_cols + 3) // 4) * 4 - return ( - col_idx % 4 - + (col_idx // 4) * (4 * 128) - + (row_idx % 32) * 16 - + ((row_idx % 128) // 32) * 4 - + (row_idx // 128) * (128 * padded_cols) - ) - - -@triton.jit -def _cutedsl_pad_q_and_sf_kernel( - q_padded_ptr, - q_ptr, - q_sf_padded_ptr, - q_sf_ptr, - num_heads, - output_heads: tl.constexpr, - packed_dim: tl.constexpr, - block_bytes: tl.constexpr, - sf_cols: tl.constexpr, - sf_cols_per_byte_block: tl.constexpr, -): - query_idx = tl.program_id(0) - byte_block = tl.program_id(1) - head_offsets = tl.arange(0, output_heads) - byte_offsets = byte_block * block_bytes + tl.arange(0, block_bytes) - head_mask = head_offsets < num_heads - byte_mask = byte_offsets < packed_dim - source_rows = query_idx * num_heads + head_offsets - destination_rows = query_idx * output_heads + head_offsets - values = tl.load( - q_ptr + source_rows[:, None] * packed_dim + byte_offsets[None, :], - mask=head_mask[:, None] & byte_mask[None, :], - other=0, - ) - tl.store( - q_padded_ptr + destination_rows[:, None] * packed_dim + byte_offsets[None, :], - values, - mask=byte_mask[None, :], - ) - sf_col_offsets = byte_block * sf_cols_per_byte_block + tl.arange(0, sf_cols_per_byte_block) - sf_col_mask = sf_col_offsets < sf_cols - source_offsets = _cutedsl_swizzled_sf_offset( - source_rows[:, None], sf_col_offsets[None, :], sf_cols - ) - destination_offsets = _cutedsl_swizzled_sf_offset( - destination_rows[:, None], sf_col_offsets[None, :], sf_cols - ) - sf_values = tl.load( - q_sf_ptr + source_offsets, - mask=head_mask[:, None] & sf_col_mask[None, :], - other=1.0, - ) - tl.store( - q_sf_padded_ptr + destination_offsets, - sf_values, - mask=sf_col_mask[None, :], - ) - - -_SM_COUNT_CACHE: dict[int, int] = {} - - -def _get_sm_count(device: torch.device) -> int: - """Return the SM (multiprocessor) count for ``device``, cached per index.""" - index = device.index if device.index is not None else torch.cuda.current_device() - count = _SM_COUNT_CACHE.get(index) - if count is None: - count = torch.cuda.get_device_properties(index).multi_processor_count - _SM_COUNT_CACHE[index] = count - return count - - -def _validate_fp4_mla_context_rope( - latent_cache: torch.Tensor, - rotary_cos_sin: torch.Tensor, - v_head_dim: int, -) -> int: - head_dim = latent_cache.shape[-1] - rope_dim = head_dim - v_head_dim - if rope_dim <= 0 or rope_dim % 2 != 0: - raise ValueError( - "FP4 MLA fused context K-RoPE requires a positive even RoPE dimension, " - f"got head_dim={head_dim}, v_head_dim={v_head_dim}." - ) - if rotary_cos_sin.device != latent_cache.device: - raise ValueError("FP4 MLA context latent cache and RoPE table must use the same device.") - if rotary_cos_sin.dtype != torch.float32: - raise TypeError( - f"FP4 MLA fused context K-RoPE requires a float32 table, got {rotary_cos_sin.dtype}." - ) - if not rotary_cos_sin.is_contiguous(): - raise ValueError("FP4 MLA fused context K-RoPE requires a contiguous RoPE table.") - table_row_size = rope_dim * 2 - if rotary_cos_sin.numel() < table_row_size or rotary_cos_sin.numel() % table_row_size != 0: - raise ValueError( - "FP4 MLA context RoPE table size must be a positive multiple of " - f"{table_row_size}, got {rotary_cos_sin.numel()}." - ) - return rope_dim - - -def _host_int_list_during_forward(value: Any, start: int, end: int) -> Optional[list[int]]: - if torch.cuda.is_current_stream_capturing(): - return None - return _host_int_list(value, start, end) - - -# FP4 MLA scale-layout helpers - - -def get_fp4_mla_v_scale_pool_size(v_head_dim: int, page_size: int) -> int: - """Return elements per page for the swizzled FP4 MLA V-scale pool. - - The PV matmul treats V as a RHS matrix shaped ``[v_head_dim, kv_tokens]``. - NVFP4 block scales therefore group along the token/K axis, not along the - latent dimension as the K-view cache does. The physical layout matches the - Triton block-scaled matmul scale layout: - ``[ceil(v_head_dim / 128), ceil(page_size / 16 / 4), 32, 16]``. - """ - return _get_fp4_mla_swizzled_scale_size(v_head_dim, page_size) - - -def _get_fp4_mla_swizzled_scale_size(rows: int, cols: int) -> int: - scale_cols = _ceil_div(cols, FP4_BLOCK_SIZE) - row_groups = _ceil_div(rows, FP4_MLA_SCALE_ROW_GROUP) - col_groups = _ceil_div(scale_cols, FP4_MLA_SCALE_COL_GROUP) - return row_groups * col_groups * 32 * 16 - - -def can_fuse_fp4_mla_q_quant( - metadata: Any, - q: torch.Tensor, - q_pe: torch.Tensor, - latent_cache: torch.Tensor, -) -> bool: - """Return whether generation can quantize Q in the fused cache update.""" - num_gen = metadata.num_seqs - metadata.num_contexts - return bool( - _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS - and _q_quant_fusion_supported() - and num_gen > 0 - and getattr(metadata, "kv_cache_manager", None) is not None - and q.shape[0] > 0 - and q.shape[0] % num_gen == 0 - and q.is_cuda - and q_pe.is_cuda - and latent_cache.is_cuda - and q.device == q_pe.device == latent_cache.device - and q.dtype == torch.bfloat16 - and q.is_contiguous() - and q.ndim == 3 - and 0 < q.shape[1] <= 128 - and q.shape[2] == FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM - and q_pe.dtype == torch.bfloat16 - and tuple(q_pe.shape) == (q.shape[0], q.shape[1], FP4_MLA_Q_RESIDUAL_DIM) - and latent_cache.dtype == torch.bfloat16 - and tuple(latent_cache.shape) == (q.shape[0], FP4_MLA_Q_PREFIX_DIM + FP4_MLA_Q_RESIDUAL_DIM) - and metadata.page_size == FP4_MLA_TOKENS_PER_BLOCK - ) - - -def _prepare_fp4_mla_q_buffers( - metadata: Any, - num_queries: int, - num_heads: int, - device: torch.device, - rows_per_seq: int = 0, -) -> tuple[torch.Tensor, torch.Tensor, int]: - """Return manager-owned, fixed-capacity packed-Q staging buffers. - - With ``rows_per_seq`` the buffers hold ``max_num_sequences`` request blocks - of that many rows (zero-filled so padded rows stay finite); otherwise one - row per (token, head) up to the token capacity. - """ - if num_queries <= 0: - raise ValueError(f"FP4 MLA Q buffer preparation needs queries, got {num_queries}.") - owner = getattr(metadata, "kv_cache_manager", None) - if owner is None: - raise RuntimeError("Fused FP4 MLA Q quantization requires a KV cache manager.") - if num_heads <= 0 or num_heads > 128: - raise ValueError(f"FP4 MLA Q buffers require 1-128 local heads, got {num_heads}.") - buffers = getattr(owner, "_fp4_mla_q_buffers", None) - if buffers is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "Cannot create FP4 MLA Q buffers while capturing a CUDA graph. " - "Run a warmup forward first." - ) - buffers = {} - setattr(owner, "_fp4_mla_q_buffers", buffers) - - max_num_tokens = int(getattr(metadata, "max_num_tokens", num_queries) or num_queries) - max_num_sequences = int( - getattr(metadata, "max_num_sequences", None) - or getattr(metadata, "max_num_requests", num_queries) - or num_queries - ) - max_query_width = 1 + int(getattr(metadata, "max_total_draft_tokens", None) or 0) - capacity = min( - max_num_tokens, - max_num_sequences * max_query_width, - _FP4_MLA_MAX_GRID_Z, - ) - if num_queries > capacity: - raise ValueError( - f"FP4 MLA active queries exceed the configured Q capacity: {num_queries} > {capacity}." - ) - - device_index = device.index if device.index is not None else torch.cuda.current_device() - canonical_device = torch.device("cuda", device_index) - if rows_per_seq > 0: - q_rows = max_num_sequences * rows_per_seq - q_key = f"q_{num_heads}_p{rows_per_seq}" - q_sf_key = f"q_sf_{num_heads}_p{rows_per_seq}" - else: - q_rows = capacity * num_heads - q_key = f"q_{num_heads}" - q_sf_key = f"q_sf_{num_heads}" - expected_q_shape = (q_rows, FP4_MLA_Q_PACKED_DIM) - expected_q_sf_shape = (_get_fp4_mla_swizzled_scale_size(q_rows, FP4_MLA_Q_LOGICAL_DIM),) - q_storage = buffers.get(q_key) - q_sf_storage = buffers.get(q_sf_key) - if q_storage is None and q_sf_storage is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "Cannot allocate FP4 MLA Q buffers while capturing a CUDA graph. " - "Run a warmup forward first." - ) - allocate = torch.zeros if rows_per_seq > 0 else torch.empty - q_storage = allocate(expected_q_shape, dtype=torch.uint8, device=canonical_device) - q_sf_storage = allocate( - expected_q_sf_shape, - dtype=torch.float8_e4m3fn, - device=canonical_device, - ) - buffers[q_key] = q_storage - buffers[q_sf_key] = q_sf_storage - if ( - q_storage is None - or q_storage.dtype != torch.uint8 - or q_storage.device != canonical_device - or tuple(q_storage.shape) != expected_q_shape - or not q_storage.is_contiguous() - or q_sf_storage is None - or q_sf_storage.dtype != torch.float8_e4m3fn - or q_sf_storage.device != canonical_device - or tuple(q_sf_storage.shape) != expected_q_sf_shape - or not q_sf_storage.is_contiguous() - ): - raise RuntimeError("FP4 MLA packed-Q buffers do not match the configured capacity.") - return q_storage, q_sf_storage, capacity - - -def _get_fp4_mla_context_start_positions(metadata: Any, num_contexts: int) -> torch.Tensor: - kv_cache_params = getattr(metadata, "kv_cache_params", None) - cached_token_lens = getattr(kv_cache_params, "num_cached_tokens_per_seq", None) - if cached_token_lens is not None: - return torch.as_tensor(cached_token_lens[:num_contexts], dtype=torch.int64, device="cpu") - - return ( - ( - metadata.kv_lens_cuda_runtime[:num_contexts] - - metadata.prompt_lens_cuda_runtime[:num_contexts] - ) - .detach() - .cpu() - ) - - -def _validate_fp4_mla_context_start_alignment( - metadata: Any, - num_contexts: int, - *, - alignment: int = HP_BLOCK_SIZE, -) -> None: - context_start_positions = _get_fp4_mla_context_start_positions(metadata, num_contexts) - bad_start = (context_start_positions < 0) | ((context_start_positions % alignment) != 0) - if bool(torch.any(bad_start).item()): - starts = context_start_positions.detach().cpu().tolist() - raise ValueError( - "FP4 MLA shared-tile context update requires every context " - f"start position to be {alignment}-token aligned, got " - f"start positions {starts}." - ) - - -def get_fp4_mla_v_scale_pool_shape( - num_layers: int, - num_pages: int, - v_head_dim: int, - page_size: int, -) -> tuple[int, int, int, int, int, int]: - """Return the logical swizzled V-scale view shape. - - The leading dimensions are ``[layer, physical_page]``. The remaining - dimensions are the preshuffled ``[N // 128, K // 16 // 4, 32, 16]`` shape - consumed by Triton block-scaled matmul for the V/PV RHS operand. - """ - token_scale_cols = _ceil_div(page_size, FP4_BLOCK_SIZE) - return ( - num_layers, - num_pages, - _ceil_div(v_head_dim, FP4_MLA_SCALE_ROW_GROUP), - _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP), - 32, - 16, - ) - - -def get_fp4_mla_v_scale_pool_view( - metadata: Any, - *, - v_head_dim: int, -) -> torch.Tensor: - """View the auxiliary MLA V-scale pool in Triton's block-scaled layout.""" - pool = getattr(metadata, "fp4_mla_v_scale_pool", None) - if pool is None: - raise RuntimeError("FP4 MLA V scale pool is not allocated.") - - elems_per_page = get_fp4_mla_v_scale_pool_size(v_head_dim, metadata.page_size) - if pool.shape[-1] < elems_per_page: - raise RuntimeError( - f"FP4 MLA V scale pool page stride is too small: got " - f"{pool.shape[-1]}, need {elems_per_page}." - ) - - token_scale_cols = _ceil_div(metadata.page_size, FP4_BLOCK_SIZE) - col_groups = _ceil_div(token_scale_cols, FP4_MLA_SCALE_COL_GROUP) - shape = get_fp4_mla_v_scale_pool_shape( - pool.shape[0], pool.shape[1], v_head_dim, metadata.page_size - ) - strides = ( - pool.stride(0), - pool.stride(1), - col_groups * 32 * 16, - 32 * 16, - 16, - 1, - ) - return torch.as_strided(pool, size=shape, stride=strides) - - -def _rebuild_fp4_mla_v_scales_from_k_scales( - sf_cache: torch.Tensor, - v_scale_pool: torch.Tensor, - page_ids: torch.Tensor, - page_valid_tokens: torch.Tensor, - *, - local_layer: int, - v_head_dim: int, - page_size: int, -) -> None: - """Bit-exactly rebuild imported MLA V scales from transferred K scales.""" - if page_ids.numel() == 0: - return - if page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - "FP4 MLA imported V-scale rebuild requires " - f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." - ) - for name, tensor in ( - ("sf_cache", sf_cache), - ("v_scale_pool", v_scale_pool), - ("page_ids", page_ids), - ("page_valid_tokens", page_valid_tokens), - ): - if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: - raise ValueError(f"{name} must be a CUDA tensor.") - if page_ids.dtype != torch.int32 or page_valid_tokens.dtype != torch.int32: - raise TypeError("FP4 MLA imported page IDs and valid-token counts must use int32.") - if page_ids.ndim != 1 or page_valid_tokens.ndim != 1: - raise ValueError("FP4 MLA imported page metadata must be one-dimensional.") - if page_ids.numel() != page_valid_tokens.numel(): - raise ValueError("FP4 MLA imported page IDs and valid-token counts must have equal length.") - if not page_ids.is_contiguous() or not page_valid_tokens.is_contiguous(): - raise ValueError("FP4 MLA imported page metadata must be contiguous.") - if not (sf_cache.device == v_scale_pool.device == page_ids.device == page_valid_tokens.device): - raise ValueError("FP4 MLA imported cache tensors must be on the same device.") - - k_sf_bytes = sf_cache.view(torch.uint8) - v_sf_bytes = v_scale_pool.view(torch.uint8) - if k_sf_bytes.ndim < 2 or v_sf_bytes.ndim < 3: - raise ValueError( - "FP4 MLA imported scale pools require per-page K storage and " - "per-layer/per-page V storage." - ) - num_layers = int(v_sf_bytes.shape[0]) - num_pages = int(v_sf_bytes.shape[1]) - if not 0 <= local_layer < num_layers: - raise IndexError( - f"local_layer={local_layer} is outside the V-scale pool with {num_layers} layers." - ) - if int(k_sf_bytes.shape[0]) != num_pages: - raise ValueError( - "FP4 MLA K/V scale pools disagree on their physical page count: " - f"{int(k_sf_bytes.shape[0])} != {num_pages}." - ) - sf_per_token = int(k_sf_bytes.shape[-1]) - required_sf_per_token = _ceil_div(v_head_dim, FP4_BLOCK_SIZE) - if sf_per_token < required_sf_per_token: - raise ValueError( - "FP4 MLA K-scale storage is too narrow for the compressed V head: " - f"{sf_per_token} < {required_sf_per_token}." - ) - required_v_page_elems = get_fp4_mla_v_scale_pool_size(v_head_dim, page_size) - if int(v_sf_bytes.shape[-1]) < required_v_page_elems: - raise ValueError( - "FP4 MLA V-scale page storage is too small for import rebuild: " - f"{int(v_sf_bytes.shape[-1])} < {required_v_page_elems}." - ) - - token_groups = page_size // HP_BLOCK_SIZE - _fp4_mla_rebuild_v_scale_from_k_scale_kernel[ - (page_ids.numel(), triton.cdiv(v_head_dim, FP4_BLOCK_SIZE)) - ]( - k_sf_bytes, - v_sf_bytes, - page_ids, - page_valid_tokens, - page_ids.numel(), - num_pages, - num_layers, - local_layer, - page_size, - k_sf_bytes.stride(0), - v_sf_bytes.stride(0), - v_sf_bytes.stride(1), - V_HEAD_D=v_head_dim, - HP_BLOCK=HP_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=token_groups, - BLOCK_TOKEN_GROUPS=triton.next_power_of_2(token_groups), - num_warps=4, - ) - - -def _stage_fp4_mla_import_page_metadata( - prompt_block_ids: list[int], - *, - prompt_len: int, - page_size: int, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - """Stage imported page metadata without synchronizing the CUDA stream.""" - if device.type != "cuda": - raise ValueError("FP4 MLA disaggregated import requires a CUDA device.") - num_prompt_pages = len(prompt_block_ids) - pin_memory = prefer_pinned() - page_ids_host = torch.tensor( - prompt_block_ids, - dtype=torch.int32, - device="cpu", - pin_memory=pin_memory, - ) - page_valid_tokens_host = torch.full( - (num_prompt_pages,), - page_size, - dtype=torch.int32, - device="cpu", - pin_memory=pin_memory, - ) - page_valid_tokens_host[-1] = prompt_len - (num_prompt_pages - 1) * page_size - # Constructing a CUDA tensor directly from a Python list synchronizes the - # current stream. During overlap scheduling that stream already waits for - # the previous forward, exposing the import rebuild as an inter-step gap. - # Explicit non-blocking copies keep the CPU free to enqueue every rebuild - # and the next forward while preserving their existing stream order. - page_ids = torch.empty_like(page_ids_host, device=device) - page_valid_tokens = torch.empty_like(page_valid_tokens_host, device=device) - page_ids.copy_(page_ids_host, non_blocking=True) - page_valid_tokens.copy_(page_valid_tokens_host, non_blocking=True) - return page_ids, page_valid_tokens - - -def rebuild_fp4_mla_disagg_imported_cache( - kv_cache_manager: Any, - request_id: int, - prompt_len: int, -) -> bool: - """Rebuild GEN-local FP4 MLA sidecars after a disaggregated KV import. - - The disaggregated payload carries the native V2 K, K-scale, and BF16 HP - roles. V scales and CuTeDSL's V-packed layout are deterministic - process-local views, so rebuilding them here avoids transfer bandwidth and - guarantees they are ready before a first decode step that may execute - through a pre-captured CUDA graph. - """ - if ( - kv_cache_manager is None - or getattr(kv_cache_manager, "dtype", None) != DataType.NVFP4 - or getattr(kv_cache_manager, "kv_factor", None) != 1 - or getattr(kv_cache_manager, "mla_v_scale_head_dim", None) is None - or not callable(getattr(kv_cache_manager, "get_fp4_mla_page_table_spec", None)) - ): - return False - if not isinstance(prompt_len, int) or prompt_len < 0: - raise ValueError( - f"FP4 MLA disaggregated import needs a nonnegative prompt_len, got {prompt_len}." - ) - # Helix assigns whole pages round-robin, so a rank may own no prompt pages. - # There are no process-local V sidecars to rebuild on that rank. - if prompt_len == 0: - return True - - page_size = int(kv_cache_manager.tokens_per_block) - if page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - "FP4 MLA disaggregated import requires " - f"tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK}, got {page_size}." - ) - pp_layers = list(getattr(kv_cache_manager, "pp_layers", ())) - num_local_layers = int(getattr(kv_cache_manager, "num_local_layers", len(pp_layers))) - if len(pp_layers) != num_local_layers: - raise RuntimeError( - "FP4 MLA disaggregated import cannot map local to global layers: " - f"{len(pp_layers)} PP layers for {num_local_layers} local layers." - ) - fp4_local_layers = list( - getattr(kv_cache_manager, "_fp4_mla_compact_to_local", range(num_local_layers)) - ) - if not fp4_local_layers: - raise RuntimeError("FP4 MLA disaggregated import found no local MLA layers.") - if any(local_layer < 0 or local_layer >= num_local_layers for local_layer in fp4_local_layers): - raise RuntimeError( - "FP4 MLA disaggregated import has invalid compact-to-local layer mapping: " - f"{fp4_local_layers}." - ) - - num_prompt_pages = _ceil_div(prompt_len, page_size) - first_attention_layer = pp_layers[fp4_local_layers[0]] - block_ids_per_seq = kv_cache_manager.get_batch_cache_indices( - [int(request_id)], layer_idx=first_attention_layer - ) - if len(block_ids_per_seq) != 1 or len(block_ids_per_seq[0]) < num_prompt_pages: - available = len(block_ids_per_seq[0]) if block_ids_per_seq else 0 - raise RuntimeError( - "FP4 MLA disaggregated import is missing prompt pages for request " - f"{request_id}: need {num_prompt_pages}, have {available}." - ) - prompt_block_ids = [int(block_id) for block_id in block_ids_per_seq[0][:num_prompt_pages]] - - v_scale_pool = kv_cache_manager.get_mla_v_scale_pool() - if not isinstance(v_scale_pool, torch.Tensor): - raise RuntimeError("FP4 MLA disaggregated import requires the manager V-scale pool.") - page_ids, page_valid_tokens = _stage_fp4_mla_import_page_metadata( - prompt_block_ids, - prompt_len=prompt_len, - page_size=page_size, - device=v_scale_pool.device, - ) - - v_scale_head_dim = int(kv_cache_manager.mla_v_scale_head_dim) - cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND - for compact_layer, local_layer in enumerate(fp4_local_layers): - layer_idx = pp_layers[local_layer] - kv_cache, sf_cache = kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) - _rebuild_fp4_mla_v_scales_from_k_scales( - sf_cache, - v_scale_pool, - page_ids, - page_valid_tokens, - local_layer=compact_layer, - v_head_dim=v_scale_head_dim, - page_size=page_size, - ) - if cutedsl_backend and not _fp4_mla_cutedsl_fused_v_transpose_enabled(): - v_head_dim = getattr(kv_cache_manager, "mla_v_head_dim", None) - if v_head_dim is None: - raise RuntimeError( - "CuTeDSL FP4 MLA disaggregated import requires a persistent V head dimension." - ) - v_packed = kv_cache_manager.get_mla_v_packed_pool(compact_layer) - if not isinstance(v_packed, torch.Tensor): - raise RuntimeError( - "CuTeDSL FP4 MLA disaggregated import requires the persistent V-packed pool." - ) - _repack_cutedsl_v_packed_cache( - v_packed, - kv_cache, - page_ids, - v_head_dim=int(v_head_dim), - page_size=page_size, - block_v=FP4_MLA_SCALE_ROW_GROUP, - ) - return True - - -# Python launch helpers - - -def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata, "_fp4_mla_kv_global_scale", None) - if ( - not isinstance(global_scale, torch.Tensor) - or global_scale.device != device - or global_scale.dtype != torch.float32 - or global_scale.numel() != 1 - ): - raise RuntimeError("FP4 MLA requires a preallocated FP32 KV global-scale tensor.") - return global_scale - - -def _get_fp4_mla_q_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata, "_fp4_mla_q_global_scale", None) - if ( - not isinstance(global_scale, torch.Tensor) - or global_scale.device != device - or global_scale.dtype != torch.float32 - or global_scale.numel() != 1 - ): - raise RuntimeError("FP4 MLA requires a preallocated FP32 Q global-scale tensor.") - return global_scale - - -def _get_fp4_mla_kv_cache_tensors( - metadata: Any, layer_idx: int -) -> tuple[torch.Tensor, torch.Tensor]: - return metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) - - -def _get_fp4_mla_hp_pool_layout( - metadata: Any, - pool: torch.Tensor, -) -> tuple[int, int]: - """Return the manager-owned HP ring size and per-token head dimension.""" - manager = getattr(metadata, "kv_cache_manager", None) - if manager is None or not hasattr(manager, "_fp4_mla_hp_pool_size"): - raise ValueError("FP4 MLA requires a V2 manager-owned HP ring.") - hp_pool_size = int(manager._fp4_mla_hp_pool_size) - if ( - hp_pool_size < HP_BLOCK_SIZE - or pool.ndim != 4 - or pool.shape[2] < 1 - or pool.shape[-1] % hp_pool_size != 0 - ): - raise ValueError( - "FP4 MLA high-precision KV pool does not match its configured " - f"ring: shape={tuple(pool.shape)}, ring_size={hp_pool_size}." - ) - return hp_pool_size, pool.shape[-1] // hp_pool_size - - -def _validate_fp4_mla_hp_generation_width( - hp_pool_size: int, - generation_len: int, -) -> None: - """Ensure one target plus rewindable drafts fit without clobbering the live tail.""" - max_rewind_len = hp_pool_size - HP_BLOCK_SIZE - if generation_len <= 0 or generation_len - 1 > max_rewind_len: - raise RuntimeError( - "FP4 MLA generation exceeds the HP ring's rewind slack: " - f"generation={generation_len}, max_rewind={max_rewind_len}." - ) - - -def _validate_fp4_mla_kv_storage_shape( - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - *, - head_dim: int, - backend: str, -) -> int: - """Validate the backend-specific physical KV and scale strides.""" - residual_dim = FP4_MLA_K_RESIDUAL_DIM if backend in _FP4_MLA_K_RESIDUAL_BACKENDS else 0 - expected_storage_head_dim = head_dim + residual_dim - storage_head_dim = kv_cache.shape[-1] * 2 - if storage_head_dim != expected_storage_head_dim: - raise RuntimeError( - "FP4 MLA KV cache storage head dimension does not match the selected backend: " - f"got {storage_head_dim}, expected {expected_storage_head_dim}. Recreate the engine " - f"after setting {FP4_MLA_ATTENTION_BACKEND_ENV}." - ) - - expected_scale_columns = expected_storage_head_dim // FP4_BLOCK_SIZE - if sf_cache.shape[-1] != expected_scale_columns: - raise RuntimeError( - "FP4 MLA KV cache scale storage does not match the contiguous data layout: " - f"got {sf_cache.shape[-1]} columns, expected {expected_scale_columns}." - ) - return storage_head_dim - - -def _scatter_fp4_mla_kv_cache_2d_context( - metadata: Any, - latent_cache: torch.Tensor, - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - v_sf: Optional[torch.Tensor], - global_scale: torch.Tensor, - rotary_cos_sin: Optional[torch.Tensor], - q_context: Optional[torch.Tensor], - q_nope_head_dim: Optional[int], - *, - token_offset: int, - local_layer: int, - v_head_dim: int, - head_dim: int, - num_tokens: int, - num_dim_blocks: int, - sf_per_token: int, - sf_per_page: int, - v_packed_base: Optional[torch.Tensor] = None, - v_page_offset: int = 0, -) -> bool: - num_contexts = metadata.num_contexts - if num_contexts > 0: - prompt_lens_cpu = metadata.prompt_lens_cpu_runtime[:num_contexts] - ctx_token_count = int(prompt_lens_cpu.sum().item()) - if num_tokens != ctx_token_count: - raise RuntimeError( - f"FP4 MLA 2D context scatter needs {ctx_token_count} context tokens, got " - f"{num_tokens}." - ) - _validate_fp4_mla_context_start_alignment(metadata, num_contexts, alignment=FP4_BLOCK_SIZE) - - apply_k_rope = rotary_cos_sin is not None - rope_dim = ( - _validate_fp4_mla_context_rope(latent_cache, rotary_cos_sin, v_head_dim) - if rotary_cos_sin is not None - else 0 - ) - rotary_cos_sin_ptr = rotary_cos_sin if rotary_cos_sin is not None else latent_cache - - apply_q_rope = q_context is not None - block_q_heads = 16 - if apply_q_rope: - if rotary_cos_sin is None or q_nope_head_dim is None: - raise ValueError("FP4 MLA fused context Q-RoPE requires a rotary table and Q layout.") - q_head_dim = q_nope_head_dim + rope_dim - if ( - q_context.dtype != torch.bfloat16 - or q_context.device != latent_cache.device - or q_context.ndim != 2 - or q_context.shape[0] != num_tokens - or q_context.shape[1] <= 0 - or q_nope_head_dim <= 0 - or q_context.shape[1] % q_head_dim != 0 - or not q_context.is_contiguous() - ): - raise ValueError( - "FP4 MLA fused context Q-RoPE requires a contiguous same-device BF16 " - f"tensor shaped [tokens, heads * ({q_nope_head_dim} + {rope_dim})]." - ) - num_q_heads = q_context.shape[1] // q_head_dim - q_context_view = q_context.view(num_tokens, num_q_heads, q_head_dim) - q_head_blocks = triton.cdiv(num_q_heads, block_q_heads) - else: - num_q_heads = 0 - q_context_view = latent_cache - q_head_blocks = 0 - - hp_pool = getattr(metadata, "high_precision_kv_pool", None) - if not isinstance(hp_pool, torch.Tensor): - raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") - if hp_pool.device != latent_cache.device: - raise ValueError("FP4 MLA latent cache and high-precision pool must share a device.") - if hp_pool.dtype != torch.bfloat16: - raise TypeError(f"FP4 MLA high-precision KV pool must use BF16, got {hp_pool.dtype}.") - hp_pool_size, pool_head_dim = _get_fp4_mla_hp_pool_layout(metadata, hp_pool) - if pool_head_dim < head_dim: - raise RuntimeError( - f"FP4 MLA HP pool head dimension is too small: got " - f"{pool_head_dim}, need at least {head_dim}." - ) - if local_layer < 0 or local_layer >= hp_pool.shape[1]: - raise ValueError( - f"FP4 MLA local layer {local_layer} is outside the HP pool's {hp_pool.shape[1]} layers." - ) - if hp_pool.stride(-1) != 1: - raise ValueError("FP4 MLA high-precision KV pool must be contiguous in head_dim.") - hp_page_ids = metadata._fp4_mla_hp_page_indices - if not isinstance(hp_page_ids, torch.Tensor): - raise RuntimeError("FP4 MLA context cache update requires HP page metadata.") - store_hp_tail = num_contexts > 0 - num_hp_pages = hp_pool.shape[0] - pool_s0 = hp_pool.stride(0) - pool_s1 = hp_pool.stride(1) - - write_v_packed = v_packed_base is not None - v_packed_output = v_packed_base if write_v_packed else kv_cache - v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 - v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 - - _fp4_mla_context_cache_update_kernel[ - ( - num_tokens, - num_dim_blocks + q_head_blocks, - ) - ]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - latent_cache, - q_context_view, - global_scale, - rotary_cos_sin_ptr, - hp_pool, - hp_page_ids, - metadata.batch_indices, - metadata.positions, - metadata.paged_kv_indices, - metadata.paged_kv_indptr, - metadata.paged_kv_indices.shape[0], - metadata.paged_kv_indptr.shape[0], - metadata.batch_indices.shape[0], - v_sf.shape[1] if v_sf is not None else kv_cache.shape[0], - v_sf.shape[0] if v_sf is not None else hp_pool.shape[1], - num_contexts, - num_hp_pages, - token_offset, - num_tokens, - local_layer, - v_page_offset if write_v_packed else 0, - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - latent_cache.stride(0), - latent_cache.stride(1), - q_context_view.stride(0), - q_context_view.stride(1) if apply_q_rope else 0, - q_context_view.stride(2) if apply_q_rope else 0, - v_sf.stride(0) if v_sf is not None else 0, - v_sf.stride(1) if v_sf is not None else 0, - v_packed_s0, - v_packed_s1, - pool_s0, - pool_s1, - HEAD_D=head_dim, - V_HEAD_D=v_head_dim, - HP_BLOCK=FP4_BLOCK_SIZE, - HP_POOL_SIZE=hp_pool_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=sf_per_page, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - STORE_K_RESIDUAL=(_fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS), - ROPE_DIM=rope_dim, - APPLY_K_ROPE=apply_k_rope, - APPLY_Q_ROPE=apply_q_rope, - NUM_DIM_BLOCKS=num_dim_blocks, - NUM_Q_HEADS=num_q_heads, - Q_NOPE_DIM=q_nope_head_dim if q_nope_head_dim is not None else 0, - BLOCK_Q_HEADS=block_q_heads, - POOL_HEAD_D=pool_head_dim, - STORE_HP_TAIL=store_hp_tail, - WRITE_V_PACKED=write_v_packed, - STORE_V_SCALE=v_sf is not None, - ) - return store_hp_tail - - -def _fp4_mla_generation_num_blocks_device(metadata: Any) -> torch.Tensor: - """Device-side scalar view holding the generation page-table capacity. - - ``paged_kv_indptr_decode[num_gen]`` is the fixed-stride generation-table - endpoint. Device kernels combine it with the live KV lengths, so inactive - slots are never consumed. - """ - num_gen = metadata.num_seqs - metadata.num_contexts - return metadata.paged_kv_indptr_decode[num_gen : num_gen + 1] - - -def _fp4_mla_uniform_generation_lengths( - metadata: Any, - num_gen_tokens: int, - num_gen: int, - plan_geometry: Optional[_PackedPlanGeometry] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Return preallocated CUDA generation KV and append lengths. - - With ``plan_geometry`` the launch that writes the lengths also writes the - step's packed split plan (one CuTe DSL kernel instead of two launches), and - a stale plan is rewritten even when the lengths are already current. - """ - if num_gen <= 0 or num_gen_tokens % num_gen != 0: - raise RuntimeError("FP4 MLA generation requires a non-empty uniform request batch.") - if plan_geometry is not None and plan_geometry.num_requests != num_gen: - raise RuntimeError("FP4 MLA packed plan geometry does not match the generation batch.") - - num_contexts = metadata.num_contexts - num_seqs = metadata.num_seqs - kv_lens_gen = metadata.kv_lens_cuda_runtime[num_contexts:num_seqs] - prompt_lens_gen = metadata.prompt_lens_cuda_runtime[num_contexts:num_seqs] - corrected_kv_lens = getattr(metadata, "fp4_mla_generation_kv_lens", None) - generation_lens = getattr(metadata, "fp4_mla_generation_append_lens", None) - tensors = ( - kv_lens_gen, - prompt_lens_gen, - corrected_kv_lens, - generation_lens, - ) - if ( - not all(isinstance(tensor, torch.Tensor) for tensor in tensors) - or corrected_kv_lens.numel() < num_gen - or generation_lens.numel() < num_gen - or not all(tensor.is_cuda for tensor in tensors) - ): - raise RuntimeError("FP4 MLA generation lengths require preallocated CUDA buffers.") - - record_for_capture = bool( - getattr(metadata, "is_cuda_graph", False) - and torch.cuda.is_current_stream_capturing() - and not getattr( - metadata, - "_fp4_mla_generation_lengths_capture_recorded", - False, - ) - ) - precomputed = ( - not record_for_capture - and metadata.fp4_mla_generation_lengths_num_tokens == num_gen_tokens - and metadata.fp4_mla_generation_lengths_num_seqs == num_gen - and metadata.fp4_mla_generation_lengths_num_contexts == num_contexts - ) - plan_pending = ( - plan_geometry is not None - and plan_geometry.key is not None - and getattr(metadata, "_fp4_mla_packed_plan_key", None) != plan_geometry.key - ) - if not precomputed: - if plan_pending: - _launch_fp4_mla_packed_plan( - plan_geometry, - kv_lens_gen, - metadata, - generation_lengths=( - prompt_lens_gen, - corrected_kv_lens[:num_gen], - generation_lens[:num_gen], - num_gen_tokens // num_gen, - ), - ) - else: - populate_fp4_mla_generation_lengths( - kv_lens_gen, - prompt_lens_gen, - corrected_kv_lens[:num_gen], - generation_lens[:num_gen], - num_gen_tokens=num_gen_tokens, - num_gen=num_gen, - ) - metadata.fp4_mla_generation_lengths_num_tokens = num_gen_tokens - metadata.fp4_mla_generation_lengths_num_seqs = num_gen - metadata.fp4_mla_generation_lengths_num_contexts = num_contexts - if record_for_capture: - metadata._fp4_mla_generation_lengths_capture_recorded = True - elif plan_pending: - _launch_fp4_mla_packed_plan(plan_geometry, corrected_kv_lens[:num_gen], metadata) - return corrected_kv_lens[:num_gen], generation_lens[:num_gen] - - -def _materialize_fp4_mla_device_page_table_for_forward( - metadata: Any, - generation_kv_lens: Optional[torch.Tensor] = None, -) -> None: - """Materialize all fixed-stride rows from final per-forward device lengths.""" - if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): - raise RuntimeError("FP4 MLA cache update requires fixed-stride device page metadata.") - if bool(getattr(metadata, "_fp4_mla_device_page_table_valid", False)): - return - - num_contexts = int(metadata.num_contexts) - num_sequences = int(metadata.num_seqs) - num_generation_sequences = num_sequences - num_contexts - if generation_kv_lens is None: - if num_generation_sequences > 0: - num_generation_tokens = int(metadata.num_tokens) - int(metadata.num_ctx_tokens) - generation_kv_lens, _ = _fp4_mla_uniform_generation_lengths( - metadata, - num_generation_tokens, - num_generation_sequences, - ) - else: - generation_kv_lens = metadata.kv_lens_cuda_runtime[num_contexts:num_sequences] - materialize_fp4_mla_device_page_table( - metadata, - metadata.kv_lens_cuda_runtime[:num_sequences], - generation_kv_lens, - ) - - -def _scatter_fp4_mla_kv_cache_2d_generation( - metadata: Any, - latent_cache: torch.Tensor, - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - v_sf: torch.Tensor, - global_scale: torch.Tensor, - *, - token_offset: int, - local_layer: int, - v_head_dim: int, - head_dim: int, - num_tokens: int, - num_dim_blocks: int, - sf_per_token: int, - sf_per_page: int, - rotary_cos_sin: torch.Tensor, - q_pe: torch.Tensor, - q_rope_out: torch.Tensor, - q_quant_input: torch.Tensor, - q_fp4_out: torch.Tensor, - q_sf_out: torch.Tensor, - v_packed_base: Optional[torch.Tensor], - v_page_offset: int, - helix_position_offsets: Optional[torch.Tensor], - helix_is_inactive_rank: Optional[torch.Tensor], -) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - num_contexts = metadata.num_contexts - num_seqs = metadata.num_seqs - num_gen = num_seqs - num_contexts - if num_gen <= 0: - return - if num_tokens < num_gen: - raise RuntimeError( - f"FP4 MLA 2D generation scatter needs at least {num_gen} generation " - f"tokens, got {num_tokens}." - ) - if num_tokens % num_gen != 0: - raise NotImplementedError( - "FP4 MLA no-dequant generation scatter requires a uniform linear MTP " - f"generation length, got {num_tokens} tokens for {num_gen} sequences." - ) - # The prompt_lens/kv_lens runtime aliases can lag at the decode anchor - # (seq_lens == 1) under CUDA graph / one-engine MTP while each generation - # sequence really appends num_tokens // num_gen tokens this step. Recover the - # true per-sequence lengths for the no-dequant kernel below (a no-op when the - # aliases already match). As the first FP4 MLA kernel of the step, the same - # launch also writes the packed split plan the decode layers reuse. - plan_geometry = None - if q_pe.ndim == 3 and q_pe.shape[1] > 0: - plan_geometry = _fp4_mla_packed_plan_geometry( - metadata, - num_heads=q_pe.shape[1], - query_len=num_tokens // num_gen, - num_requests=num_gen, - device=latent_cache.device, - helix_kv_bounds=_fp4_mla_helix_plan_bounds( - metadata, - num_tokens, - use_helix=helix_position_offsets is not None or helix_is_inactive_rank is not None, - ), - ) - kv_lens_gen, gen_lens_gen = _fp4_mla_uniform_generation_lengths( - metadata, num_tokens, num_gen, plan_geometry=plan_geometry - ) - _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens_gen) - - pool = getattr(metadata, "high_precision_kv_pool", None) - if pool is None: - raise RuntimeError("FP4 MLA 2D generation scatter requires the HP KV pool.") - try: - hp_pool_size, hp_head_dim = _get_fp4_mla_hp_pool_layout(metadata, pool) - except ValueError as error: - raise RuntimeError(str(error)) from error - if hp_head_dim < head_dim: - raise RuntimeError( - f"FP4 MLA 2D generation scatter needs at least {head_dim} HP channels, got " - f"{hp_head_dim}." - ) - hp_page_ids = _fp4_mla_generation_hp_page_ids(metadata, num_gen) - if not isinstance(hp_page_ids, torch.Tensor): - raise RuntimeError("FP4 MLA generation requires HP page metadata.") - num_hp_pages = pool.shape[0] - - max_gen_len = num_tokens // num_gen - use_helix = helix_position_offsets is not None or helix_is_inactive_rank is not None - use_helix_local_slots = use_helix and bool(getattr(metadata, "_helix_spec_tokens_valid", False)) - if use_helix: - if helix_position_offsets is None or helix_is_inactive_rank is None: - raise RuntimeError( - "FP4 MLA Helix requires both position-offset and inactive-rank metadata." - ) - if max_gen_len != 1 and not use_helix_local_slots: - raise NotImplementedError( - "FP4 MLA multi-token Helix requires speculative per-token metadata." - ) - if ( - helix_position_offsets.dtype != torch.int32 - or helix_position_offsets.device != latent_cache.device - or helix_position_offsets.ndim != 1 - or helix_position_offsets.numel() < num_tokens - or not helix_position_offsets.is_contiguous() - ): - raise ValueError( - "FP4 MLA Helix position offsets must be a contiguous same-device " - "int32 tensor covering every generation token." - ) - if ( - helix_is_inactive_rank.dtype != torch.bool - or helix_is_inactive_rank.device != latent_cache.device - or helix_is_inactive_rank.ndim != 1 - or helix_is_inactive_rank.numel() < num_gen - or not helix_is_inactive_rank.is_contiguous() - ): - raise ValueError( - "FP4 MLA Helix inactive-rank metadata must be a contiguous " - "same-device bool tensor covering every generation sequence." - ) - if use_helix_local_slots: - helix_local_slots = getattr(metadata, "helix_local_slots", None) - if ( - not isinstance(helix_local_slots, torch.Tensor) - or helix_local_slots.dtype != torch.int32 - or helix_local_slots.device != latent_cache.device - or helix_local_slots.ndim != 1 - or helix_local_slots.numel() < num_tokens - or not helix_local_slots.is_contiguous() - ): - raise ValueError( - "FP4 MLA speculative Helix local slots must be a contiguous " - "same-device int32 tensor covering every generation token." - ) - else: - helix_local_slots = helix_position_offsets - else: - helix_position_offsets = kv_lens_gen - helix_local_slots = kv_lens_gen - helix_is_inactive_rank = gen_lens_gen - _validate_fp4_mla_hp_generation_width(hp_pool_size, max_gen_len) - max_rewind_len = hp_pool_size - HP_BLOCK_SIZE - page_ids = _fp4_mla_generation_page_ids(metadata, num_gen) - rope_dim = head_dim - v_head_dim - block_q_heads = 32 - q1_kv_blocks_per_program = 1 - # Grouped Q1 kernels reuse K's packed codes for V. Keep one dimension - # block per program until their warp-specialized V quantizer is split out. - if max_gen_len == 1: - q1_kv_blocks_per_program = _fp4_mla_q1_kv_blocks_per_program(num_gen, v_head_dim) - q_prefix_block_dim = ( - FP4_MLA_Q1_PREFIX_BLOCK_DIM if max_gen_len == 1 else FP4_MLA_Q_PREFIX_BLOCK_DIM - ) - q_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // q_prefix_block_dim - q_prefix_blocks_per_program = _fp4_mla_q1_prefix_blocks_per_program( - num_gen, - q1_kv_blocks_per_program, - ) - q_work_blocks = q_prefix_blocks // q_prefix_blocks_per_program + 1 - if latent_cache.dtype != torch.bfloat16: - raise TypeError( - "Fused FP4 MLA Q/K RoPE and cache storage requires BF16 latent KV, " - f"got {latent_cache.dtype}." - ) - if rope_dim <= 0 or rope_dim % 2 != 0: - raise ValueError( - "Fused FP4 MLA K RoPE requires a positive even K tail, " - f"got head_dim={head_dim} v_head_dim={v_head_dim}." - ) - if ( - rotary_cos_sin is None - or rotary_cos_sin.device != latent_cache.device - or rotary_cos_sin.dtype != torch.float32 - or rotary_cos_sin.numel() % (rope_dim * 2) != 0 - ): - raise ValueError( - "Fused FP4 MLA K RoPE requires a same-device FP32 rotary " - f"table with rows of {rope_dim * 2} values." - ) - if ( - q_pe is None - or q_rope_out is None - or q_pe.dtype != torch.bfloat16 - or q_rope_out.dtype != torch.bfloat16 - or q_pe.device != latent_cache.device - or q_rope_out.device != latent_cache.device - or q_pe.ndim != 3 - or q_rope_out.shape != q_pe.shape - or q_pe.shape[0] != num_tokens - or q_pe.shape[1] <= 0 - or q_pe.shape[2] != rope_dim - ): - raise ValueError( - "Fused FP4 MLA Q RoPE requires same-device BF16 q_pe and " - f"q_rope_out tensors shaped [tokens, heads, {rope_dim}]." - ) - num_q_heads = q_pe.shape[1] - q_head_blocks = _ceil_div(num_q_heads, block_q_heads) - q_rows_per_seq = _fp4_mla_packed_rows_per_seq(num_q_heads, max_gen_len) - if q_rows_per_seq and q_fp4_out.shape[0] < num_gen * q_rows_per_seq: - raise RuntimeError( - "FP4 MLA packed Q buffer is smaller than the generation batch: " - f"{q_fp4_out.shape[0]} rows < {num_gen} x {q_rows_per_seq}." - ) - max_gen_tiles = _ceil_div(max_gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) - rotary_table = rotary_cos_sin - q_global_scale = _get_fp4_mla_q_global_scale(metadata, latent_cache.device) - q_pe_input = q_pe - q_rope_output = q_rope_out - q_full_input = q_quant_input - q_fp4_output = q_fp4_out - q_sf_output = q_sf_out - write_v_packed = v_packed_base is not None - v_packed_output = v_packed_base if write_v_packed else kv_cache - v_packed_s0 = v_packed_output.stride(0) if write_v_packed else 0 - v_packed_s1 = v_packed_output.stride(1) if write_v_packed else 0 - kv_work_blocks = ( - v_head_dim // FP4_BLOCK_SIZE // q1_kv_blocks_per_program + 1 - if max_gen_len == 1 - else num_dim_blocks - ) - launch_grid = ( - num_gen, - max( - kv_work_blocks, - max_gen_len * q_head_blocks * q_work_blocks, - ), - ) - store_k_residual = _fp4_mla_attention_backend() in _FP4_MLA_K_RESIDUAL_BACKENDS - - def launch_generation_update( - grid: tuple[int, ...], - *, - page_ids_len: int, - indptr_len: int, - max_gen_tiles_variant: int, - q_prefix_block_dim_variant: int, - q_prefix_blocks_variant: int, - q_prefix_blocks_per_program_variant: int, - q1_kv_blocks_per_program_variant: int, - ) -> None: - q_work_blocks_variant = q_prefix_blocks_variant // q_prefix_blocks_per_program_variant + 1 - _fp4_mla_generation_fused_qk_rope_cache_update_kernel[grid]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - pool, - latent_cache, - global_scale, - q_global_scale, - rotary_table, - q_pe_input, - q_rope_output, - q_full_input, - q_fp4_output, - q_sf_output, - kv_lens_gen, - gen_lens_gen, - helix_position_offsets, - helix_local_slots, - helix_is_inactive_rank, - page_ids, - hp_page_ids, - metadata.paged_kv_indptr_decode, - page_ids_len, - hp_page_ids.numel(), - indptr_len, - v_sf.shape[1], - num_hp_pages, - v_sf.shape[0], - local_layer, - v_page_offset if write_v_packed else 0, - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - pool.stride(0), - pool.stride(1), - v_sf.stride(0), - v_sf.stride(1), - v_packed_s0, - v_packed_s1, - q_pe_input.stride(0), - q_pe_input.stride(1) if q_pe_input.ndim > 1 else 0, - q_pe_input.stride(2) if q_pe_input.ndim > 2 else 0, - q_rope_output.stride(0), - q_rope_output.stride(1) if q_rope_output.ndim > 1 else 0, - q_rope_output.stride(2) if q_rope_output.ndim > 2 else 0, - HEAD_D=hp_head_dim, - V_HEAD_D=v_head_dim, - HP_BLOCK=FP4_BLOCK_SIZE, - HP_POOL_SIZE=hp_pool_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=sf_per_page, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - STORE_K_RESIDUAL=store_k_residual, - FUSE_ROPE_CACHE_STORE=True, - USE_HELIX=use_helix, - USE_HELIX_LOCAL_SLOTS=use_helix_local_slots, - WRITE_V_PACKED=write_v_packed, - MAX_GEN_TILES=max_gen_tiles_variant, - ROPE_DIM=rope_dim, - ROPE_PAIR_BLOCK=triton.next_power_of_2(rope_dim // 2), - NUM_DIM_BLOCKS=num_dim_blocks, - NUM_Q_HEADS=num_q_heads, - Q_HEAD_BLOCKS=max(q_head_blocks, 1), - BLOCK_Q_HEADS=block_q_heads, - Q_PREFIX_D=FP4_MLA_Q_PREFIX_DIM, - Q_PREFIX_BLOCK_D=q_prefix_block_dim_variant, - Q_PREFIX_BLOCKS=q_prefix_blocks_variant, - Q_PREFIX_BLOCKS_PER_PROGRAM=q_prefix_blocks_per_program_variant, - Q_WORK_BLOCKS=q_work_blocks_variant, - Q_SF_COLS=FP4_MLA_Q_SF_GROUPS, - WRITE_Q=True, - Q1_KV_BLOCKS_PER_PROGRAM=q1_kv_blocks_per_program_variant, - Q_ROWS_PER_SEQ=q_rows_per_seq, - maxnreg=56, - ) - - # Triton compiles and loads a CUDA module on first launch. Use runtime-zero - # work here so every reachable static tuning variant is resident before - # warmup hands the engine to serving. - if getattr(metadata, "is_warmup", False) and not torch.cuda.is_current_stream_capturing(): - configured_generation_len = int(getattr(metadata, "max_total_draft_tokens", 0) or 0) + 1 - max_preload_generation_len = max(max_gen_len, configured_generation_len) - if max_preload_generation_len - 1 > max_rewind_len: - raise NotImplementedError( - "FP4 MLA finite Triton preload exceeds the HP ring's rewind slack: " - f"max_rewind={max_rewind_len}, generation=" - f"{max_preload_generation_len}." - ) - max_num_sequences = int(getattr(metadata, "max_num_sequences", num_seqs) or num_seqs) - q1_variants = _fp4_mla_q1_preload_variants( - max_num_sequences, - v_head_dim, - ) - multi_token_tiles = tuple( - sorted( - { - _ceil_div(gen_len + FP4_BLOCK_SIZE - 1, FP4_BLOCK_SIZE) - for gen_len in range(2, max_preload_generation_len + 1) - } - ) - ) - preload_key = ( - "generation-cache-update", - str(latent_cache.device), - hp_pool_size, - hp_head_dim, - v_head_dim, - num_q_heads, - q_rows_per_seq, - metadata.page_size, - write_v_packed, - store_k_residual, - use_helix, - use_helix_local_slots, - tuple(q1_variants), - multi_token_tiles, - tuple(kv_cache.stride()), - tuple(sf_cache.stride()), - tuple(pool.stride()), - tuple(v_sf.stride()), - tuple(v_packed_output.stride()), - tuple(q_pe_input.stride()), - tuple(q_rope_output.stride()), - str(kv_cache.dtype), - str(latent_cache.dtype), - str(q_pe_input.dtype), - str(q_fp4_output.dtype), - str(q_sf_output.dtype), - ) - preload_keys = _fp4_mla_triton_preload_key_set(metadata) - if preload_key not in preload_keys: - context_page_ids = getattr(metadata, "_paged_kv_indices", None) - if not isinstance(context_page_ids, torch.Tensor): - context_page_ids = page_ids - context_indptr = getattr(metadata, "_paged_kv_indptr", None) - if not isinstance(context_indptr, torch.Tensor): - context_indptr = metadata.paged_kv_indptr_decode - context_batch_indices = getattr(metadata, "batch_indices", None) - if not isinstance(context_batch_indices, torch.Tensor): - context_batch_indices = hp_page_ids - context_positions = getattr(metadata, "positions", None) - if not isinstance(context_positions, torch.Tensor): - context_positions = hp_page_ids - _fp4_mla_context_cache_update_kernel[(1, num_dim_blocks)]( - kv_cache, - sf_cache, - v_sf, - v_packed_output, - latent_cache, - latent_cache, - global_scale, - rotary_table, - pool, - hp_page_ids, - context_batch_indices, - context_positions, - context_page_ids, - context_indptr, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - latent_cache.stride(0), - latent_cache.stride(1), - latent_cache.stride(0), - 0, - 0, - v_sf.stride(0), - v_sf.stride(1), - v_packed_s0, - v_packed_s1, - pool.stride(0), - pool.stride(1), - HEAD_D=head_dim, - V_HEAD_D=v_head_dim, - HP_BLOCK=FP4_BLOCK_SIZE, - HP_POOL_SIZE=hp_pool_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=sf_per_token, - SF_PER_PAGE=sf_per_page, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - STORE_K_RESIDUAL=store_k_residual, - ROPE_DIM=rope_dim, - APPLY_K_ROPE=True, - APPLY_Q_ROPE=False, - NUM_DIM_BLOCKS=num_dim_blocks, - NUM_Q_HEADS=0, - Q_NOPE_DIM=0, - BLOCK_Q_HEADS=16, - POOL_HEAD_D=hp_head_dim, - STORE_HP_TAIL=True, - WRITE_V_PACKED=write_v_packed, - STORE_V_SCALE=True, - ) - q1_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q1_PREFIX_BLOCK_DIM - for q1_kv_blocks, q1_prefix_blocks_per_program in q1_variants: - launch_generation_update( - (1, 1), - page_ids_len=0, - indptr_len=0, - max_gen_tiles_variant=1, - q_prefix_block_dim_variant=FP4_MLA_Q1_PREFIX_BLOCK_DIM, - q_prefix_blocks_variant=q1_prefix_blocks, - q_prefix_blocks_per_program_variant=q1_prefix_blocks_per_program, - q1_kv_blocks_per_program_variant=q1_kv_blocks, - ) - multi_prefix_blocks = FP4_MLA_Q_PREFIX_DIM // FP4_MLA_Q_PREFIX_BLOCK_DIM - for multi_token_tile in multi_token_tiles: - launch_generation_update( - (1, 1), - page_ids_len=0, - indptr_len=0, - max_gen_tiles_variant=multi_token_tile, - q_prefix_block_dim_variant=FP4_MLA_Q_PREFIX_BLOCK_DIM, - q_prefix_blocks_variant=multi_prefix_blocks, - q_prefix_blocks_per_program_variant=1, - q1_kv_blocks_per_program_variant=1, - ) - torch.cuda.synchronize(latent_cache.device) - preload_keys.add(preload_key) - - launch_generation_update( - launch_grid, - page_ids_len=page_ids.shape[0], - indptr_len=metadata.paged_kv_indptr_decode.shape[0], - max_gen_tiles_variant=max(max_gen_tiles, 1), - q_prefix_block_dim_variant=q_prefix_block_dim, - q_prefix_blocks_variant=q_prefix_blocks, - q_prefix_blocks_per_program_variant=q_prefix_blocks_per_program, - q1_kv_blocks_per_program_variant=q1_kv_blocks_per_program, - ) - return kv_lens_gen, gen_lens_gen, page_ids - - -# Public cache update and decode entry points - - -def load_fp4_mla_chunked_kv_cache( - metadata: Any, - layer_idx: int, - *, - num_ctx_cached_tokens: int, - cu_chunked_seq_len: torch.Tensor, - chunked_global_offset: torch.Tensor, - chunked_max_seq_len: int, - out_dtype: torch.dtype, - kv_lora_rank: int, - qk_rope_head_dim: int, - output_format: Literal["dense", "mxfp8"] = "dense", -) -> tuple[torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor]: - """Gather one cached-prefix partition from dense FP4 MLA V2 storage. - - ``mxfp8`` preserves BF16 reconstruction, returning ((E4M3 [M,512], - UE8M0 [pad_up(M,128)*16]), E4M3 [M,64]). Scales use R128c4 K32 - storage with K128 quantization; RoPE output uses unit FMHA scale. - """ - if output_format not in ("dense", "mxfp8"): - raise ValueError(f"Unsupported FP4 MLA gather output format: {output_format}.") - quantize_mxfp8 = output_format == "mxfp8" - if quantize_mxfp8 and (out_dtype != torch.bfloat16 or kv_lora_rank != 512): - raise ValueError("MXFP8 gather requires BF16 reconstruction and kv_lora_rank=512.") - if out_dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise TypeError(f"FP4 MLA chunk gather does not support output dtype {out_dtype}.") - if num_ctx_cached_tokens < 0: - raise ValueError("FP4 MLA chunk gather token count must be non-negative.") - if chunked_max_seq_len < 0: - raise ValueError("FP4 MLA chunk gather max sequence length must be non-negative.") - if num_ctx_cached_tokens > 0 and chunked_max_seq_len == 0: - raise ValueError("A non-empty FP4 MLA chunk gather requires a positive max length.") - if ( - kv_lora_rank <= 0 - or kv_lora_rank % FP4_BLOCK_SIZE != 0 - or qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM - ): - raise ValueError( - "FP4 MLA chunk gather requires a positive kv_lora_rank and " - f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got " - f"{kv_lora_rank} and {qk_rope_head_dim}." - ) - - num_contexts = int(metadata.num_contexts) - if num_ctx_cached_tokens > 0 and num_contexts <= 0: - raise ValueError("A non-empty FP4 MLA chunk gather requires context requests.") - tensors = (cu_chunked_seq_len, chunked_global_offset) - if any(not isinstance(tensor, torch.Tensor) or not tensor.is_cuda for tensor in tensors): - raise ValueError("FP4 MLA chunk gather metadata must use CUDA tensors.") - if ( - cu_chunked_seq_len.dtype != torch.int64 - or cu_chunked_seq_len.ndim != 1 - or cu_chunked_seq_len.numel() < num_contexts + 1 - or not cu_chunked_seq_len.is_contiguous() - ): - raise ValueError( - "FP4 MLA chunk gather requires a contiguous int64 cumulative-length tensor." - ) - if ( - chunked_global_offset.dtype != torch.int64 - or chunked_global_offset.ndim != 1 - or chunked_global_offset.numel() < num_contexts - or not chunked_global_offset.is_contiguous() - ): - raise ValueError("FP4 MLA chunk gather requires contiguous int64 global offsets.") - if cu_chunked_seq_len.device != chunked_global_offset.device: - raise ValueError("FP4 MLA chunk gather metadata tensors must share one device.") - - storage_dtype = torch.float8_e4m3fn if quantize_mxfp8 else out_dtype - compressed_kv = torch.empty( - (num_ctx_cached_tokens, kv_lora_rank), - dtype=storage_dtype, - device=cu_chunked_seq_len.device, - ) - k_pe = torch.empty( - (num_ctx_cached_tokens, qk_rope_head_dim), - dtype=storage_dtype, - device=cu_chunked_seq_len.device, - ) - if quantize_mxfp8: - compressed_sf = torch.empty( - (triton.cdiv(num_ctx_cached_tokens, 128) * 128 * (kv_lora_rank // 32),), - dtype=torch.uint8, - device=cu_chunked_seq_len.device, - ) - if num_ctx_cached_tokens == 0: - return ((compressed_kv, compressed_sf), k_pe) if quantize_mxfp8 else (compressed_kv, k_pe) - - if not bool(getattr(metadata, "_fp4_mla_device_page_table", False)): - raise RuntimeError("FP4 MLA chunk gather requires fixed-stride device page metadata.") - _materialize_fp4_mla_device_page_table_for_forward(metadata) - page_table_stride = int(metadata.fp4_mla_page_table_stride) - page_ids = metadata._paged_kv_indices - if ( - page_table_stride <= 0 - or not isinstance(page_ids, torch.Tensor) - or page_ids.dtype != torch.int32 - or not page_ids.is_cuda - or page_ids.device != cu_chunked_seq_len.device - or page_ids.numel() < num_contexts * page_table_stride - ): - raise RuntimeError("FP4 MLA chunk gather received invalid device page metadata.") - - kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) - head_dim = kv_lora_rank + qk_rope_head_dim - storage_head_dim = _validate_fp4_mla_kv_storage_shape( - kv_cache, - sf_cache, - head_dim=head_dim, - backend=_fp4_mla_attention_backend(), - ) - if storage_head_dim != head_dim + FP4_MLA_K_RESIDUAL_DIM: - raise RuntimeError( - "FP4 MLA chunk gather requires K residual storage for BF16 reconstruction." - ) - sf_cache = sf_cache.view(torch.float8_e4m3fn) - global_scale = _get_fp4_mla_global_scale(metadata, kv_cache.device) - if quantize_mxfp8: - from tensorrt_llm._torch.cute_dsl_kernels.rubin.fp4_mla_context_gather import gather_mxfp8 - - gather_mxfp8( - compressed_kv, - compressed_sf, - k_pe, - kv_cache, - sf_cache, - page_ids, - cu_chunked_seq_len, - chunked_global_offset, - global_scale, - chunked_max_seq_len, - metadata.page_size, - page_table_stride=page_table_stride, - num_contexts=num_contexts, - ) - return (compressed_kv, compressed_sf), k_pe - token_block = 16 - grid = ( - triton.cdiv(chunked_max_seq_len, token_block), - num_contexts, - triton.cdiv(head_dim, FP4_BLOCK_SIZE), - ) - _fp4_mla_chunked_cache_gather_kernel[grid]( - compressed_kv, - k_pe, - kv_cache, - sf_cache, - page_ids, - cu_chunked_seq_len, - chunked_global_offset, - global_scale, - chunked_max_seq_len, - page_table_stride, - kv_cache.shape[0], - metadata.page_size, - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - compressed_kv.stride(0), - k_pe.stride(0), - KV_LORA_RANK=kv_lora_rank, - QK_ROPE_HEAD_DIM=qk_rope_head_dim, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_TOKEN=storage_head_dim // FP4_BLOCK_SIZE, - TOKEN_BLOCK=token_block, - num_warps=4, - ) - return compressed_kv, k_pe - - -def scatter_fp4_mla_kv_cache( - metadata: Any, - latent_cache: torch.Tensor, - layer_idx: int, - *, - token_offset: int, - phase: _HPUpdatePhase, - local_layer: int, - v_head_dim: int, - rotary_cos_sin: Optional[torch.Tensor] = None, - q_pe: Optional[torch.Tensor] = None, - q_rope_out: Optional[torch.Tensor] = None, - q_quant_input: Optional[torch.Tensor] = None, - helix_position_offsets: Optional[torch.Tensor] = None, - helix_is_inactive_rank: Optional[torch.Tensor] = None, - q_context: Optional[torch.Tensor] = None, - q_nope_head_dim: Optional[int] = None, -) -> bool: - """Quantize MLA latent tokens and scatter them into the paged FP4 cache. - - Contract: this helper scatters exactly ``latent_cache.shape[0]`` tokens, - reading index metadata at ``batch_indices[token_offset : token_offset + N]`` - and ``positions[token_offset : token_offset + N]``. Callers must pass a - latent_cache pre-sliced to the current phase (context or generation) so - that ``shape[0]`` matches the number of index entries they intend to - consume. ``MLA.forward_impl`` (tensorrt_llm/_torch/modules/attention.py) - slices ``latent_cache[:num_ctx_tokens]`` for context and - ``latent_cache[num_ctx_tokens:]`` for generation before dispatching. - - Callers must pass ``phase``, ``local_layer``, and ``v_head_dim``. Context - scatter writes the final FP4 tile representation directly. Dimensions - below ``v_head_dim`` share one 16-token by 16-dim FP4 tile between K - and V, with the scale written into K's token-major and V's dim-major - layouts. Tail K-only dimensions use K's per-token 1D scales. For - exclusively owned CuTeDSL pages, context scatter also writes the - persistent packed-V sidecar. - Context scatter can rotate Q in place and rotate the K tail directly from - the unassembled latent tensor. When context Q is supplied for chunked - prefill, it also writes rotated K back for current-chunk attention. - Generation scatter rewrites each touched 16-token tile by reading - old tokens from the HP pool and new tokens from ``latent_cache``. The - static-scale generation - specialization can also rotate Q and new K tails while updating the HP pool. - The context kernel also stores the final incomplete tile in the BF16 HP - pool. When ``q_quant_input`` is supplied, the generation kernel also emits - backend-ready residual FP4 Q. The return value reports whether the current - phase updated the HP pool. - """ - if phase == "generation": - metadata._fp4_mla_prequantized_q = None - metadata._fp4_mla_prequantized_q_sf = None - metadata._fp4_mla_q_batch_capacity = None - if latent_cache.numel() == 0: - raise ValueError("FP4 MLA cache scatter requires at least one latent token.") - if q_context is not None and not latent_cache.is_contiguous(): - raise ValueError( - "FP4 MLA fused context Q/K RoPE requires contiguous latent_cache " - "storage for in-place current-K update." - ) - - latent_cache = latent_cache.reshape(latent_cache.shape[0], -1).contiguous() - num_tokens = latent_cache.shape[0] - head_dim = latent_cache.shape[-1] - if head_dim % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA KV head_dim must be divisible by {FP4_BLOCK_SIZE}, got {head_dim}." - ) - indices_len = metadata.batch_indices.shape[0] - positions_len = metadata.positions.shape[0] - if token_offset + num_tokens > indices_len or token_offset + num_tokens > positions_len: - raise RuntimeError( - f"FP4 MLA scatter would read batch_indices[{token_offset}:" - f"{token_offset + num_tokens}] / positions[{token_offset}:" - f"{token_offset + num_tokens}], but only {indices_len} / " - f"{positions_len} entries are available. This indicates " - "latent_cache was not pre-sliced to the current phase's token " - "range (see MLA.forward_impl)." - ) - - _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) - - backend = _fp4_mla_attention_backend() - global_scale = _get_fp4_mla_global_scale(metadata, latent_cache.device) - kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) - storage_head_dim = _validate_fp4_mla_kv_storage_shape( - kv_cache, - sf_cache, - head_dim=head_dim, - backend=backend, - ) - sf_per_token = storage_head_dim // FP4_BLOCK_SIZE - - if phase not in ("context", "generation"): - raise ValueError("FP4 MLA scatter requires phase='context' or 'generation'.") - context_only_cache = getattr(metadata.kv_cache_manager, "fp4_mla_context_only_cache", False) - if context_only_cache and phase != "context": - raise RuntimeError("FP4 MLA context-only cache cannot execute generation.") - if not context_only_cache and getattr(metadata, "fp4_mla_v_scale_pool", None) is None: - raise RuntimeError("FP4 MLA scatter requires the auxiliary V scale pool.") - if metadata.page_size % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA scatter requires page_size divisible by " - f"{FP4_BLOCK_SIZE}, got {metadata.page_size}." - ) - if v_head_dim > head_dim: - raise ValueError(f"FP4 MLA v_head_dim={v_head_dim} cannot exceed head_dim={head_dim}.") - if head_dim - v_head_dim != FP4_MLA_K_RESIDUAL_DIM: - raise ValueError( - "FP4 MLA K residual quantization requires the K-only tail to match " - f"the {FP4_MLA_K_RESIDUAL_DIM}-channel residual, got " - f"head_dim={head_dim} v_head_dim={v_head_dim}." - ) - if v_head_dim % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA v_head_dim must be divisible by {FP4_BLOCK_SIZE}, got {v_head_dim}." - ) - - sf_cache = sf_cache.view(torch.float8_e4m3fn) - v_sf = ( - None - if context_only_cache - else get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=v_head_dim) - ) - num_dim_blocks = triton.cdiv(head_dim, FP4_BLOCK_SIZE) - sf_per_page = metadata.page_size // FP4_BLOCK_SIZE - - generation_state = None - generation_inputs = (rotary_cos_sin, q_pe, q_rope_out, q_quant_input) - q_fp4_out = None - q_sf_out = None - if phase == "context": - if any(arg is not None for arg in (q_pe, q_rope_out, q_quant_input)): - raise ValueError("FP4 MLA context cache update does not accept generation Q tensors.") - if (q_context is None) != (q_nope_head_dim is None): - raise ValueError("FP4 MLA context Q and q_nope_head_dim must be provided together.") - hp_pool_updated = False - else: - if q_context is not None or q_nope_head_dim is not None: - raise ValueError("FP4 MLA generation cache update does not accept context Q tensors.") - if (helix_position_offsets is None) != (helix_is_inactive_rank is None): - raise ValueError( - "FP4 MLA Helix position-offset and inactive-rank metadata " - "must be provided together." - ) - if not all(arg is not None for arg in generation_inputs): - raise ValueError( - "FP4 MLA generation requires rotary_cos_sin, q_pe, q_rope_out, " - "and q_quant_input for fused RoPE, cache update, and Q quantization." - ) - if not can_fuse_fp4_mla_q_quant(metadata, q_quant_input, q_pe, latent_cache): - raise ValueError( - "Fused FP4 MLA Q quantization received an unsupported shape, dtype, " - "scale mode, or backend." - ) - num_gen_for_q = metadata.num_seqs - metadata.num_contexts - q_rows_per_seq = 0 - if num_gen_for_q > 0 and num_tokens % num_gen_for_q == 0: - q_rows_per_seq = _fp4_mla_packed_rows_per_seq( - q_quant_input.shape[1], num_tokens // num_gen_for_q - ) - q_fp4_out, q_sf_out, q_batch_capacity = _prepare_fp4_mla_q_buffers( - metadata, - num_tokens, - q_quant_input.shape[1], - q_quant_input.device, - rows_per_seq=q_rows_per_seq, - ) - metadata._fp4_mla_prequantized_q = q_fp4_out - metadata._fp4_mla_prequantized_q_sf = q_sf_out - metadata._fp4_mla_q_batch_capacity = q_batch_capacity - metadata._fp4_mla_prequantized_q_rows_per_seq = q_rows_per_seq - hp_pool_updated = True - cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND - fused_v_transpose = cutedsl_backend and _fp4_mla_cutedsl_fused_v_transpose_enabled() - kv_cache_manager = getattr(metadata, "kv_cache_manager", None) - persistent_v_packed = None - v_packed_base = None - v_page_offset = 0 - direct_v_packed_write = False - if cutedsl_backend and not fused_v_transpose and not context_only_cache: - persistent_v_packed = _get_cutedsl_persistent_v_packed_cache( - metadata, - local_layer, - kv_cache, - v_head_dim=v_head_dim, - page_size=metadata.page_size, - block_v=FP4_MLA_SCALE_ROW_GROUP, - ) - # Reused/imported pages may not carry this process-local sidecar. Write - # packed V directly only when the cache pages are exclusively owned. - # The fused generation kernel handles uniform linear-MTP batches and - # updates every 16-token tile touched by the verification window. - num_gen = metadata.num_seqs - metadata.num_contexts - block_reuse = getattr(kv_cache_manager, "enable_block_reuse", True) - direct_context_v_packed_write = ( - phase == "context" - and metadata.num_contexts > 0 - and num_tokens > 0 - and block_reuse is False - ) - direct_generation_v_packed_write = ( - phase == "generation" - and num_gen > 0 - and num_tokens >= num_gen - and num_tokens % num_gen == 0 - and block_reuse is False - ) - direct_v_packed_write = direct_context_v_packed_write or direct_generation_v_packed_write - if direct_v_packed_write: - v_packed_base = _get_fp4_mla_v_packed_pool_base(metadata) - get_v_page_offset = getattr(kv_cache_manager, "get_mla_v_packed_page_offset", None) - v_page_offset = ( - int(get_v_page_offset(local_layer)) - if callable(get_v_page_offset) - else local_layer * kv_cache.shape[0] - ) - expected_row_width = metadata.page_size // 2 - required_base_rows = (v_page_offset + kv_cache.shape[0]) * v_head_dim - if ( - not isinstance(v_packed_base, torch.Tensor) - or v_packed_base.dtype != torch.uint8 - or v_packed_base.device != kv_cache.device - or v_packed_base.ndim != 2 - or v_packed_base.shape[0] < required_base_rows - or v_packed_base.shape[1] != expected_row_width - or not v_packed_base.is_contiguous() - ): - raise RuntimeError( - "FP4 MLA direct V-packed cache update requires the " - "stable full-pool base to be a contiguous uint8 tensor " - f"with at least {required_base_rows} rows and " - f"{expected_row_width} columns on {kv_cache.device}." - ) - expected_layer_ptr = v_packed_base.data_ptr() + ( - v_page_offset * v_head_dim * expected_row_width - ) - if expected_layer_ptr != persistent_v_packed.data_ptr(): - raise RuntimeError( - "FP4 MLA V-packed layer view does not match its stable " - "full-pool base and page offset." - ) - v_pack_num_valid = None - if phase == "context": - _materialize_fp4_mla_device_page_table_for_forward(metadata) - hp_pool_updated = _scatter_fp4_mla_kv_cache_2d_context( - metadata, - latent_cache, - kv_cache, - sf_cache, - v_sf, - global_scale, - rotary_cos_sin, - q_context, - q_nope_head_dim, - token_offset=token_offset, - local_layer=local_layer, - v_head_dim=v_head_dim, - head_dim=head_dim, - num_tokens=num_tokens, - num_dim_blocks=num_dim_blocks, - sf_per_token=sf_per_token, - sf_per_page=sf_per_page, - v_packed_base=v_packed_base, - v_page_offset=v_page_offset, - ) - v_pack_page_ids = metadata.paged_kv_indices - else: - generation_state = _scatter_fp4_mla_kv_cache_2d_generation( - metadata, - latent_cache, - kv_cache, - sf_cache, - v_sf, - global_scale, - token_offset=token_offset, - local_layer=local_layer, - v_head_dim=v_head_dim, - head_dim=head_dim, - num_tokens=num_tokens, - num_dim_blocks=num_dim_blocks, - sf_per_token=sf_per_token, - sf_per_page=sf_per_page, - rotary_cos_sin=rotary_cos_sin, - q_pe=q_pe, - q_rope_out=q_rope_out, - q_quant_input=q_quant_input, - q_fp4_out=q_fp4_out, - q_sf_out=q_sf_out, - v_packed_base=v_packed_base, - v_page_offset=v_page_offset, - helix_position_offsets=helix_position_offsets, - helix_is_inactive_rank=helix_is_inactive_rank, - ) - v_pack_page_ids = _fp4_mla_generation_page_ids( - metadata, metadata.num_seqs - metadata.num_contexts - ) - if getattr(metadata, "is_cuda_graph", False): - # Frozen launch grids cannot follow the per-replay page count; - # the repack kernels stride over this device-side count instead. - v_pack_num_valid = _fp4_mla_generation_num_blocks_device(metadata) - cutedsl_repack_page_indptr = None - cutedsl_repack_kv_lens = None - cutedsl_repack_generation_lens = None - cutedsl_repack_max_touched_pages = 1 - if ( - cutedsl_backend - and not fused_v_transpose - and not direct_v_packed_write - and not context_only_cache - ): - if phase == "context": - num_contexts = metadata.num_contexts - cutedsl_v_pack_page_ids = metadata.paged_kv_indices[: metadata.num_context_blocks] - cutedsl_repack_page_indptr = metadata.paged_kv_indptr[: num_contexts + 1] - cutedsl_repack_kv_lens = metadata.kv_lens_cuda_runtime[:num_contexts] - cutedsl_repack_generation_lens = metadata.prompt_lens_cuda_runtime[:num_contexts] - cutedsl_repack_max_touched_pages = int( - metadata.fp4_mla_context_repack_max_touched_pages - ) - elif generation_state is not None: - kv_lens_gen, gen_lens_gen, generation_page_ids = generation_state - num_gen = kv_lens_gen.numel() - cutedsl_v_pack_page_ids = generation_page_ids - cutedsl_repack_page_indptr = metadata.paged_kv_indptr_decode - cutedsl_repack_kv_lens = kv_lens_gen - cutedsl_repack_generation_lens = gen_lens_gen - cutedsl_repack_max_touched_pages = _ceil_div( - num_tokens // num_gen + metadata.page_size - 1, - metadata.page_size, - ) - else: - cutedsl_v_pack_page_ids = v_pack_page_ids - _repack_cutedsl_v_packed_cache( - persistent_v_packed, - kv_cache, - cutedsl_v_pack_page_ids, - v_head_dim=v_head_dim, - page_size=metadata.page_size, - block_v=FP4_MLA_SCALE_ROW_GROUP, - page_indptr=cutedsl_repack_page_indptr, - kv_lens=cutedsl_repack_kv_lens, - generation_lens=cutedsl_repack_generation_lens, - max_touched_pages=cutedsl_repack_max_touched_pages, - ) - if context_only_cache: - return hp_pool_updated - _maybe_update_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - v_pack_page_ids, - num_queries=num_tokens, - v_head_dim=v_head_dim, - page_size=metadata.page_size, - local_layer=local_layer, - v_sf=v_sf[local_layer], - num_valid_pages=v_pack_num_valid, - ) - return hp_pool_updated - - -def _validate_fp4_mla_cache_shape(page_size: int, head_dim: int) -> None: - if page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - f"FP4 MLA KV cache requires tokens_per_block={FP4_MLA_TOKENS_PER_BLOCK} " - f"for swizzled block scales, got {page_size}." - ) - - sf_per_token = head_dim // FP4_BLOCK_SIZE - if head_dim % FP4_BLOCK_SIZE != 0 or sf_per_token % 4 != 0: - raise ValueError( - f"FP4 MLA KV head_dim must produce a scale column count divisible by 4; " - f"got head_dim={head_dim}, scale_columns={sf_per_token}." - ) - - -def _validate_fp4_mla_attention_q_shape(head_dim: int, q_residual_dim: int) -> None: - if q_residual_dim % FP4_BLOCK_SIZE != 0: - raise ValueError( - f"FP4 MLA Q residual_dim must be divisible by {FP4_BLOCK_SIZE}, got {q_residual_dim}." - ) - if q_residual_dim <= 0 or q_residual_dim > head_dim: - raise ValueError( - f"FP4 MLA Q residual_dim must be in (0, head_dim], got " - f"residual_dim={q_residual_dim}, head_dim={head_dim}." - ) - - q_head_dim = head_dim + q_residual_dim - q_sf_per_token = q_head_dim // FP4_BLOCK_SIZE - if q_head_dim % FP4_BLOCK_SIZE != 0 or q_sf_per_token % FP4_MLA_SCALE_COL_GROUP != 0: - raise ValueError( - f"FP4 MLA residual Q must produce a scale column count divisible " - f"by {FP4_MLA_SCALE_COL_GROUP}; got q_head_dim={q_head_dim}, " - f"scale_columns={q_sf_per_token}." - ) - - -def _ensure_workspace_tensor( - metadata: Any, - attr_name: str, - shape: tuple[int, ...], - *, - dtype: torch.dtype, - device: torch.device, -) -> torch.Tensor: - tensor = getattr(metadata, attr_name, None) - needs_alloc = ( - tensor is None - or tensor.dtype != dtype - or tensor.device != device - or len(tensor.shape) != len(shape) - or any(tensor.shape[idx] < dim for idx, dim in enumerate(shape)) - ) - if needs_alloc: - if torch.cuda.is_current_stream_capturing(): - raise ValueError( - f"Cannot allocate {attr_name} while capturing a CUDA graph. " - "Run a warmup prepare/forward first." - ) - tensor = torch.empty(shape, dtype=dtype, device=device) - setattr(metadata, attr_name, tensor) - - slices = tuple(slice(0, dim) for dim in shape) - return tensor[slices] - - -def _shared_v_pack_storage_enabled() -> bool: - return os.getenv("TRTLLM_FP4_MLA_SHARE_V_PACK_STORAGE", "1").lower() not in ( - "0", - "false", - "no", - "off", - ) - - -def _select_triton_block_v(num_queries: int, *, prefer_prepacked_v: bool = False) -> int: - env_block_v = _env_int("TRTLLM_FP4_MLA_BLOCK_V") - if env_block_v is not None: - return env_block_v - if prefer_prepacked_v: - return 128 - return 32 if num_queries <= 32 else 128 - - -def _v_packed_shape( - kv_cache: torch.Tensor, - v_head_dim: int, - page_size: int, - block_v: int, -) -> tuple[int, int]: - return (kv_cache.shape[0] * _ceil_div(v_head_dim, block_v) * block_v, page_size // 2) - - -def _get_fp4_mla_v_packed_pool(metadata: Any, local_layer: int) -> Optional[torch.Tensor]: - return metadata.kv_cache_manager.get_mla_v_packed_pool(local_layer) - - -def _get_fp4_mla_v_packed_pool_base(metadata: Any) -> Optional[torch.Tensor]: - return metadata.kv_cache_manager.get_mla_v_packed_pool_base() - - -def _get_fp4_mla_v_scale_pool_base(metadata: Any) -> Optional[torch.Tensor]: - return metadata.kv_cache_manager.get_mla_v_scale_pool_base() - - -def _get_cutedsl_persistent_v_packed_cache( - metadata: Any, - local_layer: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - block_v: int, -) -> torch.Tensor: - v_packed = _get_fp4_mla_v_packed_pool(metadata, local_layer) - if v_packed is None: - raise RuntimeError( - "CuTeDSL FP4 MLA requires the manager-owned persistent V-packed " - "pool; the scratch full-repack fallback has been removed." - ) - expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) - if ( - not isinstance(v_packed, torch.Tensor) - or v_packed.dtype != torch.uint8 - or v_packed.device != kv_cache.device - or tuple(v_packed.shape) != expected_shape - or not v_packed.is_contiguous() - ): - raise RuntimeError( - "FP4 MLA persistent V-packed pool must be a contiguous uint8 tensor " - f"with shape {expected_shape} on {kv_cache.device}; got " - f"{type(v_packed).__name__}, " - f"shape={getattr(v_packed, 'shape', None)}, " - f"dtype={getattr(v_packed, 'dtype', None)}, " - f"device={getattr(v_packed, 'device', None)}." - ) - return v_packed - - -def _repack_cutedsl_v_packed_cache( - v_packed: torch.Tensor, - kv_cache: torch.Tensor, - page_ids: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - block_v: int, - page_indptr: Optional[torch.Tensor] = None, - kv_lens: Optional[torch.Tensor] = None, - generation_lens: Optional[torch.Tensor] = None, - max_touched_pages: int = 1, -) -> None: - if page_ids.numel() == 0: - return - from .fp4_mla_cutedsl_v_repack import fp4_mla_repack_v_cache - - fp4_mla_repack_v_cache( - v_packed, - kv_cache, - page_ids, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - page_indptr=page_indptr, - kv_lens=kv_lens, - generation_lens=generation_lens, - max_touched_pages=max_touched_pages, - ) - - -def _v_packed_cache_tag( - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> tuple[Any, ...]: - v_sf_tag = ( - None - if v_sf is None - else ( - int(v_sf.data_ptr()), - str(v_sf.device), - str(v_sf.dtype), - tuple(int(dim) for dim in v_sf.shape), - tuple(int(stride) for stride in v_sf.stride()), - ) - ) - page_ids_tag = ( - None - if page_ids is None - else ( - int(page_ids.data_ptr()), - str(page_ids.device), - str(page_ids.dtype), - tuple(int(dim) for dim in page_ids.shape), - tuple(int(stride) for stride in page_ids.stride()), - ) - ) - return ( - int(layer_idx), - None if local_layer is None else int(local_layer), - int(kv_cache.data_ptr()), - str(kv_cache.device), - str(kv_cache.dtype), - tuple(int(dim) for dim in kv_cache.shape), - tuple(int(stride) for stride in kv_cache.stride()), - int(v_head_dim), - int(page_size), - int(block_v), - v_sf_tag, - page_ids_tag, - ) - - -def _triton_prepack_v_enabled() -> bool: - if _fp4_mla_attention_backend() != "triton": - return False - default = _env_enabled_default("TRTLLM_FP4_MLA_PREPACK_V", True) - return _env_enabled_default("TRTLLM_FP4_MLA_TRITON_PREPACK_V", default) - - -def _triton_can_prepack_v(v_head_dim: int, page_size: int, block_v: int) -> bool: - return ( - _triton_prepack_v_enabled() - and hasattr(tl, "make_tensor_descriptor") - and block_v in (32, 128) - and v_head_dim % block_v == 0 - and page_size == FP4_MLA_TOKENS_PER_BLOCK - ) - - -def _triton_v_packed_attr(layer_idx: int) -> str: - if _shared_v_pack_storage_enabled(): - return "_fp4_mla_triton_attention_v_packed_buf" - return f"_fp4_mla_triton_attention_v_packed_buf_l{layer_idx}" - - -def _triton_v_packed_valid_attr(layer_idx: int) -> str: - return f"_fp4_mla_triton_attention_v_packed_valid_l{layer_idx}" - - -def _triton_shared_v_packed_valid_attr() -> str: - return "_fp4_mla_triton_attention_v_packed_valid_tag" - - -def _triton_v_packed_cache_tag( - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> tuple[Any, ...]: - return ( - "triton", - _v_packed_cache_tag( - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ), - ) - - -def _set_triton_v_packed_cache_valid( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> None: - valid_attr = ( - _triton_shared_v_packed_valid_attr() - if _shared_v_pack_storage_enabled() - else _triton_v_packed_valid_attr(layer_idx) - ) - setattr( - metadata, - valid_attr, - _triton_v_packed_cache_tag( - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ), - ) - - -def _is_triton_v_packed_cache_valid( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> bool: - valid_attr = ( - _triton_shared_v_packed_valid_attr() - if _shared_v_pack_storage_enabled() - else _triton_v_packed_valid_attr(layer_idx) - ) - return getattr(metadata, valid_attr, None) == _triton_v_packed_cache_tag( - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ) - - -def _get_triton_v_packed_cache( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - page_ids: Optional[torch.Tensor] = None, - block_v: int = 128, -) -> Optional[torch.Tensor]: - if not _triton_can_prepack_v(v_head_dim, page_size, block_v): - return None - if not _is_triton_v_packed_cache_valid( - metadata, - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ): - return None - v_packed = getattr(metadata, _triton_v_packed_attr(layer_idx), None) - expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) - if ( - v_packed is None - or v_packed.dtype != torch.uint8 - or v_packed.device != kv_cache.device - or len(v_packed.shape) != 2 - or v_packed.shape[0] < expected_shape[0] - or v_packed.shape[1] < expected_shape[1] - ): - return None - return v_packed[: expected_shape[0], : expected_shape[1]] - - -def _update_triton_v_packed_cache( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - page_ids: torch.Tensor, - *, - v_head_dim: int, - page_size: int, - block_v: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - num_valid_pages: Optional[torch.Tensor] = None, -) -> Optional[torch.Tensor]: - if not _triton_can_prepack_v(v_head_dim, page_size, block_v): - return None - if page_ids.numel() == 0: - return None - from .fp4_mla_triton import fp4_mla_repack_v_cache - - def _tma_alloc(size: int, alignment: int, stream): - return torch.empty(size, device=kv_cache.device, dtype=torch.int8) - - triton.set_allocator(_tma_alloc) - attr_name = _triton_v_packed_attr(layer_idx) - v_packed = _ensure_workspace_tensor( - metadata, - attr_name, - _v_packed_shape(kv_cache, v_head_dim, page_size, block_v), - dtype=torch.uint8, - device=kv_cache.device, - ) - fp4_mla_repack_v_cache( - v_packed, - kv_cache, - page_ids, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - num_valid_pages=num_valid_pages, - ) - _set_triton_v_packed_cache_valid( - metadata, - layer_idx, - kv_cache, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=page_ids, - ) - return v_packed - - -def _maybe_update_triton_v_packed_cache( - metadata: Any, - layer_idx: int, - kv_cache: torch.Tensor, - page_ids: torch.Tensor, - *, - num_queries: int, - v_head_dim: int, - page_size: int, - local_layer: Optional[int] = None, - v_sf: Optional[torch.Tensor] = None, - num_valid_pages: Optional[torch.Tensor] = None, -) -> None: - block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) - _update_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - page_ids, - v_head_dim=v_head_dim, - page_size=page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - num_valid_pages=num_valid_pages, - ) - - -def _max_generation_pages(metadata: Any) -> int: - num_gen = metadata.num_seqs - metadata.num_contexts - if num_gen <= 0: - return 0 - if not getattr(metadata, "_fp4_mla_device_page_table", False): - raise RuntimeError("FP4 MLA generation requires fixed-stride device page metadata.") - max_pages = int(metadata.fp4_mla_page_table_stride) - if max_pages <= 0: - raise RuntimeError("FP4 MLA device page-table stride must be positive.") - return max_pages - - -def _fp4_mla_generation_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: - """Return the fixed-stride generation page-table view.""" - expected_num_gen = metadata.num_seqs - metadata.num_contexts - if num_gen_seqs != expected_num_gen: - raise RuntimeError( - "FP4 MLA generation sequence count does not match metadata: " - f"{num_gen_seqs} != {expected_num_gen}." - ) - max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata, "_paged_kv_indices", None) - start = metadata.num_contexts * max_pages - end = start + num_gen_seqs * max_pages - if ( - not isinstance(page_ids, torch.Tensor) - or page_ids.ndim != 1 - or page_ids.dtype != torch.int32 - or not page_ids.is_contiguous() - or page_ids.numel() < end - ): - raise RuntimeError("FP4 MLA fixed-stride generation page-table backing is invalid.") - return page_ids[start:end] - - -def _fp4_mla_generation_hp_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tensor: - """Return generation rows from the fixed-stride V2 HP page table.""" - expected_num_gen = metadata.num_seqs - metadata.num_contexts - if num_gen_seqs != expected_num_gen: - raise RuntimeError( - "FP4 MLA generation sequence count does not match metadata: " - f"{num_gen_seqs} != {expected_num_gen}." - ) - max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata, "_fp4_mla_hp_page_indices", None) - start = metadata.num_contexts * max_pages - end = start + num_gen_seqs * max_pages - if ( - not isinstance(page_ids, torch.Tensor) - or page_ids.ndim != 1 - or page_ids.dtype != torch.int32 - or not page_ids.is_contiguous() - or page_ids.numel() < end - ): - raise RuntimeError("FP4 MLA fixed-stride generation HP page-table backing is invalid.") - return page_ids[start:end] - - -def _host_int_list(value: Any, start: int, end: int) -> Optional[list[int]]: - if value is None: - return None - if isinstance(value, torch.Tensor): - if value.is_cuda: - return None - return [int(item) for item in value[start:end].tolist()] - try: - return [int(item) for item in value[start:end]] - except (TypeError, ValueError): - return None - - -def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> bool: - if getattr(metadata, "is_cuda_graph", False): - return False - - start = metadata.num_contexts - end = metadata.num_seqs - block_counts = _host_int_list(getattr(metadata, "num_blocks", None), start, end) - if block_counts is not None and ( - not block_counts or min(block_counts) != max_pages or max(block_counts) != max_pages - ): - return False - - kv_lens_cuda = getattr(metadata, "kv_lens_cuda_runtime", None) - if isinstance(kv_lens_cuda, torch.Tensor): - cache_key = ( - start, - end, - max_pages, - page_size, - tuple(block_counts) if block_counts is not None else None, - kv_lens_cuda.data_ptr(), - ) - cache = getattr(metadata, "_fp4_mla_full_pages_cache", None) - if cache is not None and cache[0] == cache_key: - return bool(cache[1]) - kv_lens = [int(item) for item in kv_lens_cuda[start:end].detach().cpu().tolist()] - result = bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size - setattr(metadata, "_fp4_mla_full_pages_cache", (cache_key, result)) - return result - - kv_cache_params = getattr(metadata, "kv_cache_params", None) - cached_token_lens = _host_int_list( - getattr(kv_cache_params, "num_cached_tokens_per_seq", None), - start, - end, - ) - seq_lens_kv = _host_int_list(getattr(metadata, "seq_lens_kv", None), start, end) - if cached_token_lens is not None and seq_lens_kv is not None: - if len(cached_token_lens) != len(seq_lens_kv): - return False - kv_lens = [ - cached_len + seq_len for cached_len, seq_len in zip(cached_token_lens, seq_lens_kv) - ] - elif kv_cache_params is None: - kv_lens = _host_int_list(getattr(metadata, "prompt_lens_cpu_runtime", None), start, end) - else: - return False - - return bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size - - -def _get_linear_mtp_query_len_per_seq( - metadata: Any, - *, - num_queries: int, - num_gen_seqs: int, -) -> int: - """Return the uniform generation query length required by linear MTP. - - Derives the length from the real query-token count (``num_queries``, taken - from the q shape) and the generation sequence count, which are reliable in - every representation. The host ``prompt_lens``/``seq_lens`` mirror can lag at - the decode anchor (== 1) under CUDA graph / one-engine MTP, so it is only - consulted to produce a precise diagnostic when the counts do not divide - evenly (a genuinely non-uniform batch, which the no-dequant path does not - support). - """ - if num_gen_seqs <= 0: - return 1 - - if num_queries % num_gen_seqs == 0: - return num_queries // num_gen_seqs - - start = metadata.num_contexts - end = metadata.num_seqs - query_lens = _host_int_list_during_forward( - getattr(metadata, "prompt_lens_cpu_runtime", None), start, end - ) - if query_lens is None: - query_lens = _host_int_list_during_forward(getattr(metadata, "seq_lens", None), start, end) - raise NotImplementedError( - "FP4 MLA no-dequant attention requires a uniform linear MTP generation " - f"query length; got {num_queries} query tokens for {num_gen_seqs} " - f"sequences (per-sequence lengths {query_lens})." - ) - - -def _run_triton_attention_decode( - *, - metadata: Any, - layer_idx: int, - local_layer: int, - q_fp4: torch.Tensor, - q_sf: torch.Tensor, - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - v_sf: torch.Tensor, - global_scale: torch.Tensor, - src_page_ids: torch.Tensor, - kv_lens: torch.Tensor, - p_fp4: torch.Tensor, - p_sf: torch.Tensor, - max_scores: torch.Tensor, - denom: torch.Tensor, - output: torch.Tensor, - num_queries: int, - num_heads: int, - head_dim: int, - kv_lora_rank: int, - q_residual_dim: int, - query_len_per_seq: int, - max_pages: int, - sm_scale: float, - q_global_scale: torch.Tensor, -) -> None: - """Dispatch the ``triton`` FP4 MLA decode pipeline. - - Mirrors the four-stage layout used by ``fp4_mla_cutile.py`` - (page-stats with packed P -> reduce-stats -> prob-scale -> PV) but - routes through the self-contained kernels in - ``fp4_mla_triton.py``. Threads through the constexpr assume flags, - TMA descriptors, occupancy/num-warps launch meta, and pipelined PV loop. - """ - from .fp4_mla_triton import ( - _fp4_mla_attention_group_reduce_stats_kernel as _attn_group_reduce_stats_kernel, - ) - from .fp4_mla_triton import _fp4_mla_attention_page_stats_kernel as _attn_page_stats_kernel - from .fp4_mla_triton import _fp4_mla_attention_prob_scale_kernel as _attn_prob_scale_kernel - from .fp4_mla_triton import _fp4_mla_attention_pv_kernel as _attn_pv_kernel - from .fp4_mla_triton import ( - _fp4_mla_attention_pv_prepacked_v_kernel as _attn_pv_prepacked_v_kernel, - ) - from .fp4_mla_triton import _fp4_mla_attention_pv_reduce_kernel as _attn_pv_reduce_kernel - from .fp4_mla_triton import _fp4_mla_attention_reduce_stats_kernel as _attn_reduce_stats_kernel - - block_h = 128 - block_t = metadata.page_size - # Adaptive BLOCK_V: the fallback PV path uses a finer V split at small batch - # on B200 (~148 SMs). PV grid = num_queries * num_head_blocks(1) * - # (kv_lora_rank / BLOCK_V). We want >= ~2*num_SMs programs so that >1 CTA - # lands per SM and hides the L1TEX scoreboard stalls. Empirically (sweep): - # bs<=32 -> BLOCK_V=32; bs>=64 -> BLOCK_V=128. - # (BLOCK_V=16 is rejected by the V TMA descriptor min-stride requirement.) - # With prepacked V, BLOCK_V=128 avoids reloading the same P tile four times - # and matches the cutile prepacked-V tile shape. - block_v = _select_triton_block_v(num_queries, prefer_prepacked_v=_triton_prepack_v_enabled()) - q_storage_head_dim = head_dim + q_residual_dim - # The virtual GEMM tail evaluates QK + Q_r K + Q K_r in one reduction. - # Q and Q_r still occupy the 640-channel interleaved physical Q buffer; - # the final Q term reuses Q's main tail groups while K_r comes from the - # contiguous 64-channel tail of the primary paged KV cache. - q_head_dim = head_dim + q_residual_dim + FP4_MLA_K_RESIDUAL_DIM - # BLOCK_K = 512 aligns the K-window with the 512-channel non-residual prefix. - block_k = 512 - full_block_end = (q_head_dim // block_k) * block_k - tail_k = q_head_dim - full_block_end - tail_block_k = 1 << (tail_k - 1).bit_length() if tail_k > 0 else block_k - q_sf_per_token = q_storage_head_dim // FP4_BLOCK_SIZE - k_sf_per_token = (head_dim + FP4_MLA_K_RESIDUAL_DIM) // FP4_BLOCK_SIZE - sf_per_page = metadata.page_size // FP4_BLOCK_SIZE - num_head_blocks = triton.cdiv(num_heads, block_h) - - assume_full_heads = num_heads % block_h == 0 - assume_full_v = kv_lora_rank % block_v == 0 - # Match the cutile path: only mark pages "full" when we can prove every - # generation sequence has the same number of cached tokens AND - # query_len_per_seq == 1 (so the kv_len adjustment is a no-op). - assume_full_pages = ( - _infer_assume_full_pages(metadata, max_pages, metadata.page_size) and query_len_per_seq == 1 - ) - # Leave validity checks on. Matches cutile's default and is correctness- - # safe. The perfect-shape PV fast path (tl.ext.make_view + load_view_tko) - # remains gated off — when measured on the TileIR backend (ENABLE_TILE=1) - # it was net-slower on the bench, so the cost of enabling it isn't worth - # the win on the FP4 MLA shapes we care about. - assume_valid_pages = False - num_gen_seqs = num_queries // query_len_per_seq - if ( - not assume_valid_pages - and assume_full_pages - and src_page_ids.numel() == num_gen_seqs * max_pages - ): - assume_valid_pages = True - # cutile checks only `make_tensor_descriptor`; on the nvt backend the - # presence of TMA descriptors implies `tl.ext.make_view` is available too. - use_tma_data_load = hasattr(triton.language, "make_tensor_descriptor") - - # Install the device-side scratch allocator on every call. Triton stores - # the allocator in a ContextVar (triton.runtime._allocation), so a single - # process-wide install is not visible from worker threads / asyncio tasks - # that run with a different Context — the kernel launch would then hit the - # default NullAllocator and raise. Matches the cutile path. - if use_tma_data_load: - - def _tma_alloc(size: int, alignment: int, stream): - return torch.empty(size, device=q_fp4.device, dtype=torch.int8) - - triton.set_allocator(_tma_alloc) - - # cutile-equivalent launch meta. occupancy=2 lets two CTAs land per SM - # which improves wave-tail efficiency at the bs=32 hot point. - # NOTE: num_stages=2 (instead of the Triton 3.6 default of 3) sidesteps - # the TritonGPUAutomaticWarpSpecialization + NVWSInsertTmemAref pass that - # ICEs on the page_stats kernel under Triton 3.6.0 / sm_100. - launch_meta = {"occupancy": 2} - # The matmul kernels (page-stats QK and PV) are register-limited: at the - # Triton default of num_warps=4 the [BLOCK_H, BLOCK_T] epilogue spills the - # register file down to ~2 CTAs/SM (12.5% occupancy), so there are too few - # warps to hide the QK/PV load latency (ncu: ~0.3 eligible warps/scheduler). - # Spreading the tile epilogue over num_warps=8 halves the per-thread - # register need and roughly doubles resident warps. Matches the cutile - # ("nvt") backend, which launches page-stats at num_warps=8. Both are - # overridable for tuning. - sm_count = _get_sm_count(q_fp4.device) - # page-stats num_warps: the full-pages fast path (uniform q_len==1 decode) - # benefits from num_warps=8 (more warps hide the QK load latency); the - # masked path (q_len>1 / ragged lengths) carries extra per-thread state and - # measured markedly faster at num_warps=4 (e.g. bs256 q_len4: 131->95ms). - page_stats_num_warps = _env_int("TRTLLM_FP4_MLA_PAGE_STATS_NUM_WARPS") - if page_stats_num_warps is None: - page_stats_num_warps = 8 if assume_full_pages else 4 - page_stats_launch_meta = {"occupancy": 2, "num_warps": page_stats_num_warps} - # PV benefits from num_warps=8 across shapes measured. - pv_num_warps = _env_int("TRTLLM_FP4_MLA_PV_NUM_WARPS") or 8 - pv_launch_meta = {"occupancy": 2, "num_warps": pv_num_warps} - # PV loop pipelining. With TMA loads, num_stages>=2 lets the next page's - # loads overlap with the current MMA via mbarrier. The PV report shows - # long_scoreboard=4.5 cycles avg on V loads at PV_LOOP_STAGES=2; bumping the - # depth pays off when the grid is small enough that occupancy can absorb - # the extra in-flight tile state — i.e. medium batch / large max_pages. - # Larger pipelines hurt at small batch (more live state, fewer dim blocks). - if num_queries <= 16 or max_pages <= 4: - pv_loop_stages = 2 - else: - pv_loop_stages = 3 - - # Page-stats kernel: per (query, head_block, page) program, does QK, - # softmax stats, and packs probs into FP4 with the per-page local-max - # scaling trick. The page-max correction is applied later by - # prob_scale_kernel via p_sf in-place rescaling. - page_stats_shape = (num_queries, max_pages, num_heads) - page_max = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_page_max_buf", - page_stats_shape, - dtype=torch.float32, - device=q_fp4.device, - ) - page_sum = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_page_sum_buf", - page_stats_shape, - dtype=torch.float32, - device=q_fp4.device, - ) - - pack_prob_in_page_stats = True - _attn_page_stats_kernel[(num_queries, num_head_blocks, max_pages)]( - page_max, - page_sum, - p_fp4, - p_sf, - q_fp4, - q_sf, - kv_cache, - sf_cache, - global_scale, - q_global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - q_fp4.stride(0), - q_fp4.stride(1), - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - sf_cache.stride(0), - page_max.stride(0), - page_max.stride(1), - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - q_fp4.shape[0], - sm_scale, - NUM_HEADS=num_heads, - Q_HEAD_D=q_head_dim, - Q_STORAGE_HEAD_D=q_storage_head_dim, - K_HEAD_D=head_dim, - Q_RESIDUAL_D=q_residual_dim, - K_RESIDUAL_D=FP4_MLA_K_RESIDUAL_DIM, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - Q_SF_PER_TOKEN=q_sf_per_token, - K_SF_PER_TOKEN=k_sf_per_token, - SF_PER_PAGE=sf_per_page, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - BLOCK_H=block_h, - BLOCK_T=block_t, - BLOCK_K=block_k, - FULL_BLOCK_END=full_block_end, - TAIL_BLOCK_K=tail_block_k, - USE_TMA_DATA_LOAD=use_tma_data_load, - PACK_PROBS=pack_prob_in_page_stats, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_VALID_PAGES=assume_valid_pages, - **page_stats_launch_meta, - ) - # Two-level softmax-stats reduction. The single-level reduce launched only - # (num_queries * num_head_blocks) CTAs, each serially walking all max_pages - # twice -- at small batch that handful of CTAs left the GPU almost idle and - # the reduce cost more than the QK matmul. Level 1 parallelizes the page - # reduction across a page-group axis (online-softmax partials, pipelined); - # level 2 reuses the existing reduce kernel to fold the few groups into the - # global (max, denom). When the (query, head) grid already fills the GPU the - # group count collapses to 1 and this degenerates to the original reduce. - seqhead_ctas = num_queries * num_head_blocks - # Aim for ~3 waves of level-1 CTAs so page loads have enough memory-level - # parallelism to hide latency, while keeping the group count small enough - # that the level-2 combine loop stays short. - target_l1_ctas = 3 * sm_count - num_reduce_groups = _ceil_div(target_l1_ctas, max(seqhead_ctas, 1)) - num_reduce_groups = max(1, min(num_reduce_groups, max_pages, 64)) - # The grouped (two-level) reduce needs an auxiliary workspace, and - # _ensure_workspace_tensor can only (re)allocate it outside CUDA graph - # capture. If a warmup forward did not already size that workspace (e.g. the - # warmup batch took the single-level path), fall back to the single-level - # reduce during capture so we never allocate mid-capture. The single-level - # reduce is numerically identical (it just launches fewer CTAs). - if num_reduce_groups > 1 and torch.cuda.is_current_stream_capturing(): - gmax = getattr(metadata, "_fp4_mla_attention_group_max_buf", None) - gsum = getattr(metadata, "_fp4_mla_attention_group_sum_buf", None) - groups_ready = ( - gmax is not None - and gsum is not None - and gmax.shape[0] >= num_queries - and gmax.shape[1] >= num_reduce_groups - and gmax.shape[2] >= num_heads - and gsum.shape[0] >= num_queries - and gsum.shape[1] >= num_reduce_groups - and gsum.shape[2] >= num_heads - ) - if not groups_ready: - num_reduce_groups = 1 - if num_reduce_groups <= 1: - _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( - max_scores, - denom, - page_max, - page_sum, - max_pages, - max_scores.stride(0), - page_max.stride(0), - page_max.stride(1), - NUM_HEADS=num_heads, - MAX_PAGES=max_pages, - BLOCK_H=block_h, - **launch_meta, - ) - else: - group_pages = _ceil_div(max_pages, num_reduce_groups) - num_reduce_groups = _ceil_div(max_pages, group_pages) - group_max = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_group_max_buf", - (num_queries, num_reduce_groups, num_heads), - dtype=torch.float32, - device=q_fp4.device, - ) - group_sum = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_group_sum_buf", - (num_queries, num_reduce_groups, num_heads), - dtype=torch.float32, - device=q_fp4.device, - ) - _attn_group_reduce_stats_kernel[(num_queries, num_head_blocks, num_reduce_groups)]( - group_max, - group_sum, - page_max, - page_sum, - max_pages, - group_max.stride(0), - group_max.stride(1), - page_max.stride(0), - page_max.stride(1), - NUM_HEADS=num_heads, - GROUP_PAGES=group_pages, - BLOCK_H=block_h, - PIPELINE_STAGES=min(group_pages, 4), - **launch_meta, - ) - _attn_reduce_stats_kernel[(num_queries, num_head_blocks)]( - max_scores, - denom, - group_max, - group_sum, - num_reduce_groups, - max_scores.stride(0), - group_max.stride(0), - group_max.stride(1), - NUM_HEADS=num_heads, - MAX_PAGES=num_reduce_groups, - BLOCK_H=block_h, - **launch_meta, - ) - _attn_prob_scale_kernel[(num_queries, num_head_blocks, max_pages)]( - p_sf, - max_scores, - denom, - page_max, - metadata.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - max_scores.stride(0), - page_max.stride(0), - page_max.stride(1), - NUM_HEADS=num_heads, - PAGE_SIZE=metadata.page_size, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - BLOCK_H=block_h, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_VALID_PAGES=assume_valid_pages, - **launch_meta, - ) - num_dim_blocks = triton.cdiv(kv_lora_rank, block_v) - v_packed = _get_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - v_head_dim=kv_lora_rank, - page_size=metadata.page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - page_ids=src_page_ids, - ) - if ( - v_packed is None - and _triton_can_prepack_v(kv_lora_rank, metadata.page_size, block_v) - and not torch.cuda.is_current_stream_capturing() - ): - v_packed = _update_triton_v_packed_cache( - metadata, - layer_idx, - kv_cache, - src_page_ids, - v_head_dim=kv_lora_rank, - page_size=metadata.page_size, - block_v=block_v, - local_layer=local_layer, - v_sf=v_sf, - ) - use_triton_v_packed_cache = v_packed is not None - - # PV page split: partition the page range across additional programs and - # reduce in a follow-up kernel. ncu showed PV at waves/SM=0.49 for bs=32 — - # PV is L1-bandwidth bound, so raising in-flight CTAs is the lever. - # BLOCK_V is bounded below by the 16-byte TMA descriptor min-stride. - # PV page split: ncu shows that with the current shape (bs=32, max_pages=256) - # the PV kernel is L1-cache-throughput bound (long_scoreboard=4.5 cycles - # avg, L1 global LD hit-rate <40%). Increasing the program count via page - # splitting reduced waves/SM idle time but did NOT improve wall-time at - # current shapes — the per-CTA L1 thrash is the limit. Gate the split off - # by default; re-enable only for very small grids where occupancy is the - # bottleneck rather than per-CTA L1 pressure. - page_split = 1 - base_grid = num_queries * num_head_blocks * num_dim_blocks - if max_pages >= 16 and base_grid < 148: - for p in (8, 4, 2): - if max_pages % p == 0 and max_pages // p >= 16 and base_grid * p <= 148 * 4: - page_split = p - break - # The page-split PV path needs a partial-output workspace, which - # _ensure_workspace_tensor can only (re)allocate outside CUDA graph capture. - # Fall back to the unsplit PV (numerically identical) during capture unless a - # warmup forward already sized that workspace, so capture never allocates. - if page_split > 1 and torch.cuda.is_current_stream_capturing(): - pbuf = getattr(metadata, "_fp4_mla_attention_pv_partial_buf", None) - partial_ready = ( - pbuf is not None - and pbuf.shape[0] >= num_queries - and pbuf.shape[1] >= page_split - and pbuf.shape[2] >= num_heads - and pbuf.shape[3] >= kv_lora_rank - ) - if not partial_ready: - page_split = 1 - if page_split > 1: - pages_per_split = max_pages // page_split - partial_out = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_pv_partial_buf", - (num_queries, page_split, num_heads, kv_lora_rank), - dtype=torch.float32, - device=q_fp4.device, - ) - if use_triton_v_packed_cache: - _attn_pv_prepacked_v_kernel[ - (num_queries, num_head_blocks, num_dim_blocks * page_split) - ]( - output, - p_fp4, - p_sf, - v_packed, - v_sf, - global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - PAGE_SPLIT=page_split, - PAGES_PER_SPLIT=pages_per_split, - PARTIAL_OUT=True, - partial_out_ptr=partial_out, - partial_s0=partial_out.stride(0), - partial_s1=partial_out.stride(1), - partial_s2=partial_out.stride(2), - partial_s3=partial_out.stride(3), - **pv_launch_meta, - ) - else: - _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks * page_split)]( - output, - p_fp4, - p_sf, - kv_cache, - kv_cache, - v_sf, - global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, - USE_PREPACKED_V=False, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - PAGE_SPLIT=page_split, - PAGES_PER_SPLIT=pages_per_split, - PARTIAL_OUT=True, - partial_out_ptr=partial_out, - partial_s0=partial_out.stride(0), - partial_s1=partial_out.stride(1), - partial_s2=partial_out.stride(2), - partial_s3=partial_out.stride(3), - **pv_launch_meta, - ) - _attn_pv_reduce_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( - output, - partial_out, - global_scale, - output.stride(0), - output.stride(1), - output.stride(2), - partial_out.stride(0), - partial_out.stride(1), - partial_out.stride(2), - partial_out.stride(3), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SPLIT=page_split, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_V=assume_full_v, - **launch_meta, - ) - else: - if use_triton_v_packed_cache: - _attn_pv_prepacked_v_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( - output, - p_fp4, - p_sf, - v_packed, - v_sf, - global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_OUT_STORE=use_tma_data_load and assume_full_heads and assume_full_v, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - **pv_launch_meta, - ) - else: - _attn_pv_kernel[(num_queries, num_head_blocks, num_dim_blocks)]( - output, - p_fp4, - p_sf, - kv_cache, - kv_cache, - v_sf, - global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode, - kv_lens, - src_page_ids.shape[0], - kv_cache.shape[0], - output.stride(0), - output.stride(1), - output.stride(2), - output.shape[0] * output.shape[1], - p_fp4.stride(0), - p_fp4.stride(1), - p_fp4.shape[0], - kv_cache.stride(0), - kv_cache.stride(2), - kv_cache.stride(4), - v_sf.stride(0), - NUM_HEADS=num_heads, - V_HEAD_D=kv_lora_rank, - PAGE_SIZE=metadata.page_size, - FP4_BLOCK=FP4_BLOCK_SIZE, - SF_PER_PAGE=sf_per_page, - QUERY_LEN_PER_SEQ=query_len_per_seq, - MAX_PAGES=max_pages, - P_GLOBAL_SCALE=FP4_MLA_P_GLOBAL_SCALE, - BLOCK_H=block_h, - BLOCK_V=block_v, - USE_TMA_P_LOAD=use_tma_data_load and assume_full_heads and assume_valid_pages, - USE_TMA_V_LOAD=use_tma_data_load and kv_lora_rank % block_v == 0, - USE_PREPACKED_V=False, - PV_LOOP_STAGES=pv_loop_stages, - ASSUME_FULL_HEADS=assume_full_heads, - ASSUME_FULL_PAGES=assume_full_pages, - ASSUME_FULL_V=assume_full_v, - ASSUME_VALID_PAGES=assume_valid_pages, - **pv_launch_meta, - ) - - -def _run_cutedsl_packed_decode( - module: Any, - metadata: Any, - q_fp4: torch.Tensor, - q_sf: torch.Tensor, - kv_cache: torch.Tensor, - sf_cache: torch.Tensor, - v_packed: torch.Tensor, - v_sf: torch.Tensor, - kv_global_scale: torch.Tensor, - q_global_scale: torch.Tensor, - src_page_ids: torch.Tensor, - page_indptr: torch.Tensor, - kv_lens: torch.Tensor, - output: torch.Tensor, - *, - num_heads: int, - query_len: int, - num_requests: int, - rows_per_seq: int, - sm_scale: float, - v_page_offset: int, - helix_kv_bounds: Optional[torch.Tensor], - softmax_stats_tensor: Optional[torch.Tensor], - plan_geometry: Optional[_PackedPlanGeometry] = None, -) -> None: - """Packed CuTe DSL decode: one kernel over (split, request) tasks. - - A single split writes the output (and Helix stats) directly; several - splits write FP32 partials that one reduction kernel folds into the - output and stats. Nothing pads, replicates, or copies Q or the output. - """ - valid_m, physical_m = module.packed_query_rows(num_heads, query_len) - if physical_m != rows_per_seq: - raise RuntimeError( - f"FP4 MLA packed Q rows {rows_per_seq} do not match the kernel geometry {physical_m}." - ) - if v_packed is None: - raise RuntimeError("FP4 MLA packed decode requires the persistent packed-V cache.") - if not output.is_contiguous() or output.numel() != num_requests * valid_m * 512: - raise RuntimeError( - "FP4 MLA packed decode requires a contiguous [tokens, heads, 512] output." - ) - device = output.device - if plan_geometry is None: - plan_geometry = _fp4_mla_packed_plan_geometry( - metadata, - num_heads=num_heads, - query_len=query_len, - num_requests=num_requests, - device=device, - helix_kv_bounds=helix_kv_bounds, - ) - if plan_geometry is None: - raise RuntimeError("FP4 MLA packed decode requires generation cache pages.") - splits = plan_geometry.splits - output_rows = output.view(num_requests * valid_m, 512) - stats_rows = None - if softmax_stats_tensor is not None: - stats_rows = softmax_stats_tensor.view(num_requests * valid_m, 2) - common = dict( - max_kv_len=plan_geometry.split_kv_len, - sm_scale=sm_scale, - num_heads=num_heads, - query_len=query_len, - num_requests=num_requests, - num_splits=splits, - q_global_scale=q_global_scale, - v_page_offset=v_page_offset, - helix_kv_bounds=helix_kv_bounds, - ) - if splits == 1: - module.run_trtllm_fp4_mla_decode_packed( - q_fp4, - q_sf, - kv_cache, - sf_cache, - v_packed, - v_sf, - kv_global_scale, - src_page_ids, - page_indptr, - kv_lens, - output_rows, - softmax_stats=stats_rows, - **common, - ) - return - from .fp4_mla_packed_reduce import reduce_packed_splits - - tasks = splits * num_requests - task_capacity = plan_geometry.task_capacity - plan = plan_geometry.plan - # The first FP4 MLA kernel of the step wrote the plan fused with the - # generation lengths; a stale key means this call plans on other bounds. - if getattr(metadata, "_fp4_mla_packed_plan_key", None) != plan_geometry.key: - _launch_fp4_mla_packed_plan(plan_geometry, kv_lens, metadata) - partial = _ensure_workspace_tensor( - metadata, - "_fp4_mla_packed_split_partial", - (task_capacity, physical_m, 512), - dtype=torch.float32, - device=device, - )[:tasks] - maxima = _ensure_workspace_tensor( - metadata, - "_fp4_mla_packed_split_maxima", - (task_capacity, physical_m), - dtype=torch.float32, - device=device, - )[:tasks] - sums = _ensure_workspace_tensor( - metadata, - "_fp4_mla_packed_split_sums", - (task_capacity, physical_m), - dtype=torch.float32, - device=device, - )[:tasks] - module.run_trtllm_fp4_mla_decode_packed( - q_fp4, - q_sf, - kv_cache, - sf_cache, - v_packed, - v_sf, - kv_global_scale, - src_page_ids, - page_indptr, - kv_lens, - output_rows, - split_kv_buffers=(partial, maxima, sums), - plan=plan, - **common, - ) - reduce_packed_splits( - partial, - maxima, - sums, - output_rows, - q_global_scale, - kv_global_scale, - num_requests=num_requests, - num_splits=splits, - valid_m=valid_m, - physical_m=physical_m, - sm_scale=sm_scale, - plan=plan, - output_scale=1.0 / module.FP4_MLA_P_GLOBAL_SCALE, - stats=stats_rows, - sm_count=_get_sm_count(device), - ) - - -def run_fp4_mla_attention_decode( - metadata: Any, - layer_idx: int, - local_layer: int, - q: torch.Tensor, - output: torch.Tensor, - *, - sm_scale: float, - kv_lora_rank: int, - qk_rope_head_dim: int, - prequantized_q: torch.Tensor, - prequantized_q_sf: torch.Tensor, - q_batch_capacity: int, - softmax_stats_tensor: Optional[torch.Tensor] = None, -) -> None: - """Run MLA decode with FP4 QK and FP4 PV tensor-core matmuls. - - Q is supplied in its assembled ``[latent, RoPE]`` layout and quantized to - FP4 directly. QK reads ``[KV-nope, K-RoPE, K-RoPE-residual]`` contiguously - from the primary cache with swizzled block scales. Softmax probabilities - are quantized to FP4 per page, and PV repacks V nibbles from the shared KV - cache while reading the auxiliary V-view scale pool. No BF16 dequantized - KV workspace is materialized on this path. Callers must supply the packed Q - and scales produced by the fused generation cache update. - """ - head_dim = kv_lora_rank + qk_rope_head_dim - if qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM: - raise ValueError( - "FP4 MLA K residual attention requires " - f"qk_rope_head_dim={FP4_MLA_K_RESIDUAL_DIM}, got {qk_rope_head_dim}." - ) - _validate_fp4_mla_cache_shape(metadata.page_size, head_dim) - if metadata.page_size != FP4_MLA_TOKENS_PER_BLOCK: - raise ValueError( - f"FP4 MLA attention decode requires page_size={FP4_MLA_TOKENS_PER_BLOCK}, " - f"got {metadata.page_size}." - ) - - if q.ndim != 3 or q.shape[-1] != head_dim: - raise ValueError( - "FP4 MLA attention Q must have shape " - f"[tokens, heads, {head_dim}], got {tuple(q.shape)}." - ) - if not q.is_contiguous(): - raise ValueError("FP4 MLA attention Q must be contiguous.") - - num_queries = q.shape[0] - if num_queries == 0: - raise ValueError("FP4 MLA attention decode requires at least one query token.") - num_gen_seqs = metadata.num_seqs - metadata.num_contexts - query_len_per_seq = _get_linear_mtp_query_len_per_seq( - metadata, - num_queries=num_queries, - num_gen_seqs=num_gen_seqs, - ) - - num_heads = q.shape[1] - if output.shape[:2] != (num_queries, num_heads): - raise ValueError("FP4 MLA attention output batch dimensions do not match.") - - backend = _fp4_mla_attention_backend() - helix_spec_tokens_valid = bool(getattr(metadata, "_helix_spec_tokens_valid", False)) - helix_kv_bounds = None - if softmax_stats_tensor is not None: - if backend != _FP4_MLA_CUTEDSL_BACKEND: - raise NotImplementedError( - "FP4 MLA Helix softmax stats require the cutedsl attention backend." - ) - if query_len_per_seq != 1 and not helix_spec_tokens_valid: - raise NotImplementedError( - "FP4 MLA multi-token Helix requires speculative per-token metadata." - ) - expected_stats_shape = (num_queries, num_heads, 2) - if ( - softmax_stats_tensor.shape != expected_stats_shape - or softmax_stats_tensor.dtype != torch.float32 - or softmax_stats_tensor.device != q.device - or not softmax_stats_tensor.is_contiguous() - ): - raise ValueError( - "FP4 MLA Helix requires contiguous same-device float32 softmax " - f"stats with shape {expected_stats_shape}." - ) - if helix_spec_tokens_valid: - helix_kv_bounds = getattr(metadata, "helix_kv_bounds", None) - if ( - not isinstance(helix_kv_bounds, torch.Tensor) - or helix_kv_bounds.dtype != torch.int32 - or helix_kv_bounds.device != q.device - or helix_kv_bounds.ndim != 1 - or helix_kv_bounds.numel() < num_queries - or not helix_kv_bounds.is_contiguous() - ): - raise ValueError( - "FP4 MLA speculative Helix KV bounds must be a contiguous " - "same-device int32 tensor covering every query token." - ) - helix_kv_bounds = helix_kv_bounds[:num_queries] - if getattr(metadata, "fp4_mla_v_scale_pool", None) is None: - raise RuntimeError( - "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." - ) - - global_scale = _get_fp4_mla_global_scale(metadata, q.device) - q_residual_dim = FP4_MLA_Q_RESIDUAL_DIM - _validate_fp4_mla_attention_q_shape(head_dim, q_residual_dim) - - if prequantized_q is None or prequantized_q_sf is None or q_batch_capacity is None: - raise RuntimeError( - "FP4 MLA decode requires Q prequantized by the fused generation cache update." - ) - - capacity = int(q_batch_capacity) - # The fused generation update records the Q layout it produced; metadata - # that never went through it (contract tests) carries the legacy layout. - recorded_rows_per_seq = getattr(metadata, "_fp4_mla_prequantized_q_rows_per_seq", None) - packed_rows_per_seq = int(recorded_rows_per_seq or 0) - if recorded_rows_per_seq is not None and packed_rows_per_seq != _fp4_mla_packed_rows_per_seq( - num_heads, query_len_per_seq - ): - raise RuntimeError( - "FP4 MLA prequantized Q layout does not match the packed decode " - f"configuration (rows per request {packed_rows_per_seq})." - ) - if packed_rows_per_seq: - active_q_rows = num_gen_seqs * packed_rows_per_seq - else: - active_q_rows = num_queries * num_heads - active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( - active_q_rows, - FP4_MLA_Q_LOGICAL_DIM, - ) - if packed_rows_per_seq: - q_contract_ok = ( - prequantized_q.ndim == 2 - and prequantized_q.shape[0] >= active_q_rows - and prequantized_q.shape[1] == FP4_MLA_Q_PACKED_DIM - and prequantized_q_sf.numel() >= active_q_sf_bytes - ) - else: - expected_q_shape = (capacity * num_heads, FP4_MLA_Q_PACKED_DIM) - expected_q_sf_shape = ( - _get_fp4_mla_swizzled_scale_size( - capacity * num_heads, - FP4_MLA_Q_LOGICAL_DIM, - ), - ) - q_contract_ok = ( - tuple(prequantized_q.shape) == expected_q_shape - and tuple(prequantized_q_sf.shape) == expected_q_sf_shape - ) - if ( - q.dtype != torch.bfloat16 - or capacity <= 0 - or num_queries > capacity - or not q_contract_ok - or prequantized_q.dtype != torch.uint8 - or not prequantized_q.is_contiguous() - or prequantized_q_sf.dtype != torch.float8_e4m3fn - or not prequantized_q_sf.is_contiguous() - ): - raise ValueError("FP4 MLA prequantized Q does not satisfy the fused-Q contract.") - q_fp4 = prequantized_q[:active_q_rows] - q_sf = prequantized_q_sf[:active_q_sf_bytes] - q_global_scale = _get_fp4_mla_q_global_scale(metadata, q.device) - - kv_cache, sf_cache = _get_fp4_mla_kv_cache_tensors(metadata, layer_idx) - _validate_fp4_mla_kv_storage_shape( - kv_cache, - sf_cache, - head_dim=head_dim, - backend=backend, - ) - sf_cache = sf_cache.view(torch.float8_e4m3fn) - - v_sf_pool = metadata.fp4_mla_v_scale_pool - v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=kv_lora_rank)[local_layer].view( - torch.float8_e4m3fn - ) - # The kv_lens runtime alias can lag at the decode anchor (seq_lens == 1) under - # CUDA graph / one-engine MTP; recover the true total per sequence so the - # per-query causal masking sees the full 1 + draft_len window (no-op when the - # alias already matches). - plan_geometry = _fp4_mla_packed_plan_geometry( - metadata, - num_heads=num_heads, - query_len=query_len_per_seq, - num_requests=num_gen_seqs, - device=q.device, - helix_kv_bounds=helix_kv_bounds, - ) - kv_lens, _ = _fp4_mla_uniform_generation_lengths( - metadata, num_queries, num_gen_seqs, plan_geometry=plan_geometry - ) - _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens) - src_page_ids = _fp4_mla_generation_page_ids(metadata, num_gen_seqs) - max_pages = _max_generation_pages(metadata) - if max_pages == 0: - raise RuntimeError("FP4 MLA attention decode requires generation cache pages.") - if backend == _FP4_MLA_CUTEDSL_BACKEND: - if get_sm_version() != 107: - raise RuntimeError( - "FP4 MLA cutedsl attention backend requires Rubin SM107; " - f"current architecture is SM{get_sm_version()}." - ) - if not _cutedsl_backend_available(): - raise RuntimeError( - "FP4 MLA cutedsl attention backend requires the Rubin CTM and " - "CuTeDSL runtime packages." - ) - if not 0 < num_heads <= 128 or kv_lora_rank != 512: - raise ValueError( - "FP4 MLA cutedsl attention requires 1-128 local heads and " - f"kv_lora_rank=512, got num_heads={num_heads}, " - f"kv_lora_rank={kv_lora_rank}." - ) - - cutedsl_kernel = _fp4_mla_cutedsl_kernel_module() - QK_LOGICAL_DIM = cutedsl_kernel.QK_LOGICAL_DIM - QK_SF_GROUPS = cutedsl_kernel.QK_SF_GROUPS - SMEM_P4_V_N_PER_CTA = cutedsl_kernel.SMEM_P4_V_N_PER_CTA - run_trtllm_fp4_mla_decode_page_native_from_raw = ( - cutedsl_kernel.run_trtllm_fp4_mla_decode_page_native_from_raw - ) - - # Keep general cases on the original path. This specialization uses - # the prepacked-V mufu16 core and linear six-query MTP only. - fold_queries = ( - not packed_rows_per_seq - and num_heads == 6 - and query_len_per_seq == 6 - and num_gen_seqs == 1 - and helix_kv_bounds is None - and softmax_stats_tensor is None - and getattr(cutedsl_kernel, "SUPPORTS_FOLDED_QUERIES", False) - and not _fp4_mla_cutedsl_fused_v_transpose_enabled() - and getattr(metadata, "_fp4_mla_fold_q", True) - and os.environ.get("TRTLLM_FP4_MLA_FOLD_Q", "1") != "0" - ) - if fold_queries: - from .fp4_mla_splitkv import select_split_count - - fold_queries = select_split_count(max_pages * metadata.page_size) > 1 - physical_heads = 128 - kernel_q = q_fp4 - kernel_q_sf = q_sf - if not packed_rows_per_seq and num_heads < physical_heads and not fold_queries: - kernel_q_storage, kernel_q_sf_storage, q_batch_capacity = _prepare_fp4_mla_q_buffers( - metadata, - num_queries, - physical_heads, - q.device, - ) - active_q_rows = num_queries * physical_heads - active_q_sf_bytes = _get_fp4_mla_swizzled_scale_size( - active_q_rows, - QK_LOGICAL_DIM, - ) - kernel_q = kernel_q_storage[:active_q_rows] - kernel_q_sf = kernel_q_sf_storage[:active_q_sf_bytes] - _cutedsl_pad_q_and_sf_kernel[(num_queries, _ceil_div(QK_LOGICAL_DIM // 2, 64))]( - kernel_q, - q_fp4, - kernel_q_sf, - q_sf, - num_heads, - output_heads=physical_heads, - packed_dim=QK_LOGICAL_DIM // 2, - block_bytes=64, - sf_cols=QK_SF_GROUPS, - sf_cols_per_byte_block=8, - ) - - fused_v_transpose = _fp4_mla_cutedsl_fused_v_transpose_enabled() - if fused_v_transpose: - # The fusion kernel reads V from the canonical KV cache and uses - # the current layer V scales directly. Keep a None placeholder so - # the mufu16 and fused-V launchers share one Python call site. - core_v_packed = None - core_v_sf = v_sf - v_page_offset = 0 - else: - v_packed = _get_cutedsl_persistent_v_packed_cache( - metadata, - local_layer, - kv_cache, - v_head_dim=kv_lora_rank, - page_size=metadata.page_size, - block_v=SMEM_P4_V_N_PER_CTA, - ) - core_v_packed = _get_fp4_mla_v_packed_pool_base(metadata) - if core_v_packed is None: - raise RuntimeError("Persistent FP4 MLA V packing requires a stable full-pool base.") - get_v_page_offset = getattr( - metadata.kv_cache_manager, "get_mla_v_packed_page_offset", None - ) - v_page_offset = ( - int(get_v_page_offset(local_layer)) - if callable(get_v_page_offset) - else local_layer * kv_cache.shape[0] - ) - page_bytes = kv_lora_rank * (metadata.page_size // 2) - expected_layer_ptr = core_v_packed.data_ptr() + v_page_offset * page_bytes - if expected_layer_ptr != v_packed.data_ptr(): - raise RuntimeError( - "Persistent FP4 MLA V-packed layer view does not match its " - "full-pool base and page offset." - ) - - v_sf_pool_base = _get_fp4_mla_v_scale_pool_base(metadata) - if v_sf_pool_base is None: - v_sf_pool_base = v_sf_pool.flatten(0, 1).view(torch.uint8) - if v_sf_pool_base.data_ptr() != v_sf_pool.data_ptr(): - raise RuntimeError( - "Persistent FP4 MLA V-scale pool must flatten without a copy." - ) - if ( - not isinstance(v_sf_pool_base, torch.Tensor) - or v_sf_pool_base.dtype != torch.uint8 - or v_sf_pool_base.device != v_sf.device - or v_sf_pool_base.ndim != 2 - or v_sf_pool_base.shape[1] != v_sf_pool.shape[-1] - or not v_sf_pool_base.is_contiguous() - ): - raise RuntimeError( - "Persistent FP4 MLA V-scale pool base must be a contiguous " - "two-dimensional uint8 tensor with the configured page stride." - ) - core_v_sf = v_sf_pool_base.view(torch.float8_e4m3fn) - get_v_sf_page_offset = getattr( - metadata.kv_cache_manager, "get_mla_v_scale_page_offset", None - ) - v_sf_page_offset = ( - int(get_v_sf_page_offset(local_layer)) - if callable(get_v_sf_page_offset) - else local_layer * kv_cache.shape[0] - ) - if v_sf_page_offset != v_page_offset: - raise RuntimeError( - "Persistent FP4 MLA V-packed and V-scale pools require " - "matching encoded layer offsets." - ) - expected_v_sf_ptr = ( - core_v_sf.data_ptr() - + v_sf_page_offset * v_sf_pool.stride(1) * v_sf_pool.element_size() - ) - if expected_v_sf_ptr != v_sf.data_ptr(): - raise RuntimeError( - "Persistent FP4 MLA V-scale layer view does not match its " - "full-pool base and page offset." - ) - - if packed_rows_per_seq: - _run_cutedsl_packed_decode( - cutedsl_kernel, - metadata, - q_fp4, - q_sf, - kv_cache, - sf_cache, - core_v_packed, - core_v_sf, - global_scale, - q_global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], - kv_lens, - output, - num_heads=num_heads, - query_len=query_len_per_seq, - num_requests=num_gen_seqs, - rows_per_seq=packed_rows_per_seq, - sm_scale=float(sm_scale), - v_page_offset=v_page_offset, - helix_kv_bounds=helix_kv_bounds, - softmax_stats_tensor=softmax_stats_tensor, - plan_geometry=plan_geometry, - ) - return - - from .fp4_mla_splitkv import run_fp4_mla_splitkv, select_split_count - - if fold_queries: - from .fp4_mla_foldq import run_fp4_mla_foldq - - run_fp4_mla_foldq( - cutedsl_kernel, - metadata, - q_fp4, - q_sf, - kv_cache, - sf_cache, - core_v_packed, - core_v_sf, - global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], - kv_lens, - output, - splits=select_split_count(max_pages * metadata.page_size), - max_kv_len=max_pages * metadata.page_size, - sm_scale=float(sm_scale), - q_global_scale=q_global_scale, - query_len_per_seq=query_len_per_seq, - v_page_offset=v_page_offset, - ) - return - - # Helix also needs per-query bounds and softmax stats from the native path. - splits = ( - select_split_count(max_pages * metadata.page_size) - if helix_kv_bounds is None and softmax_stats_tensor is None - else 1 - ) - direct_split_output = splits > 1 and getattr( - metadata, "_fp4_mla_splitkv_direct_output", True - ) - kernel_output = output - kernel_softmax_stats = None - if softmax_stats_tensor is not None: - kernel_softmax_stats = _ensure_workspace_tensor( - metadata, - "_fp4_mla_cutedsl_softmax_stats_buf", - (2, num_queries, physical_heads), - dtype=torch.float32, - device=output.device, - ) - if num_heads < physical_heads and not direct_split_output: - kernel_output = _ensure_workspace_tensor( - metadata, - "_fp4_mla_cutedsl_output_buf", - (q_batch_capacity, physical_heads, kv_lora_rank), - dtype=output.dtype, - device=output.device, - )[:num_queries] - - if splits > 1: - run_fp4_mla_splitkv( - _fp4_mla_cutedsl_kernel_module(), - metadata, - kernel_q, - kernel_q_sf, - kv_cache, - sf_cache, - core_v_packed, - core_v_sf, - global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], - kv_lens, - kernel_output, - splits=splits, - max_kv_len=max_pages * metadata.page_size, - sm_scale=float(sm_scale), - q_global_scale=q_global_scale, - query_len_per_seq=query_len_per_seq, - v_page_offset=v_page_offset, - q_batch_capacity=q_batch_capacity, - ) - else: - run_trtllm_fp4_mla_decode_page_native_from_raw( - kernel_q, - kernel_q_sf, - kv_cache, - sf_cache, - core_v_packed, - core_v_sf, - global_scale, - src_page_ids, - metadata.paged_kv_indptr_decode[: num_gen_seqs + 1], - kv_lens, - kernel_output, - max_kv_len=max_pages * metadata.page_size, - sm_scale=float(sm_scale), - num_heads=physical_heads, - q_global_scale=q_global_scale, - page_size=metadata.page_size, - query_len_per_seq=query_len_per_seq, - v_page_offset=v_page_offset, - q_batch_capacity=q_batch_capacity, - partition_runtime_valid_k=bool(getattr(metadata, "is_cuda_graph", False)), - softmax_row_max=(None if kernel_softmax_stats is None else kernel_softmax_stats[0]), - softmax_row_sum=(None if kernel_softmax_stats is None else kernel_softmax_stats[1]), - helix_kv_bounds=helix_kv_bounds, - ) - if kernel_output is not output: - output.copy_(kernel_output[:, :num_heads]) - if kernel_softmax_stats is not None: - softmax_stats_tensor[..., 0].copy_(kernel_softmax_stats[0, :, :num_heads]) - softmax_stats_tensor[..., 1].copy_(kernel_softmax_stats[1, :, :num_heads]) - return - - total_p_rows = num_queries * max_pages * num_heads - p_fp4 = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_p_buf", - (max(total_p_rows, 1), metadata.page_size // 2), - dtype=torch.uint8, - device=q.device, - )[:total_p_rows] - p_sf = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_p_sf_buf", - (max(_get_fp4_mla_swizzled_scale_size(total_p_rows, metadata.page_size), 1),), - dtype=torch.float8_e4m3fn, - device=q.device, - ) - stats_shape = (num_queries, num_heads) - max_scores = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_max_buf", - stats_shape, - dtype=torch.float32, - device=q.device, - ) - denom = _ensure_workspace_tensor( - metadata, - "_fp4_mla_attention_denom_buf", - stats_shape, - dtype=torch.float32, - device=q.device, - ) - - if backend != "triton": - raise ValueError( - f"Unsupported FP4 MLA attention backend '{backend}'. " - f"Set {FP4_MLA_ATTENTION_BACKEND_ENV} to 'triton' or " - f"'{_FP4_MLA_CUTEDSL_BACKEND}'." - ) - - # Self-contained public-Triton path: TMA-loaded QK + fused page-stats pack, - # reduce-stats, prob-scale, and PV with an optional prepacked V cache. - _run_triton_attention_decode( - metadata=metadata, - layer_idx=layer_idx, - local_layer=local_layer, - q_fp4=q_fp4, - q_sf=q_sf.contiguous().view(-1), - kv_cache=kv_cache, - sf_cache=sf_cache, - v_sf=v_sf, - global_scale=global_scale, - src_page_ids=src_page_ids, - kv_lens=kv_lens, - p_fp4=p_fp4, - p_sf=p_sf, - max_scores=max_scores, - denom=denom, - output=output, - num_queries=num_queries, - num_heads=num_heads, - head_dim=head_dim, - kv_lora_rank=kv_lora_rank, - q_residual_dim=q_residual_dim, - query_len_per_seq=query_len_per_seq, - max_pages=max_pages, - sm_scale=float(sm_scale), - q_global_scale=q_global_scale, - ) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py deleted file mode 100644 index b9d2d9adc58f..000000000000 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ /dev/null @@ -1,473 +0,0 @@ -# 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 management for MiniMax-M3 sparse attention. - -Provides: - * :class:`MiniMaxM3SparseIndexCache` — plain-tensor side cache used by - algorithm-only unit tests (no pyexecutor dependency required for - construction). - * :class:`MiniMaxM3KVCacheManagerV2` — :class:`KVCacheManagerV2` - subclass that registers a per-sparse-layer ``Role.INDEX_KEY`` paged - buffer alongside the standard K/V buffers. -""" - -from __future__ import annotations - -from typing import List, Optional, Sequence - -import torch - -from tensorrt_llm._torch.disaggregation.resource.page import MapperKind -from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor -from tensorrt_llm.bindings import DataType -from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode -from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX -from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole - -from ....pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role - - -class MiniMaxM3SparseIndexCache: - """Plain-tensor side cache for the M3 sparse index branch. - - Slot layout matches the main ``KVCacheManagerV2`` paged buffer - geometry so the same ``req_to_token`` mapping addresses both - caches. One ``[num_slots, 1, sparse_index_dim]`` index-K buffer is - allocated per sparse layer; index-V is allocated only for layers - not listed in ``disable_index_value_layer_ids``. - """ - - def __init__( - self, - *, - num_layers: int, - sparse_layer_ids: List[int], - disable_index_value_layer_ids: List[int], - num_slots: int, - sparse_index_dim: int, - dtype: torch.dtype, - device: torch.device, - ): - if num_slots <= 0: - raise ValueError(f"num_slots must be > 0, got {num_slots}") - if sparse_index_dim <= 0: - raise ValueError(f"sparse_index_dim must be > 0, got {sparse_index_dim}") - self.num_layers = int(num_layers) - self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids) - self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids) - self.num_slots = int(num_slots) - self.sparse_index_dim = int(sparse_index_dim) - self.dtype = dtype - self.device = device - - self._index_k: dict[int, torch.Tensor] = {} - self._index_v: dict[int, torch.Tensor] = {} - for layer_idx in self.sparse_layer_ids: - if not (0 <= layer_idx < self.num_layers): - raise ValueError(f"sparse layer_idx {layer_idx} outside [0, {self.num_layers})") - self._index_k[layer_idx] = torch.zeros( - (self.num_slots, 1, self.sparse_index_dim), - dtype=dtype, - device=device, - ) - if layer_idx not in self.disable_index_value_layer_ids: - self._index_v[layer_idx] = torch.zeros( - (self.num_slots, 1, self.sparse_index_dim), - dtype=dtype, - device=device, - ) - - def has_index_value(self, layer_idx: int) -> bool: - return layer_idx in self._index_v - - def get_index_k_buffer(self, layer_idx: int) -> torch.Tensor: - if layer_idx not in self._index_k: - raise KeyError( - f"layer_idx {layer_idx} is not a sparse layer; " - f"sparse layers: {self.sparse_layer_ids}" - ) - return self._index_k[layer_idx] - - def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._index_v.get(layer_idx) - - def set_index_k(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_k: torch.Tensor) -> None: - """Write ``idx_k`` into the index-K cache at ``out_cache_loc``.""" - buf = self.get_index_k_buffer(layer_idx) - if idx_k.shape[1] != 1: - raise ValueError( - f"index K is replicated single-head; expected shape " - f"[N, 1, {self.sparse_index_dim}], got {tuple(idx_k.shape)}" - ) - buf.index_copy_(0, out_cache_loc.to(torch.long), idx_k.to(buf.dtype)) - - def set_index_v(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_v: torch.Tensor) -> None: - """Write ``idx_v`` into the index-V cache (only when allocated).""" - buf = self.get_index_v_buffer(layer_idx) - if buf is None: - raise RuntimeError( - f"layer {layer_idx} has disable_index_value=True; index V is not allocated" - ) - if idx_v.shape[1] != 1: - raise ValueError( - f"index V is replicated single-head; expected shape " - f"[N, 1, {self.sparse_index_dim}], got {tuple(idx_v.shape)}" - ) - buf.index_copy_(0, out_cache_loc.to(torch.long), idx_v.to(buf.dtype)) - - -class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): - """KVCacheManagerV2 subclass with a V2-managed paged index-K cache - per sparse layer. - - Each sparse local layer registers a ``Role.INDEX_KEY`` - :class:`BufferConfig` via :meth:`_extra_buffers_per_layer` so the - index-K cache participates in the V2 paged lifecycle (allocation, - free, prefix reuse) and shares ``life_cycle_id`` with main K/V. - - The optional index-V branch is kept as a plain CUDA tensor for the - rare ``disable_index_value=False`` test path. M3 production sets - ``disable_index_value=True`` everywhere, so ``_index_v_buffers`` - stays empty in production. - - Constructor extras (forwarded kwargs go to :class:`KVCacheManagerV2`): - * ``sparse_layer_ids`` — layer indices using sparse attention. - * ``disable_index_value_layer_ids`` — subset whose index-V is - omitted. - * ``sparse_index_dim`` — width of the index-K/V vectors. - """ - - def __init__( - self, - *args, - sparse_layer_ids=None, - disable_index_value_layer_ids=None, - sparse_index_dim: Optional[int] = None, - **kwargs, - ): - # Resolve M3 sparse-layer metadata from explicit kwargs first, - # then from ``sparse_attn_config``, then from the M3 checkpoint - # convention (layers 0..2 dense, 3..N-1 sparse, - # disable_index_value=True, sparse_index_dim=128). - sparse_attn_config = kwargs.get("sparse_attn_config") - num_layers = kwargs.get("num_layers") - - if sparse_index_dim is None: - sparse_index_dim = int(getattr(sparse_attn_config, "sparse_index_dim", 0) or 0) or 128 - if sparse_layer_ids is None: - if num_layers is not None: - sparse_layer_ids = list(range(3, int(num_layers))) - else: - sparse_layer_ids = [] - if disable_index_value_layer_ids is None: - disable_index_value_layer_ids = list(sparse_layer_ids) - - # Must be set BEFORE super().__init__ — the base - # ``_build_base_config`` invokes ``_extra_buffers_per_layer`` - # which reads these attributes. - self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids) - self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids) - self.sparse_index_dim = int(sparse_index_dim) - - super().__init__(*args, **kwargs) - - index_v_layer_ids = set(self.sparse_layer_ids) - self.disable_index_value_layer_ids - if self.is_disagg and index_v_layer_ids: - raise ValueError( - "MiniMax M3 disaggregated serving requires disable_index_value=True " - "for every sparse layer because the optional test-only index-V cache " - "is not managed or transferred by KVCacheManagerV2; enabled layers=" - f"{sorted(index_v_layer_ids)}" - ) - - # Optional plain-tensor index-V cache for non-disabled sparse - # layers (test-only; production has disable_index_value=True - # on every sparse layer). - num_total_slots = self._compute_num_total_slots() - torch_dtype = self._torch_dtype_for_index_cache() - device = torch.device("cuda") - self._index_v_buffers: dict[int, torch.Tensor] = {} - for layer_idx in self.sparse_layer_ids: - if layer_idx not in self.layer_offsets: - continue - if layer_idx not in self.disable_index_value_layer_ids: - self._index_v_buffers[layer_idx] = torch.zeros( - (num_total_slots, 1, self.sparse_index_dim), - dtype=torch_dtype, - device=device, - ) - - def _extra_buffers_per_layer(self, *, tokens_per_block): - """Register a per-sparse-layer ``Role.INDEX_KEY`` :class:`BufferConfig`. - - ``size`` is bytes per **block**: ``1 * sparse_index_dim * - elem_bytes * tokens_per_block``. Keyed by **local** layer id — - the base ``_build_base_config`` iterates local ids, so keying - by global ids would silently skip registration on non-trivial - PP ranks. - """ - torch_dtype = self._torch_dtype_for_index_cache() - elem_bytes = torch.tensor([], dtype=torch_dtype).element_size() - bytes_per_token = 1 * self.sparse_index_dim * elem_bytes - size_per_block = bytes_per_token * tokens_per_block - return { - self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=size_per_block)] - for layer_id in self.sparse_layer_ids - if layer_id in self.layer_offsets - } - - def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: - """Declare MiniMax M3's token-major K/V and replicated index-K.""" - return { - Role.ALL: MapperKind.NHD, - Role.INDEX_KEY: MapperKind.REPLICATED, - } - - def _compute_num_total_slots(self) -> int: - """Total token slots across all blocks in the main K pool. - - Sizes the plain-tensor index-V cache only; the V2-managed - index-K cache pulls its slot count from ``Role.INDEX_KEY``'s - page upper bound. - """ - if not self.layer_offsets: - return int(self.max_batch_size * self.max_seq_len) - any_layer_offset = next(iter(self.layer_offsets.values())) - page_upper = self.impl.get_page_index_upper_bound(any_layer_offset, Role.KEY) - kv_factor = 1 if self.kv_cache_type == CacheTypeCpp.SELFKONLY else 2 - return int((page_upper // kv_factor) * self.tokens_per_block) - - def _torch_dtype_for_index_cache(self) -> torch.dtype: - """Match the main cache dtype where possible, fall back to bf16.""" - if self.dtype == DataType.HALF: - return torch.float16 - if self.dtype == DataType.FLOAT: - return torch.float32 - return torch.bfloat16 - - def get_index_k_buffer(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: - """Return the V2-managed paged index-K view for ``layer_idx``. - - NHD shape is ``[num_pages, tokens_per_block, 1, sparse_index_dim]``; - HND shape is ``[num_pages, 1, tokens_per_block, sparse_index_dim]``. - """ - return super().get_index_k_buffer( - layer_idx, - num_heads=1, - head_dim=self.sparse_index_dim, - dtype=self._torch_dtype_for_index_cache(), - kv_layout=kv_layout, - ) - - def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]: - """Plain-tensor index-V cache for non-disabled sparse layers.""" - return self._index_v_buffers.get(layer_idx) - - def has_index_value(self, layer_idx: int) -> bool: - return layer_idx in self._index_v_buffers - - def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: - """Return a paged K+V view with strides spanning the coalesced pool. - - The base :meth:`KVCacheManagerV2.get_buffers` produces a - ``[num_pages, kv_factor, ...]`` view with contiguous strides - that assume the slot holds exactly one layer's K+V. In M3's - pool the slot packs K+V for *all* layers of the group - (``scale >= 2 * num_layers_in_group``), so the base view's - dim-0 stride does not reach the next slot's K for this layer. - (When INDEX_KEY's per-block size coincides with K/V's, it is - coalesced into the same pool and contributes to ``scale`` too.) - - The override builds a ``[num_slots, scale, ...]`` view rooted - at K's base, then slices ``[:, :2]`` to extract K+V. The slice - preserves the dim-0 stride (``scale * page_stride``), so - ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``. - """ - if kv_layout not in ("NHD", "HND"): - raise ValueError(f"Unsupported kv_layout: {kv_layout}") - if self.kv_cache_type == CacheTypeCpp.SELFKONLY: - raise NotImplementedError( - "MiniMaxM3KVCacheManagerV2.get_buffers does not support SELFKONLY cache type" - ) - - layer_offset = self.layer_offsets[layer_idx] - addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY) - addr_value = self.impl.get_mem_pool_base_address(layer_offset, Role.VALUE) - page_stride_key = self.impl.get_page_stride(layer_offset, Role.KEY) - page_stride_value = self.impl.get_page_stride(layer_offset, Role.VALUE) - # V2 always lays V immediately after K within the per-layer - # contribution to a slot. The slice ``[:, :2]`` depends on this. - assert addr_key + page_stride_value == addr_value, ( - f"MiniMaxM3 get_buffers requires addr_K + page_stride " - f"== addr_V (V immediately after K in slot); got " - f"addr_K={addr_key} page_stride_V={page_stride_value} " - f"addr_V={addr_value} for layer {layer_idx}." - ) - assert page_stride_key == page_stride_value, ( - f"MiniMaxM3 get_buffers requires equal K and V page " - f"strides; got K={page_stride_key} V=" - f"{page_stride_value}." - ) - - converter = self.impl.get_page_index_converter(layer_offset, Role.KEY) - scale = int(converter.scale) - layer_offset_pages = int(converter.layer_offset) - page_upper_K = self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) - num_slots_total = page_upper_K + layer_offset_pages - assert num_slots_total % scale == 0, ( - f"V2 storage inconsistency: page_upper_K + " - f"layer_offset_pages = {num_slots_total} is not " - f"divisible by scale = {scale}." - ) - num_slots = num_slots_total // scale - - element_per_container = 1 - dtype = self.dtype - if dtype == DataType.NVFP4: - element_per_container = 2 - torch_dtype = torch.int8 - else: - torch_dtype = binding_to_torch_dtype(dtype) - - layer_head_dim = self.head_dim_per_layer[layer_offset] - num_kv_heads = self.num_kv_heads_per_layer[layer_offset] - - if kv_layout == "NHD": - full_slot_shape = [ - num_slots, - scale, - self.tokens_per_block, - num_kv_heads, - layer_head_dim // element_per_container, - ] - else: - full_slot_shape = [ - num_slots, - scale, - num_kv_heads, - self.tokens_per_block, - layer_head_dim // element_per_container, - ] - - full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) - return full_view[:, :2] - - def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> int: - """Pool-mapping offset from the layer's physical position in its pool. - - The base formula ``exact_div(addr_offset, key_bytes * kv_factor * - tokens_per_block)`` assumes each layer contributes exactly K+V to - its pool slot. When index-K coalesces into the K/V pool the layer - stride is non-uniform (sparse layers add an INDEX_KEY sub-page), - so no uniform-stride offset exists. The M3 forward path uses - :meth:`get_buffers` / :meth:`get_index_k_buffer` rather than this - mapping, so the offset just needs to be a consistent per-layer - position. Rank the group's layers by their K base address instead - of by ``layer_grouping`` iteration order: the ordering of - ``layer_grouping`` is not a V2 API contract, while the address - rank always reflects the physical slot layout (and keeps the - NVFP4 ``block_scale_offset == offset`` cross-check in the base - pool-mapping loop meaningful). - """ - layers_by_addr = sorted( - self.impl.layer_grouping[int(layer_group_id)], - key=lambda lid: self.impl.get_mem_pool_base_address( - lid, Role.KEY, PageIndexMode.SHARED - ), - ) - return layers_by_addr.index(int(layer_id)) - - def _get_batch_cache_indices_by_pool_id( - self, - request_ids, - *, - pool_id: int = 0, - is_kv_aggregate: bool = True, - num_blocks_per_seq: Optional[Sequence[int]] = None, - index_scale: Optional[int] = None, - kv_factor: Optional[int] = None, - ): - """Return page indices; padded entries remain ``BAD_PAGE_INDEX`` (-1). - - The base method converts slot ids to V1-style block ids via - ``base_idx * index_scales[pool_id] // kv_factor``, which is - only correct when each layer contributes exactly K+V. M3's slot - packs K+V for all layers of the group, so the scale breaks the - V1 conversion and produces out-of-bounds block ids during V2 - warmup. - - Bypass the conversion: the M3 forward path indexes paged - views (built by :meth:`get_buffers` / - :meth:`get_index_k_buffer`) directly by slot id. - ``BAD_PAGE_INDEX`` slots remain ``-1`` here because disaggregation's - :class:`KVRegionExtractorV1` filters ``region_ids >= 0``. - :meth:`get_block_ids_per_seq` maps them to zero for the attention - metadata's padded tensor. - - Args: - request_ids: Request IDs whose page-index rows are returned. - pool_id: V2 pool whose page indices are requested. - is_kv_aggregate: Kept for compatibility with the base virtual method. - num_blocks_per_seq: Optional per-request truncation limits. When - omitted, preserve the full padded width required by MiniMax - CUDA-graph metadata initialization. - index_scale: Kept for compatibility with the base virtual method; - M3 bypasses the V1 block-id conversion entirely, so any - caller-supplied scale is ignored alongside ``index_scales``. - """ - res = [] - for req_idx, req_id in enumerate(request_ids): - kv_cache = self.kv_cache_map[req_id] - base_page_indices = kv_cache.get_base_page_indices(pool_id) - if num_blocks_per_seq is not None: - num_blocks = min(kv_cache.num_blocks, num_blocks_per_seq[req_idx]) - base_page_indices = base_page_indices[:num_blocks] - res.append(list(base_page_indices)) - return res - - def get_block_ids_per_seq(self, request_ids): - """Return per-request slot ids matching the per-layer paged view's dim-0. - - Drops the base's final ``i // num_local_layers`` step (paired - with the base ``index_scales`` multiplication that's also - bypassed here). Pads with ``0`` to preserve shape. - """ - block_ids_per_seq = self.get_batch_cache_indices(request_ids) - block_ids_per_seq_tensors = [ - torch.tensor( - [i if i != BAD_PAGE_INDEX else 0 for i in sublist], - dtype=torch.int, - ) - for sublist in block_ids_per_seq - ] - padded_tensor = torch.nn.utils.rnn.pad_sequence( - block_ids_per_seq_tensors, - batch_first=True, - padding_value=0, - ) - return padded_tensor - - -def get_minimax_m3_kv_cache_manager_cls(): - """Backward-compatible accessor; prefer importing the class directly.""" - return MiniMaxM3KVCacheManagerV2 - - -__all__ = [ - "MiniMaxM3KVCacheManagerV2", - "MiniMaxM3SparseIndexCache", - "get_minimax_m3_kv_cache_manager_cls", -] diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index 1a0614d0c101..a149d781204e 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -16,6 +16,7 @@ from ..._utils import get_sm_version from ...logger import logger from ...math_utils import ceil_div, pad_up +from ..attention.kernels.blackwell.dspark.attention import DSparkAttention from ..autotuner import ( AutoTuner, ConstraintSpec, @@ -24,11 +25,14 @@ TunableRunner, TuningConfig, ) -from ..cute_dsl_kernels.blackwell.dspark.attention import DSparkAttention -from ..cute_dsl_kernels.blackwell.utils import make_ptr -from ..cute_dsl_utils import IS_CUTLASS_DSL_INTERNAL_AVAILABLE +from ..cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from ..kernels.blackwell.utils import make_ptr from ..utils import fp8_scale_infer_shape, mxfp8_scale_infer_shape, next_positive_power_of_2 -from .dspark_rmsnorm_rope_custom_op import _get_dspark_arch_str +from .dspark_rmsnorm_rope_custom_op import ( + _get_dspark_arch_str, + cute_dsl_dspark_rmsnorm_rope_cache_write, + cute_dsl_dspark_rmsnorm_rope_draft_block, +) _DSPARK_NUM_HEADS = 128 _DSPARK_HEAD_DIM = 512 @@ -442,7 +446,7 @@ def _( ) -if IS_CUTLASS_DSL_INTERNAL_AVAILABLE: +if IS_CUTLASS_DSL_RUBIN_AVAILABLE: from .cute_dsl_custom_ops import CuteDSLFp8RubinBmmQuantRunner, CuteDSLMXFP8RubinLinear @functools.cache @@ -588,7 +592,7 @@ def _forward_tvm_ffi_prepared( bias: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: - compiled_gemm, swap_ab = launch_plan + compiled_gemm, swap_ab, split_k = launch_plan if len(inputs) == 5: a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, alpha_tensor = inputs output_tensor = None @@ -606,7 +610,11 @@ def _forward_tvm_ffi_prepared( n = b_tensor.shape[0] partition_id = kwargs.get("partition_id", -1) if output_tensor is not None: - c_tensor = output_tensor[:, partition_id * n : (partition_id + 1) * n] + c_tensor = ( + output_tensor[:, partition_id * n : (partition_id + 1) * n] + if partition_id >= 0 + else output_tensor + ) elif self.to_userbuffers: c_tensor = torch.ops.trtllm.create_userbuffers_tensor([m, n], self.output_dtype) else: @@ -619,6 +627,9 @@ def _forward_tvm_ffi_prepared( if swap_ab: c_tensor = c_tensor.permute(1, 0) + if split_k > 1: + c_tensor.zero_() + real_k = k * self.packed_k_multiplier sf_m = pad_up(m, 128) sf_k = pad_up(real_k // self.sf_vec_size, 4) @@ -642,7 +653,7 @@ def _forward_tvm_ffi_prepared( kernel_b_tensor = b_tensor kernel_b_sf_tensor = b_sf_tensor - c_ld = output_tensor.shape[1] if output_tensor is not None else 0 + c_ld = output_tensor.shape[1] if output_tensor is not None and partition_id >= 0 else 0 compiled_gemm( kernel_m, kernel_n, @@ -974,3 +985,244 @@ def _( shape = list(mat_a.shape) shape[-1] = mat_b.shape[-2] return mat_a.new_empty(shape, dtype=torch.bfloat16) + + +def _log_unsupported(reason: str, key: str) -> bool: + logger.debug_once( + f"Falling back from fused DSV4 DSpark attention: {reason}", + key=("fused_dsv4_dspark_attention_unsupported", key), + ) + return False + + +def is_dsv4_dspark_attention_config_supported( + block_size: int, + num_heads: int, + head_dim: int, + window_size: int, +) -> bool: + """Return whether static model geometry matches the DSV4 specialization.""" + return ( + block_size in _DSV4_DSPARK_BLOCK_SIZES + and num_heads in _DSV4_DSPARK_NUM_HEADS + and head_dim == _DSV4_DSPARK_HEAD_DIM + and window_size == _DSV4_DSPARK_WINDOW_SIZE + ) + + +def is_fused_dsv4_dspark_attention_supported( + q: torch.Tensor, + kv_cache: torch.Tensor, + valid_len: torch.Tensor, + attn_sink: torch.Tensor, + inverse_rope_freqs: torch.Tensor, +) -> bool: + """Return whether tensors can use the fused DSV4 DSpark attention kernel.""" + if _get_dspark_arch_str() is None: + return _log_unsupported(f"SM {get_sm_version()} is not SM100 or SM103", "sm_version") + if not all(t.is_cuda for t in (q, kv_cache, valid_len, attn_sink, inverse_rope_freqs)): + return _log_unsupported("all inputs must be CUDA tensors", "device") + if q.dtype != torch.bfloat16: + return _log_unsupported(f"q dtype must be BF16, got {q.dtype}", "q_dtype") + if kv_cache.dtype != q.dtype: + return _log_unsupported( + f"kv_cache dtype must match q, got {kv_cache.dtype}", "kv_cache_dtype" + ) + if valid_len.dtype != torch.int64: + return _log_unsupported( + f"valid_len must use INT64, got {valid_len.dtype}", "valid_len_dtype" + ) + if attn_sink.dtype != torch.float32 or inverse_rope_freqs.dtype != torch.float32: + return _log_unsupported( + "attn_sink and inverse_rope_freqs must be FP32; " + f"got {attn_sink.dtype} and {inverse_rope_freqs.dtype}", + "aux_dtype", + ) + if q.ndim != 4 or kv_cache.ndim != 3: + return _log_unsupported( + f"expected q/kv_cache ranks 4/3, got {q.ndim}/{kv_cache.ndim}", "tensor_ranks" + ) + if q.shape[1] not in _DSV4_DSPARK_BLOCK_SIZES: + return _log_unsupported(f"draft block size must be 5 or 6, got {q.shape[1]}", "block_size") + if q.shape[2] not in _DSV4_DSPARK_NUM_HEADS or q.shape[3] != _DSV4_DSPARK_HEAD_DIM: + return _log_unsupported( + "q must have 128 heads and head_dim 512; " + f"got {q.shape[2]} heads and head_dim {q.shape[3]}", + "q_shape", + ) + if kv_cache.shape[1:] != (_DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM): + return _log_unsupported( + f"kv_cache must have trailing shape (128, 512), got {tuple(kv_cache.shape[1:])}", + "kv_cache_shape", + ) + if ( + kv_cache.stride(0) <= 0 + or kv_cache.stride(0) % 8 != 0 + or kv_cache.stride(1) != _DSV4_DSPARK_HEAD_DIM + or kv_cache.stride(2) != 1 + ): + return _log_unsupported( + "kv_cache strides must have a positive 16-byte-aligned page stride " + f"and trailing strides (512, 1), got {kv_cache.stride()}", + "kv_cache_layout", + ) + if valid_len.shape != (q.shape[0],): + return _log_unsupported( + f"valid_len shape must be {(q.shape[0],)}, got {tuple(valid_len.shape)}", + "valid_len_shape", + ) + if attn_sink.shape != (q.shape[2],): + return _log_unsupported( + f"attn_sink shape must be {(q.shape[2],)}, got {tuple(attn_sink.shape)}", + "attn_sink_shape", + ) + expected_freqs_shape = (q.shape[0], q.shape[1], _DSV4_DSPARK_ROPE_DIM // 2, 2) + if inverse_rope_freqs.shape != expected_freqs_shape: + return _log_unsupported( + f"inverse_rope_freqs shape must be {expected_freqs_shape}, " + f"got {tuple(inverse_rope_freqs.shape)}", + "inverse_rope_freqs_shape", + ) + for name, tensor in ( + ("q", q), + ("valid_len", valid_len), + ("attn_sink", attn_sink), + ("inverse_rope_freqs", inverse_rope_freqs), + ): + if not tensor.is_contiguous(): + return _log_unsupported(f"{name} must be contiguous", f"{name}_layout") + return True + + +def _is_fused_dsv4_dspark_attention_input_supported( + q: torch.Tensor, + draft_block: torch.Tensor, + kv_cache: torch.Tensor, + slots_i32: torch.Tensor, + cache_seqs: torch.Tensor, + valid_len: torch.Tensor, + attn_sink: torch.Tensor, + inverse_rope_freqs: torch.Tensor, +) -> bool: + if not is_fused_dsv4_dspark_attention_supported( + q, kv_cache, valid_len, attn_sink, inverse_rope_freqs + ): + return False + if not all(t.is_cuda for t in (draft_block, slots_i32, cache_seqs)): + return _log_unsupported("all fused-op inputs must be CUDA tensors", "input_device") + if draft_block.dtype != q.dtype: + return _log_unsupported( + f"draft_block dtype must match q, got {draft_block.dtype}", "draft_block_dtype" + ) + if slots_i32.dtype != torch.int32 or cache_seqs.dtype != torch.int32: + return _log_unsupported( + "slots_i32 and cache_seqs must use INT32; " + f"got {slots_i32.dtype} and {cache_seqs.dtype}", + "index_dtype", + ) + expected_draft_block_shape = ( + q.shape[0], + _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, + _DSV4_DSPARK_HEAD_DIM, + ) + if draft_block.shape != expected_draft_block_shape: + return _log_unsupported( + f"draft_block shape must be {expected_draft_block_shape}, " + f"got {tuple(draft_block.shape)}", + "draft_block_shape", + ) + expected_batch_shape = (q.shape[0],) + if slots_i32.shape != expected_batch_shape or cache_seqs.shape != expected_batch_shape: + return _log_unsupported( + f"slots_i32 and cache_seqs shapes must be {expected_batch_shape}; " + f"got {tuple(slots_i32.shape)} and {tuple(cache_seqs.shape)}", + "index_shape", + ) + for name, tensor in ( + ("draft_block", draft_block), + ("slots_i32", slots_i32), + ("cache_seqs", cache_seqs), + ): + if not tensor.is_contiguous(): + return _log_unsupported(f"{name} must be contiguous", f"{name}_layout") + return True + + +def warmup_fused_dsv4_dspark_attention(block_size: int, eps: float) -> None: + """Best-effort prewarm of the production DSV4 DSpark fused-op path.""" + if _get_dspark_arch_str() is None or block_size not in _DSV4_DSPARK_BLOCK_SIZES: + return + + batch = 1 + device = torch.device("cuda") + try: + with torch.inference_mode(): + q = torch.zeros( + (batch, block_size, 128, _DSV4_DSPARK_HEAD_DIM), + dtype=torch.bfloat16, + device=device, + ) + main_x = torch.zeros( + (batch, 1, _DSV4_DSPARK_HEAD_DIM), dtype=torch.bfloat16, device=device + ) + block_x = torch.zeros( + (batch, block_size, _DSV4_DSPARK_HEAD_DIM), + dtype=torch.bfloat16, + device=device, + ) + weight = torch.ones((_DSV4_DSPARK_HEAD_DIM,), dtype=torch.bfloat16, device=device) + main_freqs = torch.zeros( + (batch, _DSV4_DSPARK_ROPE_DIM // 2, 2), + dtype=torch.float32, + device=device, + ) + block_freqs = torch.zeros( + (batch * block_size, _DSV4_DSPARK_ROPE_DIM // 2, 2), + dtype=torch.float32, + device=device, + ) + kv_cache = torch.zeros( + (batch, _DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM), + dtype=torch.bfloat16, + device=device, + ) + slots = torch.zeros((batch,), dtype=torch.int64, device=device) + start_pos = torch.zeros((batch,), dtype=torch.int64, device=device) + slots_i32, cache_seqs = cute_dsl_dspark_rmsnorm_rope_cache_write( + main_x, + weight, + main_freqs, + kv_cache, + slots, + start_pos, + eps, + ) + draft_block = cute_dsl_dspark_rmsnorm_rope_draft_block( + block_x, + weight, + block_freqs, + eps, + ) + fused_dsv4_dspark_attention( + q, + draft_block, + kv_cache, + slots_i32[:, None], + cache_seqs.to(torch.int64), + torch.zeros((batch,), dtype=torch.int64, device=device), + torch.full((batch,), _DSV4_DSPARK_WINDOW_SIZE, dtype=torch.int64, device=device), + torch.zeros((128,), dtype=torch.float32, device=device), + block_freqs.view( + batch, + block_size, + _DSV4_DSPARK_ROPE_DIM // 2, + 2, + ), + 1.0, + ) + torch.cuda.synchronize() + except RuntimeError as e: + logger.warning( + "DSV4 DSpark CuTe DSL attention prewarm failed; the op will " + f"self-JIT on first use. {type(e).__name__}: {e}" + ) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 31e8ee1f7275..0fc68189f14c 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -10,9 +10,10 @@ import torch from ..._utils import get_sm_version -from ..cute_dsl_kernels.blackwell.dspark_rmsnorm_rope import ( +from ..attention.kernels.blackwell.dspark_rmsnorm_rope import ( DSparkConcatFP8QuantKernel, DSparkRMSNormFP8QuantKernel, + DSparkRMSNormRoPECacheWriteKernel, DSparkRMSNormRoPEDraftBlockKernel, DSparkRMSNormRoPEKernel, DSparkRMSNormRoPEPageWriteKernel, @@ -20,12 +21,20 @@ _DSV4_DSPARK_HEAD_DIM = 512 _DSV4_DSPARK_ROPE_DIM = 64 +_DSV4_DSPARK_WINDOW_SIZE = 128 _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE = 8 - +_DSV4_DSPARK_BLOCK_SIZES = (5, 6) _DSV4_DSPARK_ARCH_BY_SM = { 100: "sm_100", 103: "sm_103", } +# The RMSNorm/RoPE kernels are plain SIMT CuTe kernels with no architecture-specific +# instructions, so they run on more GPUs than the DSV4 DSpark attention kernel gated +# by _DSV4_DSPARK_ARCH_BY_SM. +_DSPARK_RMSNORM_ROPE_SM_VERSIONS = frozenset((100, 103, 107)) + + +_dspark_fp8_input_compile_allowed = False @functools.cache @@ -36,6 +45,14 @@ def _get_dspark_arch_str(sm_version: int | None = None) -> str | None: return _DSV4_DSPARK_ARCH_BY_SM.get(sm_version) +@functools.cache +def _is_dspark_rmsnorm_rope_sm_supported(sm_version: int | None = None) -> bool: + """Return whether the fused RMSNorm/RoPE kernels support this GPU.""" + if sm_version is None: + sm_version = get_sm_version() + return sm_version in _DSPARK_RMSNORM_ROPE_SM_VERSIONS + + def _has_regular_row_stride(x: torch.Tensor) -> bool: """Return whether leading dimensions flatten to non-overlapping strided rows.""" if x.is_contiguous(): @@ -106,6 +123,70 @@ def is_fused_dspark_rmsnorm_rope_supported( ) +def _is_dspark_cache_write_supported( + x: torch.Tensor, + weight: torch.Tensor, + freqs: torch.Tensor, + kv_cache: torch.Tensor, + slots: torch.Tensor, + start_pos: torch.Tensor, +) -> bool: + batch = x.shape[0] if x.ndim == 3 else -1 + return ( + x.shape == (batch, 1, _DSV4_DSPARK_HEAD_DIM) + and is_fused_dspark_rmsnorm_rope_supported( + x, weight, freqs, num_heads=1, rope_dim=_DSV4_DSPARK_ROPE_DIM + ) + and kv_cache.is_cuda + and kv_cache.dtype == x.dtype + and kv_cache.ndim == 3 + and kv_cache.shape[1:] == (_DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM) + and kv_cache.stride(1) == _DSV4_DSPARK_HEAD_DIM + and kv_cache.stride(2) == 1 + and slots.is_cuda + and slots.dtype == torch.int64 + and slots.shape == (batch,) + and slots.is_contiguous() + and start_pos.is_cuda + and start_pos.dtype == torch.int64 + and start_pos.shape == (batch,) + and start_pos.is_contiguous() + ) + + +def _is_dspark_draft_block_supported( + x: torch.Tensor, + weight: torch.Tensor, + freqs: torch.Tensor, +) -> bool: + block_size = x.shape[1] if x.ndim == 3 else -1 + return ( + block_size in _DSV4_DSPARK_BLOCK_SIZES + and x.shape[-1:] == (_DSV4_DSPARK_HEAD_DIM,) + and is_fused_dspark_rmsnorm_rope_supported( + x, weight, freqs, num_heads=1, rope_dim=_DSV4_DSPARK_ROPE_DIM + ) + ) + + +def is_fused_dspark_attention_preparation_supported( + main_x: torch.Tensor, + block_x: torch.Tensor, + weight: torch.Tensor, + main_freqs: torch.Tensor, + block_freqs: torch.Tensor, + kv_cache: torch.Tensor, + slots: torch.Tensor, + start_pos: torch.Tensor, +) -> bool: + """Return whether RMSNorm/RoPE can directly prepare the attention inputs.""" + return ( + main_x.shape[0] == block_x.shape[0] + and _is_dspark_cache_write_supported(main_x, weight, main_freqs, kv_cache, slots, start_pos) + and _is_dspark_draft_block_supported(block_x, weight, block_freqs) + ) + + @functools.cache def _compile_fused_dspark_rmsnorm_rope( hidden_dim: int, @@ -218,6 +299,105 @@ def _compile_fused_dspark_rope_into( ) +@functools.cache +def _compile_dspark_rmsnorm_rope_cache_write(eps: float): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "DSpark RMSNorm/RoPE cache-write must be warmed up before CUDA graph capture" + ) + rows = cute.sym_int() + pages = cute.sym_int() + x_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (rows, _DSV4_DSPARK_HEAD_DIM), stride_order=(1, 0) + ) + weight_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (_DSV4_DSPARK_HEAD_DIM,), stride_order=(0,) + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (rows, _DSV4_DSPARK_ROPE_DIM // 2, 2), + stride_order=(2, 1, 0), + ) + cache_fake = cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (pages, _DSV4_DSPARK_WINDOW_SIZE, _DSV4_DSPARK_HEAD_DIM), + stride=(cute.sym_int64(), _DSV4_DSPARK_HEAD_DIM, 1), + ) + slots_fake = cute.runtime.make_fake_compact_tensor(cutlass.Int64, (rows,), stride_order=(0,)) + start_pos_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int64, (rows,), stride_order=(0,) + ) + slots_i32_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (rows,), stride_order=(0,) + ) + cache_seqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (rows,), stride_order=(0,) + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkRMSNormRoPECacheWriteKernel( + _DSV4_DSPARK_HEAD_DIM, + _DSV4_DSPARK_ROPE_DIM, + eps, + _DSV4_DSPARK_WINDOW_SIZE, + ) + return cute.compile( + kernel, + x_fake, + weight_fake, + freqs_fake, + cache_fake, + slots_fake, + start_pos_fake, + slots_i32_fake, + cache_seqs_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + +@functools.cache +def _compile_dspark_rmsnorm_rope_draft_block(block_size: int, eps: float, device: int): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "DSpark RMSNorm/RoPE draft-block must be warmed up before CUDA graph capture" + ) + rows = cute.sym_int() + batch = cute.sym_int() + x_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (rows, _DSV4_DSPARK_HEAD_DIM), stride_order=(1, 0) + ) + weight_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, (_DSV4_DSPARK_HEAD_DIM,), stride_order=(0,) + ) + freqs_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + (rows, _DSV4_DSPARK_ROPE_DIM // 2, 2), + stride_order=(2, 1, 0), + ) + output_fake = cute.runtime.make_fake_compact_tensor( + cutlass.BFloat16, + (batch, _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM), + stride_order=(2, 1, 0), + ) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = DSparkRMSNormRoPEDraftBlockKernel( + _DSV4_DSPARK_HEAD_DIM, + _DSV4_DSPARK_ROPE_DIM, + eps, + block_size, + _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, + ) + return cute.compile( + kernel, + x_fake, + weight_fake, + freqs_fake, + output_fake, + stream_fake, + options="--opt-level 2 --enable-tvm-ffi", + ) + + @torch.library.custom_op( "trtllm::cute_dsl_dspark_rmsnorm_rope", mutates_args=(), @@ -357,6 +537,93 @@ def _( return None +@torch.library.custom_op( + "trtllm::cute_dsl_dspark_rmsnorm_rope_cache_write", + mutates_args=("kv_cache",), + device_types="cuda", +) +def cute_dsl_dspark_rmsnorm_rope_cache_write( + x: torch.Tensor, + weight: torch.Tensor, + freqs: torch.Tensor, + kv_cache: torch.Tensor, + slots: torch.Tensor, + start_pos: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Prepare attention metadata after the caller validates the fused contract.""" + compiled = _compile_dspark_rmsnorm_rope_cache_write(eps) + slots_i32 = torch.empty_like(slots, dtype=torch.int32) + cache_seqs = torch.empty_like(start_pos, dtype=torch.int32) + compiled( + x.view(-1, _DSV4_DSPARK_HEAD_DIM), + weight, + freqs, + kv_cache, + slots, + start_pos, + slots_i32, + cache_seqs, + ) + return slots_i32, cache_seqs + + +@torch.library.register_fake("trtllm::cute_dsl_dspark_rmsnorm_rope_cache_write") +def _( + x: torch.Tensor, + weight: torch.Tensor, + freqs: torch.Tensor, + kv_cache: torch.Tensor, + slots: torch.Tensor, + start_pos: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del x, weight, freqs, kv_cache, eps + return ( + torch.empty_like(slots, dtype=torch.int32), + torch.empty_like(start_pos, dtype=torch.int32), + ) + + +@torch.library.custom_op( + "trtllm::cute_dsl_dspark_rmsnorm_rope_draft_block", + mutates_args=(), + device_types="cuda", +) +def cute_dsl_dspark_rmsnorm_rope_draft_block( + x: torch.Tensor, + weight: torch.Tensor, + freqs: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Prepare a zero-padded draft block after validating the fused contract.""" + _validate_dspark_preparation(x, weight, freqs) + block_size = x.shape[1] + if block_size not in (5, 6): + raise ValueError("DSpark draft block size must be 5 or 6") + with torch.cuda.device(x.device): + compiled = _compile_dspark_rmsnorm_rope_draft_block(block_size, eps, x.device.index) + output = x.new_empty((x.shape[0], _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM)) + compiled( + x.view(-1, _DSV4_DSPARK_HEAD_DIM), + weight, + freqs, + output, + ) + return output + + +@torch.library.register_fake("trtllm::cute_dsl_dspark_rmsnorm_rope_draft_block") +def _( + x: torch.Tensor, + weight: torch.Tensor, + freqs: torch.Tensor, + eps: float, +) -> torch.Tensor: + del weight, freqs, eps + return x.new_empty((x.shape[0], _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM)) + + def _validate_dspark_preparation( x: torch.Tensor, weight: torch.Tensor, freqs: torch.Tensor ) -> None: @@ -503,91 +770,6 @@ def _( return None -@functools.cache -def _compile_dspark_rmsnorm_rope_draft_block(block_size: int, eps: float, device: int): - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "DSpark RMSNorm/RoPE draft-block must be warmed up before CUDA graph capture" - ) - rows = cute.sym_int() - batch = cute.sym_int() - x_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (rows, _DSV4_DSPARK_HEAD_DIM), stride_order=(1, 0) - ) - weight_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, (_DSV4_DSPARK_HEAD_DIM,), stride_order=(0,) - ) - freqs_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, - (rows, _DSV4_DSPARK_ROPE_DIM // 2, 2), - stride_order=(2, 1, 0), - ) - output_fake = cute.runtime.make_fake_compact_tensor( - cutlass.BFloat16, - (batch, _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM), - stride_order=(2, 1, 0), - ) - stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - kernel = DSparkRMSNormRoPEDraftBlockKernel( - _DSV4_DSPARK_HEAD_DIM, - _DSV4_DSPARK_ROPE_DIM, - eps, - block_size, - _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, - ) - return cute.compile( - kernel, - x_fake, - weight_fake, - freqs_fake, - output_fake, - stream_fake, - options="--opt-level 2 --enable-tvm-ffi", - ) - - -@torch.library.custom_op( - "trtllm::cute_dsl_dspark_rmsnorm_rope_draft_block", - mutates_args=(), - device_types="cuda", -) -def cute_dsl_dspark_rmsnorm_rope_draft_block( - x: torch.Tensor, - weight: torch.Tensor, - freqs: torch.Tensor, - eps: float, -) -> torch.Tensor: - """Prepare a zero-padded draft block after validating the fused contract.""" - _validate_dspark_preparation(x, weight, freqs) - block_size = x.shape[1] - if block_size not in (5, 6): - raise ValueError("DSpark draft block size must be 5 or 6") - with torch.cuda.device(x.device): - compiled = _compile_dspark_rmsnorm_rope_draft_block(block_size, eps, x.device.index) - output = x.new_empty((x.shape[0], _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM)) - compiled( - x.view(-1, _DSV4_DSPARK_HEAD_DIM), - weight, - freqs, - output, - ) - return output - - -@torch.library.register_fake("trtllm::cute_dsl_dspark_rmsnorm_rope_draft_block") -def _( - x: torch.Tensor, - weight: torch.Tensor, - freqs: torch.Tensor, - eps: float, -) -> torch.Tensor: - del weight, freqs, eps - return x.new_empty((x.shape[0], _DSV4_DSPARK_DRAFT_BLOCK_STORAGE_SIZE, _DSV4_DSPARK_HEAD_DIM)) - - -_dspark_fp8_input_compile_allowed = False - - def _r128c4_scale_numel(rows: int, hidden_dim: int) -> int: return ((rows + 127) // 128 * 128) * (hidden_dim // 32) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index bcd58da099c1..87b1daca257c 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -630,7 +630,7 @@ def block_bytes_per_group(page_table: KVCachePageTable) -> list[int | None]: only once. Non-attention groups retain a ``None`` placeholder so the result remains aligned with receive-request layer-group indices. """ - from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup + from tensorrt_llm._torch.disaggregation.resource.page import CacheKind from tensorrt_llm._torch.disaggregation.resource.utils import ( get_physical_pool, get_pool_view_num_layers, diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index 58689e5d8a74..8cace14f89d6 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -28,10 +28,14 @@ from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.page import ( AttentionLayerGroup, + CacheKind, KVCachePageTable, MapperKind, ) -from tensorrt_llm._torch.disaggregation.resource.utils import get_layer_to_layer_group +from tensorrt_llm._torch.disaggregation.resource.utils import ( + find_replicated_role_mismatch, + get_layer_to_layer_group, +) from tensorrt_llm._utils import nvtx_range @@ -562,10 +566,21 @@ def _tpb_check(self, local: int, peer: int, peer_ri: RankInfo) -> bool: @staticmethod def validate_peer_compatible( - self_page_table: KVCachePageTable | None, - peer_page_table: KVCachePageTable | None, + self_page_table: Optional[KVCachePageTable], + peer_page_table: Optional[KVCachePageTable], ) -> None: - """Validate standalone draft storage on overlapping layers.""" + """Reject a peer whose PAGED groups declare different replicated roles. + + Pool matching drops a view with no counterpart silently, so an + index-key side cache the peer never declares would leave the receiver + holding zeroed state instead of raising. Raises ``ValueError``. + """ + differing = find_replicated_role_mismatch(self_page_table, peer_page_table, CacheKind.PAGED) + if differing: + raise ValueError( + "AttentionPolicy.validate_peer_compatible: replicated roles differ on " + f"overlapping layers: {differing}" + ) if self_page_table is not None and peer_page_table is not None: local_layers = get_layer_to_layer_group(self_page_table) peer_layers = get_layer_to_layer_group(peer_page_table) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 647099dc6fab..035a5314ed8f 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -15,7 +15,7 @@ from collections import Counter from dataclasses import dataclass, field, replace -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple, Union import numpy as np @@ -26,7 +26,12 @@ from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, MapperKind, PoolView +from tensorrt_llm._torch.disaggregation.resource.page import ( + AttentionLayerGroup, + CacheKind, + MapperKind, + PoolView, +) from tensorrt_llm._torch.disaggregation.resource.utils import ( get_layer_byte_ranges, get_layer_to_layer_group, @@ -184,6 +189,29 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) + if self._ri.num_draft_layers != peer_ri.num_draft_layers: + # The layer-count tolerance below would transfer only the target + # layers and leave the drafter's prompt KV unset on the receiver. + raise ValueError( + "PeerRegistrar: unified draft KV layer count differs " + f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " + "both peers must use the same speculative decoding KV cache layout; " + "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." + ) + if self._ri.num_draft_layers: + if self._ri.tp_size != peer_ri.tp_size: + raise ValueError( + "PeerRegistrar: a shared draft KV cache needs equal tp_size " + f"(local={self._ri.tp_size}, peer={peer_ri.tp_size}); " + "unset TRTLLM_DISAGG_SHARED_DRAFT_KV or use the same TP on both workers." + ) + if self._ri.attention.enable_attention_dp != peer_ri.attention.enable_attention_dp: + raise ValueError( + "PeerRegistrar: a shared draft KV cache needs the same attention DP " + f"setting (local={self._ri.attention.enable_attention_dp}, " + f"peer={peer_ri.attention.enable_attention_dp}); unset " + "TRTLLM_DISAGG_SHARED_DRAFT_KV or match it on both workers." + ) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) @@ -404,12 +432,34 @@ def get_kv_map( pool_idx=peer_pi, ) - policy = self._attention_policy + # Polymorphic dispatch: attention vs mamba policy + policy = self._get_policy(self_lg.kind) + + # For mamba, compute ptr-array layer offsets for partial PP overlap. + # extract_slot returns ptrs for ALL local_layer_ids in sorted order; + # the mapper must slice only the overlapping subset. + extra_kwargs = {} + if self_lg.kind == CacheKind.STATE: + # Positions index the ptr array ``extract_slot`` builds, which is + # ordered by view offset. + self_gid_to_pos = {gid: i for i, gid in enumerate(self_global_ids)} + peer_gid_to_pos = {gid: i for i, gid in enumerate(peer_global_ids)} + self_overlap_positions = [self_gid_to_pos[gid] for gid in overlapping_layers] + peer_overlap_positions = [peer_gid_to_pos[gid] for gid in overlapping_layers] + # Under contiguous PP partitioning, overlapping layers form a + # contiguous block in the ptrs array. + extra_kwargs["src_layer_off"] = ( + self_overlap_positions[0] if self_overlap_positions else 0 + ) + extra_kwargs["dst_layer_off"] = ( + peer_overlap_positions[0] if peer_overlap_positions else 0 + ) + mapper_peer_ri = peer_ri - if self_lg.cp_as_tp: + if getattr(self_lg, "cp_as_tp", False): policy = AttentionPolicy(self._group_rank_info(self._ri, self_lg)) mapper_peer_ri = self._group_rank_info(peer_ri, peer_lg) - mapper = policy.build_kv_mapper( + mapper = policy.build_mapper( peer_ri=mapper_peer_ri, mapper_kind=self_pv.mapper_kind, self_layer_offsets=self_layer_offsets, @@ -573,7 +623,7 @@ def should_send_pool( """ layer_group = self._self_ext_cache.page_table.layer_groups[layer_group_id] pool_view = layer_group.pool_views[pool_idx] - if layer_group.cp_as_tp: + if getattr(layer_group, "cp_as_tp", False): peer_lg_id, _ = self.get_pool_mapping(peer_rank_info)[(layer_group_id, pool_idx)] peer_lg = peer_rank_info.page_table.layer_groups[peer_lg_id] local = self._group_rank_info(self._ri, layer_group) @@ -585,9 +635,10 @@ def should_send_pool( return False duplicate, _ = AttentionPolicy(local).duplicate_head_factors(peer) return duplicate <= 1 or peer.dp_rank % duplicate == local.tp_rank % duplicate - if pool_view.mapper_kind == MapperKind.REPLICATED: - return self._owns_tp_fan_in(peer_rank_info) - return self.should_send_kv(peer_overlap, peer_rank_info) + # Delegate to the policy's should_send. None means fan-in election. + policy = self._get_policy(layer_group.kind) + result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) + return self._owns_tp_fan_in(peer_rank_info) if result is None else result @staticmethod def _group_rank_info(rank_info: RankInfo, layer_group: AttentionLayerGroup) -> RankInfo: diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 87731d303204..358b1b576412 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -49,6 +49,9 @@ class RankInfo: aux_meta: Optional[AuxBufferMeta] = None page_table: Optional[KVCachePageTable] = None bounded_replay_on_generation: bool = False + # Layers a unified KV cache manager owns for the drafter (layer specs with + # a DRAFT owner). They ride the target's pages, so both peers must agree. + num_draft_layers: int = 0 @property def tp_size_per_dp_group(self) -> int: @@ -114,6 +117,9 @@ def from_kv_cache_manager( aux_meta=aux_buffer_meta, page_table=build_page_table_from_manager(kvm), bounded_replay_on_generation=getattr(kvm, "bounded_replay_on_generation", False), + num_draft_layers=( + len(kvm.draft_layer_ids) if getattr(kvm, "layer_specs", None) is not None else 0 + ), ) @classmethod diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 807c2417a19c..aefe805f2faa 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -15,16 +15,17 @@ from __future__ import annotations +import math import os import queue import struct import threading import time import weakref -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, List, Optional, Union +from typing import TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, Optional, Union import msgpack import numpy as np @@ -45,9 +46,7 @@ TransferOp, TransferRequest, ) -from tensorrt_llm._torch.disaggregation.base.region import IndexRange, KVRegionSpec from tensorrt_llm._torch.disaggregation.base.transfer import ( - KVSlice, ReceiverBase, RxSessionBase, SenderBase, @@ -56,27 +55,49 @@ TxSessionBase, WaitResult, ) +from tensorrt_llm._torch.disaggregation.lifecycle.retirement import ( + QuiescenceFatalEvent, + RetirementDeadline, + RetirementWatchdog, +) from tensorrt_llm._torch.disaggregation.native.auxiliary import ( AuxBuffer, build_aux_transfer_layout, get_non_empty_aux_indices, ) +from tensorrt_llm._torch.disaggregation.native.extent import Chunk from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message -from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy +from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( + MambaPolicy, + mamba_receiver_payload_bytes, +) +from tensorrt_llm._torch.disaggregation.native.ownership import ( + PhysicalOperationState as _PhysicalOperationState, # noqa: F401 -- native diagnostic alias +) +from tensorrt_llm._torch.disaggregation.native.ownership import ( + ReceiveOperationOwner as _ReceiveOperationOwner, +) +from tensorrt_llm._torch.disaggregation.native.ownership import SendOperationOwner +from tensorrt_llm._torch.disaggregation.native.ownership import ( + TransferNotSubmittedError as _TransferNotSubmittedError, +) from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable, MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import CUASSERT, nvtx_range from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle if TYPE_CHECKING: + from tensorrt_llm._torch.disaggregation.base.agent import TransferStatus + from .bounce import Config AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType @@ -94,65 +115,20 @@ _FALLBACK_TX_OVERALL_TIMEOUT_S = 60.0 -class _AgentOperationGate: - """Reader/writer gate for NIXL transfers and remote-metadata updates.""" - - def __init__(self): - self._condition = threading.Condition() - self._active_transfers = 0 - self._metadata_active = False - self._metadata_waiters = 0 - - @contextmanager - def transfer(self): - with self._condition: - # Prefer a waiting metadata update so a continuous transfer stream - # cannot starve first-use peer registration. - while self._metadata_active or self._metadata_waiters > 0: - self._condition.wait() - self._active_transfers += 1 - try: - yield - finally: - with self._condition: - self._active_transfers -= 1 - if self._active_transfers == 0: - self._condition.notify_all() - - @contextmanager - def metadata(self): - with self._condition: - self._metadata_waiters += 1 - try: - while self._metadata_active or self._active_transfers > 0: - self._condition.wait() - self._metadata_active = True - finally: - self._metadata_waiters -= 1 - try: - yield - finally: - with self._condition: - self._metadata_active = False - self._condition.notify_all() - - @dataclass class RecvReqInfo: sender_req_id: int instance_name: str instance_rank: int - block_ids_per_layer_groups: list[ - np.ndarray - ] # Block IDs per layer group, each np.ndarray(dtype=np.int64) + # Positional pool slots per layer group (see Chunk): entry i is the + # receiver's slot for block ordinal i, or -1 where it accepts nothing. + block_ids_per_layer_groups: list[np.ndarray] unique_rid: int - # Block-aligned token offset where the receiver's block list starts. - # None means "end-of-range suffix" — sender derives it from len(blocks). + # Optional extra lower bound (block-aligned token offset) below which the + # receiver accepts nothing, on top of the -1 holes already in its table. dst_start_token: Optional[int] = None aux_slot: Optional[int] = None - mamba_state_index: Optional[int] = None slice_id: Optional[int] = None - need_aux: bool = False bounce_dst_base: Optional[int] = None def to_bytes(self) -> bytes: @@ -167,9 +143,7 @@ def to_bytes(self) -> bytes: "unique_rid": self.unique_rid, "dst_start_token": self.dst_start_token, "aux_slot": self.aux_slot, - "mamba_state_index": self.mamba_state_index, "slice_id": self.slice_id, - "need_aux": self.need_aux, "bounce_dst_base": self.bounce_dst_base, } ) @@ -208,12 +182,11 @@ class WriteMeta: sizes: np.ndarray # dtype=np.int64 dst_device_id: Optional[int] = None slice_id: Optional[int] = None + # Peer task index; slice_id identifies the local sender task. + receiver_slice_id: int = 0 is_last_slice: bool = False meta_type: WriteMetaType = WriteMetaType.KV bounce_dst_base: Optional[int] = None - # Set by Sender._enqueue while the work item owns the Tx session's - # queued-or-active lifetime reservation. Worker completion releases it. - owner_session: Optional["TxSession"] = None class MessageType: @@ -224,6 +197,7 @@ class MessageType: REGISTER_RANK_INFO = b"REGISTER_RANK_INFO" AUX_AGENT_RESULT = b"AUX_AGENT_RESULT" CANCEL_SESSION = b"CANCEL_SESSION" + SESSION_QUIESCED = b"SESSION_QUIESCED" class PeerIncompatibleError(ValueError): @@ -243,21 +217,159 @@ class TaskStatus(Enum): ERROR = "ERROR" +@dataclass(frozen=True) +class _LogicalOutcome: + """A committed delivery decision, independent of physical/report progress.""" + + status: SessionStatus + reason: str = "" + by_peer: bool = False + + +class _LogicalOutcomes: + """Arbitrate task completion against session failure/cancellation at event time. + + A failed or cancelled session ends only pending tasks. Already committed + results, including their original cause, never change. No task, request or + backend handle is retained here: memory ownership remains a separate concern. + """ + + def __init__(self, retirement: Optional[RetirementDeadline] = None) -> None: + self._lock = threading.Lock() if retirement is None else retirement.lock + self._retirement = retirement + self._terminal: Optional[_LogicalOutcome] = None + self._results: list[Optional[_LogicalOutcome]] = [] + self._sealed = False + if retirement is not None: + retirement.bind_timeout_outcome(self._timeout_locked) + + def _timeout_locked(self) -> None: + """Commit timeout metadata without entering session, backend or CUDA code.""" + self._end_locked(_LogicalOutcome(SessionStatus.ERROR, reason="KV transfer timed out")) + + def _end_locked(self, outcome: _LogicalOutcome) -> None: + """End pending results once, preserving every previously committed outcome.""" + if self._terminal is None: + self._terminal = outcome + for index, result in enumerate(self._results): + if result is None: + self._results[index] = outcome + + @property + def terminal(self) -> Optional[_LogicalOutcome]: + """Return the stable terminal decision, including an elapsed request timeout.""" + with self._lock: + if self._retirement is not None: + self._retirement.check() + return self._terminal + + def add_task(self) -> int: + with self._lock: + if self._sealed: + raise RuntimeError("transfer session already registered its final piece") + index = len(self._results) + self._results.append(self._terminal) + return index + + def seal(self) -> None: + """Declare that every expected KV and AUX result has been registered.""" + with self._lock: + self._sealed = True + self._complete_session_locked() + + def _complete_session_locked(self) -> None: + if ( + self._retirement is not None + and self._sealed + and self._results + and all( + result is not None and result.status is SessionStatus.TRANSFERRED + for result in self._results + ) + ): + self._retirement.complete_pieces() + + def get(self, index: int) -> Optional[_LogicalOutcome]: + with self._lock: + if self._retirement is not None: + self._retirement.check() + return self._results[index] + + def complete(self, index: int) -> None: + with self._lock: + if self._retirement is not None: + self._retirement.check() + if self._results[index] is None: + self._results[index] = _LogicalOutcome(SessionStatus.TRANSFERRED) + self._complete_session_locked() + + def fail(self, error: Exception) -> None: + self._end(_LogicalOutcome(SessionStatus.ERROR, reason=str(error))) + + def cancel(self, by_peer: bool) -> None: + self._end(_LogicalOutcome(SessionStatus.CANCELLED, by_peer=by_peer)) + + def _end(self, outcome: _LogicalOutcome) -> None: + with self._lock: + if self._retirement is not None: + self._retirement.check() + if self._terminal is not None: + return + self._end_locked(outcome) + if self._retirement is not None: + self._retirement.request_drain(outcome.reason or outcome.status.value) + + +class _LogicalTask: + def __init__(self) -> None: + self._logical_outcomes = _LogicalOutcomes() + self._logical_index = self._logical_outcomes.add_task() + self._retirement: Optional[RetirementDeadline] = None + + def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: + """Join the session before this task is exposed to a worker or caller.""" + self._logical_outcomes = outcomes + self._logical_index = outcomes.add_task() + self._retirement = outcomes._retirement + + @property + def logical_outcome(self) -> Optional[_LogicalOutcome]: + return self._logical_outcomes.get(self._logical_index) + + class AgentResult(Enum): SUCCESS = "SUCCESS" FAILED = "FAILED" + IN_DOUBT = "IN_DOUBT" + FAILED_QUIESCED = "FAILED_QUIESCED" # KV_AGENT_RESULT prefix in one struct frame (was ascii frames serialized/parsed under the # GIL per slice per writer): instance_rank, unique_rid, slice_id, is_last, status, # transfer_size. The optional bounce tail follows at message[2:]. _KV_RESULT_PREFIX = struct.Struct(" None: - with self._state_lock: - if self._started: - return - if self._stop.is_set(): - raise RuntimeError(f"{self._owner_name} message outbox is stopped") - self._started = True - self._thread.start() - - def submit(self, endpoint: str, message: list[bytes]) -> bool: - if not endpoint: - raise ValueError(f"{self._owner_name}: peer endpoint is empty") - with self._state_lock: - if self._stop.is_set(): - logger.error( - f"Cannot queue reliable message for endpoint {endpoint}: " - f"{self._owner_name} outbox is shutting down" - ) - return False - # Copy the multipart list so callers may safely reuse their object. - self._queue.put(_QueuedMessage(endpoint, list(message))) - return True - - def _get_or_connect_dealer(self, endpoint: str) -> ZMQMessenger: - if endpoint not in self._dealers: - self._dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) - return self._dealers[endpoint] - - def _drop_dealer(self, endpoint: str) -> None: - dealer = self._dealers.pop(endpoint, None) - if dealer is None: - return - try: - dealer.stop() - except Exception as exc: - logger.warning( - f"{self._owner_name} outbox failed to close dealer for endpoint {endpoint}: {exc}" - ) - - def _try_send(self, pending: _QueuedMessage) -> bool: - attempts = self._IMMEDIATE_ATTEMPTS if pending.initial_attempt else 1 - for attempt in range(1, attempts + 1): - try: - self._get_or_connect_dealer(pending.endpoint).send(pending.message) - return True - except Exception as exc: - self._drop_dealer(pending.endpoint) - logger.warning( - f"{self._owner_name} reliable send failed for endpoint " - f"{pending.endpoint} (attempt {attempt}/{attempts}): {exc}" - ) - return False - - def process_once(self) -> bool: - """Process one queued item; return False for the shutdown sentinel.""" - pending = self._queue.get() - if pending is None: - return False - if not self._try_send(pending): - pending.initial_attempt = False - if not self._stop.is_set(): - self._queue.put(pending) - self._stop.wait(self._RETRY_INTERVAL_S) - return True +def _make_aux_result_msg(instance_rank: int, unique_rid: int, result: AgentResult) -> list[bytes]: + return [ + MessageType.AUX_AGENT_RESULT, + str(instance_rank).encode("ascii"), + str(unique_rid).encode("ascii"), + result.value.encode("ascii"), + ] - def _run(self) -> None: - try: - while self.process_once(): - pass - finally: - for endpoint in list(self._dealers): - self._drop_dealer(endpoint) - def shutdown(self, timeout: float = 5.0) -> None: - with self._state_lock: - if self._stop.is_set(): - return - self._stop.set() - if not self._started: - return - # FIFO ordering gives every already-queued message one final send - # attempt before the sentinel terminates the worker. - self._queue.put(None) - self._thread.join(timeout=timeout) - if self._thread.is_alive(): - logger.warning( - f"{self._owner_name} reliable-message outbox did not terminate within {timeout:g}s" - ) +@dataclass +class _PendingSettlement: + write_meta: WriteMeta + initial_report: Optional[list[bytes]] + send_slot_id: Optional[int] = None + backend_done: bool = False + report_error_logged: bool = False -class SendTaskBase: +class SendTaskBase(_LogicalTask): def __init__(self, params: DisaggregatedParams): + super().__init__() self.status = TaskStatus.INIT self._event = threading.Event() self._exception: Optional[Exception] = None @@ -407,13 +418,27 @@ def __init__(self, params: DisaggregatedParams): self._params = params self._unique_rid: Optional[int] = params.disagg_request_id self._perf_timer = PerfTimer() if perf_log_manager.enabled else None + self._physical_owner = SendOperationOwner() + # Preserve the task's diagnostic view of its owner's exact roots. + self._physical_operations = self._physical_owner._physical_operations + + def bind_logical_outcomes(self, outcomes: _LogicalOutcomes) -> None: + """Bind session outcome and deadline before exposing the task to workers. + + Args: + outcomes: The session's stable logical-result arbiter. + """ + super().bind_logical_outcomes(outcomes) + self._physical_owner.bind_retirement(self._retirement) def fail(self, exc: Exception) -> None: + self._logical_outcomes.fail(exc) self._exception = exc self.status = TaskStatus.ERROR self._event.set() def complete(self) -> None: + self._logical_outcomes.complete(self._logical_index) self.status = TaskStatus.TRANSFERRED self._event.set() @@ -425,6 +450,43 @@ def wait(self, timeout: Optional[float] = None) -> bool: def is_done(self) -> bool: return self._event.is_set() + def begin_physical_operation(self, peer_rank: int) -> bool: + """Admit a participant once in the native physical owner.""" + return self._physical_owner.begin_physical_operation(peer_rank) + + def begin_backend_submission(self, peer_rank: int, request: TransferRequest) -> None: + """Root the request before the native backend submission escapes.""" + self._physical_owner.begin_backend_submission(peer_rank, request) + + def record_backend_submission(self, peer_rank: int, status: TransferStatus) -> None: + """Retain the exact status returned by the native backend.""" + self._physical_owner.record_backend_submission(peer_rank, status) + + def mark_physical_operation_in_doubt(self, peer_rank: int) -> None: + """Retain ambiguous access independently of the logical outcome.""" + self._physical_owner.mark_physical_operation_in_doubt(peer_rank) + + def retire_unsubmitted_physical_operation(self, peer_rank: int) -> None: + """Record that a participant never submitted a backend operation.""" + self._physical_owner.retire_unsubmitted_physical_operation(peer_rank) + + def retire_backend_done_physical_operation(self, peer_rank: int) -> bool: + """Settle backend-defined completion through the common deadline arbiter.""" + return self._physical_owner.retire_backend_done_physical_operation(peer_rank) + + def poll_in_doubt_physical_operation(self, peer_rank: int) -> bool: + """Query and settle only the exact retained ambiguous operation.""" + return self._physical_owner.poll_in_doubt_physical_operation(peer_rank) + + def has_started_physical_operation(self, peer_rank: int) -> bool: + """Whether this task has admitted the participant.""" + return self._physical_owner.has_started_physical_operation(peer_rank) + + @property + def resources_drained(self) -> bool: + """Whether operation evidence and the session deadline permit retirement.""" + return self._physical_owner.resources_drained + def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int): if self._perf_timer is None: return @@ -448,18 +510,35 @@ def __init__(self, params: DisaggregatedParams, slot: Optional[int]): class KVSendTask(SendTaskBase): def __init__( self, - kv_slice: KVSlice, + chunk: Chunk, params: DisaggregatedParams, slice_id: int, prompt_len: Optional[int] = None, - beam_width: int = 1, ): super().__init__(params) self.slice_id = slice_id self.transferred_count = 0 - self._slice = kv_slice + self.expected_transfers = 0 + self._chunk = chunk self._prompt_len = prompt_len - self._beam_width = beam_width + + @property + def reports_outstanding(self) -> bool: + """Whether any of this piece's writes has still to finish. + + One peer failing ends the task while the writes to its siblings run on, so the task's own + status cannot answer this. Zero expected means nothing was ever submitted. + """ + return self.transferred_count < self.expected_transfers + + +@dataclass(frozen=True) +class _SessionQuiescence: + """A FIFO marker for one cancelled request's no-future-access acknowledgment.""" + + unique_rid: int + peer_rank: int + endpoint: Optional[str] class Sender(SenderBase): @@ -468,20 +547,21 @@ class Sender(SenderBase): # RecvReqInfo that never gets consumed. Entries older than this # are evicted during periodic sweeps. _STALE_REQ_INFO_TTL_S = 120.0 - # Bound synchronous worker retries before handing delivery to the shared - # reliable outbox. - _AGENT_RESULT_IMMEDIATE_ATTEMPTS = 2 def __init__( self, peer_registrar: PeerRegistrar, agent: BaseTransferAgent, bounce=None, - ): + enforce_physical_ownership: bool = False, + retirement_watchdog: Optional[RetirementWatchdog] = None, + ) -> None: self._registrar = peer_registrar self._device_id = peer_registrar.self_rank_info.device_id self._agent = agent self._bounce = bounce + self._enforce_physical_ownership = enforce_physical_ownership + self._retirement_watchdog = retirement_watchdog self._peer_requests: dict = {} self._peer_requests_timestamps: dict[int, float] = {} # unique_rid -> insert time self._peer_requests_lock = threading.Lock() @@ -489,31 +569,33 @@ def __init__( self._dealers = {} # used by listener thread only (single-threaded path) self._thread_local = threading.local() # per-thread DEALER cache for worker threads self._sessions = {} # unique_rid -> TxSession - self._sessions_lock = threading.Lock() # Protects sessions and cancel tombstones - self._pre_cancelled_rids: set[int] = set() - self._pre_cancelled_timestamps: dict[int, float] = {} - # A pre-cancelled sender may receive duplicate/late REQUEST_DATA. - # Remember the terminal replies already emitted so each receiver-side - # peer reservation is resolved at most once by this path. - self._cancel_failed_result_keys: set[tuple[int, str, int, int, bool]] = set() + self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids + # unique_rid -> whether the peer asked, for cancels that arrive before the session exists. + # Who asked is not recoverable from the id alone, and the two differ: a peer's cancel is a + # transfer error, our own side's is an ordinary end. + # Ownership mode also fences existing sessions: keep one entry per + # cancelled immutable ID until Sender teardown, not merely metadata TTL. + self._pre_cancelled_rids: dict[int, bool] = {} + self._pending_session_quiescence: dict[tuple[int, int], _SessionQuiescence] = {} self._shutdown = False + self._shutdown_requested = False self._instance_rank = self._registrar.self_rank_info.instance_rank # Guards concurrent add() from the listener thread. self._loaded_remote_agents: set[str] = set() self._loaded_remote_agents_lock = threading.Lock() - # NIXL remote-metadata mutation is not safe while getXferStatus() is - # polling an active request. The C++ wrapper's shared lock covers - # submit_transfer_requests(), but the returned status outlives that - # lock. Serialize lazy peer loads with the complete submit+wait - # interval so a first request from another peer cannot invalidate an - # unrelated transfer already in flight. - self._agent_operation_gate = _AgentOperationGate() + self._ownership_poisoned: Optional[Exception] = None + self._ownership_poison_lock = ( + threading.Lock() if retirement_watchdog is None else retirement_watchdog.lock + ) self._num_threads = KV_TRANSFER_NUM_THREADS self._send_task_queues: List[queue.Queue] = [ queue.Queue() for _ in range(self._num_threads) ] - # Agent results and CANCEL_SESSION share one reliable, fair outbox. - self._reliable_message_outbox = _ReliableMessageOutbox("Sender") + # Each dictionary belongs to its worker, preserving that peer stream's + # result ordering and ZMQ socket affinity even after the first failure. + self._pending_settlements: list[dict[tuple[SendTaskBase, int], _PendingSettlement]] = [ + {} for _ in range(self._num_threads) + ] self._worker_threads: List[threading.Thread] = [ threading.Thread(target=self._process_task_queue, args=(i,), daemon=True) for i in range(self._num_threads) @@ -531,25 +613,6 @@ def __init__( def endpoint(self): return self._messenger.endpoint - def _load_remote_agent_serialized(self, name: str, agent_desc: bytes) -> None: - with self._agent_operation_gate.metadata(): - self._agent.load_remote_agent(name, agent_desc) - - def _submit_and_wait_agent_transfer(self, request: TransferRequest): - with self._agent_operation_gate.transfer(): - status = self._agent.submit_transfer_requests(request) - try: - completed = status.wait() - last_status = None - if not completed: - last_status = getattr(status, "last_status_str", lambda: "")() - finally: - # The pybind status destructor calls releaseXferReq(). Drop - # the final local reference before releasing the reader gate, - # so a lazy remote-metadata load cannot race that NIXL call. - del status - return completed, last_status - def _add_req_info(self, unique_rid: int, instance_rank: int, req_info: RecvReqInfo): with self._peer_requests_lock: if unique_rid not in self._peer_requests: @@ -581,48 +644,69 @@ def _remove_req_info(self, unique_rid: int): self._peer_requests_timestamps.pop(unique_rid, None) def sweep_stale_req_infos(self): - """Evict stale request infos and cancel tombstones with no TxSession. + """Evict RecvReqInfo entries that have no matching TxSession and exceed the TTL. Called opportunistically from the listener thread when a new REQUEST_DATA arrives. With gen-first ADP broadcast, non-assigned DP ranks accumulate entries that are never consumed; this sweep prevents unbounded growth. """ now = time.monotonic() - with self._sessions_lock, self._peer_requests_lock: - stale_req_rids = [ + with self._peer_requests_lock: + stale_rids = [ rid for rid, ts in self._peer_requests_timestamps.items() if now - ts > self._STALE_REQ_INFO_TTL_S ] - stale_cancel_rids = [ - rid - for rid, ts in self._pre_cancelled_timestamps.items() - if now - ts > self._STALE_REQ_INFO_TTL_S - ] - for rid in stale_req_rids: + if not stale_rids: + return + for rid in stale_rids: + with self._sessions_lock, self._peer_requests_lock: if rid not in self._sessions and rid in self._peer_requests: self._peer_requests.pop(rid, None) self._peer_requests_timestamps.pop(rid, None) logger.debug(f"Swept stale RecvReqInfo for rid={rid}") - for rid in stale_cancel_rids: - if rid in self._sessions: - continue - self._pre_cancelled_rids.discard(rid) - self._pre_cancelled_timestamps.pop(rid, None) - self._cancel_failed_result_keys = { - key for key in self._cancel_failed_result_keys if key[0] != rid - } - logger.debug(f"Swept stale pre-cancel tombstone for rid={rid}") def setup_session(self, tx_session: "TxSession"): unique_rid = tx_session.disagg_request_id pre_cancel = False + cancelled_by_peer = False with self._sessions_lock: - self._sessions[unique_rid] = weakref.ref(tx_session) + if self._shutdown_requested: + raise RuntimeError("cannot create a TxSession after Sender shutdown") + if self._enforce_physical_ownership and unique_rid in self._sessions: + raise RuntimeError("cannot replace an owned TxSession") if unique_rid in self._pre_cancelled_rids: pre_cancel = True + cancelled_by_peer = self._pre_cancelled_rids[unique_rid] + if not self._enforce_physical_ownership: + self._pre_cancelled_rids.pop(unique_rid) + if not (pre_cancel and self._enforce_physical_ownership): + self._sessions[unique_rid] = ( + tx_session if self._enforce_physical_ownership else weakref.ref(tx_session) + ) if pre_cancel: - tx_session.cancel() + if self._enforce_physical_ownership: + tx_session.cancel_local(by_peer=cancelled_by_peer, report_unsubmitted_aux=False) + # Make the cancelled session visible only after its terminal + # state and the previously saved requests are captured. Late + # creation cannot reopen admission while the listener + # acknowledges the generation-first cancellation. + with self._sessions_lock: + if self._shutdown_requested: + raise RuntimeError("cannot create a TxSession after Sender shutdown") + if unique_rid in self._sessions: + raise RuntimeError("cannot replace an owned TxSession") + # Keep the fence for this immutable request ID after + # session cleanup, so delayed creation cannot reopen it. + with self._peer_requests_lock: + req_infos = list(self._peer_requests.get(unique_rid, {}).values()) + self._sessions[unique_rid] = tx_session + for info in req_infos: + self._queue_session_quiescence( + unique_rid, info.instance_name, info.instance_rank + ) + else: + tx_session.cancel(by_peer=cancelled_by_peer) return req_info = self._get_first_req_info(unique_rid) @@ -638,37 +722,43 @@ def setup_session(self, tx_session: "TxSession"): return def _get_session(self, unique_rid: Optional[int]) -> Optional["TxSession"]: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: + session_entry = self._sessions.get(unique_rid) + if session_entry is None: return None - session = session_ref() - if session is None: - logger.warning(f"TxSession {unique_rid} has been garbage collected") + if isinstance(session_entry, weakref.ReferenceType): + session = session_entry() + if session is None: + logger.warning(f"TxSession {unique_rid} has been garbage collected") + return session + return session_entry + + def _begin_task_operation(self, task: SendTaskBase, peer_rank: int) -> Optional[bool]: + """Return True for admission, None for a duplicate, or False for rejection.""" + if not task.begin_physical_operation(peer_rank): return None - return session + with self._ownership_poison_lock: + if self._ownership_poisoned is not None: + task.retire_unsubmitted_physical_operation(peer_rank) + return False + with self._sessions_lock: + session = self._get_session(task._unique_rid) + if session is None: + task.retire_unsubmitted_physical_operation(peer_rank) + return False + with session.lock: + if session._closed or session.has_failed() or task.is_done: + task.retire_unsubmitted_physical_operation(peer_rank) + return False + if task.status == TaskStatus.INIT: + task.status = TaskStatus.TRANSFERRING + return True def _enqueue(self, write_meta: WriteMeta): # Route by (unique_rid, peer_rank) so that: # - Same peer's slices stay ordered on one thread (is_last_slice correctness) # - Different peers can run on different threads (better load balancing) thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads - with self._sessions_lock: - session = self._get_session(write_meta.unique_rid) - if session is not None: - with session.lock: - # A terminal/closed session no longer exposes source buffers. - # Still enqueue the metadata so the worker sends FAILED and - # resolves the receiver's already-dispatched peer reservation. - if session._terminal_status is None and not session._closed: - session._begin_nixl_op_locked() - write_meta.owner_session = session - try: - self._send_task_queues[thread_idx].put(write_meta) - except Exception: - if write_meta.owner_session is not None: - write_meta.owner_session._finish_nixl_op() - write_meta.owner_session = None - raise + self._send_task_queues[thread_idx].put(write_meta) def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger: """Get or create a per-thread DEALER socket via threading.local(). @@ -684,52 +774,63 @@ def _get_or_connect_thread_dealer(self, endpoint: Optional[str]) -> ZMQMessenger dealers[endpoint] = ZMQMessenger(mode="DEALER", endpoint=endpoint) return dealers[endpoint] - def _drop_thread_dealer(self, endpoint: str) -> None: - """Discard the current thread's failed socket before reconnecting.""" - dealers = getattr(self._thread_local, "dealers", None) - if not dealers: - return - dealer = dealers.pop(endpoint, None) - if dealer is None: - return - try: - dealer.stop() - except Exception as exc: - logger.warning(f"Failed to close agent-result dealer for endpoint {endpoint}: {exc}") + def _get_result_dealer( + self, endpoint: Optional[str], *, legacy_shared: bool = False + ) -> ZMQMessenger: + """Keep legacy early-failure routing outside the ownership bridge.""" + if legacy_shared and not self._enforce_physical_ownership: + return self._get_or_connect_dealer(endpoint) + return self._get_or_connect_thread_dealer(endpoint) - def _try_send_agent_result( + def _submit_transfer( self, - endpoint: str, - message: list[bytes], - *, - attempts: int, - ) -> bool: - """Try an agent-result send, rebuilding the thread-local socket on error.""" - for attempt in range(1, attempts + 1): - try: - self._get_or_connect_thread_dealer(endpoint).send(message) - return True - except Exception as exc: - self._drop_thread_dealer(endpoint) - logger.warning( - f"Agent-result send failed for endpoint {endpoint} " - f"(attempt {attempt}/{attempts}): {exc}" - ) - return False - - def _send_agent_result(self, endpoint: str, message: list[bytes]) -> bool: - """Send on the worker fast path, queueing only transport failures. - - Delivery is at-least-once. Receiver result processing is idempotent, so - retrying after an ambiguous transport failure is safe. - """ - if self._try_send_agent_result( - endpoint, - message, - attempts=self._AGENT_RESULT_IMMEDIATE_ATTEMPTS, - ): - return True - return self._reliable_message_outbox.submit(endpoint, message) + task: SendTaskBase, + peer_rank: int, + request: TransferRequest, + ) -> tuple[bool, Optional[str]]: + if not self._enforce_physical_ownership: + status = self._agent.submit_transfer_requests(request) + completed = status.wait() + detail = None if completed else getattr(status, "last_status_str", lambda: None)() + return completed, detail + with self._ownership_poison_lock: + if self._ownership_poisoned is not None: + error = self._ownership_poisoned + task.retire_unsubmitted_physical_operation(peer_rank) + raise _TransferNotSubmittedError( + "NIXL sender rejected transfer before backend submission" + ) from error + task.begin_backend_submission(peer_rank, request) + try: + # SUBMITTING already retains the request. Neither backend admission + # nor completion may hold the lock needed by deadline containment. + status = self._agent.submit_transfer_requests(request) + task.record_backend_submission(peer_rank, status) + except Exception as error: + with self._ownership_poison_lock: + self._ownership_poisoned = error + task.mark_physical_operation_in_doubt(peer_rank) + return False, str(error) + try: + if not status.wait(): + # A non-success query is a logical transfer failure, but the + # current binding does not expose a backend-defined proof that + # the request handle has released every memory accessor. Treat + # it as physically unresolved until that contract is available. + detail = getattr(status, "last_status_str", lambda: "")() + error = RuntimeError(f"NIXL transfer outcome is ambiguous: {detail}") + with self._ownership_poison_lock: + self._ownership_poisoned = error + task.mark_physical_operation_in_doubt(peer_rank) + return False, detail + except Exception as error: + with self._ownership_poison_lock: + self._ownership_poisoned = error + task.mark_physical_operation_in_doubt(peer_rank) + return False, str(error) + if not task.retire_backend_done_physical_operation(peer_rank): + return False, "quiescence deadline expired before backend completion" + return True, None def _process_task_queue(self, thread_idx: int): device_id = self._device_id @@ -739,13 +840,42 @@ def _process_task_queue(self, thread_idx: int): task_queue = self._send_task_queues[thread_idx] try: while True: - write_meta = task_queue.get() + if self._pending_settlements[thread_idx]: + self._poll_in_doubt_transfers(thread_idx) + if any( + pending.initial_report is not None + for pending in self._pending_settlements[thread_idx].values() + ): + # A later queued FAILED may describe only an unsubmitted + # chunk. It must never overtake the earlier IN_DOUBT. + time.sleep(0.01) + continue + try: + write_meta = task_queue.get(timeout=0.01) + except queue.Empty: + continue + else: + write_meta = task_queue.get() if write_meta is None: break + if isinstance(write_meta, _SessionQuiescence): + if not self._send_session_quiesced(write_meta): + task_queue.put(write_meta) + time.sleep(0.01) + continue + if isinstance(write_meta, tuple): + endpoint, message = write_meta + try: + self._get_or_connect_thread_dealer(endpoint).send(message) + except Exception as e: + logger.warning(f"failed to send transfer rejection: {e}") + continue try: if write_meta.meta_type == WriteMetaType.AUX: logger.debug( - f"_process_task_queue[{thread_idx}]: delivering aux task to agent: {write_meta}" + "_process_task_queue[%d]: delivering aux task to agent: %s", + thread_idx, + write_meta, ) self._deliver_aux_to_agent(write_meta) else: @@ -756,10 +886,6 @@ def _process_task_queue(self, thread_idx: int): f"unique_rid={write_meta.unique_rid}: {e}" ) write_meta.task.fail(e) - finally: - if write_meta.owner_session is not None: - write_meta.owner_session._finish_nixl_op() - write_meta.owner_session = None finally: # Clean up this thread's DEALER sockets. threading.local storage # is only accessible from the owning thread, so shutdown must @@ -776,6 +902,71 @@ def _process_task_queue(self, thread_idx: int): ) dealers.clear() + def _retain_in_doubt_transfer( + self, + write_meta: WriteMeta, + initial_report: list[bytes], + send_slot_id: Optional[int] = None, + ) -> None: + thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads + key = (write_meta.task, write_meta.peer_rank) + if key in self._pending_settlements[thread_idx]: + return + pending = _PendingSettlement(write_meta, initial_report, send_slot_id) + self._pending_settlements[thread_idx][key] = pending + try: + self._get_result_dealer(write_meta.peer_endpoint).send(initial_report) + except Exception as error: + logger.warning(f"Failed to report ambiguous transfer; retaining evidence: {error}") + pending.report_error_logged = True + else: + pending.initial_report = None + + def _poll_in_doubt_transfers(self, thread_idx: int) -> None: + pending_transfers = self._pending_settlements[thread_idx] + for key, pending in list(pending_transfers.items()): + meta = pending.write_meta + try: + dealer = self._get_result_dealer(meta.peer_endpoint) + if pending.initial_report is not None: + dealer.send(pending.initial_report) + pending.initial_report = None + if not pending.backend_done: + if not meta.task.poll_in_doubt_physical_operation(meta.peer_rank): + continue + pending.backend_done = True + if pending.send_slot_id is not None: + self._bounce.release_send(pending.send_slot_id) + pending.send_slot_id = None + if meta.meta_type == WriteMetaType.AUX: + message = _make_aux_result_msg( + self._instance_rank, + meta.unique_rid, + AgentResult.FAILED_QUIESCED, + ) + else: + message = _make_kv_result_msg( + self._instance_rank, + meta.unique_rid, + meta.receiver_slice_id, + True, + AgentResult.FAILED_QUIESCED, + ) + dealer.send(message) + except Exception as error: + # Sending may have escaped before raising. A duplicate settlement + # is idempotent; losing the only remaining report is not safe. + if not pending.report_error_logged: + logger.warning(f"Failed to report physical settlement; will retry: {error}") + pending.report_error_logged = True + continue + with meta.task.lock: + if meta.meta_type == WriteMetaType.AUX: + meta.task._transfer_count += 1 + else: + meta.task.transferred_count += 1 + del pending_transfers[key] + @staticmethod @nvtx_range("_make_agent_request") def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferRequest": @@ -795,7 +986,11 @@ def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferReque f"_make_agent_request: dst_device_id is None for KV transfer " f"unique_rid={write_meta.unique_rid}" ) - src_dev, dst_dev, mem_type = device_id, write_meta.dst_device_id, MemoryType.VRAM + src_dev, dst_dev, mem_type = ( + device_id, + write_meta.dst_device_id, + MemoryType.VRAM, + ) if n == 0: src_memory_descs = MemoryDescs(mem_type, []) @@ -812,7 +1007,11 @@ def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferReque # After this call, src_memory_descs and dst_memory_descs are in a moved-from # state and must NOT be accessed again. return TransferRequest( - TransferOp.WRITE, src_memory_descs, dst_memory_descs, write_meta.peer_name, None + TransferOp.WRITE, + src_memory_descs, + dst_memory_descs, + write_meta.peer_name, + None, ) @nvtx_range("_deliver_kv_to_agent") @@ -821,7 +1020,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"WriteMeta ptr/size mismatch for unique_rid={write_meta.unique_rid}" ) - assert write_meta.slice_id is not None + owned = self._enforce_physical_ownership with self._sessions_lock: session = self._get_session(write_meta.unique_rid) if session is None: @@ -829,21 +1028,11 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"_deliver_kv_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" ) logger.error(msg) + if owned: + write_meta.task.retire_unsubmitted_physical_operation(write_meta.peer_rank) write_meta.task.fail(RuntimeError(msg)) - # REQUEST_DATA was already accepted by this sender. Resolve the - # receiver reservation even though cancellation/cleanup removed - # the source session before this queued item reached a worker. - self._send_agent_result( - write_meta.peer_endpoint, - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - write_meta.is_last_slice, - AgentResult.FAILED, - ), - ) return + assert write_meta.slice_id is not None task = session.kv_tasks[write_meta.slice_id] timer = task._perf_timer if timer: @@ -851,13 +1040,13 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): # Hold session.lock to serialize the INIT→TRANSFERRING transition with # cancel(): prevents cancel_request() from freeing KV pages while a # worker is about to write into them. - active_nixl_op = write_meta.src_ptrs.size > 0 with session.lock: status = session.status if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): should_abort = True else: - task.status = TaskStatus.TRANSFERRING + if not owned: + task.status = TaskStatus.TRANSFERRING should_abort = False if should_abort: @@ -867,18 +1056,19 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): ) # Task may have been enqueued after cancel() already iterated kv_tasks, # so its future was never set by cancel(). Set it here as a fallback. + if owned: + task.retire_unsubmitted_physical_operation(write_meta.peer_rank) task.fail( RuntimeError(f"session {write_meta.unique_rid} {status.value}, transfer aborted") ) - self._send_agent_result( - write_meta.peer_endpoint, + self._get_result_dealer(write_meta.peer_endpoint, legacy_shared=True).send( _make_kv_result_msg( self._instance_rank, write_meta.unique_rid, - write_meta.slice_id, + write_meta.receiver_slice_id, True, # is_last_slice — ensures receiver resolves its task future AgentResult.FAILED, - ), + ) ) return @@ -886,18 +1076,46 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): agent_result = AgentResult.SUCCESS send_slot_id = None - try: - if active_nixl_op: + submitted_and_completed = False + if write_meta.src_ptrs.size > 0: + try: request, send_slot_id = build_send_request( self._bounce, write_meta, lambda: Sender._make_agent_request(write_meta, device_id=self._device_id), ) - if timer: - timer.record_transfer_start(write_meta.peer_rank) - completed, last_status = self._submit_and_wait_agent_transfer(request) - if not completed: - agent_result = AgentResult.FAILED + except Exception as e: + # Don't let a gather fault escape: without a result the receiver would hang and its + # region leak. Tell the receiver it failed and fail the local task instead. + logger.error( + f"_deliver_kv_to_agent: failed to build the KV send request for " + f"{write_meta.unique_rid} slice={write_meta.slice_id}: {e}" + ) + if owned: + task.retire_unsubmitted_physical_operation(write_meta.peer_rank) + task.fail(RuntimeError(f"build_send_request failed: {e}")) + self._get_result_dealer(write_meta.peer_endpoint, legacy_shared=True).send( + _make_kv_result_msg( + self._instance_rank, + write_meta.unique_rid, + write_meta.receiver_slice_id, + True, # is_last_slice — ensures receiver resolves its task future + AgentResult.FAILED, + ) + ) + return + if timer: + timer.record_transfer_start(write_meta.peer_rank) + transfer_finished = False + try: + transfer_finished, last_status = self._submit_transfer( + task, write_meta.peer_rank, request + ) + if transfer_finished: + submitted_and_completed = True + del request + if not transfer_finished: + agent_result = AgentResult.IN_DOUBT if owned else AgentResult.FAILED agent_name = getattr(self._agent, "name", "") detail = ( f"KV transfer agent failed: " @@ -912,47 +1130,53 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"nixl_status={last_status} agent={agent_name}" ) logger.error(detail) - task.fail(RuntimeError(detail)) - except Exception as exc: - agent_result = AgentResult.FAILED - task.fail(exc) - logger.error( - f"KV transfer raised for rid={write_meta.unique_rid}, " - f"peer_rank={write_meta.peer_rank}: {exc}" - ) - finally: - if send_slot_id is not None: - self._bounce.release_send(send_slot_id) - if timer: - try: - timer.record_transfer_end(write_meta.peer_rank) - except Exception as exc: - logger.warning( - f"Failed to record KV transfer end for rid={write_meta.unique_rid}: {exc}" - ) + error = RuntimeError(detail) + task.fail(error) + except _TransferNotSubmittedError as error: + # Backend admission failed, so local resources are safe to retire. + transfer_finished = True + agent_result = AgentResult.FAILED + task.fail(error) + finally: + if send_slot_id is not None and (not owned or transfer_finished): + self._bounce.release_send(send_slot_id) + elif owned: + task.retire_unsubmitted_physical_operation(write_meta.peer_rank) + if timer: + timer.record_transfer_end(write_meta.peer_rank) - ## TODO: just last slice need to send task state? + # Report every chunk so failures reach the receiver immediately. tail = ( encode_result_tail(write_meta) if send_slot_id is not None and agent_result == AgentResult.SUCCESS else None ) - transfer_size = timer.get_transfer_size(write_meta.peer_rank) if timer else 0 + # Attested written bytes: only what was actually submitted to the + # agent AND completed. The receiver checks the per-slice sum against + # the byte total of its published destinations before admitting the + # transferred range to KV block reuse (verified-range admission), so + # this must reflect the completed submission, never intent — gate on + # submitted_and_completed, not on agent_result, so a SUCCESS message + # that did not come from a completed submit attests zero bytes. + # Matches the perf timer's transfer-size accounting (kv_sizes.sum() + # recorded in _build_kv_write_meta) but is populated even with perf + # logging off. + transfer_size = 0 + if submitted_and_completed and agent_result == AgentResult.SUCCESS: + transfer_size = int(write_meta.sizes.sum()) result_msg = _make_kv_result_msg( self._instance_rank, write_meta.unique_rid, - write_meta.slice_id, + write_meta.receiver_slice_id, write_meta.is_last_slice, agent_result, transfer_size=transfer_size, tail=tail, ) - self._send_agent_result(write_meta.peer_endpoint, result_msg) - - if timer: - timer.record_task_end(write_meta.peer_rank) - ri = self._registrar.self_rank_info - task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) + if agent_result == AgentResult.IN_DOUBT: + self._retain_in_doubt_transfer(write_meta, result_msg, send_slot_id) + return + self._get_result_dealer(write_meta.peer_endpoint).send(result_msg) with task.lock: task.transferred_count += 1 @@ -973,6 +1197,17 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): if all(t.status == TaskStatus.TRANSFERRED for t in session.kv_tasks): session.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() + # Diagnostics must not delay authoritative completion or turn it into failure. + try: + if timer: + timer.record_task_end(write_meta.peer_rank) + ri = self._registrar.self_rank_info + task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) + except Exception as error: + logger.warning( + f"KV send perf logging failed for request {write_meta.unique_rid}: {error}" + ) + logger.debug( f"deliver_kv_to_agent completed: unique_rid={write_meta.unique_rid}, " f"slice_id={write_meta.slice_id}, agent_result={agent_result}" @@ -980,20 +1215,14 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): @nvtx_range("_deliver_aux_to_agent") def _deliver_aux_to_agent(self, write_meta: WriteMeta): + owned = self._enforce_physical_ownership session = self._get_session(write_meta.unique_rid) if session is None: msg = f"_deliver_aux_to_agent: TxSession {write_meta.unique_rid} not found or already GC'd" logger.error(msg) + if owned: + write_meta.task.retire_unsubmitted_physical_operation(write_meta.peer_rank) write_meta.task.fail(RuntimeError(msg)) - self._send_agent_result( - write_meta.peer_endpoint, - [ - MessageType.AUX_AGENT_RESULT, - str(self._instance_rank).encode("ascii"), - str(write_meta.unique_rid).encode("ascii"), - AgentResult.FAILED.value.encode("ascii"), - ], - ) return aux_task = session.aux_task assert aux_task is not None, f"aux_task is None for session {write_meta.unique_rid}" @@ -1001,73 +1230,70 @@ def _deliver_aux_to_agent(self, write_meta: WriteMeta): if timer: timer.record_push_end(write_meta.peer_rank) - active_nixl_op = write_meta.src_ptrs.size > 0 - with session.lock: - status = session.status - if status in (SessionStatus.ERROR, SessionStatus.CANCELLED): - should_abort = True - else: - aux_task.status = TaskStatus.TRANSFERRING - should_abort = False - - if should_abort: - aux_task.fail( - RuntimeError( - f"session {write_meta.unique_rid} {status.value}, aux transfer aborted" + agent_result = AgentResult.SUCCESS + if write_meta.src_ptrs.size > 0: + try: + request = Sender._make_agent_request(write_meta, device_id=self._device_id) + except Exception as error: + logger.error( + "_deliver_aux_to_agent: failed to build the auxiliary send " + f"request for {write_meta.unique_rid}: {error}" + ) + if owned: + aux_task.retire_unsubmitted_physical_operation(write_meta.peer_rank) + aux_task.fail(RuntimeError(f"build aux transfer request failed: {error}")) + # An admitted operation normally makes this worker the only + # eligible claimant. Keep the claim at the send boundary so + # duplicate dispatch cannot publish contradictory evidence. + if not owned or session._claim_aux_terminal_result(write_meta.peer_rank): + self._get_result_dealer(write_meta.peer_endpoint).send( + _make_aux_result_msg( + self._instance_rank, + write_meta.unique_rid, + AgentResult.FAILED, + ) + ) + return + if timer: + timer.record_transfer_start(write_meta.peer_rank) + try: + transfer_finished, last_status = self._submit_transfer( + aux_task, write_meta.peer_rank, request + ) + if transfer_finished: + del request + if not transfer_finished: + agent_result = AgentResult.IN_DOUBT if owned else AgentResult.FAILED + session.set_exception(f"aux transfer agent request failed: {last_status}") + except _TransferNotSubmittedError as error: + agent_result = AgentResult.FAILED + aux_task.fail(error) + if timer: + timer.record_transfer_end(write_meta.peer_rank) + elif owned: + aux_task.retire_unsubmitted_physical_operation(write_meta.peer_rank) + + # An admitted operation normally makes this worker the only eligible + # claimant. Keep the claim at the send boundary so future cleanup-path + # changes cannot publish contradictory evidence. + if not owned or session._claim_aux_terminal_result(write_meta.peer_rank): + if agent_result == AgentResult.IN_DOUBT: + self._retain_in_doubt_transfer( + write_meta, + _make_aux_result_msg(self._instance_rank, write_meta.unique_rid, agent_result), + ) + return + self._get_result_dealer(write_meta.peer_endpoint).send( + _make_aux_result_msg( + self._instance_rank, + write_meta.unique_rid, + agent_result, ) ) - self._send_agent_result( - write_meta.peer_endpoint, - [ - MessageType.AUX_AGENT_RESULT, - str(self._instance_rank).encode("ascii"), - str(write_meta.unique_rid).encode("ascii"), - AgentResult.FAILED.value.encode("ascii"), - ], - ) + + if agent_result == AgentResult.IN_DOUBT: return - agent_result = AgentResult.SUCCESS - try: - if active_nixl_op: - request = Sender._make_agent_request(write_meta, device_id=self._device_id) - if timer: - timer.record_transfer_start(write_meta.peer_rank) - completed, _ = self._submit_and_wait_agent_transfer(request) - if not completed: - agent_result = AgentResult.FAILED - session.set_exception("aux transfer agent request failed") - except Exception as exc: - agent_result = AgentResult.FAILED - aux_task.fail(exc) - logger.error( - f"Aux transfer raised for rid={write_meta.unique_rid}, " - f"peer_rank={write_meta.peer_rank}: {exc}" - ) - finally: - if timer: - try: - timer.record_transfer_end(write_meta.peer_rank) - except Exception as exc: - logger.warning( - f"Failed to record aux transfer end for rid={write_meta.unique_rid}: {exc}" - ) - - self._send_agent_result( - write_meta.peer_endpoint, - [ - MessageType.AUX_AGENT_RESULT, - str(self._instance_rank).encode("ascii"), - str(write_meta.unique_rid).encode("ascii"), - agent_result.value.encode("ascii"), - ], - ) - - if timer: - timer.record_task_end(write_meta.peer_rank) - ri = self._registrar.self_rank_info - aux_task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) - with aux_task.lock: aux_task._transfer_count += 1 count = aux_task._transfer_count @@ -1083,106 +1309,64 @@ def _deliver_aux_to_agent(self, write_meta: WriteMeta): f"aux task received more than {write_meta.expected_transfers} transfers" ) - @staticmethod - def _align_kv_blocks( - src_block_ids: np.ndarray, - dst_block_ids: np.ndarray, - src_token_start: int, - dst_token_start: int, - tokens_per_block: int, - minimum_token_start: int = 0, - ) -> tuple[np.ndarray, np.ndarray]: - """Align src/dst block arrays using explicit token-start positions. - - Both src_token_start and dst_token_start must be block-aligned - (multiples of tokens_per_block), which is always true for prefix-cache - boundaries in current KV cache managers. minimum_token_start applies a - shared lower bound, such as an SWA stale boundary, and trims both arrays - instead of relabeling their first blocks. - - Returns the (src, dst) sub-arrays that cover the shared token overlap. - Returns a pair of empty arrays when there is no overlap (i.e. this - context slice is entirely within generation's already-cached prefix). - - This handles four cases without special-casing: - 1. No prefix cache on either side → identity (start_token == 0 both) - 2. Context prefix cache (src starts later than 0) → trim dst head - 3. Generation prefix cache (dst starts later than 0) → trim src head - 4. Chunked context (each slice has its own token_range) → correct - overlap even when the slice is entirely before dst_token_start - """ - overlap_start = max(src_token_start, dst_token_start, minimum_token_start) - src_skip = (overlap_start - src_token_start) // tokens_per_block - dst_skip = (overlap_start - dst_token_start) // tokens_per_block - n_transfer = min(src_block_ids.size - src_skip, dst_block_ids.size - dst_skip) - if n_transfer <= 0: - return np.array([], dtype=np.int64), np.array([], dtype=np.int64) - return ( - src_block_ids[src_skip : src_skip + n_transfer], - dst_block_ids[dst_skip : dst_skip + n_transfer], - ) - - @staticmethod - def _beam0_block_count(block_ids: np.ndarray, total_blocks: int, beam_width: int) -> int: - """Return the number of beam-0 blocks in a packed 1-D beam layout.""" - if beam_width <= 1 or block_ids.size <= total_blocks: - return block_ids.size - return max(0, block_ids.size - (beam_width - 1)) - - @staticmethod - def _helix_windowed_block_ids( - src_block_ids: np.ndarray, - total_blocks: int, - beam_width: int, - cp_rank: int, - cp_size: int, - ) -> np.ndarray: - """Select the blocks of a windowed (suffix) list owned by one Helix rank. - - The list covers global ordinals ``[total_blocks - k, total_blocks)``; - rank ``cp_rank`` owns every ordinal congruent to it modulo ``cp_size``. - """ - k = Sender._beam0_block_count(src_block_ids, total_blocks, beam_width) - first_ordinal = total_blocks - k - keep = [i for i in range(k) if (first_ordinal + i) % cp_size == cp_rank] - return src_block_ids[keep] + try: + if timer: + timer.record_task_end(write_meta.peer_rank) + ri = self._registrar.self_rank_info + aux_task.print_perf_info(write_meta.peer_rank, ri.instance_name, ri.instance_rank) + except Exception as error: + logger.warning( + f"AUX send perf logging failed for request {write_meta.unique_rid}: {error}" + ) @staticmethod - def _trim_receiver_window_head( + def _pair_ordinals( src_block_ids: np.ndarray, dst_block_ids: np.ndarray, - peer_window_size: Optional[int], - beam_width: int, - ) -> np.ndarray: - """Drop the receiver's extra leading blocks for a windowed layer group. - - A windowed receiver keeps a larger window when only it runs speculative - decoding, so its suffix starts earlier and the extra blocks are at the - head. Both starts are derived from list length, so trimming the tail - instead would shift every block one position early. - - Non-windowed lists are both trimmed to ceil(prompt_len / tpb) in - _create_kv_slice, so there dst must not exceed src. A smaller dst - (generation prefix-cache reuse) is handled via dst_start. + dst_start_block: int = 0, + ) -> tuple[np.ndarray, np.ndarray]: + """Pair src/dst positional block tables ordinal by ordinal. + + Both tables are indexed by block ordinal (see Chunk); an entry is -1 + where that side has nothing to offer/accept. A block moves only where + both sides hold it, so a sender that has not evicted yet, a receiver + that already reused a prefix, a receiver with a wider speculative + window, or a sender chunk that blanks everything outside its range all + fall out of the same intersection without any per-case arithmetic. + Ordinals below ``dst_start_block`` are excluded as well. """ - block_diff = dst_block_ids.size - src_block_ids.size - if block_diff <= 0: - return dst_block_ids - if peer_window_size is None or beam_width > 1: + if src_block_ids.size != dst_block_ids.size: raise ValueError( - f"src/dst block count mismatch: {src_block_ids.size} vs " - f"{dst_block_ids.size} (dst must not exceed src)" + f"positional block tables differ in length: src={src_block_ids.size} " + f"dst={dst_block_ids.size}; both sides must span the same prompt" ) - return dst_block_ids[block_diff:] + mask = (src_block_ids >= 0) & (dst_block_ids >= 0) + if dst_start_block > 0: + mask[:dst_start_block] = False + idx = np.flatnonzero(mask) + return src_block_ids[idx], dst_block_ids[idx] @nvtx_range("_build_kv_write_meta") def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> WriteMeta: peer_ri = self._registrar.get_peer_rank_info(req_info.instance_name, req_info.instance_rank) + self_ri = self._registrar.self_rank_info + token_range = task._chunk.token_range + # Only a piece short of the whole prompt runs the chunk projection below, and that is the + # part with no answer under CP; a sole piece covering everything takes the untouched path. + is_partial = token_range.start > 0 or token_range.end < task._prompt_len + if is_partial and (self_ri.cp_size > 1 or peer_ri.cp_size > 1): + raise ValueError( + "enable_pipelined_transfer is not supported with context parallelism " + f"(sender cp_size={self_ri.cp_size}, receiver cp_size={peer_ri.cp_size})" + ) timer = task._perf_timer if timer: timer.record_prepare_args_start(peer_ri.instance_rank) targets = self._registrar.get_peer_overlap(peer_ri, peer_ri.dp_rank) expected_transfers = len(targets.ranks) + # Recorded on the task as well: the count on the write meta settles the task, while the + # handle has to answer how many writes are still out after one of them has already failed. + task.expected_transfers = expected_transfers # Aggregate fragment pointers from all matching pool pairs. # Each pool pair produces one or more region pairs (rp), each containing @@ -1201,109 +1385,58 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write ) pool_mapping = self._registrar.get_pool_mapping(peer_ri) dst_block_ids_per_groups = req_info.block_ids_per_layer_groups - src_block_ids_per_groups = task._slice.block_ids_per_layer_groups + src_block_ids_per_groups = task._chunk.block_ids_per_layer_groups # Aggregate fragments from all matching pools using numpy concatenation. # Send ownership is per pool: replicated pools elect one fan-in # owner, sharded pools keep head-duplication routing. for (self_lg, self_pi), (peer_lg, peer_pi) in pool_mapping.items(): + if (self_lg, self_pi) in task._chunk.excluded_pool_views: + continue if not self._registrar.should_send_pool(targets, peer_ri, self_lg, self_pi): continue + + lg_info = extractor.page_table.layer_groups[self_lg] src_block_ids = src_block_ids_per_groups[self_lg] dst_block_ids = dst_block_ids_per_groups[peer_lg] - tpb = extractor.page_table.tokens_per_block - token_range = task._slice.token_range - lg_info = extractor.page_table.layer_groups[self_lg] - window_size = getattr(lg_info, "sliding_window_size", None) - if ( - peer_ri.cp_size > 1 - and self._registrar.self_rank_info.cp_size == 1 - and not getattr(lg_info, "cp_as_tp", False) - ): - # Helix: the receiver owns global blocks [cp_rank::cp_size] - # (same protocol as partition_context_for_helix); block reuse - # is rejected under helix. Windowed lists are a suffix of the - # global ordinals, so route them by ordinal, not by position. - if window_size is None: + # Extract regions: STATE = per-layer ptrs, PAGED = per-block ptrs + if lg_info.kind == CacheKind.STATE: + if src_block_ids.size == 0 or dst_block_ids.size == 0: + continue + if not MambaPolicy.is_paired(self._registrar.self_rank_info, peer_ri): + continue + src_region = extractor.extract_slot(int(src_block_ids[0]), self_lg, self_pi) + dst_region = peer_extractor.extract_slot(int(dst_block_ids[0]), peer_lg, peer_pi) + else: + src_block_ids = np.asarray(src_block_ids, dtype=np.int64) + dst_block_ids = np.asarray(dst_block_ids, dtype=np.int64) + if ( + peer_ri.cp_size > 1 + and self_ri.cp_size == 1 + and not getattr(extractor.page_table.layer_groups[self_lg], "cp_as_tp", False) + ): + # Helix: the receiver owns global blocks [cp_rank::cp_size] + # (same protocol as partition_context_for_helix), so its + # table is the strided subset of ours; block reuse is + # rejected under helix. src_block_ids = src_block_ids[peer_ri.cp_rank :: peer_ri.cp_size] - else: - slice_end = token_range.end if token_range is not None else 0 - src_block_ids = Sender._helix_windowed_block_ids( - src_block_ids, - total_blocks=(slice_end + tpb - 1) // tpb, - beam_width=task._beam_width, - cp_rank=peer_ri.cp_rank, - cp_size=peer_ri.cp_size, - ) - - # Non-windowed lists are both trimmed to ceil(prompt_len / tpb) in - # _create_kv_slice, so dst must not exceed src. A smaller dst - # (generation prefix-cache reuse) is handled via dst_start below. - block_diff = dst_block_ids.size - src_block_ids.size - if block_diff > 0: - # A windowed receiver keeps a larger window when only it runs - # speculative decoding, so its suffix starts earlier. Drop the - # extra head: dst_start is derived from len(dst), so trimming - # the tail instead would shift every block by block_diff. - peer_lg_info = peer_extractor.page_table.layer_groups[peer_lg] - peer_window = getattr(peer_lg_info, "sliding_window_size", None) - if peer_window is not None and task._beam_width <= 1: - dst_block_ids = dst_block_ids[block_diff:] - else: - raise ValueError( - f"src/dst block count mismatch: {src_block_ids.size} vs " - f"{dst_block_ids.size} (dst must not exceed src)" - ) - - # Block lists are the suffix of [..., slice_end); cached prefix - # is implicit in their size. token_start = (total_blocks - n) * tpb. - slice_end = token_range.end if token_range is not None else 0 - total_blocks = (slice_end + tpb - 1) // tpb - src_beam0_blocks = Sender._beam0_block_count( - src_block_ids, total_blocks, task._beam_width - ) - dst_beam0_blocks = Sender._beam0_block_count( - dst_block_ids, total_blocks, task._beam_width - ) - assert src_beam0_blocks <= total_blocks, ( - f"src beam-0 block list ({src_beam0_blocks}) exceeds total slice " - f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" - ) - assert dst_beam0_blocks <= total_blocks, ( - f"dst beam-0 block list ({dst_beam0_blocks}) exceeds total slice " - f"blocks ({total_blocks}); slice_end={slice_end}, tpb={tpb}" - ) - src_start = (total_blocks - src_beam0_blocks) * tpb - dst_start = (total_blocks - dst_beam0_blocks) * tpb - minimum_start = req_info.dst_start_token or 0 - if window_size is not None: - # SWA stale_end uses the request prompt_len (not slice_end — - # they differ for non-final slices). prompt_len must be plumbed - # via the session; falling back to slice_end is wrong on - # non-final slices. - assert task._prompt_len is not None, ( - "SWA layer requires session.prompt_len; " - "set TxSession(prompt_len=request.prompt_len)." + tpb = extractor.page_table.tokens_per_block + dst_start_block = ( + req_info.dst_start_token // tpb if req_info.dst_start_token is not None else 0 ) - stale_end = max(0, (task._prompt_len + 1 - window_size) // tpb) - minimum_start = max(minimum_start, stale_end * tpb) - src_block_ids, dst_block_ids = Sender._align_kv_blocks( - src_block_ids, - dst_block_ids, - src_token_start=src_start, - dst_token_start=dst_start, - tokens_per_block=tpb, - minimum_token_start=minimum_start, - ) + # Both tables are positional (index == block ordinal, -1 == not + # held), so pairing is a plain intersection: SWA eviction, prefix + # reuse, and pipelined chunk bounds are all already encoded as + # holes by whichever side owns that state. + src_block_ids, dst_block_ids = Sender._pair_ordinals( + src_block_ids, dst_block_ids, dst_start_block=dst_start_block + ) + if src_block_ids.size == 0: + continue + src_region = extractor.extract(src_block_ids, self_lg, self_pi) + dst_region = peer_extractor.extract(dst_block_ids, peer_lg, peer_pi) - src_region = extractor.extract(src_block_ids, layer_group_id=self_lg, pool_idx=self_pi) - dst_region = peer_extractor.extract( - dst_block_ids, layer_group_id=peer_lg, pool_idx=peer_pi - ) - if token_range is not None: - token_spec = KVRegionSpec(tokens=IndexRange(token_range.start, token_range.end - 1)) - src_region = src_region._replace(spec=token_spec) - dst_region = dst_region._replace(spec=token_spec) + # Map and collect fragments (same for all layer types) mapper = self._registrar.get_kv_map(peer_ri, (self_lg, self_pi), (peer_lg, peer_pi)) region_pair = mapper.map(src_region, dst_region) region_pairs = region_pair if isinstance(region_pair, list) else [region_pair] @@ -1326,20 +1459,6 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write dst_frags = np.array([], dtype=np.int64) kv_sizes = np.array([], dtype=np.int64) - # handle mamba fragments - m_src, m_dst, m_sizes = MambaPolicy.collect_frags( - self_page_table=extractor.page_table, - peer_page_table=peer_extractor.page_table, - src_slot=task._slice.mamba_state_index, - dst_slot=req_info.mamba_state_index, - self_ri=self._registrar.self_rank_info, - peer_ri=peer_ri, - ) - if m_src: - src_frags = np.concatenate([src_frags, np.array(m_src, dtype=np.int64)]) - dst_frags = np.concatenate([dst_frags, np.array(m_dst, dtype=np.int64)]) - kv_sizes = np.concatenate([kv_sizes, np.array(m_sizes, dtype=np.int64)]) - if timer: timer.record_prepare_args_end(peer_ri.instance_rank) timer.record_transfer_sizes(peer_ri.instance_rank, int(kv_sizes.sum()), dst_frags.size) @@ -1356,7 +1475,8 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write peer_endpoint=peer_ri.self_endpoint, unique_rid=task._unique_rid, slice_id=task.slice_id, - is_last_slice=task._slice.is_last_slice, + receiver_slice_id=req_info.slice_id if req_info.slice_id is not None else 0, + is_last_slice=task._chunk.is_last, bounce_dst_base=req_info.bounce_dst_base, ) @@ -1393,7 +1513,9 @@ def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> Wri if timer: timer.record_prepare_args_end(peer_ri.instance_rank) timer.record_transfer_sizes( - peer_ri.instance_rank, int(sizes.sum()) if sizes.size > 0 else 0, src_ptrs.size + peer_ri.instance_rank, + int(sizes.sum()) if sizes.size > 0 else 0, + src_ptrs.size, ) return WriteMeta( @@ -1418,33 +1540,44 @@ def dispatch_task( # critical section small. When not provided, we fetch it here (legacy / standalone path). if req_info_snapshot is None: req_info_snapshot = dict(self._get_req_info(task._unique_rid) or {}) - req_infos = list(req_info_snapshot.values()) - for info_idx, info in enumerate(req_infos): + first_error: Optional[Exception] = None + for info in req_info_snapshot.values(): try: - if task._perf_timer is not None: - task._perf_timer.record_task_start(info.instance_rank) - if isinstance(task, KVSendTask): - trans_meta = self._build_kv_write_meta(task, info) - else: - trans_meta = self._build_aux_write_meta(task, info) - if task._perf_timer is not None: - task._perf_timer.record_push_start(trans_meta.peer_rank) - self._enqueue(trans_meta) - except Exception as exc: - task.fail(exc) - # The receiver reserved one peer transfer before publishing - # every RecvReqInfo in this snapshot. Entries before info_idx - # already own an enqueued WriteMeta and will report their own - # terminal result. This entry and every later one have not - # been dispatched, so resolve each reservation exactly once - # before propagating the build/enqueue failure. - for undispatched_info in req_infos[info_idx:]: - self._send_failed_result_to_receiver( - undispatched_info, - slice_id=(task.slice_id if isinstance(task, KVSendTask) else None), - aux=isinstance(task, AuxSendTask), - ) - raise + self._dispatch_task_to_peer(task, info) + except Exception as error: + if not self._enforce_physical_ownership: + raise + if first_error is None: + first_error = error + if first_error is not None: + raise first_error + + def _dispatch_task_to_peer(self, task: KVSendTask | AuxSendTask, info: RecvReqInfo) -> None: + owned = False + if self._enforce_physical_ownership: + admission = self._begin_task_operation(task, info.instance_rank) + if admission is None: + return + if not admission: + self._send_failed_task_result_to_receiver(task, info) + return + owned = True + try: + if task._perf_timer is not None: + task._perf_timer.record_task_start(info.instance_rank) + trans_meta = ( + self._build_kv_write_meta(task, info) + if isinstance(task, KVSendTask) + else self._build_aux_write_meta(task, info) + ) + if task._perf_timer is not None: + task._perf_timer.record_push_start(trans_meta.peer_rank) + self._enqueue(trans_meta) + except Exception: + if owned: + task.retire_unsubmitted_physical_operation(info.instance_rank) + self._send_failed_task_result_to_receiver(task, info) + raise def _start_listener(self): def handle_message(messages: list[bytes]): @@ -1474,19 +1607,24 @@ def handle_message(messages: list[bytes]): self._messenger.start_listener(handle_message) def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): - # Skip late messages so we don't race shutdown's invalidate loop. - if self._shutdown: + if self._shutdown_requested: return torch.cuda.set_device(self._device_id) CUASSERT(cudart.cudaSetDevice(self._device_id)) ri: RankInfo = RankInfo.from_bytes(message[1]) - - self._registrar.register(ri.instance_name, ri.instance_rank, ri) - agent_name = ri.instance_name + str(ri.instance_rank) logger.debug(f"Loading remote transfer agent descriptor for peer '{agent_name}'") - self._load_remote_agent_serialized(agent_name, ri.transfer_engine_info) + # NIXL serializes descriptor load against submit internally. This lock + # only closes the local load-vs-shutdown window; active transfers do + # not need to drain before a new peer is registered. with self._loaded_remote_agents_lock: + if self._shutdown_requested: + return + self._registrar.register(ri.instance_name, ri.instance_rank, ri) + self._agent.load_remote_agent( + agent_name, + ri.transfer_engine_info, + ) self._loaded_remote_agents.add(agent_name) logger.debug( f"Completed handling REGISTER_RANK_INFO for instance='{ri.instance_name}', rank={ri.instance_rank}" @@ -1495,201 +1633,235 @@ def _register_peer_rank(self, _send_id: bytes, message: list[bytes]): def _handle_cancel_session(self, message: list[bytes]): unique_rid = int(message[1]) session = None - pre_cancel_infos: list[RecvReqInfo] = [] with self._sessions_lock: - self._pre_cancelled_rids.add(unique_rid) - self._pre_cancelled_timestamps[unique_rid] = time.monotonic() - session_ref = self._sessions.get(unique_rid) - if session_ref is None: - with self._peer_requests_lock: - req_info_map = self._peer_requests.pop(unique_rid, None) - self._peer_requests_timestamps.pop(unique_rid, None) - if req_info_map: - pre_cancel_infos = list(req_info_map.values()) + session = self._get_session(unique_rid) + if session is None or self._enforce_physical_ownership: + # Parked with the peer as who asked: the session built afterwards reads its + # cancellation from here, and an id on its own would read as a local cancel. + if self._enforce_physical_ownership: + self._pre_cancelled_rids.setdefault(unique_rid, True) + else: + self._pre_cancelled_rids[unique_rid] = True + if session is not None: + if self._enforce_physical_ownership: + session.cancel_local(by_peer=True) else: - session = session_ref() - if session is None: - with self._peer_requests_lock: - req_info_map = self._peer_requests.pop(unique_rid, None) - self._peer_requests_timestamps.pop(unique_rid, None) - if req_info_map: - pre_cancel_infos = list(req_info_map.values()) + session.cancel(by_peer=True) + if self._enforce_physical_ownership and len(message) == 4: + # Routing is carried by CANCEL itself, not by expiring REQUEST_DATA. + self._queue_session_quiescence(unique_rid, message[2].decode(), int(message[3])) + + def _queue_session_quiescence( + self, unique_rid: int, instance_name: str, instance_rank: int + ) -> None: + """Queue proof on the same worker/socket as this peer's piece results.""" + endpoint = self._registrar.get_peer_rank_info(instance_name, instance_rank).self_endpoint + key = (unique_rid, instance_rank) + with self._sessions_lock: + if self._shutdown_requested or key in self._pending_session_quiescence: + return + marker = _SessionQuiescence(unique_rid, instance_rank, endpoint) + self._pending_session_quiescence[key] = marker + self._send_task_queues[hash(key) % self._num_threads].put(marker) + + def _send_session_quiesced(self, marker: _SessionQuiescence) -> bool: + """Acknowledge a fenced, drained session; return False to retry later. + + Called only by the ordered worker, after earlier queued piece reports. + Session absence is safe only with the persistent admission fence. An + existing session must also prove that all KV and AUX accesses drained. + """ + key = (marker.unique_rid, marker.peer_rank) + with self._sessions_lock: + if marker.unique_rid not in self._pre_cancelled_rids: + return False + session = self._get_session(marker.unique_rid) if session is not None: - session.cancel() with session.lock: - need_kv_result = not session.kv_tasks - need_aux_result = session._need_aux and session.aux_task is None - if need_kv_result or need_aux_result: - # cancel() has made future send()/send_aux() calls fail, so no - # worker can be created for the missing task type. Remove the - # saved requests and resolve those reservations directly. - with self._peer_requests_lock: - req_info_map = self._peer_requests.pop(unique_rid, None) - self._peer_requests_timestamps.pop(unique_rid, None) - direct_infos = list(req_info_map.values()) if req_info_map else [] - for info in direct_infos: - self._send_pre_cancel_failed_results( - info, - send_kv=need_kv_result, - send_aux=need_aux_result, - ) - else: - # These REQUEST_DATA messages were accepted, but no TxSession will - # ever enqueue a worker result on this (possibly non-selected ADP) - # rank. Explicitly resolve the receiver's per-peer reservations. - for info in pre_cancel_infos: - self._send_pre_cancel_failed_results(info) + if not session.has_failed() or not session.resources_drained(): + return False + thread_idx = hash(key) % self._num_threads + if any( + pending.write_meta.unique_rid == marker.unique_rid + and pending.write_meta.peer_rank == marker.peer_rank + for pending in self._pending_settlements[thread_idx].values() + ): + return False + try: + self._get_or_connect_thread_dealer(marker.endpoint).send( + [ + MessageType.SESSION_QUIESCED, + str(self._instance_rank).encode("ascii"), + str(marker.unique_rid).encode("ascii"), + ] + ) + except Exception as error: + logger.warning(f"Failed to acknowledge fenced session {marker.unique_rid}: {error}") + return False + with self._sessions_lock: + self._pending_session_quiescence.pop(key, None) + return True @nvtx_range("_respond_with_kv") def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): # _sessions_lock prevents a race between session lookup and req_info save. - # session.lock serializes _enqueue calls from both paths. - self.sweep_stale_req_infos() + # session.lock atomically saves peer info and snapshots tasks against send(). info: RecvReqInfo = RecvReqInfo.from_bytes(message[1]) - pre_cancelled = False with self._sessions_lock: session = self._get_session(info.unique_rid) - if session is None: - pre_cancelled = info.unique_rid in self._pre_cancelled_rids - if not pre_cancelled: - self._save_peer_req_info(info) - return - if pre_cancelled: - # CANCEL_SESSION may race ahead of REQUEST_DATA. Do not retain a - # request that can never acquire a source session; acknowledge it - # exactly once instead. - self._send_pre_cancel_failed_results(info) - return - with session.lock: - terminal = session.status in ( - SessionStatus.ERROR, - SessionStatus.CANCELLED, + fenced = ( + self._enforce_physical_ownership and info.unique_rid in self._pre_cancelled_rids ) - if terminal: - with self._peer_requests_lock: - existing = self._peer_requests.get(info.unique_rid, {}).get(info.instance_rank) - already_known = existing is not None and existing.slice_id == info.slice_id - tasks = list(session.kv_tasks) - aux_task_exists = session.aux_task is not None - else: + if session is None and not fenced: self._save_peer_req_info(info) - tasks = list(session.kv_tasks) - if session.aux_task is not None: - tasks.append(session.aux_task) - already_known = False - aux_task_exists = False - if terminal: - # Do not enqueue duplicate work into a cancelled/error session. - # A previously known request with an existing task already has a - # queued/active worker which will report its own terminal result; - # otherwise this late request needs a direct per-peer reply. - self._send_pre_cancel_failed_results( - info, - send_kv=not already_known or not tasks, - send_aux=(info.need_aux and (not already_known or not aux_task_exists)), - ) + return + if fenced: + self._queue_session_quiescence(info.unique_rid, info.instance_name, info.instance_rank) return + assert session is not None + with session.lock: + with self._sessions_lock: + if self._get_session(info.unique_rid) is not session or session._closed: + tasks = None + terminal = True + include_aux = bool(session._claim_unsubmitted_aux_failures_locked((info,))) + else: + self._save_peer_req_info(info) + tasks = list(session.kv_tasks) + terminal = session.has_failed() + include_aux = terminal and bool( + session._claim_unsubmitted_aux_failures_locked((info,)) + ) + if tasks is None: + self._send_failed_result_to_receiver(info, include_aux=include_aux) + return + if terminal: + self._send_failed_result_to_receiver(info, include_aux=include_aux) + return for task in tasks: - try: - if task._perf_timer is not None: - task._perf_timer.record_task_start(info.instance_rank) - trans_meta = ( - self._build_kv_write_meta(task, info) - if isinstance(task, KVSendTask) - else self._build_aux_write_meta(task, info) - ) - if task._perf_timer is not None: - task._perf_timer.record_push_start(trans_meta.peer_rank) - self._enqueue(trans_meta) - except Exception as exc: - task.fail(exc) - session.set_exception(str(exc)) - is_aux = isinstance(task, AuxSendTask) - self._send_failed_result_to_receiver( - info, slice_id=-1 if is_aux else task.slice_id, aux=is_aux - ) - raise - - def _cancel_failed_result_key( - self, info: RecvReqInfo, *, aux: bool - ) -> tuple[int, str, int, int, bool]: - slice_id = -1 if aux else (info.slice_id if info.slice_id is not None else 0) - return ( - info.unique_rid, - info.instance_name, - info.instance_rank, - slice_id, - aux, - ) - - def _send_pre_cancel_failed_result_once(self, info: RecvReqInfo, *, aux: bool) -> bool: - key = self._cancel_failed_result_key(info, aux=aux) - with self._peer_requests_lock: - if key in self._cancel_failed_result_keys: - return True - self._cancel_failed_result_keys.add(key) - sent = self._send_failed_result_to_receiver( - info, - slice_id=info.slice_id, - aux=aux, - ) - if not sent: - # A later duplicate CANCEL/REQUEST_DATA may retry a transport - # failure; successful sends remain idempotent. - with self._peer_requests_lock: - self._cancel_failed_result_keys.discard(key) - return sent - - def _send_pre_cancel_failed_results( - self, - info: RecvReqInfo, - *, - send_kv: bool = True, - send_aux: Optional[bool] = None, - ) -> bool: - if send_aux is None: - send_aux = info.need_aux - sent = True - if send_kv: - sent = self._send_pre_cancel_failed_result_once(info, aux=False) - if send_aux: - sent = self._send_pre_cancel_failed_result_once(info, aux=True) and sent - return sent + self._dispatch_task_to_peer(task, info) def _send_failed_result_to_receiver( self, info: RecvReqInfo, - slice_id: Optional[int] = None, - aux: bool = False, - ) -> bool: + *, + include_aux: bool = False, + defer_to_worker: bool = False, + ): try: - peer_ri = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) - if aux: - message = [ - MessageType.AUX_AGENT_RESULT, - str(self._instance_rank).encode("ascii"), - str(info.unique_rid).encode("ascii"), - AgentResult.FAILED.value.encode("ascii"), - ] - else: - resolved_slice_id = ( - slice_id - if slice_id is not None - else (info.slice_id if info.slice_id is not None else 0) - ) - message = _make_kv_result_msg( + endpoint = self._get_result_receiver_endpoint(info) + slice_id = info.slice_id if info.slice_id is not None else 0 + messages = [ + _make_kv_result_msg( self._instance_rank, info.unique_rid, - resolved_slice_id, + slice_id, True, # is_last_slice AgentResult.FAILED, ) - return self._send_agent_result(peer_ri.self_endpoint, message) + ] + if include_aux: + messages.append( + _make_aux_result_msg( + self._instance_rank, + info.unique_rid, + AgentResult.FAILED, + ) + ) + self._route_result_messages_to_receiver( + info, + endpoint, + messages, + defer_to_worker=defer_to_worker, + ) except Exception as e: logger.warning( f"_respond_with_kv: failed to abort receiver for rid={info.unique_rid}: {e}" ) - return False + + def _send_aux_result_to_receiver( + self, + info: RecvReqInfo, + result: AgentResult, + *, + defer_to_worker: bool = True, + ) -> None: + try: + endpoint = self._get_result_receiver_endpoint(info) + message = _make_aux_result_msg( + self._instance_rank, + info.unique_rid, + result, + ) + self._route_result_messages_to_receiver( + info, + endpoint, + [message], + defer_to_worker=defer_to_worker, + ) + except Exception as error: + logger.warning( + f"failed to report auxiliary transfer result for rid={info.unique_rid}: {error}" + ) + + def _send_failed_task_result_to_receiver(self, task: SendTaskBase, info: RecvReqInfo) -> None: + if not task.is_done: + task.fail(RuntimeError("transfer rejected before backend submission")) + try: + endpoint = self._get_result_receiver_endpoint(info) + if isinstance(task, KVSendTask): + receiver_slice_id = info.slice_id if info.slice_id is not None else 0 + message = _make_kv_result_msg( + self._instance_rank, + info.unique_rid, + receiver_slice_id, + True, + AgentResult.FAILED, + ) + else: + with self._sessions_lock: + session = self._get_session(info.unique_rid) + if session is not None and not session._claim_aux_terminal_result( + info.instance_rank + ): + return + message = _make_aux_result_msg( + self._instance_rank, + info.unique_rid, + AgentResult.FAILED, + ) + self._route_result_messages_to_receiver( + info, + endpoint, + [message], + defer_to_worker=True, + ) + except Exception as error: + logger.warning(f"failed to settle rejected transfer {info.unique_rid}: {error}") + + def _get_result_receiver_endpoint(self, info: RecvReqInfo) -> Optional[str]: + return self._registrar.get_peer_rank_info( + info.instance_name, info.instance_rank + ).self_endpoint + + def _route_result_messages_to_receiver( + self, + info: RecvReqInfo, + endpoint: Optional[str], + messages: Iterable[list[bytes]], + *, + defer_to_worker: bool, + ) -> None: + if defer_to_worker or self._enforce_physical_ownership: + # Listener-side rejection must share the worker's ordered stream; + # it cannot overtake an ambiguous write's pending IN_DOUBT report. + thread_idx = hash((info.unique_rid, info.instance_rank)) % self._num_threads + for message in messages: + self._send_task_queues[thread_idx].put((endpoint, message)) + return + dealer = self._get_or_connect_dealer(endpoint) + for message in messages: + dealer.send(message) def _get_or_connect_dealer(self, endpoint: Optional[str]): if endpoint is None: @@ -1737,65 +1909,52 @@ def send_cancel_to_receivers(self, unique_rid: int) -> None: peer_ri = self._registrar.get_peer_rank_info( req_info.instance_name, req_info.instance_rank ) - message = [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] - outbox = getattr(self, "_reliable_message_outbox", None) - accepted = ( - outbox.submit(peer_ri.self_endpoint, message) - if outbox is not None - else self._send_agent_result(peer_ri.self_endpoint, message) + self._get_or_connect_dealer(peer_ri.self_endpoint).send( + [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] ) - if not accepted: - logger.error(f"send_cancel_to_receivers: outbox rejected rid={unique_rid}") except Exception as e: logger.warning(f"send_cancel_to_receivers: failed for rid={unique_rid}: {e}") def shutdown(self): - if getattr(self, "_shutdown", False): - return - self._shutdown = True + watchdog = getattr(self, "_retirement_watchdog", None) + if watchdog is not None: + watchdog.request_shutdown() + watchdog.require_retired() + with self._sessions_lock: + if self._shutdown: + return + if self._enforce_physical_ownership and self._sessions: + raise RuntimeError("Sender refuses shutdown while transfer ownership is active") + if any(self._pending_settlements): + raise RuntimeError("Sender refuses shutdown while settlement reports are pending") + self._shutdown_requested = True - # Quiesce listener before invalidate to avoid set/map mutation races. - messenger = getattr(self, "_messenger", None) - if messenger is not None: - messenger.stop() + self._messenger.stop() - for q in getattr(self, "_send_task_queues", []): + for q in self._send_task_queues: q.put(None) - for t in getattr(self, "_worker_threads", []): + for t in self._worker_threads: t.join(timeout=5) - # Workers cannot enqueue any more terminal results after joining. Let - # the shared result/cancel outbox make one final attempt for every - # queued message, then stop. - reliable_outbox = getattr(self, "_reliable_message_outbox", None) - if reliable_outbox is not None: - reliable_outbox.shutdown() - # Snapshot under lock as defense in depth. - loaded_agents: list[str] = [] - loaded_remote_agents_lock = getattr(self, "_loaded_remote_agents_lock", None) - if loaded_remote_agents_lock is not None: - with loaded_remote_agents_lock: - loaded_remote_agents = getattr(self, "_loaded_remote_agents", set()) - loaded_agents = list(loaded_remote_agents) - loaded_remote_agents.clear() + with self._loaded_remote_agents_lock: + loaded_agents = list(self._loaded_remote_agents) + self._loaded_remote_agents.clear() # Invalidate all loaded remote agents to release fabric/POSIX FD resources. - agent = getattr(self, "_agent", None) for agent_name in loaded_agents: try: - if agent is not None: - agent.invalidate_remote_agent(agent_name) + self._agent.invalidate_remote_agent(agent_name) except Exception as e: logger.warning( f"Failed to invalidate remote agent '{agent_name}' during shutdown: {e}" ) - dealers = getattr(self, "_dealers", {}) - for dealer in dealers.values(): + for dealer in self._dealers.values(): try: dealer.stop() except Exception as e: logger.warning(f"Failed to stop dealer during Sender shutdown: {e}") - dealers.clear() + self._dealers.clear() + self._shutdown = True def __del__(self): try: @@ -1811,6 +1970,9 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): class TxSession(TxSessionBase): + cancelled_by_peer: bool = False + """Only meaningful once the session is CANCELLED.""" + def __init__( self, request_id: int, @@ -1819,34 +1981,43 @@ def __init__( aux_buffer: Optional[AuxBuffer] = None, timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, - beam_width: int = 1, overall_timeout_s: Optional[float] = None, + *, + retirement_watchdog: Optional[RetirementWatchdog] = None, ): - super().__init__( - sender, - SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), - ) + super().__init__(sender, SessionArgsBase(params, prompt_len=prompt_len)) self._timeout_s = timeout_s self._overall_timeout_s = overall_timeout_s self._deadline_monotonic_s: Optional[float] = None self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST or ( aux_buffer is not None and aux_buffer.has_draft_history ) + self._enforce_physical_ownership = getattr(sender, "_enforce_physical_ownership", False) self._sender: Sender # narrow base class type for Pylance self.request_id = request_id + self._retirement = None + if retirement_watchdog is not None: + if overall_timeout_s is None: + raise ValueError("deadline retirement requires the overall transfer timeout") + self._retirement = retirement_watchdog.create_owner( + request_id, "send", overall_timeout_s + ) self._aux_buffer = aux_buffer - self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None + self.aux_slot = None + try: + if aux_buffer is not None: + self.aux_slot = aux_buffer.alloc_slot().id + except Exception: + if self._retirement is not None: + self._retirement.close() + raise self.receiver_ready: bool = False self.kv_tasks = [] self.aux_task = None + self._reported_aux_peer_ranks: set[int] = set() + self._has_last_slice = False self.lock = threading.Lock() - # TaskStatus is a logical aggregate status and cannot represent - # concurrent writes to multiple peers: one failed peer may set ERROR - # while another peer still owns the registered source addresses. - # Track queued-or-active peer writes independently under ``self.lock`` - # so cancellation cannot unregister source addresses before a queued - # worker either finishes NIXL or reports FAILED to the receiver. - self._active_nixl_ops = 0 + self._logical_outcomes = _LogicalOutcomes(self._retirement) self._exception: Optional[Exception] = None self._closed = False @@ -1855,7 +2026,20 @@ def __init__( self.transfer_end_time = None # Must be last: makes session visible to listener thread, # so all attributes above must be initialized first. - self._sender.setup_session(self) + try: + self._sender.setup_session(self) + except Exception: + with self._sender._sessions_lock: + entry = self._sender._sessions.get(self.disagg_request_id) + if entry is self or (isinstance(entry, weakref.ReferenceType) and entry() is self): + self._sender._sessions.pop(self.disagg_request_id, None) + if self._retirement is not None: + self._retirement.close() + if self._aux_buffer is not None and self.aux_slot is not None: + self._aux_buffer.free_slot(self.aux_slot) + self.aux_slot = None + self._closed = True + raise @property def disagg_request_id(self) -> int: @@ -1871,31 +2055,40 @@ def disagg_request_id(self) -> int: @property def status(self) -> SessionStatus: + if getattr(self, "_retirement", None) is not None: + terminal = self._logical_outcomes.terminal + if terminal is not None: + return terminal.status if self._terminal_status is not None: return self._terminal_status - kv_all_transferred = bool(self.kv_tasks) and all( - t.status == TaskStatus.TRANSFERRED for t in self.kv_tasks + # A chunk may fail before the final task is created. + if self._exception is not None or any(t.status == TaskStatus.ERROR for t in self.kv_tasks): + return SessionStatus.ERROR + kv_all_transferred = ( + self._has_last_slice + and bool(self.kv_tasks) + and all(t.status == TaskStatus.TRANSFERRED for t in self.kv_tasks) + ) + aux_done = not self._need_aux or ( + self.aux_task is not None and self.aux_task.status == TaskStatus.TRANSFERRED ) - if kv_all_transferred: - if self.aux_task is not None and self.aux_task.status == TaskStatus.TRANSFERRED: - return SessionStatus.FULLY_TRANSFERRED - return SessionStatus.KV_TRANSFERRED - if self.kv_tasks and any(t.status == TaskStatus.TRANSFERRING for t in self.kv_tasks): + if kv_all_transferred and aux_done: + retirement = getattr(self, "_retirement", None) + if retirement is not None and not retirement.is_complete: + return SessionStatus.TRANSFERRING + return SessionStatus.TRANSFERRED + if kv_all_transferred or ( + self.kv_tasks and any(t.status == TaskStatus.TRANSFERRING for t in self.kv_tasks) + ): return SessionStatus.TRANSFERRING return SessionStatus.READY if self.receiver_ready else SessionStatus.INIT - def send(self, slice: KVSlice) -> None: + def send(self, chunk: Chunk) -> None: if self.transfer_start_time is None: self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() with self.lock: - if self._terminal_status in ( - SessionStatus.ERROR, - SessionStatus.CANCELLED, - ): - raise RuntimeError( - f"TxSession {self.disagg_request_id} is " - f"{self._terminal_status.value}; refusing new KV work" - ) + if self._retirement is not None and self._has_last_slice: + raise RuntimeError("cannot send another KV piece after the final slice") if not self.kv_tasks: overall_timeout_s = self._overall_timeout_s if overall_timeout_s is None or overall_timeout_s <= 0: @@ -1904,39 +2097,95 @@ def send(self, slice: KVSlice) -> None: params = self._base_args.params slice_id = len(self.kv_tasks) task = KVSendTask( - slice, + chunk, params, slice_id, prompt_len=self._base_args.prompt_len, ) task._unique_rid = self.disagg_request_id + task.bind_logical_outcomes(self._logical_outcomes) self.kv_tasks.append(task) + self._has_last_slice |= chunk.is_last + if self._retirement is not None and self._has_last_slice: + if not self._need_aux or self.aux_task is not None: + self._logical_outcomes.seal() req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) self._sender.dispatch_task(task, req_info_snapshot) def send_aux(self) -> AuxSendTask: + aux_failures: list[RecvReqInfo] = [] + terminal_error: Optional[Exception] = None + task: Optional[AuxSendTask] = None with self.lock: - if self._terminal_status in ( - SessionStatus.ERROR, - SessionStatus.CANCELLED, - ): - raise RuntimeError( - f"TxSession {self.disagg_request_id} is " - f"{self._terminal_status.value}; refusing new aux work" + if self._retirement is not None and self.aux_task is not None: + return self.aux_task + # In the ownership bridge's generation-first no-retry profile, + # prefill waits for every REQUEST_DATA, so this peer snapshot is sealed. + req_info_snapshot = dict(self._sender._get_req_info(self.disagg_request_id) or {}) + if self._closed or self._terminal_status is not None: + terminal_error = self._exception or RuntimeError( + f"TxSession {self.disagg_request_id} became terminal before auxiliary submission" ) - if not self._need_aux: - raise RuntimeError( - "Aux transfer requires GENERATION_FIRST scheduling so the " - "receiver reserves its destination buffer before dispatch" + aux_failures = self._claim_unsubmitted_aux_failures_locked( + req_info_snapshot.values() ) - params = self._base_args.params - task = AuxSendTask(params, self.aux_slot) - task._unique_rid = self.disagg_request_id - self.aux_task = task - req_info_snapshot = dict(self._sender._get_req_info(task._unique_rid) or {}) + else: + params = self._base_args.params + task = AuxSendTask(params, self.aux_slot) + task._unique_rid = self.disagg_request_id + task.bind_logical_outcomes(self._logical_outcomes) + self.aux_task = task + if self._retirement is not None and self._has_last_slice: + self._logical_outcomes.seal() + self._report_unsubmitted_aux_failures(aux_failures) + if terminal_error is not None: + raise terminal_error + assert task is not None self._sender.dispatch_task(task, req_info_snapshot) return task + def _claim_unsubmitted_aux_failures_locked( + self, req_infos: Iterable[RecvReqInfo] + ) -> list[RecvReqInfo]: + if not self._enforce_physical_ownership or not self._need_aux: + return [] + claimed = [] + for info in req_infos: + peer_rank = info.instance_rank + if self.aux_task is not None and self.aux_task.has_started_physical_operation( + peer_rank + ): + continue + if not self._claim_aux_terminal_result_locked(peer_rank): + continue + claimed.append(info) + return claimed + + def _claim_aux_terminal_result_locked(self, peer_rank: int) -> bool: + """Claim the only terminal AUX result report permitted for one peer. + + Terminal cleanup and worker completion both use this claim. Exactly one + path may publish terminal evidence so the receiver never sees + contradictory results for the same operation. The caller must hold + ``lock``. + """ + if peer_rank in self._reported_aux_peer_ranks: + return False + self._reported_aux_peer_ranks.add(peer_rank) + return True + + def _claim_aux_terminal_result(self, peer_rank: int) -> bool: + with self.lock: + return self._claim_aux_terminal_result_locked(peer_rank) + + def _claim_unsubmitted_aux_failure(self, info: RecvReqInfo) -> bool: + with self.lock: + return bool(self._claim_unsubmitted_aux_failures_locked((info,))) + + def _report_unsubmitted_aux_failures(self, req_infos: list[RecvReqInfo]) -> None: + for info in req_infos: + self._sender._send_aux_result_to_receiver(info, AgentResult.FAILED) + def pack_aux(self, request: LlmRequest) -> None: """Fill the aux buffer slot with token data from the given request.""" assert self._aux_buffer is not None, "No aux_buffer set for this session" @@ -1945,10 +2194,9 @@ def pack_aux(self, request: LlmRequest) -> None: def is_completed(self) -> bool: """Non-blocking check: has the transfer completed successfully?""" - status = self.status - if self._need_aux: - return status == SessionStatus.FULLY_TRANSFERRED - return status in (SessionStatus.KV_TRANSFERRED, SessionStatus.FULLY_TRANSFERRED) + return self.status is SessionStatus.TRANSFERRED and ( + not getattr(self, "_enforce_physical_ownership", False) or self.resources_drained() + ) def has_failed(self) -> bool: """Non-blocking check: has the transfer failed or been cancelled?""" @@ -1958,45 +2206,73 @@ def has_failed(self) -> bool: return True return self.aux_task is not None and self.aux_task.status == TaskStatus.ERROR - def cancel(self) -> None: - """Cancel the session and notify the remote receiver. + def resources_drained(self) -> bool: + if not getattr(self, "_enforce_physical_ownership", False): + return not any(task.status == TaskStatus.TRANSFERRING for task in self.kv_tasks) + retirement = getattr(self, "_retirement", None) + if retirement is not None and not retirement.can_retire(): + return False + tasks = self.kv_tasks + ([self.aux_task] if self.aux_task is not None else []) + return all(task.resources_drained for task in tasks) + + def _failed_wait_result(self) -> Optional[WaitResult]: + return ( + None + if getattr(self, "_enforce_physical_ownership", False) and not self.resources_drained() + else WaitResult.FAILED + ) + + def _completed_wait_result(self) -> Optional[WaitResult]: + if getattr(self, "_retirement", None) is not None and not self.is_completed(): + return None + return WaitResult.COMPLETED + + def cancel(self, by_peer: bool = False) -> bool: + """Cancel the session and notify the remote receiver. ``True`` if the ask went out. Safe to call multiple times. TRANSFERRING tasks keep running (mid-write). Only INIT tasks have their events signalled immediately. The lock serializes with _deliver_kv_to_agent() so has_transferring_tasks() is accurate the moment this returns. """ + if not self.cancel_local(by_peer=by_peer): + return False + self._sender.send_cancel_to_receivers(self.disagg_request_id) + return True + + def cancel_local(self, by_peer: bool = False, *, report_unsubmitted_aux: bool = True) -> bool: + aux_failures: list[RecvReqInfo] = [] + retirement = getattr(self, "_retirement", None) + if retirement is not None: + retirement.request_drain("peer cancellation" if by_peer else "cancellation") with self.lock: + # A later cancellation must still notify peers after logical failure. + # The logical arbiter preserves whichever outcome already committed. if self._terminal_status == SessionStatus.CANCELLED: - return + return False + self._logical_outcomes.cancel(by_peer) self._terminal_status = SessionStatus.CANCELLED + # Who asked is not recoverable later, and the two differ: the peer asking is a transfer + # error, our own side asking is an ordinary end. + self.cancelled_by_peer = by_peer exc = RuntimeError(f"TxSession {self.disagg_request_id} cancelled") for task in self.kv_tasks: if task.status == TaskStatus.INIT: task.fail(exc) if self.aux_task is not None and self.aux_task.status == TaskStatus.INIT: self.aux_task.fail(exc) - # Send outside the lock to avoid holding it during I/O. - self._sender.send_cancel_to_receivers(self.disagg_request_id) + req_infos = list((self._sender._get_req_info(self.disagg_request_id) or {}).values()) + if report_unsubmitted_aux: + aux_failures = self._claim_unsubmitted_aux_failures_locked(req_infos) + self._report_unsubmitted_aux_failures(aux_failures) + return True def has_transferring_tasks(self) -> bool: - """True while queued/active peer work still owns session buffers. + """True if any KV or auxiliary task is currently mid-write. cancel_request() must return False while this is True. """ - with self.lock: - return self._active_nixl_ops > 0 - - def _begin_nixl_op_locked(self) -> None: - """Reserve source-address lifetime; caller must hold ``self.lock``.""" - self._active_nixl_ops += 1 - - def _finish_nixl_op(self) -> None: - with self.lock: - if self._active_nixl_ops <= 0: - logger.error(f"TxSession {self.disagg_request_id} NIXL operation counter underflow") - return - self._active_nixl_ops -= 1 + return not self.resources_drained() def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: """Poll or block until KV (and optionally aux) transfer finishes. @@ -2010,14 +2286,14 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: or aux is not yet done. """ if self.has_failed(): - return WaitResult.FAILED + return self._failed_wait_result() if not self.kv_tasks: return None if not blocking: has_pending = False for task in self.kv_tasks: if task.status == TaskStatus.ERROR: - return WaitResult.FAILED + return self._failed_wait_result() if task.status != TaskStatus.TRANSFERRED: has_pending = True if has_pending: @@ -2026,10 +2302,12 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: if self.aux_task is None: return None if self.aux_task.status == TaskStatus.ERROR: - return WaitResult.FAILED + return self._failed_wait_result() if self.aux_task.status != TaskStatus.TRANSFERRED: return None - return WaitResult.COMPLETED + return ( + self._failed_wait_result() if self.has_failed() else self._completed_wait_result() + ) # send() normally anchors this once, at the first dispatched KV task. # Keep direct/internal TxSession construction bounded as well, and never @@ -2049,14 +2327,14 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]: if wait_slice_s is None or wait_slice_s <= 0: wait_slice_s = _FALLBACK_TX_WAIT_SLICE_S - def wait_for_task(task: SendTaskBase) -> WaitResult: + def wait_for_task(task: SendTaskBase) -> Optional[WaitResult]: while True: # A task/session terminal state observed at the deadline # boundary takes precedence over TIMEOUT. if self.has_failed(): - return WaitResult.FAILED + return self._failed_wait_result() if task.status == TaskStatus.TRANSFERRED: - return WaitResult.COMPLETED + return self._failed_wait_result() if self.has_failed() else WaitResult.COMPLETED remaining_s = None if self._deadline_monotonic_s is not None: @@ -2066,9 +2344,13 @@ def wait_for_task(task: SendTaskBase) -> WaitResult: # checks above and the clock read. Preserve boundary # precedence before classifying this as a timeout. if self.has_failed(): - return WaitResult.FAILED + return self._failed_wait_result() if task.status == TaskStatus.TRANSFERRED: - return WaitResult.COMPLETED + return ( + self._failed_wait_result() + if self.has_failed() + else WaitResult.COMPLETED + ) return WaitResult.TIMEOUT timeout_s = wait_slice_s if remaining_s is None else min(wait_slice_s, remaining_s) task.wait(timeout=timeout_s) @@ -2081,7 +2363,7 @@ def wait_for_task(task: SendTaskBase) -> WaitResult: if self._need_aux: if self.aux_task is None: # _finalize_send() installs the aux task synchronously before - # publishing the request to _send_reqs. Once every KV task is + # returning from respond_and_send_async(). Once every KV task is # terminal, a missing required aux task is an invariant error, # not an asynchronously pending transfer. with self.lock: @@ -2092,44 +2374,59 @@ def wait_for_task(task: SendTaskBase) -> WaitResult: self._exception = RuntimeError( "required auxiliary transfer was not dispatched" ) + self._logical_outcomes.fail(self._exception) self._terminal_status = SessionStatus.ERROR return WaitResult.FAILED result = wait_for_task(self.aux_task) if result != WaitResult.COMPLETED: return result return ( - WaitResult.FAILED + self._failed_wait_result() if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) - else WaitResult.COMPLETED + else self._completed_wait_result() ) - def set_exception(self, reason: str = ""): + def set_exception(self, reason: str = "") -> None: msg = f"TxSession {self.disagg_request_id} exception" if reason: msg += f": {reason}" + aux_failures: list[RecvReqInfo] = [] with self.lock: - self._exception = RuntimeError(msg) - self._terminal_status = SessionStatus.ERROR + if self._exception is None: + self._exception = RuntimeError(msg) + self._logical_outcomes.fail(self._exception) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR for task in self.kv_tasks: if not task.is_done: task.fail(self._exception) if self.aux_task is not None and not self.aux_task.is_done: self.aux_task.fail(self._exception) + req_infos = list((self._sender._get_req_info(self.disagg_request_id) or {}).values()) + aux_failures = self._claim_unsubmitted_aux_failures_locked(req_infos) + self._report_unsubmitted_aux_failures(aux_failures) @property def exception(self) -> Optional[Exception]: return self._exception - def close(self): - if getattr(self, "_closed", False): - return - self._closed = True + def close(self) -> bool: + with self.lock: + if self._closed: + return True + if getattr(self, "_enforce_physical_ownership", False) and not self.resources_drained(): + return False + retirement = getattr(self, "_retirement", None) + if retirement is not None and not retirement.close(): + return False + self._closed = True if self._aux_buffer is not None and self.aux_slot is not None: self._aux_buffer.free_slot(self.aux_slot) self.aux_slot = None # Unregister from Sender; keep fields alive for in-flight worker threads. if self._sender is not None: self._sender.clear_session(self.disagg_request_id) + return True def __enter__(self): return self @@ -2144,53 +2441,169 @@ def __del__(self): logger.warning(f"TxSession.__del__: exception during close: {e}") -class KVRecvTask: +class KVRecvTask(_LogicalTask): def __init__( self, unique_rid: Optional[int], - kv_slice: KVSlice, + chunk: Chunk, slice_id: int, params: DisaggregatedParams, aux_slot: Optional[int], + expected_write_bytes: Optional[int] = None, ): + super().__init__() self._event = threading.Event() self.slice_id = slice_id self.status = TaskStatus.INIT self.expected_transfers = 0 - self.last_slice_count = 0 - # Peers for which REQUEST_DATA was successfully dispatched but no - # KV_AGENT_RESULT has been received yet. This deliberately tracks - # actual dispatches rather than expected topology peers so a partial - # dispatch failure cannot leave a phantom active write forever. - self.pending_peer_ranks: set[int] = set() - self.resolved_peer_ranks: set[int] = set() - self.aux_reserved_peer_ranks: set[int] = set() - self.is_broadcast_dispatch = False + # One terminal result per writer rank, keyed by the rank the result frame carries. + self._writer_reports: dict[int, bool] = {} + # Verified-range admission accounting: the receiver-computed byte + # total the writers must cover for this piece, and the bytes the + # writers attested as actually submitted-and-completed. None means + # "unknown" and the task can never become write-verified. + self.expected_write_bytes = expected_write_bytes + self.verified_write_bytes = 0 self._unique_rid = unique_rid - self._kv_slice = kv_slice + self._chunk = chunk self._params = params self._exception: Optional[Exception] = None self._aux_slot = aux_slot self._perf_timer = PerfTimer() if perf_log_manager.enabled else None + self._physical_owner: Optional[_ReceiveOperationOwner] = None + self._ownership_state_lock: Optional[threading.Lock] = None def fail(self, exc: Exception) -> None: - self._exception = exc - self.status = TaskStatus.ERROR - self._event.set() + self._logical_outcomes.fail(exc) + if self._ownership_state_lock is None: + self._exception = exc + self.status = TaskStatus.ERROR + self._event.set() + return + with self._ownership_state_lock: + self._exception = exc + self.status = TaskStatus.ERROR + self._event.set() def complete(self) -> None: - self.status = TaskStatus.TRANSFERRED - self._event.set() + if self._ownership_state_lock is None: + self._logical_outcomes.complete(self._logical_index) + self.status = TaskStatus.TRANSFERRED + self._event.set() + return + with self._ownership_state_lock: + if self.status == TaskStatus.ERROR: + return + self._logical_outcomes.complete(self._logical_index) + self.status = TaskStatus.TRANSFERRED + self._event.set() def wait(self, timeout: Optional[float] = None) -> bool: """Block until terminal state. Returns True if done, False on timeout.""" return self._event.wait(timeout=timeout) + @property + def write_verified(self) -> bool: + """True when every published destination byte has attested write coverage. + + Strict equality: an over-count (more attested bytes than published) + indicates duplicated or misrouted writes and is as disqualifying as a + missing write. + """ + return ( + self.status == TaskStatus.TRANSFERRED + and self.expected_write_bytes is not None + and self.verified_write_bytes == self.expected_write_bytes + ) + @property def is_done(self) -> bool: return self._event.is_set() + def note_writer_report(self, peer_rank: int, succeeded: bool) -> tuple[bool, bool]: + """Record one writer's terminal word: whether this rank had yet to speak, and whether every + expected writer has now succeeded. + + The first answer is an edge, not a level. A rank keeps the first result it reports, and a + repeat is refused outright, so the caller settles the piece once rather than once per frame. + Callers hold ``RxSession.lock``, which serializes the writes. + """ + # TODO: a report from a rank outside the authorized writer set is counted, not rejected. + if peer_rank in self._writer_reports: + previous = self._writer_reports[peer_rank] + if previous != succeeded: + logger.warning( + f"KV writer {peer_rank} reported success={succeeded} for slice {self.slice_id} " + f"after reporting success={previous}; the first result stands" + ) + return False, False + self._writer_reports[peer_rank] = succeeded + all_succeeded = len(self._writer_reports) == self.expected_transfers and all( + self._writer_reports.values() + ) + return True, all_succeeded + + @property + def reports_outstanding(self) -> bool: + """Whether a writer's terminal word about this task is still missing. + + A logical outcome is not an answer to this: one writer's failure settles the task while + its siblings write on. The ownership owner tracks its own cohort and answers for itself. + """ + if self._physical_owner is not None: + return not self._physical_owner.all_writers_reported + return len(self._writer_reports) < self.expected_transfers + + def begin_publication(self) -> None: + if self._physical_owner is None: + self._physical_owner = _ReceiveOperationOwner(self._retirement) + self._ownership_state_lock = threading.Lock() + self._physical_owner.begin_publication() + + def _get_physical_owner(self) -> _ReceiveOperationOwner: + if self._physical_owner is None: + raise RuntimeError("receive operation owner was not initialized") + return self._physical_owner + + def seal_writer_cohort(self, writer_cohort: Optional[set[int]] = None) -> None: + self._get_physical_owner().seal_writer_cohort(self.expected_transfers, writer_cohort) + + def finish_publication(self) -> None: + self._get_physical_owner().finish_publication() + + def abort_publication(self, published_writers: set[int]) -> None: + self._get_physical_owner().abort_publication(published_writers) + + def cancel_unpublished(self) -> bool: + return self._get_physical_owner().cancel_unpublished() + + def record_writer_in_doubt(self, peer_rank: int) -> bool: + return self._get_physical_owner().record_writer_in_doubt(peer_rank) + + def record_writer_settlement(self, peer_rank: int) -> bool: + return self._get_physical_owner().record_writer_settlement(peer_rank) + + def record_writer_result( + self, + peer_rank: int, + succeeded: bool, + *, + wait_for_local_completion: bool, + ) -> tuple[bool, bool]: + return self._get_physical_owner().record_writer_result( + peer_rank, + succeeded, + wait_for_local_completion=wait_for_local_completion, + ) + + def finish_local_completion(self) -> None: + self._get_physical_owner().finish_local_completion() + + @property + def resources_drained(self) -> bool: + return self._get_physical_owner().resources_drained + def print_perf_info(self, peer_rank: int, instance_name: str, instance_rank: int): if self._perf_timer is None: return @@ -2210,10 +2623,16 @@ def __init__( peer_registrar: PeerRegistrar, agent: BaseTransferAgent, bounce=None, + enforce_physical_ownership: bool = False, + retirement_watchdog: Optional[RetirementWatchdog] = None, ): self._registrar = peer_registrar self._agent = agent self._bounce = bounce + # Internal component gate only. Production wiring is added after the + # complete Phase 1 ownership path is available and qualified. + self._enforce_physical_ownership = enforce_physical_ownership + self._retirement_watchdog = retirement_watchdog self._dealers = {} self._sender_ep_instance_map = {} # info_endpoint -> diagnostic message for peers that failed the @@ -2225,9 +2644,13 @@ def __init__( self._messenger = ZMQMessenger(mode="ROUTER") self._sessions = {} # unique_rid -> RxSession self._sessions_lock = threading.Lock() # Protects _sessions and _pre_cancelled_rids - self._pre_cancelled_rids: set[int] = set() + # unique_rid -> whether the peer asked, for cancels that arrive before the session exists. + # Who asked is not recoverable from the id alone, and the two differ: a peer's cancel is a + # transfer error, our own side's is an ordinary end. + self._pre_cancelled_rids: dict[int, bool] = {} self._shutdown = False - self._reliable_message_outbox = _ReliableMessageOutbox("Receiver") + self._ownership_admission_lock = threading.Lock() + self._ownership_poisoned: Optional[Exception] = None self._start_listener() logger.info(f"Receiver init with endpoint: {self._messenger.endpoint}") @@ -2237,49 +2660,96 @@ def endpoint(self): return self._messenger.endpoint def shutdown(self): - if getattr(self, "_shutdown", False): - return - self._shutdown = True - # Quiesce listener callbacks before putting the outbox sentinel behind - # all already-accepted CANCEL_SESSION messages. - messenger = getattr(self, "_messenger", None) - if messenger is not None: - messenger.stop() - reliable_outbox = getattr(self, "_reliable_message_outbox", None) - if reliable_outbox is not None: - reliable_outbox.shutdown() - dealers = getattr(self, "_dealers", {}) - for dealer in dealers.values(): + watchdog = getattr(self, "_retirement_watchdog", None) + if watchdog is not None: + watchdog.request_shutdown() + watchdog.require_retired() + with self._sessions_lock: + if self._shutdown: + return + if self._enforce_physical_ownership and self._sessions: + raise RuntimeError("Receiver refuses shutdown while transfer ownership is active") + self._shutdown = True + for dealer in self._dealers.values(): try: dealer.stop() except Exception as e: logger.warning(f"Failed to stop dealer during Receiver shutdown: {e}") - dealers.clear() + self._dealers.clear() + self._messenger.stop() def clear_session(self, unique_rid: int): with self._sessions_lock: self._sessions.pop(unique_rid, None) + def _get_ownership_admission_lock(self) -> threading.Lock: + lock = getattr(self, "_ownership_admission_lock", None) + if lock is None: + lock = threading.Lock() + self._ownership_admission_lock = lock + return lock + + def _acquire_ownership_admission(self) -> Callable[[], None]: + lock = self._get_ownership_admission_lock() + lock.acquire() + watchdog = getattr(self, "_retirement_watchdog", None) + if watchdog is not None and watchdog.fatal is not None: + lock.release() + raise RuntimeError("receiver retirement admission is closed") + error = getattr(self, "_ownership_poisoned", None) + if error is not None: + lock.release() + raise RuntimeError("receiver ownership admission is quarantined") from error + return lock.release + + def _poison_ownership_locked(self, error: Exception) -> None: + if getattr(self, "_ownership_poisoned", None) is None: + self._ownership_poisoned = error + def setup_session(self, rx_session: RxSessionBase): + release_admission = None + if getattr(rx_session, "_enforce_physical_ownership", False): + release_admission = self._acquire_ownership_admission() pre_cancel = False - with self._sessions_lock: - self._sessions[rx_session.disagg_request_id] = weakref.ref(rx_session) - if rx_session.disagg_request_id in self._pre_cancelled_rids: - pre_cancel = True - self._pre_cancelled_rids.discard(rx_session.disagg_request_id) + cancelled_by_peer = False + try: + with self._sessions_lock: + if self._shutdown: + raise RuntimeError("cannot create an RxSession after Receiver shutdown") + if getattr(rx_session, "_enforce_physical_ownership", False): + self._sessions[rx_session.disagg_request_id] = rx_session + else: + self._sessions[rx_session.disagg_request_id] = weakref.ref(rx_session) + if rx_session.disagg_request_id in self._pre_cancelled_rids: + pre_cancel = True + cancelled_by_peer = self._pre_cancelled_rids.pop(rx_session.disagg_request_id) + finally: + if release_admission is not None: + release_admission() if pre_cancel: - rx_session.cancel() + if getattr(rx_session, "_enforce_physical_ownership", False): + rx_session.cancel_local(by_peer=cancelled_by_peer) + else: + rx_session.cancel(by_peer=cancelled_by_peer) def _get_session(self, unique_rid: Optional[int]) -> Optional["RxSession"]: with self._sessions_lock: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: - return None - session = session_ref() - if session is None: - logger.warning(f"RxSession {unique_rid} has been garbage collected") + session_entry = self._sessions.get(unique_rid) + return self._resolve_session_entry(unique_rid, session_entry) + + @staticmethod + def _resolve_session_entry( + unique_rid: Optional[int], + session_entry: Optional[Union["RxSession", weakref.ReferenceType["RxSession"]]], + ) -> Optional["RxSession"]: + if session_entry is None: return None - return session + if isinstance(session_entry, weakref.ReferenceType): + session = session_entry() + if session is None: + logger.warning(f"RxSession {unique_rid} has been garbage collected") + return session + return session_entry def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: self_ri = self._registrar.self_rank_info @@ -2302,13 +2772,11 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: sender_req_id=sender_req_id, instance_name=self_ri.instance_name, instance_rank=self_ri.instance_rank, - block_ids_per_layer_groups=task._kv_slice.block_ids_per_layer_groups, + block_ids_per_layer_groups=task._chunk.block_ids_per_layer_groups, unique_rid=task._unique_rid, dst_start_token=None, aux_slot=task._aux_slot, - mamba_state_index=task._kv_slice.mamba_state_index, slice_id=task.slice_id, - need_aux=(task._params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST), ) @staticmethod @@ -2319,8 +2787,8 @@ def _fanin_bounce_safe( ) -> bool: """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. The split assumes every writer contributes the same size, which holds when: - * duplicate_head_factor == 1 -- else some ranks don't send KV (should_send_kv) yet still - count in expected_transfers, so the live writers overflow their slots; + * duplicate_head_factor == 1 -- else some ranks don't send KV yet still count in + expected_transfers, so the live writers overflow their slots; * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; for PP fan-in, every overlapping stage must hold the same number of layers (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage @@ -2340,19 +2808,6 @@ def _fanin_bounce_safe( # a masked PP stage may advertise no replicated view even though another # stage owns one that is visible in the receiver's page table. if len(overlap.ranks) > 1: - if peer_ri.page_table is not None: - for layer_group in peer_ri.page_table.layer_groups: - if not getattr(layer_group, "cp_as_tp", False): - continue - draft_tp, _ = MambaPolicy._mamba_tp(peer_ri) - if ( - layer_group.total_kv_head_num is None - or layer_group.kv_head_num_per_rank * draft_tp - > layer_group.total_kv_head_num - ): - # Draft head replicas may elect fewer writers than the - # target's heads, so their payload shares are unequal. - return False for page_table in (peer_ri.page_table, receiver_page_table): if page_table is None: continue @@ -2362,6 +2817,18 @@ def _fanin_bounce_safe( return False return True + def _get_mamba_slot(self, req_info: RecvReqInfo) -> Optional[int]: + """Extract the mamba slot index from block_ids_per_layer_groups, or None.""" + pt = self._registrar.self_extractor.page_table + if pt is None: + return None + for lg_idx, lg in enumerate(pt.layer_groups): + if lg.kind == CacheKind.STATE: + block_ids = req_info.block_ids_per_layer_groups[lg_idx] + if block_ids.size > 0: + return int(block_ids[0]) + return None + def dispatch_task(self, task: KVRecvTask) -> None: params = task._params logger.debug( @@ -2382,11 +2849,17 @@ def dispatch_task(self, task: KVRecvTask) -> None: # validation and may still run a transfer that is then # discarded. The task is already in the session's _kv_tasks, no # bounce reservation exists yet, and session._sender_endpoints - # is still empty, so no cleanup is needed here. + # is still empty, so the unpublished owner can close locally. logger.error( "dispatch_task: context peer incompatible, failing request " f"unique_rid={task._unique_rid}: {e}" ) + if self._enforce_physical_ownership: + session = self._get_session(task._unique_rid) + if session is None: + task.cancel_unpublished() + else: + session.cancel_unpublished_task(task) task.fail(e) return @@ -2418,7 +2891,6 @@ def dispatch_task(self, task: KVRecvTask) -> None: task.expected_transfers = len(peer_overlap.ranks) else: task.expected_transfers = len(dp0_overlap.ranks) - task.is_broadcast_dispatch = True # TP fan-in splits ONE region equally, so allow it only for a uniform writer set: # _fanin_bounce_safe() (TP-by-head / even-PP), and never under ADP broadcast (sender_dp_rank # None), where the real writer count exceeds expected_transfers and would overflow the slot. @@ -2436,21 +2908,18 @@ def dispatch_task(self, task: KVRecvTask) -> None: self._registrar.self_extractor.page_table, ) ) - # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks (the sender - # appends its MambaPolicy fragments in _build_kv_write_meta), so the bounce region must be - # sized for those bytes too or the write would overrun into the neighboring slot. + # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks, + # so the bounce region must be sized for those bytes too. extra_bytes = 0 - if receiver_req.mamba_state_index is not None: + mamba_dst_slot = self._get_mamba_slot(receiver_req) + if mamba_dst_slot is not None: if peer_infos.page_table is None: allow_bounce = False # cannot size the recurrent-state payload else: - # Receiver-local sizing: RankInfoServer runs on ctx rank 0 - # only, so a sender-side simulation under-reserves for every - # receiver whose paired sender is not rank 0. - extra_bytes = MambaPolicy.receiver_payload_bytes( + extra_bytes = mamba_receiver_payload_bytes( sender_page_table=peer_infos.page_table, receiver_page_table=self._registrar.self_extractor.page_table, - dst_slot=receiver_req.mamba_state_index, + dst_slot=mamba_dst_slot, ) bounced = allow_bounce and self._bounce.reserve( receiver_req, task.expected_transfers, extra_bytes=extra_bytes @@ -2461,39 +2930,65 @@ def dispatch_task(self, task: KVRecvTask) -> None: f"dispatch_task: RxSession {task._unique_rid} not found; " "session may have been closed before dispatch" ) + sender_endpoints = {peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks} # Fan-in: each sender gets its own sub-region base (writers must not overwrite); else serialize once. fanin_bounce = bounced and task.expected_transfers > 1 key = (receiver_req.unique_rid, receiver_req.slice_id) receiver_req_bytes = None if fanin_bounce else receiver_req.to_bytes() - successful_dispatches = 0 - for i, rank in enumerate(peer_overlap.ranks): - endpoint = peer_infos.sender_endpoints[rank] - if not session.begin_peer_transfer(task.slice_id, rank, endpoint): - exc = RuntimeError( - f"dispatch_task: RxSession {task._unique_rid} was cancelled " - f"before dispatch to peer_rank={rank}" - ) - task.fail(exc) - raise exc - if task._perf_timer is not None: - task._perf_timer.record_task_start(rank) + if not self._enforce_physical_ownership: + session.mark_transferring(task.slice_id) + # Cache sender endpoints so cancel() can send CANCEL_SESSION to them. + session._sender_endpoints.update(sender_endpoints) + for i, rank in enumerate(peer_overlap.ranks): + if task._perf_timer is not None: + task._perf_timer.record_task_start(rank) + if fanin_bounce: + receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) + receiver_req_bytes = receiver_req.to_bytes() + self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) + return + + peer_ranks = list(peer_overlap.ranks) + serialized_requests: list[tuple[int, bytes]] = [] + for i, rank in enumerate(peer_ranks): if fanin_bounce: receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) - receiver_req_bytes = receiver_req.to_bytes() - try: - self._request_sender_data(endpoint, receiver_req_bytes) - except Exception as exc: - # The synchronous send did not complete, so this peer cannot - # own the destination. Keep only earlier successful peer - # reservations alive for the async cancellation/drain path. - session.abort_peer_transfer(task.slice_id, rank) - task.fail(exc) - logger.error( - f"Receiver.dispatch_task failed for rid={task._unique_rid}, " - f"peer_rank={rank}, successful_dispatches={successful_dispatches}: {exc}" - ) - raise - successful_dispatches += 1 + payload = receiver_req.to_bytes() + else: + assert receiver_req_bytes is not None + payload = receiver_req_bytes + serialized_requests.append((rank, payload)) + + published_writers: set[int] = set() + + def publish_requests() -> None: + for rank, payload in serialized_requests: + if task._perf_timer is not None: + task._perf_timer.record_task_start(rank) + published_writers.add(rank) + try: + self._request_sender_data(peer_infos.sender_endpoints[rank], payload) + except Exception: + published_writers.discard(rank) + raise + + # Gen-first ADP publishes the destination to every eligible DP group, + # but the qualified immutable-request, no-retry/no-reroute profile + # admits exactly one group. Track that group's distinct terminal + # writers by count instead of sealing the full broadcast union. + writer_cohort = ( + set(peer_ranks) if sender_dp_rank is not None or peer_infos.dp_size == 1 else None + ) + if not session.try_begin_transfer( + task.slice_id, + sender_endpoints, + writer_cohort, + writer_candidates=set(peer_ranks), + publish=publish_requests, + published_writers=published_writers, + ): + self._bounce.release_idle_reservation(key) + task.cancel_unpublished() return @staticmethod @@ -2532,7 +3027,7 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: finally: messenger.stop() - # Recurrent-state (Mamba/KDA) layout gate on the receiver side. + # Peer compatibility gates on the receiver side. # The sender-side check (PeerRegistrar.register) runs in the # sender's listener thread, where exceptions are only logged, so # reject here — before REGISTER_RANK_INFO is even sent, so no @@ -2541,12 +3036,17 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: # (handled in dispatch_task) so only requests targeting this peer # fail, and cached so later requests fail fast. try: + self._registrar.validate_bounded_replay_compatible(sender_info) MambaPolicy.validate_peer_compatible( self._registrar.self_rank_info, sender_info, self._registrar.self_extractor.page_table, sender_info.page_table, ) + AttentionPolicy.validate_peer_compatible( + self._registrar.self_extractor.page_table, + sender_info.page_table, + ) except ValueError as e: msg = ( f"context peer at '{info_endpoint}' is incompatible: {e} " @@ -2569,18 +3069,25 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: else: return self._sender_ep_instance_map[info_endpoint] - def send_cancel_to_senders(self, unique_rid: int, sender_endpoints: set[str]) -> None: - """Notify all senders involved in this session to cancel.""" + def send_cancel_to_senders(self, unique_rid: int, sender_endpoints: set[str]) -> set[str]: + """Notify senders and return the endpoints whose cancel message was sent.""" + message = [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")] + if self._enforce_physical_ownership: + rank_info = self._registrar.self_rank_info + message.extend( + [ + rank_info.instance_name.encode(), + str(rank_info.instance_rank).encode("ascii"), + ] + ) + sent = set() for endpoint in sender_endpoints: try: - accepted = self._reliable_message_outbox.submit( - endpoint, - [MessageType.CANCEL_SESSION, str(unique_rid).encode("ascii")], - ) - if not accepted: - logger.error(f"send_cancel_to_senders: outbox rejected rid={unique_rid}") + self._get_or_connect_dealer(endpoint).send(message) + sent.add(endpoint) except Exception as e: logger.warning(f"send_cancel_to_senders: failed for rid={unique_rid}: {e}") + return sent def _start_listener(self): def handle_message(messages: list[bytes]) -> bool: @@ -2599,6 +3106,11 @@ def handle_message(messages: list[bytes]) -> bool: self._process_aux_agent_result(send_id, msg) except Exception as e: logger.error(f"Receiver: error handling AUX_AGENT_RESULT: {e}") + case MessageType.SESSION_QUIESCED: + try: + self._process_session_quiesced(msg) + except Exception as e: + logger.error(f"Receiver: error handling SESSION_QUIESCED: {e}") case MessageType.CANCEL_SESSION: try: self._handle_cancel_session(msg) @@ -2610,19 +3122,30 @@ def handle_message(messages: list[bytes]) -> bool: self._messenger.start_listener(handle_message) + def _process_session_quiesced(self, message: list[bytes]) -> None: + """Dispatch no-future-access proof only to the matching live request.""" + if not self._enforce_physical_ownership: + return + peer_rank, unique_rid = int(message[1]), int(message[2]) + session = self._get_session(unique_rid) + if session is not None: + session.process_session_quiesced(peer_rank) + def _handle_cancel_session(self, message: list[bytes]): unique_rid = int(message[1]) session = None with self._sessions_lock: - session_ref = self._sessions.get(unique_rid) - if session_ref is None: - self._pre_cancelled_rids.add(unique_rid) - else: - session = session_ref() - if session is None: - self._pre_cancelled_rids.add(unique_rid) + session_entry = self._sessions.get(unique_rid) + session = self._resolve_session_entry(unique_rid, session_entry) + if session is None: + # Parked with the peer as who asked: the session built afterwards reads its + # cancellation from here, and an id on its own would read as a local cancel. + self._pre_cancelled_rids[unique_rid] = True if session is not None: - session.cancel() + if self._enforce_physical_ownership: + session.cancel_local(by_peer=True) + else: + session.cancel(by_peer=True) def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): if message[0] != MessageType.KV_AGENT_RESULT: @@ -2630,9 +3153,14 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): f"_process_kv_agent_result: unexpected msg_type={message[0]!r}, expected KV_AGENT_RESULT" ) return - peer_rank, unique_rid, sender_slice_id, is_last_slice, status_code, transfer_size = ( - _KV_RESULT_PREFIX.unpack(message[1]) - ) + ( + peer_rank, + unique_rid, + receiver_slice_id, + is_last_slice, + status_code, + transfer_size, + ) = _KV_RESULT_PREFIX.unpack(message[1]) from .bounce import decode_result_tail dst_ptrs, sizes, src_base = decode_result_tail(message) @@ -2644,7 +3172,7 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): return session.process_kv_agent_result( peer_rank, - sender_slice_id, + receiver_slice_id, is_last_slice, _AGENT_RESULT_BY_CODE[status_code], dst_ptrs=dst_ptrs, @@ -2685,6 +3213,9 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): class RxSession(RxSessionBase): + cancelled_by_peer: bool = False + """Only meaningful once the session is CANCELLED.""" + def __init__( self, request_id: int, @@ -2693,20 +3224,31 @@ def __init__( aux_buffer: Optional[AuxBuffer] = None, timeout_s: Optional[float] = None, prompt_len: Optional[int] = None, - beam_width: int = 1, + *, + retirement_watchdog: Optional[RetirementWatchdog] = None, ): - super().__init__( - receiver, - SessionArgsBase(params, prompt_len=prompt_len, beam_width=beam_width), - ) + super().__init__(receiver, SessionArgsBase(params, prompt_len=prompt_len)) self._timeout_s = timeout_s self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST or ( aux_buffer is not None and aux_buffer.has_draft_history ) + self._enforce_physical_ownership = getattr(receiver, "_enforce_physical_ownership", False) self._receiver: Receiver # narrow base class type for Pylance self.request_id = request_id + self._retirement = None + if retirement_watchdog is not None: + if timeout_s is None: + raise ValueError("deadline retirement requires the transfer timeout") + self._retirement = retirement_watchdog.create_owner(request_id, "receive", timeout_s) self._aux_buffer = aux_buffer - self.aux_slot = aux_buffer.alloc_slot().id if aux_buffer is not None else None + self.aux_slot = None + try: + if aux_buffer is not None: + self.aux_slot = aux_buffer.alloc_slot().id + except Exception: + if self._retirement is not None: + self._retirement.close() + raise self._exception: Optional[Exception] = None self._closed = False self._terminal_status: Optional[SessionStatus] = None @@ -2716,18 +3258,45 @@ def __init__( self._kv_tasks: list[KVRecvTask] = [] self._aux_count = 0 self._aux_status: TaskStatus = TaskStatus.INIT - self._pending_aux_peer_ranks: set[int] = set() - self._resolved_aux_peer_ranks: set[int] = set() + self._publication_may_have_escaped = False + self._aux_physical_owner: Optional[_ReceiveOperationOwner] = None + if self._enforce_physical_ownership and self._need_aux: + self._aux_physical_owner = _ReceiveOperationOwner(self._retirement) + self._aux_physical_owner.begin_publication() self._sender_endpoints: set[str] = set() - # Error-triggered ADP cancellation is endpoint-idempotent. Sender - # notifications are claimed under ``lock`` and sent after releasing it. - self._failure_cancelled_sender_endpoints: set[str] = set() + self._cancelled_sender_endpoints: set[str] = set() + self._cancel_notification_started = False + # Serialize REQUEST_DATA publication with cancellation notification + # without holding the session state lock across a potentially blocking + # network send. The ordering is publication -> cancellation whenever + # publication wins the state transition. + self._publication_lock = threading.Lock() self.lock = threading.Lock() - # See TxSession._active_nixl_ops. On the receiver side each - # successfully dispatched peer request conservatively owns one - # destination operation until its agent-result message arrives. - self._active_nixl_ops = 0 - self._receiver.setup_session(self) + self._logical_outcomes = _LogicalOutcomes(self._retirement) + self._aux_logical_index = ( + self._logical_outcomes.add_task() + if self._retirement is not None and self._need_aux + else None + ) + try: + self._receiver.setup_session(self) + except Exception: + if self._aux_physical_owner is not None: + self._aux_physical_owner.cancel_unpublished() + if self._retirement is not None: + self._retirement.close() + with self._receiver._sessions_lock: + session_entry = self._receiver._sessions.get(self.disagg_request_id) + if ( + self._receiver._resolve_session_entry(self.disagg_request_id, session_entry) + is self + ): + self._receiver._sessions.pop(self.disagg_request_id, None) + if self._aux_buffer is not None and self.aux_slot is not None: + self._aux_buffer.free_slot(self.aux_slot) + self.aux_slot = None + self._closed = True + raise @property def disagg_request_id(self) -> int: @@ -2741,126 +3310,239 @@ def disagg_request_id(self) -> int: return params.ctx_request_id return self.request_id + @contextmanager + def _ownership_evidence_guard(self) -> Iterator[None]: + get_lock = getattr(self._receiver, "_get_ownership_admission_lock", None) + if not self._enforce_physical_ownership or get_lock is None: + yield + return + with get_lock(): + yield + + def _poison_receiver_ownership(self, error: Exception) -> None: + poison = getattr(self._receiver, "_poison_ownership_locked", None) + if poison is not None: + poison(error) + + def _record_ownership_evidence_error(self, error: Exception) -> None: + """Record a fatal ownership-evidence error and close receiver admission.""" + self._logical_outcomes.fail(error) + self._exception = error + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + self._poison_receiver_ownership(error) + @property def status(self) -> SessionStatus: + if getattr(self, "_retirement", None) is not None: + terminal = self._logical_outcomes.terminal + if terminal is not None: + return terminal.status if self._terminal_status is not None: return self._terminal_status if self._exception is not None or any(t.status == TaskStatus.ERROR for t in self._kv_tasks): return SessionStatus.ERROR if self._kv_tasks: kv_all_transferred = all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks) - if kv_all_transferred and self._aux_status == TaskStatus.TRANSFERRED: - return SessionStatus.FULLY_TRANSFERRED - if kv_all_transferred: - return SessionStatus.KV_TRANSFERRED - if any(t.status == TaskStatus.TRANSFERRING for t in self._kv_tasks): + aux_done = not self._need_aux or self._aux_status == TaskStatus.TRANSFERRED + if kv_all_transferred and aux_done: + retirement = getattr(self, "_retirement", None) + if retirement is not None and not retirement.is_complete: + return SessionStatus.TRANSFERRING + return SessionStatus.TRANSFERRED + if kv_all_transferred or any( + t.status == TaskStatus.TRANSFERRING for t in self._kv_tasks + ): return SessionStatus.TRANSFERRING return SessionStatus.INIT - def begin_peer_transfer(self, slice_id: int, peer_rank: int, sender_endpoint: str) -> bool: - """Reserve destination lifetime before dispatching REQUEST_DATA. - - Returns False when cancellation/error won the race, in which case the - caller must not contact the sender. - """ - with self.lock: - if self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED): - return False - task = self._kv_tasks[slice_id] - if peer_rank in task.pending_peer_ranks: - raise RuntimeError( - f"duplicate peer dispatch for request {self.request_id}, " - f"slice={slice_id}, peer_rank={peer_rank}" - ) - task.pending_peer_ranks.add(peer_rank) - reserve_aux = ( - self._need_aux - and peer_rank not in self._pending_aux_peer_ranks - and peer_rank not in self._resolved_aux_peer_ranks - ) - if reserve_aux: - self._pending_aux_peer_ranks.add(peer_rank) - task.aux_reserved_peer_ranks.add(peer_rank) - task.status = TaskStatus.TRANSFERRING - self._sender_endpoints.add(sender_endpoint) - self._active_nixl_ops += 1 + int(reserve_aux) + def try_begin_transfer( + self, + slice_id: int, + sender_endpoints: set[str], + writer_cohort: Optional[set[int]] = None, + *, + writer_candidates: Optional[set[int]] = None, + publish: Optional[Callable[[], None]] = None, + published_writers: Optional[set[int]] = None, + ) -> bool: + """Seal and publish a writer cohort unless cancellation won first.""" + if publish is not None and published_writers is None: + raise ValueError("publication must track successfully queued writers") + with self._publication_lock: + # Only metadata and publication reservation use the admission gate. + # Network I/O must not hold the arbiter needed by the watchdog. + retirement = getattr(self, "_retirement", None) + with ( + self._ownership_evidence_guard(), + self.lock, + retirement.lock if retirement is not None else nullcontext(), + ): + if self._closed or self._terminal_status is not None: + return False + error = getattr(self._receiver, "_ownership_poisoned", None) + if error is not None: + raise RuntimeError("receiver ownership admission is quarantined") from error + task = self._kv_tasks[slice_id] + owners = [task._get_physical_owner()] + if self._aux_physical_owner is not None: + owners.append(self._aux_physical_owner) + unpublished = [owner for owner in owners if owner._expected_writers is None] + try: + for owner in owners: + owner.seal_writer_cohort( + task.expected_transfers, + writer_cohort, + published_writers=writer_candidates, + ) + if retirement is not None and not retirement.expose(*owners): + raise RuntimeError("destination retirement admission is closed") + except Exception: + # No send has run for this reservation. Roll back only new + # owners; AUX may already belong to an earlier published slice. + for owner in unpublished: + owner.abort_publication(set()) + raise + task.status = TaskStatus.TRANSFERRING + self._sender_endpoints.update(sender_endpoints) + if publish is not None: + assert published_writers is not None + self._publication_may_have_escaped = True + try: + publish() + expected_publication = ( + writer_candidates if writer_candidates is not None else writer_cohort + ) + if ( + expected_publication is not None + and published_writers != expected_publication + ): + raise RuntimeError("publication did not queue the complete writer cohort") + except Exception: + # ZMQ delivers multipart messages atomically. A successful + # send transfers responsibility to that writer; a failed + # send does not. Retain only the successfully queued prefix + # and wait for each of those writers to settle. + self._receiver._bounce.abort_publication( + (self.disagg_request_id, task.slice_id), + published_writers, + ) + task.abort_publication(published_writers) + if self._aux_physical_owner is not None: + self._aux_physical_owner.abort_publication(published_writers) + raise + task.finish_publication() + if self._aux_physical_owner is not None: + self._aux_physical_owner.finish_publication() return True - def abort_peer_transfer(self, slice_id: int, peer_rank: int) -> None: - """Undo a reservation whose REQUEST_DATA send raised synchronously.""" - with self.lock: - task = self._kv_tasks[slice_id] - if peer_rank not in task.pending_peer_ranks: - return - task.pending_peer_ranks.remove(peer_rank) - op_count = 1 - if peer_rank in task.aux_reserved_peer_ranks: - task.aux_reserved_peer_ranks.remove(peer_rank) - if peer_rank in self._pending_aux_peer_ranks: - self._pending_aux_peer_ranks.remove(peer_rank) - op_count += 1 - if self._active_nixl_ops < op_count: - logger.error(f"RxSession {self.disagg_request_id} NIXL operation counter underflow") - return - self._active_nixl_ops -= op_count - - def _discard_broadcast_kv_phantoms_locked(self, task: KVRecvTask) -> None: - """Drop non-selected ADP broadcast peers after all real responders reply.""" - if ( - not task.is_broadcast_dispatch - or self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED) - or len(task.resolved_peer_ranks) < task.expected_transfers - ): - return - phantom_count = len(task.pending_peer_ranks) - if phantom_count == 0: + def mark_transferring(self, slice_id: int, writer_cohort: Optional[set[int]] = None) -> None: + if not self._enforce_physical_ownership: + with self.lock: + self._kv_tasks[slice_id].status = TaskStatus.TRANSFERRING return - task.pending_peer_ranks.clear() - self._active_nixl_ops = max(0, self._active_nixl_ops - phantom_count) + if writer_cohort is None: + raise ValueError("ownership-mode publication requires an explicit writer cohort") + if not self.try_begin_transfer(slice_id, set(), writer_cohort): + raise RuntimeError( + f"RxSession {self.disagg_request_id} became terminal before publication" + ) - def _discard_broadcast_aux_phantoms_locked(self, task: KVRecvTask) -> None: - if ( - not task.is_broadcast_dispatch - or self._terminal_status in (SessionStatus.ERROR, SessionStatus.CANCELLED) - or len(self._resolved_aux_peer_ranks) < task.expected_transfers - ): - return - phantom_count = len(self._pending_aux_peer_ranks) - if phantom_count == 0: - return - phantom_ranks = set(self._pending_aux_peer_ranks) - self._pending_aux_peer_ranks.clear() - for kv_task in self._kv_tasks: - kv_task.aux_reserved_peer_ranks.difference_update(phantom_ranks) - self._active_nixl_ops = max(0, self._active_nixl_ops - phantom_count) - - def _claim_broadcast_failure_cancel_locked(self, task: KVRecvTask) -> set[str]: - """Claim sender endpoints that still need ERROR-path cancellation.""" - if not task.is_broadcast_dispatch or self._terminal_status != SessionStatus.ERROR: - return set() - endpoints = self._sender_endpoints.difference(self._failure_cancelled_sender_endpoints) - self._failure_cancelled_sender_endpoints.update(endpoints) - return endpoints - - def receive(self, slice: KVSlice) -> None: + def receive(self, chunk: Chunk, expected_write_bytes: Optional[int] = None) -> None: if self.transfer_start_time is None: self.transfer_start_time = tensorrt_llm.bindings.global_steady_clock_now() + if self._enforce_physical_ownership: + task = self.prepare_receive(chunk, expected_write_bytes=expected_write_bytes) + if task is not None: + self.dispatch_prepared_receive(task) + return params = self._base_args.params slice_id = len(self._kv_tasks) task = KVRecvTask( self.disagg_request_id, - slice, + chunk, slice_id, params, aux_slot=self.aux_slot, + expected_write_bytes=expected_write_bytes, ) + task.bind_logical_outcomes(self._logical_outcomes) self._kv_tasks.append(task) self._receiver.dispatch_task(task) + def prepare_receive( + self, chunk: Chunk, expected_write_bytes: Optional[int] = None + ) -> Optional[KVRecvTask]: + """Create an unpublished task under the cancellation linearization lock.""" + with self.lock: + if self._closed or self._terminal_status is not None: + return None + params = self._base_args.params + task = KVRecvTask( + self.disagg_request_id, + chunk, + len(self._kv_tasks), + params, + aux_slot=self.aux_slot, + expected_write_bytes=expected_write_bytes, + ) + task.bind_logical_outcomes(self._logical_outcomes) + task.begin_publication() + self._kv_tasks.append(task) + if self._retirement is not None and chunk.is_last: + self._logical_outcomes.seal() + return task + + def kv_write_verified(self) -> bool: + """True when every KV task's published destination bytes are attested written. + + This is the admission predicate for verified-range KV block reuse: a + session whose writers merely reported SUCCESS without covering the + published byte range (skipped, partial, or misrouted writes) is not + verified, and the received range must not be committed to the reuse + tree. An empty session (nothing received) is never verified. + """ + with self.lock: + if not self._kv_tasks: + return False + return all(task.write_verified for task in self._kv_tasks) + + def dispatch_prepared_receive(self, task: KVRecvTask) -> None: + try: + self._receiver.dispatch_task(task) + except Exception as error: + if self.cancel_unpublished_task(task): + self._receiver._bounce.release_idle_reservation( + (self.disagg_request_id, task.slice_id) + ) + self.fail_admission(error) + raise + + def fail_admission(self, error: Exception) -> None: + """Fail logical admission without releasing possibly published destinations.""" + with self.lock: + self._logical_outcomes.fail(error) + self._exception = error + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + for task in self._kv_tasks: + if task.status == TaskStatus.INIT: + if self._enforce_physical_ownership: + task.cancel_unpublished() + task.fail(error) + if self._aux_physical_owner is not None and not self._publication_may_have_escaped: + self._aux_physical_owner.cancel_unpublished() + + def cancel_unpublished_task(self, task: KVRecvTask) -> bool: + cancelled = task.cancel_unpublished() + if self._aux_physical_owner is not None: + self._aux_physical_owner.cancel_unpublished() + return cancelled + def process_kv_agent_result( self, peer_rank: int, - sender_slice_id: int, + receiver_slice_id: int, is_last_slice: bool, status: AgentResult, dst_ptrs=None, @@ -2868,167 +3550,255 @@ def process_kv_agent_result( src_base=None, transfer_size: int = 0, ): - cancel_endpoints: set[str] = set() - with self.lock: - self.kv_cache_size_bytes += transfer_size - assert sender_slice_id < len(self._kv_tasks), ( - f"Receiver got slice_id={sender_slice_id} from sender but only has " + with self._ownership_evidence_guard(), self.lock: + if self._retirement is not None and self._closed: + return + assert receiver_slice_id < len(self._kv_tasks), ( + f"Receiver got receiver_slice_id={receiver_slice_id} but only has " f"{len(self._kv_tasks)} receive task(s) for request {self.request_id}. " - f"Sender/receiver slice count mismatch." + f"The sender echoes this index back from RecvReqInfo, so it must " + f"address a task this receiver posted." ) - task = self._kv_tasks[sender_slice_id] - if peer_rank not in task.pending_peer_ranks: - kind = "duplicate" if peer_rank in task.resolved_peer_ranks else "unexpected" - logger.warning( - f"RxSession {self.request_id} ignoring {kind} KV result for " - f"slice={sender_slice_id}, peer_rank={peer_rank}" + task = self._kv_tasks[receiver_slice_id] + if status not in ( + AgentResult.SUCCESS, + AgentResult.FAILED, + AgentResult.IN_DOUBT, + AgentResult.FAILED_QUIESCED, + ): + raise ValueError( + f"Session {self.request_id} received unknown task status: {status.value}" ) - return - try: - if status == AgentResult.SUCCESS: - from .bounce import scatter_write_result - - on_done = None - if is_last_slice: - task.last_slice_count += 1 - if task.last_slice_count == task.expected_transfers: - request_id = self.request_id - - def on_done( - success, - task=task, - peer_rank=peer_rank, - sender_slice_id=sender_slice_id, - request_id=request_id, - ): - if not success: - task.fail( - RuntimeError( - f"KV bounce scatter failed for request {request_id} " - f"slice={sender_slice_id}" - ) - ) - return - if task.status == TaskStatus.ERROR: - return - task.complete() - # Transfer end for perf/time-sync is valid only - # after the bounce data has landed in the real KV pool. - if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): - self.transfer_end_time = ( - tensorrt_llm.bindings.global_steady_clock_now() - ) - logger.debug( - f"KV transfer complete for request {request_id} " - f"slice={sender_slice_id}" - ) - if task._perf_timer is not None: - task._perf_timer.record_task_end(peer_rank) - ri = self._receiver._registrar.self_rank_info - task.print_perf_info(peer_rank, ri.instance_name, ri.instance_rank) - - scatter_write_result( - self._receiver._bounce, - (self.disagg_request_id, task.slice_id), - peer_rank, - dst_ptrs, - sizes, - src_base, - on_done, + if status == AgentResult.IN_DOUBT: + if not self._enforce_physical_ownership: + raise RuntimeError( + "received IN_DOUBT evidence without physical ownership enabled" ) - elif status == AgentResult.FAILED: - detail = ( - f"KV transfer failed for request {self.request_id} slice={sender_slice_id} " - f"peer_rank={peer_rank} is_last_slice={is_last_slice} " - f"(reported by remote agent; see sender-side log for nixl_status)" + try: + accepted = task.record_writer_in_doubt(peer_rank) + except Exception as error: + task.fail(error) + self._record_ownership_evidence_error(error) + raise + if not accepted: + return + error = RuntimeError( + f"KV transfer quiescence is unproven for request {self.request_id} " + f"slice={receiver_slice_id} peer_rank={peer_rank}" + ) + task.fail(error) + self._record_ownership_evidence_error(error) + return + if status == AgentResult.FAILED_QUIESCED: + if not self._enforce_physical_ownership: + raise RuntimeError("received physical settlement without ownership enabled") + try: + accepted = task.record_writer_settlement(peer_rank) + except Exception as error: + self._record_ownership_evidence_error(error) + raise + if accepted: + self._receiver._bounce.record_failure( + (self.disagg_request_id, task.slice_id), peer_rank ) - logger.error(detail) - task.fail(RuntimeError(detail)) - if self._terminal_status is None: # Don't overwrite CANCELLED with ERROR - self._terminal_status = SessionStatus.ERROR + return + if self._enforce_physical_ownership: + if status == AgentResult.FAILED or is_last_slice: + try: + accepted, all_succeeded = task.record_writer_result( + peer_rank, + status == AgentResult.SUCCESS, + wait_for_local_completion=(status == AgentResult.SUCCESS), + ) + except Exception as error: + task.fail(error) + self._record_ownership_evidence_error(error) + raise + if not accepted: + return else: - raise ValueError( - f"Session {self.request_id} received unknown task status: {status.value}" - ) - finally: - # A result is sent only after the remote submit/wait path has - # stopped using our destination. Resolve this peer exactly - # once even when result processing itself raises. - task.pending_peer_ranks.remove(peer_rank) - task.resolved_peer_ranks.add(peer_rank) - if self._active_nixl_ops <= 0: - logger.error( - f"RxSession {self.disagg_request_id} NIXL operation counter underflow" + all_succeeded = False + else: + all_succeeded = False + # A failure and a last-slice success are both one writer's terminal word. + if status == AgentResult.FAILED or is_last_slice: + accepted, all_succeeded = task.note_writer_report( + peer_rank, status == AgentResult.SUCCESS ) - else: - self._active_nixl_ops -= 1 - self._discard_broadcast_kv_phantoms_locked(task) - cancel_endpoints = self._claim_broadcast_failure_cancel_locked(task) - if cancel_endpoints: - # Never perform messaging I/O while holding the session lock. - self._receiver.send_cancel_to_senders(self.disagg_request_id, cancel_endpoints) + if not accepted: + return + self.kv_cache_size_bytes += transfer_size + if status == AgentResult.SUCCESS: + # Verified-range accounting: transfer_size is the sender's + # attestation of destination bytes actually submitted and + # completed for this chunk (0 unless the write completed). + task.verified_write_bytes += transfer_size + + from .bounce import scatter_write_result + + on_done = None + if is_last_slice and all_succeeded: + # Completing message: defer task.complete()+perf until the scatter has actually + # landed. scatter_write_result fires this inline for the non-bounced path, or on + # the scatter worker (after cudaStreamSynchronize) for the bounced path, so the + # gen consumer never observes completion before the KV is scattered into place. + request_id = self.request_id + ri = self._receiver._registrar.self_rank_info + instance_name, instance_rank = ri.instance_name, ri.instance_rank + + def on_done( + success, + task=task, + peer_rank=peer_rank, + receiver_slice_id=receiver_slice_id, + request_id=request_id, + instance_name=instance_name, + instance_rank=instance_rank, + ): + # Runs on the scatter worker thread for the bounced path. Do not acquire + # RxSession.lock here: the non-bounced path invokes this callback inline + # while already holding it. Task outcome/ownership locks are independent. + # complete() sets status before _event for wait_complete's status-first poll. + if self._enforce_physical_ownership: + task.finish_local_completion() + if not success: + task.fail( + RuntimeError( + f"KV bounce scatter failed for request {request_id} " + f"slice={receiver_slice_id}" + ) + ) + return + if task.status == TaskStatus.ERROR: + return # a concurrent FAILED writer already failed it; don't un-fail + task.complete() + # Record completion before best-effort diagnostics can block the worker. + if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): + self.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() + try: + if task._perf_timer is not None: + task._perf_timer.record_task_end(peer_rank) + task.print_perf_info(peer_rank, instance_name, instance_rank) + except Exception as e: # perf is best-effort; never block completion + logger.warning( + f"KV transfer perf logging failed for request {request_id} " + f"slice={receiver_slice_id}: {e}" + ) + logger.debug( + f"KV transfer complete for request {request_id} " + f"slice={receiver_slice_id}" + ) + + scatter_write_result( + self._receiver._bounce, + (self.disagg_request_id, task.slice_id), + peer_rank, + dst_ptrs, + sizes, + src_base, + on_done, + ) + elif status == AgentResult.FAILED: + detail = ( + f"KV transfer failed for request {self.request_id} " + f"slice={receiver_slice_id} " + f"peer_rank={peer_rank} is_last_slice={is_last_slice} " + f"(reported by remote agent; see sender-side log for nixl_status)" + ) + logger.error(detail) + # Drain-before-release: record this writer FAILED; the owner frees the shared region + # only once every fan-in writer is terminal (freeing now could race a sibling's RMA). + self._receiver._bounce.record_failure( + (self.disagg_request_id, task.slice_id), peer_rank + ) + task.fail(RuntimeError(detail)) + if self._terminal_status is None: # Don't overwrite CANCELLED with ERROR + self._terminal_status = SessionStatus.ERROR def process_aux_agent_result(self, peer_rank: int, status: AgentResult): # Aux is session-level (not per-slice); expected_transfers is identical # across all kv_tasks, so any task provides the right count. - cancel_endpoints: set[str] = set() - with self.lock: - if peer_rank not in self._pending_aux_peer_ranks: - kind = "duplicate" if peer_rank in self._resolved_aux_peer_ranks else "unexpected" - logger.warning( - f"RxSession {self.request_id} ignoring {kind} aux result " - f"from peer_rank={peer_rank}" - ) + with self._ownership_evidence_guard(), self.lock: + if self._retirement is not None and self._closed: return if not self._kv_tasks: logger.warning( f"Aux result received before any KV tasks for request {self.request_id}" ) - self._pending_aux_peer_ranks.remove(peer_rank) - self._resolved_aux_peer_ranks.add(peer_rank) - if self._active_nixl_ops > 0: - self._active_nixl_ops -= 1 return task = self._kv_tasks[0] - try: - if status == AgentResult.SUCCESS: - self._aux_count += 1 - - if self._aux_count == task.expected_transfers: - self._aux_status = TaskStatus.TRANSFERRED - elif self._aux_count > task.expected_transfers: - self._aux_status = TaskStatus.ERROR - self._exception = RuntimeError( - f"Session {self.request_id} received too many aux transfers" - ) - if self._terminal_status is None: - self._terminal_status = SessionStatus.ERROR - logger.error(str(self._exception)) - elif status == AgentResult.FAILED: + all_succeeded = False + if status == AgentResult.IN_DOUBT: + if self._aux_physical_owner is None: + raise RuntimeError( + "received auxiliary IN_DOUBT evidence without physical ownership enabled" + ) + try: + accepted = self._aux_physical_owner.record_writer_in_doubt(peer_rank) + except Exception as error: self._aux_status = TaskStatus.ERROR - self._exception = RuntimeError(f"Session {self.request_id} aux transfer failed") - if self._terminal_status is None: - self._terminal_status = SessionStatus.ERROR - else: - raise ValueError( - f"Session {self.request_id} received unknown aux send status: {status}" + self._record_ownership_evidence_error(error) + raise + if not accepted: + return + error = RuntimeError( + f"auxiliary transfer quiescence is unproven for request " + f"{self.request_id} peer_rank={peer_rank}" + ) + self._aux_status = TaskStatus.ERROR + self._record_ownership_evidence_error(error) + return + if status == AgentResult.FAILED_QUIESCED: + if self._aux_physical_owner is None: + raise RuntimeError("received auxiliary settlement without ownership enabled") + try: + self._aux_physical_owner.record_writer_settlement(peer_rank) + except Exception as error: + self._record_ownership_evidence_error(error) + raise + return + if self._aux_physical_owner is not None: + try: + accepted, all_succeeded = self._aux_physical_owner.record_writer_result( + peer_rank, + status == AgentResult.SUCCESS, + wait_for_local_completion=False, ) - finally: - self._pending_aux_peer_ranks.remove(peer_rank) - self._resolved_aux_peer_ranks.add(peer_rank) - for kv_task in self._kv_tasks: - kv_task.aux_reserved_peer_ranks.discard(peer_rank) - if self._active_nixl_ops <= 0: - logger.error( - f"RxSession {self.disagg_request_id} NIXL operation counter underflow" + except Exception as error: + self._aux_status = TaskStatus.ERROR + self._record_ownership_evidence_error(error) + raise + if not accepted: + return + if status == AgentResult.SUCCESS: + self._aux_count += 1 + + if (self._aux_physical_owner is not None and all_succeeded) or ( + self._aux_physical_owner is None and self._aux_count == task.expected_transfers + ): + self._aux_status = TaskStatus.TRANSFERRED + aux_logical_index = getattr(self, "_aux_logical_index", None) + if aux_logical_index is not None: + self._logical_outcomes.complete(aux_logical_index) + elif self._aux_count > task.expected_transfers: + self._aux_status = TaskStatus.ERROR + self._exception = RuntimeError( + f"Session {self.request_id} received too many aux transfers" ) - else: - self._active_nixl_ops -= 1 - self._discard_broadcast_aux_phantoms_locked(task) - cancel_endpoints = self._claim_broadcast_failure_cancel_locked(task) - if cancel_endpoints: - # Never perform messaging I/O while holding the session lock. - self._receiver.send_cancel_to_senders(self.disagg_request_id, cancel_endpoints) + self._logical_outcomes.fail(self._exception) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + logger.error(str(self._exception)) + elif status == AgentResult.FAILED: + self._aux_status = TaskStatus.ERROR + self._exception = RuntimeError(f"Session {self.request_id} aux transfer failed") + self._logical_outcomes.fail(self._exception) + if self._terminal_status is None: + self._terminal_status = SessionStatus.ERROR + else: + raise ValueError( + f"Session {self.request_id} received unknown aux send status: {status}" + ) @property def exception(self) -> Optional[Exception]: @@ -3055,58 +3825,131 @@ def unpack_aux(self, request: LlmRequest) -> None: }, } - def unpack_draft_history(self, request: LlmRequest) -> None: - """Read standalone history without changing context-first token fields.""" - if self._aux_buffer is None or self.aux_slot is None: - raise ValueError("Standalone draft transfer requires an auxiliary buffer slot") - request.py_draft_transfer_history = self._aux_buffer.get_slot_draft_history(self.aux_slot) - def is_completed(self) -> bool: - """Non-blocking check: has the transfer completed successfully?""" - status = self.status - if self._need_aux: - return status == SessionStatus.FULLY_TRANSFERRED - return status in (SessionStatus.KV_TRANSFERRED, SessionStatus.FULLY_TRANSFERRED) + """Non-blocking check: has the transfer completed successfully? + + One read of ``status`` answers both halves at one instant. Asking the KV handles and the + auxiliary handle separately would compose two instants, and retirement feeds a cross-rank + consensus that must not see a combination which never held. + """ + return self.status is SessionStatus.TRANSFERRED and ( + not self._enforce_physical_ownership or self.resources_drained() + ) def has_failed(self) -> bool: """Non-blocking check: has the transfer failed or been cancelled?""" return self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED) - def cancel(self) -> None: - """Cancel the session and notify the remote sender. - - Safe to call multiple times. TRANSFERRING tasks keep running (mid-write). - Only INIT tasks have their events signalled immediately. - The lock serializes with process_kv_agent_result() / process_aux_agent_result(). - """ + def resources_drained(self) -> bool: + if not self._enforce_physical_ownership: + return not any(task.status == TaskStatus.TRANSFERRING for task in self._kv_tasks) + retirement = getattr(self, "_retirement", None) + if retirement is not None and not retirement.can_retire(): + return False + kv_drained = all(task.resources_drained for task in self._kv_tasks) + aux_drained = self._aux_physical_owner is None or self._aux_physical_owner.resources_drained + return kv_drained and aux_drained + + def cancel_local(self, by_peer: bool = False) -> bool: + """Commit cancellation under the same lock used for publication.""" + retirement = getattr(self, "_retirement", None) + if retirement is not None: + retirement.request_drain("peer cancellation" if by_peer else "cancellation") with self.lock: + # A later cancellation must still notify peers after logical failure. + # The logical arbiter preserves whichever outcome already committed. if self._terminal_status == SessionStatus.CANCELLED: + return False + self._logical_outcomes.cancel(by_peer) + self._terminal_status = SessionStatus.CANCELLED + # Who asked is not recoverable later, and the two differ: the peer asking is a transfer + # error, our own side asking is an ordinary end. + self.cancelled_by_peer = by_peer + exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") + for task in self._kv_tasks: + rid_slice = (self.disagg_request_id, task.slice_id) + if task.status == TaskStatus.INIT: + if self._enforce_physical_ownership: + task.cancel_unpublished() + self._receiver._bounce.release_idle_reservation(rid_slice) + task.fail(exc) + else: + has_active_access = ( + not task.resources_drained + if self._enforce_physical_ownership + else task.status == TaskStatus.TRANSFERRING + ) + if has_active_access: + self._receiver._bounce.orphan_reservation(rid_slice) + if self._aux_physical_owner is not None and not self._publication_may_have_escaped: + self._aux_physical_owner.cancel_unpublished() + return True + + def capture_cancel_targets(self) -> set[str]: + with self.lock: + return set(self._sender_endpoints) + + def notify_cancel(self, sender_endpoints: Optional[set[str]] = None) -> None: + with self._publication_lock: + if sender_endpoints is None: + sender_endpoints = self.capture_cancel_targets() + if self._enforce_physical_ownership: + with self.lock: + self._cancel_notification_started = True + sender_endpoints = sender_endpoints - self._cancelled_sender_endpoints + sent = self._receiver.send_cancel_to_senders(self.disagg_request_id, sender_endpoints) + if self._enforce_physical_ownership and sent: + self._cancelled_sender_endpoints.update(sent) + + def process_session_quiesced(self, peer_rank: int) -> None: + """Settle missing pieces only after every published candidate fences writes.""" + with self._ownership_evidence_guard(), self.lock: + if ( + not self._enforce_physical_ownership + or self._closed + or not self._cancel_notification_started + or not self.has_failed() + ): return - if self._terminal_status == SessionStatus.ERROR: - # Preserve the diagnostically useful ERROR while still - # notifying any endpoints not covered by failure fan-out. - endpoints = self._sender_endpoints.difference( - self._failure_cancelled_sender_endpoints - ) - self._failure_cancelled_sender_endpoints.update(endpoints) - else: - self._terminal_status = SessionStatus.CANCELLED - exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") - for task in self._kv_tasks: - if task.status == TaskStatus.INIT: - task.fail(exc) - endpoints = set(self._sender_endpoints) - # Send outside the lock to avoid holding it during I/O. - if endpoints: - self._receiver.send_cancel_to_senders(self.disagg_request_id, endpoints) + owners = [task._get_physical_owner() for task in self._kv_tasks] + if self._aux_physical_owner is not None: + owners.append(self._aux_physical_owner) + try: + for owner in owners: + if owner._expected_writers is not None and not owner._cancelled_unpublished: + owner.record_session_quiesced(peer_rank) + except Exception as error: + self._record_ownership_evidence_error(error) + raise + + def cancel(self, by_peer: bool = False) -> bool: + """Cancel locally, then notify every writer without holding the lock. + + ``True`` if the ask went out; ``False`` means the session had already ended. + """ + if not self.cancel_local(by_peer=by_peer): + return False + self.notify_cancel() + return True def has_transferring_tasks(self) -> bool: - """True while any dispatched peer may still own destination buffers. + """True while a destination accessor may remain active. cancel_request() must return False while this is True. """ - with self.lock: - return self._active_nixl_ops > 0 + return not self.resources_drained() + + def _failed_wait_result(self) -> Optional[WaitResult]: + return ( + WaitResult.FAILED + if not self._enforce_physical_ownership or self.resources_drained() + else None + ) + + def _completed_wait_result(self) -> Optional[WaitResult]: + if getattr(self, "_retirement", None) is not None and not self.is_completed(): + return None + return WaitResult.COMPLETED def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: """Poll or block until transfer completes. @@ -3116,6 +3959,8 @@ def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: With blocking=True: waits up to _timeout_s for each task. Returns WaitResult.COMPLETED on full success, WaitResult.FAILED on error/timeout. """ + if self.has_failed() and self._enforce_physical_ownership: + return self._failed_wait_result() if not blocking: # Use task.status instead of task.wait(timeout=0): task.complete() # sets status before event, so a GIL switch between the two steps @@ -3124,32 +3969,60 @@ def wait_complete(self, blocking: bool = False) -> Optional[WaitResult]: if task.status == TaskStatus.TRANSFERRED: continue if task.status == TaskStatus.ERROR: - return WaitResult.FAILED + return self._failed_wait_result() return None # task not yet done; re-poll next cycle else: timeout = self._timeout_s for task in self._kv_tasks: done = task.wait(timeout=timeout) if not done: - return WaitResult.FAILED # timeout + return self._failed_wait_result() if task.status == TaskStatus.ERROR: - return WaitResult.FAILED + return self._failed_wait_result() if self._need_aux: while True: status = self.status - if status == SessionStatus.FULLY_TRANSFERRED: - return WaitResult.COMPLETED + if status == SessionStatus.TRANSFERRED: + return ( + self._failed_wait_result() + if self.has_failed() + else self._completed_wait_result() + ) elif status in (SessionStatus.ERROR, SessionStatus.CANCELLED): - return WaitResult.FAILED + return self._failed_wait_result() if not blocking: return None # KV done, aux still in flight; re-poll next cycle time.sleep(0.001) - return WaitResult.COMPLETED + return self._failed_wait_result() if self.has_failed() else self._completed_wait_result() - def close(self): - if getattr(self, "_closed", False): - return - self._closed = True + def close(self) -> bool: + if self._enforce_physical_ownership: + with self.lock: + if self._closed: + return True + retirement = getattr(self, "_retirement", None) + if retirement is not None and retirement.controller.shutdown_requested: + # Publication also holds this session lock. Only unsealed + # owners can be cancelled here; a reserved send remains live + # even before its network callback has actually started. + cancelled = [task for task in self._kv_tasks if task.cancel_unpublished()] + aux_cancelled = ( + self._aux_physical_owner is not None + and self._aux_physical_owner.cancel_unpublished() + ) + if cancelled or aux_cancelled: + self._logical_outcomes.cancel(by_peer=False) + for task in cancelled: + task.fail(RuntimeError("receiver shutdown before publication")) + if not self.resources_drained(): + return False + if retirement is not None and not retirement.close(): + return False + self._closed = True + else: + if getattr(self, "_closed", False): + return True + self._closed = True if self._aux_buffer is not None and self.aux_slot is not None: self._aux_buffer.free_slot(self.aux_slot) self.aux_slot = None @@ -3160,22 +4033,41 @@ def close(self): for task in self._kv_tasks: self._receiver._bounce.orphan_reservation((self.disagg_request_id, task.slice_id)) self._receiver.clear_session(self.disagg_request_id) + return True def __enter__(self): return self def __exit__(self, _exc_type, _exc, _tb): - self.close() + if not self.close(): + logger.warning( + f"RxSession {self.disagg_request_id} refused close while resources remain active" + ) def __del__(self): try: - self.close() + if not self.close(): + logger.warning( + f"RxSession {self.disagg_request_id} refused destructor close while " + "resources remain active" + ) except Exception as e: logger.warning(f"RxSession.__del__: exception during close: {e}") + def unpack_draft_history(self, request: LlmRequest) -> None: + """Read standalone history without changing context-first token fields.""" + if self._aux_buffer is None or self.aux_slot is None: + raise ValueError("Standalone draft transfer requires an auxiliary buffer slot") + request.py_draft_transfer_history = self._aux_buffer.get_slot_draft_history(self.aux_slot) + class RankInfoServer: - def __init__(self, rank_info: RankInfo, addr: Optional[str] = None, port: Optional[int] = None): + def __init__( + self, + rank_info: RankInfo, + addr: Optional[str] = None, + port: Optional[int] = None, + ): self._rank_info = rank_info self._shutdown = False # must be set before _start_listener() so __del__ is safe if addr is None and port is None: @@ -3230,7 +4122,13 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): self.shutdown() -def _create_nixl_agent(name: str, rank: int, world_size: int) -> NixlTransferAgent: +def _create_nixl_agent( + name: str, + rank: int, + world_size: int, + agent_buffer_size_mb: int = 0, + agent_bounce_params: Optional[Dict[str, str]] = None, +) -> NixlTransferAgent: num_threads = int(os.environ.get("TRTLLM_NIXL_NUM_THREADS", "8")) kwargs = {} if "TRTLLM_NIXL_SPLIT_BATCH_SIZE" in os.environ: @@ -3241,14 +4139,16 @@ def _create_nixl_agent(name: str, rank: int, world_size: int) -> NixlTransferAge num_threads=num_threads, rank=rank, world_size=world_size, + agent_buffer_size_mb=agent_buffer_size_mb, + agent_bounce_params=agent_bounce_params, **kwargs, ) def _make_aux_buffer( - kvm: KVCacheManager, max_slots: int, max_draft_len: Optional[int] = None + kvm: KVCacheManagerV2, max_slots: int, max_draft_len: Optional[int] = None ) -> Optional[AuxBuffer]: - draft_history = bool(getattr(kvm, "draft_layer_ids", ())) + draft_history = bool(getattr(kvm, "_draft_layer_sizes", ())) if max_slots <= 0: if draft_history: raise ValueError("Standalone draft transfer requires auxiliary buffer slots") @@ -3266,7 +4166,7 @@ def _make_aux_buffer( @dataclass class TransferWorkerConfig: - kv_cache_manager: KVCacheManager + kv_cache_manager: KVCacheManagerV2 device_id: int instance_name: str max_concurrent_sessions: int = 0 @@ -3275,11 +4175,33 @@ class TransferWorkerConfig: rx_timeout_s: Optional[float] = None bounce: Optional["Config"] = None tx_overall_timeout_s: Optional[float] = None + enforce_physical_ownership: bool = False + # Internal activation gate: only the qualified containment policy supplies it. + quiescence_fatal_callback: Optional[Callable[[QuiescenceFatalEvent], None]] = None + # Transfer-agent staging (bounce v2) buffer size in MiB; 0 disables the + # fast path. Derived upstream from CacheTransceiverConfig: it equals + # kv_cache_bounce_size_mb when agent_bounce_buffer_enable is set (in which + # case `bounce` above is off) and 0 otherwise. + agent_buffer_size_mb: int = 0 + # Expert bounce knobs (TRTLLM_NIXL_BOUNCE_* names without the prefix and the + # trailing _BYTES, lowercased); dict > env > default. Ignored when + # agent_buffer_size_mb == 0. + agent_bounce_params: Optional[Dict[str, str]] = None class TransferWorker: def __init__(self, config: TransferWorkerConfig): self._config = config + self._retirement_watchdog = None + if config.quiescence_fatal_callback is not None: + if not config.enforce_physical_ownership: + raise ValueError("deadline retirement requires physical ownership") + for timeout in (config.tx_overall_timeout_s, config.rx_timeout_s): + if timeout is None or not math.isfinite(timeout) or timeout <= 0: + raise ValueError( + "deadline retirement requires finite positive transfer timeouts" + ) + self._retirement_watchdog = RetirementWatchdog(self._on_quiescence_fatal) kvm = config.kv_cache_manager self._aux_buffer = _make_aux_buffer( kvm, config.max_concurrent_sessions, config.max_draft_len @@ -3292,6 +4214,24 @@ def __init__(self, config: TransferWorkerConfig): ) self._setup_peer_infrastructure(kvm) self._setup_transfer_engine() + if self._retirement_watchdog is not None: + self._retirement_watchdog.start() + + def _on_quiescence_fatal(self, event: QuiescenceFatalEvent) -> None: + """Forward fatal evidence while rooting all worker-owned resources. + + Args: + event: Immutable deadline decision; never authorizes resource release. + """ + callback = self._config.quiescence_fatal_callback + assert callback is not None + callback(event) + + def request_shutdown(self) -> None: + """Close admission and start drain before any potentially blocking teardown.""" + watchdog = getattr(self, "_retirement_watchdog", None) + if watchdog is not None: + watchdog.request_shutdown() def populate_instance_and_rank_info(self, endpoints: list[str], layer_num_per_pp: list[int]): assert self._rank_info is not None @@ -3308,8 +4248,8 @@ def create_tx_session(self, request: LlmRequest) -> TxSession: aux_buffer=self._aux_buffer, timeout_s=self._config.tx_timeout_s, prompt_len=request.prompt_len, - beam_width=request.py_beam_width, overall_timeout_s=self._config.tx_overall_timeout_s, + retirement_watchdog=getattr(self, "_retirement_watchdog", None), ) def create_rx_session(self, request: LlmRequest) -> RxSession: @@ -3322,7 +4262,7 @@ def create_rx_session(self, request: LlmRequest) -> RxSession: aux_buffer=self._aux_buffer, timeout_s=self._config.rx_timeout_s, prompt_len=request.prompt_len, - beam_width=request.py_beam_width, + retirement_watchdog=getattr(self, "_retirement_watchdog", None), ) def has_all_peer_req_infos_for_send(self, unique_rid: int) -> bool: @@ -3332,7 +4272,7 @@ def sweep_stale_req_infos(self): """Forward to Sender to evict orphaned RecvReqInfo from ADP broadcast.""" self._sender.sweep_stale_req_infos() - def _setup_peer_infrastructure(self, kvm: KVCacheManager): + def _setup_peer_infrastructure(self, kvm: KVCacheManagerV2): self._rank_info_server = RankInfoServer(self._rank_info) if kvm.mapping.rank == 0 else None self._kv_extractor = KVRegionExtractorV1(kvm) self._peer_registrar = PeerRegistrar(self._rank_info, self._kv_extractor) @@ -3345,7 +4285,22 @@ def _setup_transfer_engine(self): self._rank_info.instance_name + str(self._rank_info.instance_rank), rank=mapping.rank, world_size=mapping.world_size, + agent_buffer_size_mb=self._config.agent_buffer_size_mb, + agent_bounce_params=self._config.agent_bounce_params, ) + # The pure-Python agent has no bounce_enabled attribute and already warns on its own. + if ( + self._config.agent_buffer_size_mb > 0 + and hasattr(self._agent, "bounce_enabled") + and not self._agent.bounce_enabled + ): + logger.warning( + f"TransferWorker: the C++ transfer-agent bounce was requested " + f"(agent_bounce_buffer_enable=True, kv_cache_bounce_size_mb=" + f"{self._config.agent_buffer_size_mb}) but is inactive on agent " + f"{self._agent.name} (init failed); standard per-descriptor NIXL " + "transfers are in use." + ) self._registered_mem: list = [] try: self._register_kv_cache() @@ -3362,8 +4317,20 @@ def _setup_transfer_engine(self): # The bounce object owns its own NIXL descriptors (deregistered by Transport.close in # shutdown), so they are NOT added to _registered_mem — single-source to avoid a # double deregister. - self._sender = Sender(self._peer_registrar, self._agent, bounce=self._bounce) - self._receiver = Receiver(self._peer_registrar, self._agent, bounce=self._bounce) + self._sender = Sender( + self._peer_registrar, + self._agent, + bounce=self._bounce, + enforce_physical_ownership=self._config.enforce_physical_ownership, + retirement_watchdog=self._retirement_watchdog, + ) + self._receiver = Receiver( + self._peer_registrar, + self._agent, + bounce=self._bounce, + enforce_physical_ownership=self._config.enforce_physical_ownership, + retirement_watchdog=self._retirement_watchdog, + ) self._rank_info.transfer_engine_info = bytes(self._agent.get_local_agent_desc()) self._rank_info.self_endpoint = self._receiver.endpoint except Exception: @@ -3417,20 +4384,25 @@ def page_table(self): def shutdown(self): if getattr(self, "_shutdown", False): return - self._shutdown = True + self.request_shutdown() + watchdog = getattr(self, "_retirement_watchdog", None) + if watchdog is not None: + watchdog.require_retired() + watchdog.stop() # Use getattr guards: __init__ may have failed partway, leaving some # attributes unset. Without them, __del__ -> shutdown() raises # AttributeError and ZMQ resources from already-created sub-objects # are never cleaned up. rank_info_server = getattr(self, "_rank_info_server", None) - if rank_info_server is not None: - rank_info_server.shutdown() sender = getattr(self, "_sender", None) if sender is not None: sender.shutdown() receiver = getattr(self, "_receiver", None) if receiver is not None: receiver.shutdown() + if rank_info_server is not None: + rank_info_server.shutdown() + self._shutdown = True # Close the bounce transport (stops its scatter thread, deregisters its own descriptors and # frees the VMM buffers) once the receiver listener is stopped so no new scatter is enqueued. # No-op for the non-bounced path (NoBounceTransport.close); without this the daemon thread + fabric @@ -3445,6 +4417,14 @@ def shutdown(self): # (e.g. when the KV cache manager is recreated after profiling). agent = getattr(self, "_agent", None) if agent is not None: + try: + if getattr(agent, "bounce_enabled", False): + logger.info( + f"TransferWorker.shutdown: agent {agent.name} bounce_submit_count=" + f"{agent.bounce_submit_count} bounce_reject_count={agent.bounce_reject_count}" + ) + except Exception as e: # observability only; never skip deregister/shutdown below + logger.debug(f"TransferWorker.shutdown: bounce counters unavailable: {e}") registered = getattr(self, "_registered_mem", []) while registered: desc = registered.pop(0) diff --git a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py index 5dea7ba2eb26..7945d3395435 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py +++ b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py @@ -19,9 +19,8 @@ import numpy as np -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode from .page import AttentionLayerGroup @@ -145,9 +144,8 @@ def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 # V2 already returns per-cache-level pool slot indices (not logical block # IDs), and active sequences GPU-lock their pages (_UniqPageLock enforces # cache_level==GPU), so the slot_ids yielded here are already the right - # offsets for primary-pool pointer arithmetic. No translation is needed, - # unlike V1 (see _CacheReuseAdapterV1.get_block_ids). - if lg.cp_as_tp: + # offsets for primary-pool pointer arithmetic. No translation is needed. + if getattr(lg, "cp_as_tp", False): ordinals = self.get_block_ordinals(req, group_idx, lg) return ordinals[ordinals >= 0] return np.fromiter( diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 1209490b31eb..629af5872c8c 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -43,15 +43,10 @@ get_layer_byte_ranges, get_physical_pool, ) -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( - MambaHybridCacheManager, - MambaHybridCacheManagerV2, -) -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm._utils import get_size_in_bytes, nvtx_range -from tensorrt_llm.bindings import DataType -from tensorrt_llm.runtime.kv_cache_manager_v2 import PoolGroupDesc, SlotDescVariant +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole, PoolGroupDesc, SlotDescVariant # Mapper kinds a V2 manager may declare via get_disagg_role_mapper_kinds(). # A physical pool may mix kinds (V2 storage coalesces buffers purely by @@ -224,7 +219,7 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: # matching this ID is required to be consistent for. inverse = {} for (model_layer, attn_type), layer_id in manager._layer_attn_to_layer_id.items(): - inverse[layer_id] = (model_layer, attn_type.value) + inverse.setdefault(layer_id, (model_layer, attn_type.value)) for draft_layer in getattr(manager, "draft_layer_ids", ()): inverse[manager.layer_offsets[draft_layer]] = (draft_layer, 0) @@ -505,6 +500,19 @@ def _window_size_for_layer(internal_layer_id: int): total_kv_head_num=manager._draft_total_num_kv_heads, ) continue + # A life cycle is recurrent when its layers are the manager's + # mamba layers. Recurrent and attention layers never share a life + # cycle: their buffer configs differ, so V2 groups them apart. + is_recurrent = isinstance(manager, MambaHybridCacheManagerV2) and any( + manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids + ) + if is_recurrent and not all( + manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids + ): + raise ValueError( + f"Layer group {layer_group_id} mixes recurrent and attention layers" + ) + cache_kind = CacheKind.STATE if is_recurrent else CacheKind.PAGED # Bucket buffer entries by (pool, mapper kind). One PoolView is # emitted per bucket and spans every layer of that role class, diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 850257d8f5ab..837ce31ce20d 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -293,8 +293,6 @@ class AttentionLayerGroup(LayerGroup): kv_head_num_per_rank: int = 0 sliding_window_size: Optional[int] = None - local_layers: List[LocalLayer] = field(default_factory=list) - pool_views: List[PoolView] = field(default_factory=list) # Full-sequence KV whose heads use the flattened TP x CP rank grid. cp_as_tp: bool = False total_kv_head_num: Optional[int] = None @@ -317,6 +315,8 @@ def from_dict(cls, data: dict) -> "AttentionLayerGroup": pool_group_idx=int(data["pool_group_idx"]), local_layers=[LocalLayer.from_dict(x) for x in data.get("local_layers", [])], pool_views=[PoolView.from_dict(pv) for pv in data.get("pool_views", [])], + kv_head_num_per_rank=int(data["kv_head_num_per_rank"]), + sliding_window_size=data.get("sliding_window_size"), cp_as_tp=bool(data.get("cp_as_tp", False)), total_kv_head_num=data.get("total_kv_head_num"), ) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 84c3014471a8..51ae7033ae8c 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -58,12 +58,10 @@ ) from tensorrt_llm._torch.disaggregation.resource.utils import get_pool_view_slot_bytes from tensorrt_llm._torch.distributed.communicator import Distributed -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( - MambaHybridCacheManager, - MambaHybridCacheManagerV2, +from tensorrt_llm._torch.pyexecutor import hang_detector +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + BlockReusePolicy, + KVCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -635,6 +633,7 @@ def _create_cache_extent(self, req: LlmRequest) -> CacheExtent: registered by.""" rid = get_unique_rid(req) assert rid is not None + self._pack_draft_history(req) return CacheExtent(name=rid, local=self._describe_local(req)) def _describe_local(self, req: LlmRequest) -> Chunk: @@ -654,11 +653,18 @@ def _describe_local(self, req: LlmRequest) -> Chunk: tpb = adapter.tokens_per_block assert self._page_table is not None layer_groups = self._page_table.layer_groups - prompt_len = self._global_prompt_len(req) - prompt_blocks = (prompt_len + tpb - 1) // tpb - manager = self._kv_cache_manager - cp_size = getattr(manager, "_helix_cp_size", 1) - cp_rank = getattr(manager, "_helix_cp_rank", 0) + remote_tail_start = getattr(req, "py_csa2_remote_tail_start", None) + transfer_len = ( + remote_tail_start if remote_tail_start is not None else self._global_prompt_len(req) + ) + prompt_blocks = (transfer_len + tpb - 1) // tpb + remote_tail_split = getattr(req, "py_csa2_remote_tail_split", None) + excluded_pool_views = set() + if remote_tail_split is not None: + exclude = getattr(self._kv_cache_manager, "remote_tail_excluded_pool_views", None) + if exclude is None: + raise ValueError("remote-tail replay requires a CSA2 cache manager") + excluded_pool_views = exclude(self._page_table, remote_tail_split) is_gen_only = req.is_generation_only_request cached_per_lg = ( @@ -667,104 +673,26 @@ def _describe_local(self, req: LlmRequest) -> Chunk: else [0] * len(layer_groups) ) - unified_draft = isinstance(manager, KVCacheManagerV2) and bool(manager.draft_layer_ids) - transfer_length = prompt_len if unified_draft else req.prompt_len - token_range = None - if transfer_length > 0: - # end must match the trimmed block list below (ceil(prompt_len / tpb) - # blocks). num_extra_kv_tokens slots (speculative decoding) are not - # transferred. In the previously added support for ctx disabling - # speculative decoding while gen enables it, both sides currently - # use prompt_len as the transfer range, so the ranges stay - # consistent. - # TODO: the accuracy impact of not transferring num_extra_kv_tokens - # on MTP and other speculative decoding paths is currently unclear; - # revisit whether these extra KV slots need to be transferred. - token_range = TokenRange(start=0, end=transfer_length) - + empty = np.array([], dtype=np.int64) groups = [] kinds = [lg.kind for lg in layer_groups] for idx, lg in enumerate(layer_groups): - if isinstance(lg, MambaLayerGroup): - groups.append(np.array([], dtype=np.int64)) - continue - if unified_draft: - # KVSlice carries a contiguous retained suffix in this protocol. - local_prompt_blocks = ( - prompt_blocks if lg.cp_as_tp else len(range(cp_rank, prompt_blocks, cp_size)) - ) - stale_end = 0 - if lg.sliding_window_size is not None: - stale_end = max(0, (prompt_len + 1 - lg.sliding_window_size) // tpb) - if not lg.cp_as_tp: - stale_end = len(range(cp_rank, stale_end, cp_size)) - ordinals = adapter.get_block_ordinals(req, idx, lg) - pages = ordinals[stale_end:local_prompt_blocks] - if pages.size != local_prompt_blocks - stale_end or np.any(pages < 0): - raise ValueError("Missing allocated prompt pages for unified draft KV transfer") - cache_skip = max(0, cached_per_lg[idx] // tpb - stale_end) - groups.append(pages[cache_skip:]) - continue - block_ids = adapter.get_block_ids(req, idx, lg) - # Limit to prompt_len blocks, matching C++ cacheFormatter behavior. - total_blocks = (req.prompt_len + tpb - 1) // tpb - window_size = lg.sliding_window_size - - if window_size is not None: - draft_reserve_tokens = get_draft_token_length(req) if is_gen_only else 0 - allocated_blocks = ( - req.prompt_len - + draft_reserve_tokens - + self._kv_cache_manager.num_extra_kv_tokens - + tpb - - 1 - ) // tpb - beam0_block_ids, tail_block_ids = self._split_packed_beam_block_ids( - block_ids, - req.py_beam_width, - allocated_blocks, - ) - if beam0_block_ids.size > allocated_blocks: - beam0_block_ids = beam0_block_ids[:allocated_blocks] - block_ids = ( - np.concatenate([beam0_block_ids, tail_block_ids]) - if tail_block_ids.size > 0 - else beam0_block_ids - ) - # Current PyExecutor cache managers disable KV-cache token sinks, - # so SWA block lists contain an evictable prompt prefix followed - # by the speculative scratch tail. If token sinks are enabled, - # this must use block-ordinal metadata to preserve the sink prefix. - # Remove scratch before trimming stale prompt blocks; otherwise a - # boundary-crossing allocation can displace initialized prompt KV. - scratch_blocks = max(0, allocated_blocks - total_blocks) - if scratch_blocks > 0 and getattr(self._kv_cache_manager, "_has_cp_helix", False): - # Helix ledgers expose one windowed page per - # tokens_per_block * cp_size tokens, so the per-page scratch - # arithmetic above would empty a one-entry HP-tail list. - scratch_blocks = 0 - if scratch_blocks > 0: - if req.py_beam_width != 1: - raise ValueError("speculative scratch blocks require beam_width == 1") - block_ids = ( - block_ids[:-scratch_blocks] - if scratch_blocks < block_ids.size - else np.array([], dtype=np.int64) - ) - # Drop stale blocks the manager may still expose (V1 pre-eviction). - stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb) - expected_valid = max(0, total_blocks - stale_end) - # Stale prefix already pruned above; skip reuse-hit blocks that - # land inside the window. Clamp to 0: ctx side has cached_per_lg - # synthetically 0, and a reuse hit may fall entirely inside the - # stale region (those blocks were already pruned, no extra skip). - cache_skip = max(0, cached_per_lg[idx] // tpb - stale_end) + if lg.kind == CacheKind.STATE: + slot = self._get_mamba_slot_for_request(req) + group = np.array([slot], dtype=np.int64) if slot is not None else empty else: # Block lists carry beam 0 only (beam-search attention reads # prompt positions through beam 0's block table), so the # positional path applies to every beam width. ordinals = adapter.get_block_ordinals(req, idx, lg) - group = self._positional_window(ordinals, prompt_blocks, cached_per_lg[idx] // tpb) + local_prompt_blocks = prompt_blocks + if not getattr(lg, "cp_as_tp", False): + cp_size = getattr(self._kv_cache_manager, "_helix_cp_size", 1) + cp_rank = getattr(self._kv_cache_manager, "_helix_cp_rank", 0) + local_prompt_blocks = len(range(cp_rank, prompt_blocks, cp_size)) + group = self._positional_window( + ordinals, local_prompt_blocks, cached_per_lg[idx] // tpb + ) groups.append(group) return Chunk( @@ -793,7 +721,7 @@ def _positional_window( def _need_aux_transfer(self, req: LlmRequest) -> bool: params = req.py_disaggregated_params - return bool(getattr(self._kv_cache_manager, "draft_layer_ids", ())) or ( + return bool(getattr(self._kv_cache_manager, "_draft_layer_sizes", ())) or ( params is not None and params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST ) @@ -824,7 +752,7 @@ def _validate_draft_history_range(self, req: LlmRequest, history: dict) -> None: def _pack_draft_history(self, req: LlmRequest) -> None: manager = getattr(self, "_kv_cache_manager", None) - if not getattr(manager, "draft_layer_ids", ()): + if not getattr(manager, "_draft_layer_sizes", ()): return history = self._kv_cache_manager.export_draft_history(req.py_request_id) self._validate_draft_history_range(req, history) @@ -833,7 +761,7 @@ def _pack_draft_history(self, req: LlmRequest) -> None: def _restore_draft_history(self, req: LlmRequest) -> None: history = getattr(req, "py_draft_transfer_history", None) manager = getattr(self, "_kv_cache_manager", None) - has_draft = bool(getattr(manager, "draft_layer_ids", ())) + has_draft = bool(getattr(manager, "_draft_layer_sizes", ())) if history is None: if has_draft: raise ValueError( @@ -1158,7 +1086,6 @@ def _apply_aux(self, session, req: LlmRequest): """Unpack aux tokens from session into request's context_phase_params.""" params = req.py_disaggregated_params if params is not None and params.schedule_style != DisaggScheduleStyle.GENERATION_FIRST: - # Context-first tokens and usage already arrived in the context response. session.unpack_draft_history(req) return session.unpack_aux(req) @@ -1291,28 +1218,20 @@ def respond_and_send_async(self, req: LlmRequest) -> None: rid = get_unique_rid(req) assert rid is not None session = self._get_or_create_send_session(req) - req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS - # Register before dispatch: dispatch_task may synchronously fail after - # enqueueing work for an earlier peer. Status polling must always have - # the matching request while that partial dispatch drains. - self._send_reqs[rid] = req - try: - self._pack_draft_history(req) - session.send(self._create_kv_slice(req)) - self._finalize_send(req, session) - except Exception as exc: - req.py_disagg_send_setup_failed = True - session.cancel() - if session.has_transferring_tasks(): - # At least one worker already entered NIXL. Keep the session, - # request, and pinned pages registered; - # normal status polling will classify the safe cancellation - # as a setup failure once every active operation has drained. - req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS - logger.error( - "CTX send setup failed after transfer started; " - f"retaining rid={rid} until NIXL leaves mid-write: {exc}" - ) + if session is None: + return + bridge_enabled = getattr(self, "_fp4_mla_bridge_enabled", False) + if bridge_enabled: + # Root the request before backend admission so its source pages + # outlive every ownership-enabled NIXL operation. + self._send_reqs[rid] = req + req.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + # A remote cancellation can win before the local TxSession is + # created. Sender.setup_session() records that terminal state and + # acknowledges the fenced session to every known receiver; + # leave retirement to the normal status path without attempting + # to publish KV or auxiliary memory afterward. + if session.has_failed(): return try: if self.pipeline_transfer_enabled: @@ -1368,9 +1287,11 @@ def request_and_receive_sync(self, req: LlmRequest) -> None: # blocking path has no cross-rank outcome consensus). req.py_kv_transfer_verified = session.kv_write_verified() # KV-transfer timing setters deferred to #15871 (clock-source consistency); size only. - req.set_kv_cache_size(self._slice_num_bytes(kv_slice) * self._kv_size_rank_factor) - self._prepare_received_history(session, req) - self._rebuild_fp4_mla_disagg_imported_cache(req) + req.set_kv_cache_size(self._request_kv_bytes(req, paged_bytes, state_bytes)) + if self._need_aux_transfer(req): + self._apply_aux(session, req) + self._assert_disagg_history_declared(req) + self._restore_draft_history(req) req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE else: req.state = LlmRequestState.DISAGG_TRANS_ERROR @@ -1561,6 +1482,7 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]) -> GenT 0 if need_progress else wait_num, block_all, ) + completed, failed, cancelled = [], [], [] for rid in to_process: session = self._recv_sessions[rid] @@ -1627,49 +1549,20 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]) -> GenT for rid in completed: session = self._recv_sessions[rid] req = self._recv_reqs[rid] - completed_reqs[rid] = req - try: - self._prepare_received_history(session, req) - self._rebuild_fp4_mla_disagg_imported_cache(req) - except Exception as exc: - local_postprocess_failed.append(rid) - req.state = LlmRequestState.DISAGG_TRANS_ERROR - logger.error( - f"Disagg gen postprocess failed rank={self._dist.rank} rid={rid}: {exc}" - ) - finally: - # NIXL is complete at this point, so unregistering is safe even - # if aux application or an FP4 MLA sidecar rebuild fails. - session.close() - del self._recv_reqs[rid] - del self._recv_sessions[rid] - - # ``completed`` is globally consistent, so guarding this collective is - # itself symmetric and avoids another allgather on empty polling cycles. - if completed: - all_postprocess_failed = self._allgather_or_passthrough( - local_postprocess_failed, - self._gen_allgather, - self._gen_need_sync, - ) - global_postprocess_failed = self._union(all_postprocess_failed) - else: - global_postprocess_failed = set() - postprocess_failed = [rid for rid in completed if rid in global_postprocess_failed] - postprocess_failed_set = set(postprocess_failed) - for rid in completed: - req = completed_reqs[rid] - if rid in postprocess_failed_set: - # This also covers remote-only failure: a locally successful - # rank must not publish TRANS_COMPLETE when another - # participating rank failed the same request. - req.state = LlmRequestState.DISAGG_TRANS_ERROR - else: - req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE - - completed = [rid for rid in completed if rid not in postprocess_failed_set] - transfer_failed = failed - failed = transfer_failed + postprocess_failed + # Stamp the write-verification verdict on the request before the + # session is closed and dropped: commit_blocks_for_reuse runs + # later (batch preparation) and gates reuse admission on it. + req.py_kv_transfer_verified = rid in verified + # transfer_end already stamped at completion detection above. + req.set_kv_cache_size(getattr(req, "py_kv_cache_xfer_bytes", 0)) + if self._need_aux_transfer(req): + self._apply_aux(session, req) + self._assert_disagg_history_declared(req) + self._restore_draft_history(req) + self._close_session_or_raise(session, rid, "completed") + req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + del self._recv_reqs[rid] + del self._recv_sessions[rid] if failed: logger.warning( f"Disagg gen transfer FAILED rank={self._dist.rank} " diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 8fb53462c6d1..f41b868d9e72 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -68,7 +68,8 @@ def dspark_layer_window_size( def declares_lilicorr(config: PretrainedConfig) -> bool: """Whether the checkpoint requires the LiLiCorr candidate scorer.""" return is_lilicorr_checkpoint( - getattr(config, "architectures", None), getattr(config, "dflash_config", None) or {} + getattr(config, "architectures", None), + getattr(config, "dflash_config", None) or {}, ) @@ -135,7 +136,8 @@ def dflash2_score_edges( top_k = candidate_ids.shape[-1] successors = successor_codebook[candidate_ids].float() predecessor_ids = torch.cat( - (anchor_token_ids[:, None, None].expand(-1, 1, top_k), candidate_ids[:, :-1]), dim=1 + (anchor_token_ids[:, None, None].expand(-1, 1, top_k), candidate_ids[:, :-1]), + dim=1, ) predecessors = predecessor_codebook[predecessor_ids].float() return unary_logits.float()[:, :, None] + torch.einsum( @@ -222,7 +224,11 @@ def __init__( self.kernel_projection.weight.data.copy_(kernel_projection_weight) def _convolve( - self, hidden_states: torch.Tensor, delta: torch.Tensor, side: int, block_size: int + self, + hidden_states: torch.Tensor, + delta: torch.Tensor, + side: int, + block_size: int, ) -> torch.Tensor: return dflash2_grouped_conv( hidden_states, delta, self.base_kernel[side], block_size, self.group_size @@ -465,7 +471,6 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): ) self.target_layer_ids = dflash_config.get("target_layer_ids", None) - # Upstream keeps the dflash_config override; rubin-advance dropped it. self.block_size = dflash_config.get( "block_size", getattr(pretrained_config, "block_size", None) ) @@ -835,7 +840,10 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): # Laguna DFlash checkpoints may ship a fused self_attn.qkv_proj; the draft # loader expects split q/k/v (a fused key is silently dropped otherwise). if any(k.endswith("self_attn.qkv_proj.weight") for k in weights): - for attr in ("num_attention_heads_per_layer", "num_key_value_heads_per_layer"): + for attr in ( + "num_attention_heads_per_layer", + "num_key_value_heads_per_layer", + ): per_layer = getattr(self.config, attr, None) if per_layer is not None and len(set(per_layer)) > 1: raise ValueError( @@ -845,7 +853,9 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): "is required for this checkpoint." ) head_dim = getattr( - self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads + self.config, + "head_dim", + self.config.hidden_size // self.config.num_attention_heads, ) num_kv_heads = getattr( self.config, "num_key_value_heads", self.config.num_attention_heads @@ -936,7 +946,11 @@ def _load_target_projection(self, weights: Dict) -> None: ) weight = weight[:, self._fc_input_slice] self.fc = nn.Linear( - weight.shape[1], weight.shape[0], bias=False, device="cuda", dtype=weight.dtype + weight.shape[1], + weight.shape[0], + bias=False, + device="cuda", + dtype=weight.dtype, ) self.fc.weight.data.copy_(weight) @@ -961,7 +975,10 @@ def take(key: str) -> torch.Tensor: attention_convs = nn.ModuleList() mlp_convs = nn.ModuleList() for layer_idx in range(len(self.model.layers)): - for name, convs in (("attention_conv", attention_convs), ("mlp_conv", mlp_convs)): + for name, convs in ( + ("attention_conv", attention_convs), + ("mlp_conv", mlp_convs), + ): prefix = f"layers.{layer_idx}.{name}." convs.append( DFlash2BlockConv( @@ -1636,6 +1653,20 @@ def dflash_forward( append_pos_long = append_positions.long() fp8_pages = block_tables[append_rows, append_pos_long // page_size].long() fp8_offsets = append_pos_long % page_size + else: # VANILLA + # Slot and column that each block K/V row occupies + self._prepare_dflash_index_buffers( + ctx_k_cache.device, + ctx_k_cache.size(0), + block_size, + need_batch_indices=False, + ) + append_batch_indices = ( + cache_batch_idx_i32.view(-1, 1).expand(-1, block_size).reshape(-1).long() + ) + append_positions = ( + (cache_seqlens_i32.view(-1, 1) + self._dflash_block_offsets).reshape(-1).long() + ) # Flatten query positions once for the fused QK-norm-RoPE kernel. query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) @@ -2082,7 +2113,11 @@ def load_weights(self, weights, weight_mapper=None, **kwargs): for k in aux_keys: w = weights.pop(k) norm = nn.RMSNorm( - w.shape[0], eps=eps, device="cuda", elementwise_affine=True, dtype=w.dtype + w.shape[0], + eps=eps, + device="cuda", + elementwise_affine=True, + dtype=w.dtype, ) norm.weight.data.copy_(w) norms.append(norm) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index dc3ee32f041b..e3e0c47251ca 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -16,8 +16,12 @@ # DSpark backbone ported from the DeepSeek-V4-Pro-DSpark reference # (`inference/model.py`: DSparkBlock / Transformer.forward_spec). # -# Captured-context attention reads managed KV pages and the noncausal draft -# block, with a learnable attention sink contributing only to the denominator. +# The captured-context attention primitives are ported from DeepSeek's DeepSpec +# reference ``inference/kernel.py`` (``sparse_attn``) and ``inference/model.py`` +# (``get_dspark_topk_idxs``). The reference computes these with a TileLang +# kernel; this is a functional-first pure-PyTorch port with the same math +# (index-gather + online softmax + a learnable attention sink that contributes +# only to the softmax denominator). # # The draft I/O stages are ported from the same reference # (`inference/model.py`: DSparkBlock.forward_embed / forward_head). @@ -55,8 +59,9 @@ class or the edge would become a cycle. ``hc_head`` — the block-draft output head. The heads themselves are shared with the standalone path and live in :mod:`modeling_speculative`. -2. **Captured-context attention** — projections and rotary helpers feeding - paged attention over the retained context and the noncausal draft block. +2. **Captured-context attention primitives** — the dense sliding-window MLA the + draft block attends with (``get_dspark_topk_idxs`` / ``dspark_sparse_attn`` + and the rotary helpers). Hardware-agnostic and unit-testable in isolation. 3. **Block draft I/O** — ``build_draft_input_ids`` (the ``[bonus_token, noise, ...]`` block input) and ``dspark_propose`` (Markov @@ -78,7 +83,7 @@ class or the edge would become a cycle. import os import re from functools import lru_cache -from typing import Dict, List, NamedTuple, Optional +from typing import Callable, Dict, List, NamedTuple, Optional import torch import torch.nn.functional as F @@ -91,9 +96,11 @@ class or the edge would become a cycle. from tensorrt_llm.quantization.mode import QuantAlgo from ..._utils import get_sm_version, is_sm_100f -from ..attention_backend.dspark import write_dspark_context -from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_INTERNAL_AVAILABLE -from ..distributed import AllReduceParams +from ...models.modeling_utils import QuantConfig +from ..attention.backends.dspark import write_dspark_context +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE +from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.gated_mlp import GatedMLP @@ -104,7 +111,7 @@ class or the edge would become a cycle. from ..speculative.interface import SpeculativeDecodingMode from ..speculative.spec_tail_kernels import DraftBlockPlacement from ..speculative.vocab_parallel import markov_chain_logits -from ..utils import AuxStreamType +from ..utils import AuxStreamType, swizzle_sf, unswizzle_sf from .modeling_deepseekv4 import ( DeepseekV4DecoderLayer, DeepseekV4WeightLoader, @@ -122,6 +129,13 @@ class or the edge would become a cycle. ) from .modeling_utils import DecoderModel, register_auto_model, register_draft_model +_DSPARK_LINEAR_BACKEND_ENV = "TRTLLM_DSPARK_LINEAR_BACKEND" + +if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + from ..custom_ops.dspark_attention_custom_op import ( + _run_cute_dsl_dspark_mxfp8_gemm_rubin_prevalidated, + ) + if IS_CUTLASS_DSL_AVAILABLE: from ..custom_ops import dspark_attention_custom_op # noqa: F401 from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( @@ -134,6 +148,17 @@ class or the edge would become a cycle. is_fused_dspark_rmsnorm_rope_supported, ) + +if IS_CUTLASS_DSL_AVAILABLE: + from ..custom_ops.dspark_attention_custom_op import ( + fused_dsv4_dspark_attention, + is_fused_dsv4_dspark_attention_supported, + ) + from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( + cute_dsl_dspark_rmsnorm_rope_cache_write, + is_fused_dspark_attention_preparation_supported, + ) + # ---------------------------------------------------------------------------- # Captured-context attention primitives. # ---------------------------------------------------------------------------- @@ -210,6 +235,148 @@ def apply_dspark_rotary_batched( return out.to(orig_dtype) +@lru_cache(maxsize=64) +def _topk_matrix(window_size: int, block_size: int, start_pos: int) -> torch.Tensor: + # [min(window, start_pos+1)] context positions in the rolling KV window, + # followed by [block_size] positions for the current block's own K/V (which + # the caller appends to the window at offset ``window_size``). + ctx = torch.arange(min(window_size, start_pos + 1)) + blk = window_size + torch.arange(block_size) + return torch.cat([ctx, blk]).int() + + +def get_dspark_topk_idxs( + window_size: int, + bsz: int, + block_size: int, + start_pos: int, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Per-query attended-position indices for the DSpark draft block. + + Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` + query positions attends to the same set — the ``min(window_size, start_pos+1)`` + most-recent context positions in the rolling KV window, then the + ``block_size`` positions of the current block (stored at offset + ``window_size`` in the concatenated KV). Note this is *non-causal* within the + block (every position sees every block position), matching the reference. + + Args: + window_size: sliding-window length of the captured-context KV cache. + bsz: batch size. + block_size: number of draft positions per request. + start_pos: absolute decode position (must be > 0); bounds the context. + device: device for the returned index tensor. + + Returns: + int32 tensor ``[bsz, block_size, topk]`` with + ``topk = min(window_size, start_pos+1) + block_size``. + """ + assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" + matrix = _topk_matrix(int(window_size), int(block_size), int(start_pos)).to(device) + return matrix.view(1, 1, -1).expand(bsz, block_size, -1).contiguous() + + +def get_dspark_topk_idxs_batched( + window_size: int, + block_size: int, + start_pos: torch.Tensor, + valid_len: torch.Tensor | None = None, +) -> torch.Tensor: + """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. + + Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width + ``min(window_size, start_pos+1) + block_size`` depends on the host int + ``start_pos``), this always returns the **fixed** width ``window_size + + block_size`` and masks the unfilled context slots with ``-1``. The masked + slots are excluded by :func:`dspark_sparse_attn` exactly as if they were + absent while the shape remains CUDA-graph safe. + + Every query attends to the actually written circular-window suffix, followed + by the current-block positions. Without ``valid_len`` this preserves the + legacy ``start_pos``-only behavior. + + Args: + window_size: sliding-window length of the captured-context KV cache. + block_size: number of draft positions per request. + start_pos: ``[G]`` int tensor of per-request absolute decode positions. + valid_len: optional ``[G]`` count of actually written rolling-window + entries. When omitted, preserve the legacy ``start_pos`` mask. + + Returns: + int32 tensor ``[G, block_size, window_size + block_size]``. + """ + device = start_pos.device + g = start_pos.shape[0] + ctx_cols = torch.arange(window_size, device=device) # [win] + if valid_len is None: + valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] + else: + # The valid entries are the contiguous logical suffix ending at + # start_pos, but their physical slots wrap modulo window_size. + valid_len = valid_len.clamp(min=0, max=window_size) + age = torch.remainder(start_pos.unsqueeze(1) - ctx_cols.unsqueeze(0), window_size) + valid = age < valid_len.unsqueeze(1) + ctx_idx = torch.where( + valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) + ) + blk_idx = window_size + torch.arange(block_size, device=device) # [block] + blk_idx = blk_idx.unsqueeze(0).expand(g, -1) # [G, block] + row = torch.cat([ctx_idx, blk_idx], dim=1).to(torch.int32) # [G, win+block] + return row.unsqueeze(1).expand(g, block_size, -1).contiguous() + + +def dspark_sparse_attn( + q: torch.Tensor, + kv: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Index-gathered multi-query attention with an attention sink. + + Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For + each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by + ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled + dot-product softmax over them, and adds a per-head learnable *sink* logit that + participates only in the softmax denominator (i.e. an "attend-to-nothing" + option with a zero value vector). KV is shared across query heads (MQA). + + Args: + q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). + kv: ``[b, n, d]`` keys/values (shared across heads). + attn_sink: ``[h]`` per-head sink logits (fp32). + topk_idxs: ``[b, m, topk]`` int gather indices into ``kv`` (``-1`` masks). + softmax_scale: scalar applied to the q·k scores (``head_dim ** -0.5``). + + Returns: + ``[b, m, h, d]`` attention output, in ``q.dtype``. + """ + b, m, h, d = q.shape + idx = topk_idxs.long() # [b, m, topk] + valid = idx >= 0 + safe = idx.clamp(min=0) + + # Invalid slots read kv[0, :] (via safe.clamp), but masked_fill below + # zeros their softmax probs, so the einsum nullifies them. + kv_exp = kv.unsqueeze(1).expand(b, m, kv.shape[1], d) + gathered = torch.gather(kv_exp, 2, safe.unsqueeze(-1).expand(b, m, safe.shape[-1], d)).float() + + # Scores [b, m, h, topk]; mask invalid slots to -inf before the softmax. + scores = torch.einsum("bmhd,bmkd->bmhk", q.float(), gathered) * softmax_scale + scores = scores.masked_fill(~valid.unsqueeze(2), float("-inf")) + + # Online-softmax max is taken over gathered positions only (the sink is added + # to the denominator afterwards), matching the kernel's reduce order. + smax = scores.max(dim=-1, keepdim=True).values # [b, m, h, 1] + smax = torch.where(torch.isinf(smax), torch.zeros_like(smax), smax) + probs = torch.exp(scores - smax) # masked slots -> exp(-inf) = 0 + sink = torch.exp(attn_sink.to(torch.float32).view(1, 1, h) - smax.squeeze(-1)) + denom = probs.sum(dim=-1) + sink # [b, m, h] + out = torch.einsum("bmhk,bmkd->bmhd", probs, gathered) / denom.unsqueeze(-1) + return out.to(q.dtype) + + def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: """RMSNorm matching the DeepSpec reference (fp32 reduce, then * weight).""" dtype = x.dtype @@ -218,6 +385,15 @@ def _rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: return (weight.float() * xf).to(dtype) +def _rope_last_dims( + t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply RoPE to the last ``rope_head_dim`` dims; pass the rest through.""" + nope = t[..., :-rope_head_dim] + rope = apply_dspark_rotary(t[..., -rope_head_dim:], freqs_cis, inverse=inverse) + return torch.cat([nope, rope], dim=-1) + + def _rope_last_dims_batched( t: torch.Tensor, rope_head_dim: int, freqs_cis: torch.Tensor, inverse: bool = False ) -> torch.Tensor: @@ -310,14 +486,11 @@ def _rmsnorm_rope_batched( return t -def dspark_attention_forward( +def _dspark_attention_forward_legacy( x: torch.Tensor, main_x: torch.Tensor, - start_pos: torch.Tensor, + start_pos: int, kv_cache: torch.Tensor, - block_tables: torch.Tensor, - capacities: torch.Tensor, - valid_len: torch.Tensor, *, wq_a: torch.Tensor, q_norm_w: torch.Tensor, @@ -336,6 +509,103 @@ def dspark_attention_forward( eps: float, softmax_scale: float, freqs_cis: torch.Tensor, + persist: bool = False, + q_b_norm_enabled: bool = True, +) -> torch.Tensor: + """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). + + Functional port of DeepSpec ``DSparkAttention.forward`` for the dense + (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) + with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys + gathered from a rolling captured-context window (``kv_cache``, into which the + projected ``main_x`` context is written at ``start_pos % window_size``) plus the + block's own positions, attention-sink softmax, inverse-RoPE on the output, and a + grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). + + Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / + dequantized projection weights); ``wo_a`` is the raw grouped weight matrix + ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is + ``[b, window_size, head_dim]`` and is updated functionally (cloned). + + Returns: + ``[b, block_size, dim]`` attention output (residual stream contribution). + """ + assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" + b, block, _ = x.shape + rd = rope_head_dim + # Host-int slices run off the end of the table silently, unlike the batched + # path's gathers: a short blk_freqs drops RoPE from the block tail and shows + # up only as lower acceptance. Fail loudly instead. + assert start_pos + block < freqs_cis.shape[0], ( + f"DSpark RoPE table holds {freqs_cis.shape[0]} entries, short of the " + f"{start_pos + block + 1} the block at start_pos {start_pos} needs" + ) + main_freqs = freqs_cis[start_pos : start_pos + 1] + blk_freqs = freqs_cis[start_pos + 1 : start_pos + 1 + block] + + # Captured-context K/V from main_x (MQA, shared across heads). + main_kv = _rmsnorm(F.linear(main_x, wkv), kv_norm_w, eps) # [b, 1, head_dim] + main_kv = _rope_last_dims(main_kv, rd, main_freqs) + + # Query: low-rank + per-head RMS + RoPE. + q = _rmsnorm(F.linear(x, wq_a), q_norm_w, eps) + q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [b, block, h, head_dim] + # Per-head RMS in the query dtype (matches the reference inline normalization, + # which is NOT the fp32 RMSNorm path). + if q_b_norm_enabled: + q = q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + q = _rope_last_dims(q, rd, blk_freqs) + + # Block K/V. + kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] + kv = _rope_last_dims(kv, rd, blk_freqs) + + # Write the context K/V into the rolling window, then attend over + # [window context | block] with the sink. ``persist=True`` writes through + # to the caller's buffer (cross-step decode, worker-owned window); the + # default clones so single-shot callers (golden / unit tests) stay pure. + cache = kv_cache if persist else kv_cache.clone() + cache[:, start_pos % window_size] = main_kv.squeeze(1) + kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] + topk = get_dspark_topk_idxs(window_size, b, block, start_pos, device=x.device) + o = dspark_sparse_attn(q, kv_full, attn_sink, topk, softmax_scale) # [b, block, h, head_dim] + o = _rope_last_dims(o, rd, blk_freqs, inverse=True) + + # Grouped low-rank O projection. + o = o.reshape(b, block, n_groups, -1) + wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) + o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) + return F.linear(o.flatten(2), wo_b) + + +def dspark_attention_forward( + x: torch.Tensor, + main_x: torch.Tensor, + start_pos: torch.Tensor, + kv_cache: torch.Tensor, + block_tables: Optional[torch.Tensor] = None, + capacities: Optional[torch.Tensor] = None, + valid_len: Optional[torch.Tensor] = None, + *, + wq_a: Optional[torch.Tensor], + q_norm_w: torch.Tensor, + wq_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + wkv: Optional[torch.Tensor], + kv_norm_w: torch.Tensor, + wo_a: Optional[torch.Tensor], + wo_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + attn_sink: torch.Tensor, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_groups: int, + o_lora_rank: int, + window_size: int, + eps: float, + softmax_scale: float, + freqs_cis: torch.Tensor, + q_b_norm_enabled: bool = True, + persist: bool = False, qkv_proj: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, o_a_proj: Optional[torch.Tensor] = None, o_a_scale: Optional[torch.Tensor] = None, @@ -352,6 +622,34 @@ def dspark_attention_forward( contains RoPE frames (logical token positions + 1); ``valid_len`` counts initialized context tokens. Zero capacities disable synthetic page access. """ + if block_tables is None: + return _dspark_attention_forward_legacy( + x, + main_x, + start_pos, + kv_cache, + wq_a=wq_a, + q_norm_w=q_norm_w, + wq_b=wq_b, + wkv=wkv, + kv_norm_w=kv_norm_w, + wo_a=wo_a, + wo_b=wo_b, + attn_sink=attn_sink, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + n_groups=n_groups, + o_lora_rank=o_lora_rank, + window_size=window_size, + eps=eps, + softmax_scale=softmax_scale, + freqs_cis=freqs_cis, + q_b_norm_enabled=q_b_norm_enabled, + persist=persist, + ) + if capacities is None or valid_len is None: + raise ValueError("Paged DSpark attention requires capacities and valid lengths") if window_size != 128 or rope_head_dim != 64: raise ValueError("Embedded DSpark attention requires a 128-token window and 64 RoPE dims") g, block, _ = x.shape @@ -361,24 +659,32 @@ def dspark_attention_forward( blk_pos = start_pos.unsqueeze(1) + 1 + torch.arange(block, device=x.device) # [G, block] blk_freqs = freqs_cis[blk_pos] # [G, block, rd//2] - # Keep the two GEMM outputs live until dispatch: the fused path lets - # their RMSNorm/RoPE kernels store directly into attention's physical inputs. - main_kv_input = F.linear(main_x, wkv) + q_latent, block_kv_input, main_kv_input = _qkv_projections( + x, + main_x, + wq_a=wq_a, + wkv=wkv, + q_norm_w=q_norm_w, + qkv_proj=qkv_proj, + use_fused_fp8_input=use_fused_qkv_input_projection, + ) # Query: low-rank + per-head RMS + RoPE. q = _q_projection( q_latent, q_norm_w, eps, wq_b, use_fused_q_input_projection, blk_freqs ).unflatten(-1, (n_heads, head_dim)) - q = _rmsnorm_rope_batched( - q, - None, - eps, - rd, - blk_freqs, - num_heads=n_heads, - apply_weight=False, - apply_rmsnorm=q_b_norm_enabled, - ) + if q_b_norm_enabled: + q = _rmsnorm_rope_batched( + q, + None, + eps, + rd, + blk_freqs, + num_heads=n_heads, + apply_weight=False, + ) + else: + q = _rope_last_dims(q, rd, blk_freqs) if IS_CUTLASS_DSL_AVAILABLE and _get_dspark_arch_str() is not None and block in (5, 6): main_rope_freqs = torch.view_as_real(main_freqs).reshape(-1, rd // 2, 2).contiguous() @@ -473,6 +779,195 @@ def dspark_attention_forward( o = torch.einsum("bsht,btd->bshd", probs, kv.float()).to(q.dtype) o = _rope_last_dims_batched(o, rd, blk_freqs, inverse=True) + return _output_projection( + o, + wo_a=wo_a, + wo_b=wo_b, + n_groups=n_groups, + o_lora_rank=o_lora_rank, + o_a_proj=o_a_proj, + o_a_scale=o_a_scale, + o_a_dequant=o_a_dequant, + use_fused_output_projection=use_fused_output_projection, + ) + + +def dspark_attention_forward_batched( + x: torch.Tensor, + main_x: torch.Tensor, + start_pos: torch.Tensor, + kv_cache: torch.Tensor, + slots: torch.Tensor, + valid_len: torch.Tensor | None = None, + *, + wq_a: torch.Tensor, + q_norm_w: torch.Tensor, + wq_b: torch.Tensor, + wkv: torch.Tensor, + kv_norm_w: torch.Tensor, + wo_a: torch.Tensor, + wo_b: torch.Tensor, + attn_sink: torch.Tensor, + n_heads: int, + head_dim: int, + rope_head_dim: int, + n_groups: int, + o_lora_rank: int, + window_size: int, + eps: float, + softmax_scale: float, + freqs_cis: torch.Tensor, + persist: bool = False, + q_b_norm_enabled: bool = True, +) -> torch.Tensor: + """Batched, CUDA-graph-safe captured-context DSpark draft attention. + + This path belongs to the embedded DeepSeek-V4-Pro draft. Standalone DSpark + drafters (including Kimi-K3-DSpark) use the DFlash paged-KV decode instead. + + Numerically identical, per request, to :func:`dspark_attention_forward`, but + free of host syncs and data-dependent shapes so it can be captured into a CUDA + graph (the one-engine drafter runs inside the target's graph). The differences + from the scalar path are purely mechanical: + + * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen + request) instead of a python int; RoPE phases are *gathered* per request from + the fixed ``freqs_cis`` table rather than sliced. + * the rolling-window context K/V is written/read through the ``slots`` index + into a shared ``kv_cache`` (``persist=True`` writes through to the caller's + worker-owned buffer; otherwise a clone is used), instead of mutating a + per-request cache in place. + * the attended-position list has the fixed width ``window_size + block_size`` + with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). + + Args: + x: ``[G, block, dim]`` block layer input (per gen request). + main_x: ``[G, 1, hidden]`` projected captured context. + start_pos: ``[G]`` int tensor of absolute decode positions (> 0). + kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows + (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). + slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. + valid_len: optional ``[G]`` count of actually written context entries; + masks holes left when absolute positions are bootstrapped without + receiving the corresponding DSpark rolling-window state. + freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; + must satisfy ``maxlen > start_pos.max() + block_size``. + + Returns: + ``[G, block, dim]`` attention output (residual stream contribution). + """ + g, block, _ = x.shape + if kv_cache.shape[1] != window_size: + raise ValueError( + f"kv_cache window extent {kv_cache.shape[1]} does not match window_size {window_size}" + ) + rd = rope_head_dim + # Per-request RoPE phases gathered from the fixed table (no host-int slicing). + main_freqs = freqs_cis[start_pos].unsqueeze(1) # [G, 1, rd//2] + blk_pos = start_pos.unsqueeze(1) + 1 + torch.arange(block, device=x.device) # [G, block] + blk_freqs = freqs_cis[blk_pos] # [G, block, rd//2] + + # Keep the two GEMM outputs live until dispatch: the fused path lets + # their RMSNorm/RoPE kernels store directly into attention's physical inputs. + main_kv_input = F.linear(main_x, wkv) + + # Query: low-rank + per-head RMS + RoPE. + q = _rmsnorm_rope_batched(F.linear(x, wq_a), q_norm_w, eps, 0, blk_freqs) + q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [G, block, h, head_dim] + q = _rmsnorm_rope_batched( + q, + None, + eps, + rd, + blk_freqs, + num_heads=n_heads, + apply_weight=False, + apply_rmsnorm=q_b_norm_enabled, + ) + + block_kv_input = F.linear(x, wkv) + + # ``persist=True`` writes through to the worker-owned rolling window; + # single-shot callers keep the original functional clone behavior. + write_target = kv_cache if persist else kv_cache.clone() + main_rope_freqs = torch.view_as_real(main_freqs).reshape(-1, rd // 2, 2).contiguous() + inverse_rope_freqs = torch.view_as_real(blk_freqs).contiguous() + block_rope_freqs = inverse_rope_freqs.reshape(-1, rd // 2, 2) + + # The DSV4 fused kernel requires explicit physical-window validity; None + # keeps the reference path. + use_fused_dsv4_dspark_attention = ( + valid_len is not None + and IS_CUTLASS_DSL_AVAILABLE + and is_fused_dsv4_dspark_attention_supported( + q, write_target, valid_len, attn_sink, inverse_rope_freqs + ) + and is_fused_dspark_attention_preparation_supported( + main_kv_input, + block_kv_input, + kv_norm_w, + main_rope_freqs, + block_rope_freqs, + write_target, + slots, + start_pos, + ) + ) + if use_fused_dsv4_dspark_attention: + slots_i32, cache_seqs = cute_dsl_dspark_rmsnorm_rope_cache_write( + main_kv_input, + kv_norm_w, + main_rope_freqs, + write_target, + slots, + start_pos, + eps, + ) + draft_block = cute_dsl_dspark_rmsnorm_rope_draft_block( + block_kv_input, + kv_norm_w, + block_rope_freqs, + eps, + ) + # Two logical pages alias the same ring. Shifting the exclusive end + # into the second page preserves wraparound and the initialized suffix. + o = fused_dsv4_dspark_attention( + q, + draft_block, + write_target, + torch.stack((slots_i32, slots_i32), dim=1), + cache_seqs.to(torch.int64).remainder(window_size) + window_size + 1, + valid_len.to(torch.int64), + torch.full_like(start_pos, 2 * window_size, dtype=torch.int64), + attn_sink, + inverse_rope_freqs, + softmax_scale, + ) + else: + main_kv = _rmsnorm_rope_batched(main_kv_input, kv_norm_w, eps, rd, main_freqs) + kv = _rmsnorm_rope_batched(block_kv_input, kv_norm_w, eps, rd, blk_freqs) + main_kv_flat = main_kv.squeeze(1).to(write_target.dtype) + + slot_pos = start_pos % window_size # [G] + write_target[slots, slot_pos] = main_kv_flat + cache_rows = write_target[slots] # [G, window, head_dim] + kv_full = torch.cat([cache_rows, kv], dim=1) # [G, window + block, head_dim] + topk = get_dspark_topk_idxs_batched(window_size, block, start_pos, valid_len) + o = dspark_sparse_attn( + q, kv_full, attn_sink, topk, softmax_scale + ) # [G, block, h, head_dim] + o = _rmsnorm_rope_batched( + o, + kv_norm_w, + eps, + rd, + blk_freqs, + num_heads=n_heads, + apply_weight=False, + apply_rmsnorm=False, + inverse_rope=True, + ) + # Grouped low-rank O projection. o = o.reshape(g, block, n_groups, -1) wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) @@ -813,15 +1308,20 @@ def __init__( dtype=config.torch_dtype, quant_config=self._capture_quant_config(model_config), skip_create_weights_in_init=model_config.skip_create_weights_in_init, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, ) self.main_norm = RMSNorm( - hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, ) # Last stage: the block-draft output heads + mHC head + final norm. if self.has_heads: self.norm = RMSNorm( - hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, ) if self.uses_hc_head: self.hc_head = HCHead(config.hc_mult, config.hidden_size) @@ -1002,45 +1502,54 @@ def __init__( # Fallback: _dspark_freqs_table prefers _runtime_position_ceiling. self._freqs_cap = _runtime_position_cap(model_config, config, self.block_size + 2) self._freqs_table_cache: Dict = {} + self._context_wkv_fp8 = None + self._use_bf16_linear = _use_dspark_bf16_linear(model_config.mapping) + + def post_load_weights(self, weights: Optional[Dict] = None) -> None: + """Transform draft linears and bind the selected attention backend.""" + # DSparkForCausalLM.load_weights supplies the complete raw draft + # checkpoint before the generic zero-argument post-load walk. Only that + # explicit draft load owns these derived projections; a target-model + # reload does not reload the separate draft model. + # The separately constructed draft config does not receive the target + # model's SM107 CuTe block-scaling default. Select the same backend for + # the draft attention linears before their one-time weight transform; + # otherwise they materialize DeepGEMM int32 scales and cannot consume + # the fused Rubin attention/projection layouts. + if weights is not None and not self._use_bf16_linear and get_sm_version() == 107: + for stage in self.mtp_layers: + for module in stage.self_attn.modules(): + if isinstance(module, Linear): + module.use_cute_dsl_blockscaling_mm = True - def post_load_weights(self) -> None: - """Run the one-shot post-load transforms for the draft's quant linears. - - The fp8 UE8M0 linears we invoke as modules (``main_proj``, shared experts, - the heads) need ``resmooth_to_fp8_e8m0`` + ``transform_sf_into_required_layout`` - before the first forward, or the kernel reads raw scales and emits NaNs. - ``Linear.transform_weights`` is idempotent; the routed-expert MoE packs - itself in its own ``load_weights``. - - The bf16 captured-context attention does NOT use the MLA module's forward — - it runs ``dspark_attention_forward`` on dequantized bf16 weights cached via - :meth:`cache_attn_weights_from_checkpoint` — so the MLA projection linears are - skipped here (they would otherwise be transformed into the deep_gemm layout we - don't consume). - """ attn_linear_ids = set() - for stage in self.mtp_layers: - for m in stage.self_attn.modules(): - if isinstance(m, Linear): - attn_linear_ids.add(id(m)) + if weights is None or self._use_bf16_linear: + for stage in self.mtp_layers: + for module in stage.self_attn.modules(): + if isinstance(module, Linear): + attn_linear_ids.add(id(module)) for module in self.modules(): if isinstance(module, Linear) and id(module) not in attn_linear_ids: module.transform_weights() + if weights is None: + return + if self._use_bf16_linear: + self._bind_bf16_attn_projections(weights) + else: + self._bind_attn_projections() @staticmethod - def _block_dequant(w_fp8: torch.Tensor, scale: torch.Tensor, block: int = 128) -> torch.Tensor: - """DeepSeek ``block``×``block`` block-scale dequant → bf16: ``real = fp8 * scale``. - - ``scale`` (possibly UE8M0) is broadcast over each ``block``×``block`` tile. - Pure-torch (matches the golden-validated reference dequant), robust to the - e8m0 scale dtype. - """ - wf = w_fp8.float() - out, inn = wf.shape - s = scale.float() - s_full = s.repeat_interleave(block, 0)[:out].repeat_interleave(block, 1)[:, :inn] - return (wf * s_full).to(torch.bfloat16) + def _block_dequant( + weight_fp8: torch.Tensor, scale: torch.Tensor, block: int = 128 + ) -> torch.Tensor: + """Dequantize a block-scaled FP8 checkpoint weight to BF16.""" + weight = weight_fp8.float() + output_features, input_features = weight.shape + scale = scale.float() + scale = scale.repeat_interleave(block, 0)[:output_features] + scale = scale.repeat_interleave(block, 1)[:, :input_features] + return (weight * scale).to(torch.bfloat16) def _cache_attn_weights(self, src: Dict) -> None: """Populate each stage's ``_dspark_attn`` from a dict of raw ``mtp.{s}.attn.*`` @@ -1053,6 +1562,8 @@ def _cache_attn_weights(self, src: Dict) -> None: ``q_a``+``kv`` and stores the scale interleaved). This mirrors the golden-validated dequant exactly. """ + self._use_bf16_linear = True + self._context_wkv_fp8 = None for s, stage in enumerate(self.mtp_layers): pref = f"mtp.{s}.attn." dev = stage.input_layernorm.weight.device @@ -1256,48 +1767,110 @@ def _draft_quant_config_dict(model_config, base: int, num_stages: int): changed = True return new_qcd if changed else None - def _project_context_wkv(self, main_x: torch.Tensor) -> Optional[torch.Tensor]: - """Quantize once and project every stage's WKV with one stacked GEMM.""" - if self._context_wkv_fp8 is None: - return None - rows = main_x.reshape(-1, main_x.shape[-1]) - input_fp8, input_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(rows) - weight, weight_scale = self._context_wkv_fp8 - output = torch.ops.trtllm.cute_dsl_dspark_context_wkv_gemm_rubin( - input_fp8, weight, input_scale, weight_scale - ) - head_dim = self._attn_params["head_dim"] - return output.view(*main_x.shape[:-1], len(self.mtp_layers), head_dim) + @torch.inference_mode() + def write_context_windows( + self, + main_hidden: torch.Tensor, + positions: torch.Tensor, + stage_windows: torch.Tensor, + ) -> None: + """Write captured-context ``main_kv`` into the rolling per-stage KV windows. + + Replicates exactly the per-position context write that + :func:`dspark_attention_forward` performs each generation step + (``main_kv = RoPE_pos(rmsnorm(wkv @ main_x))`` written at + ``window[pos % window_size]``), but for an arbitrary set of + ``(captured-hidden, position)`` pairs. Used to (a) seed a request's + window from its prompt at prefill and (b) back-fill the intermediate + accepted tokens of a multi-accept step — both of which the per-step + generation path would otherwise leave as holes, starving the draft + attention of context (acceptance-rate only; verified decoding keeps + output correctness regardless). - def write_context_pages( + Args: + main_hidden: ``[M, num_capture * hidden]`` captured target hiddens. + positions: ``[M]`` absolute window positions (used for BOTH the RoPE + phase and the slot ``pos % window_size``). By the generation-path + convention this is ``committed_position + 1``. Must hold at most + ``window_size`` entries with distinct slots (the caller passes a + contiguous, deduplicated range) so the scatter is well defined. + stage_windows: ``[num_stages, window_size, head_dim]`` window for one + request's slot; updated in place. + """ + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + return + M = int(main_hidden.shape[0]) + if M == 0: + return + win = int(self._attn_params["window_size"]) + rd = int(self._attn_params["rope_head_dim"]) + eps = float(self._attn_params["eps"]) + positions = positions.to(main_hidden.device).long() + # main_x is stage-invariant (stage 0's projection), matching forward_embed. + stage0 = self.mtp_layers[0] + main_x = stage0.main_norm(stage0.main_proj(main_hidden)) # [M, hidden] + freqs = self._dspark_freqs_table(main_x.device)[positions] + slots = positions % win # [M] + mx = main_x.unsqueeze(0) # [1, M, hidden] for the per-position RoPE layout + for s, stage in enumerate(self.mtp_layers): + a = stage._dspark_attn + kv = _rmsnorm(F.linear(mx, a["wkv"]), a["kv_norm_w"], eps) # [1, M, head_dim] + kv = _rope_last_dims(kv, rd, freqs) # [1, M, head_dim] + stage_windows[s, slots] = kv[0].to(stage_windows.dtype) + + def write_context_windows_batched( self, main_hidden: torch.Tensor, positions: torch.Tensor, + slots: torch.Tensor, mask: torch.Tensor, - pools: tuple[torch.Tensor, ...], - block_tables: torch.Tensor, - capacities: torch.Tensor, + kv_windows: torch.Tensor, ) -> None: - """Project captured context and store only valid tokens in managed pages. + """CUDA-graph-safe batched + masked variant of :meth:`write_context_windows`. - Inputs have fixed [batch, tokens] leading dimensions. Positions are RoPE - frames (logical token position + 1), including for chunked prefill. + Back-fills the intermediate accepted tokens of a multi-accept step into the + rolling per-stage KV windows for ALL gen requests at once, with a fixed + ``[G, M]`` shape (``M = max interim per request``) and a validity mask + (invalid entries are no-ops via a read-modify-write), so nothing depends on + the per-request accept count. Same per-position math as the scalar version + (``RoPE_pos(rmsnorm(wkv @ main_x))`` written at ``window[pos % window]``), + but indexed/scattered through ``slots`` into the shared persistent buffer. + + Args: + main_hidden: ``[G, M, num_capture * hidden]`` captured target hiddens + (rows beyond a request's interim count are masked). + positions: ``[G, M]`` absolute window positions (RoPE phase + slot). + slots: ``[G]`` row index of each request into ``kv_windows``. + mask: ``[G, M]`` bool — which ``(g, m)`` entries are real interim writes. + kv_windows: ``[N, num_stages, window_size, head_dim]`` persistent buffer; + updated in place. """ - if main_hidden.numel() == 0 or getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: return + G, M = positions.shape + if G == 0 or M == 0: + return + win = int(self._attn_params["window_size"]) rd = int(self._attn_params["rope_head_dim"]) eps = float(self._attn_params["eps"]) + positions = positions.long() + slots = slots.long() safe_positions = torch.where(mask, positions, torch.zeros_like(positions)) - freqs = self._dspark_freqs_table(main_hidden.device)[safe_positions] + freqs = self._dspark_freqs_table(main_hidden.device)[safe_positions] # [G, M, rd//2] + cols = positions % win # [G, M] + rows = slots[:, None].expand(-1, M) # [G, M] + mask3 = mask.unsqueeze(-1) # [G, M, 1] stage0 = self.mtp_layers[0] - main_x = stage0.main_norm(stage0.main_proj(main_hidden)) - head_dim = int(self._attn_params["head_dim"]) - projected_wkv = self._project_context_wkv(main_x) - for s, (stage, pool) in enumerate(zip(self.mtp_layers, pools)): + main_x = stage0.main_norm(stage0.main_proj(main_hidden)) # [G, M, hidden] + for s, stage in enumerate(self.mtp_layers): a = stage._dspark_attn kv = _rmsnorm(F.linear(main_x, a["wkv"]), a["kv_norm_w"], eps) # [G, M, head_dim] kv = _rope_last_dims_batched(kv, rd, freqs) # [G, M, head_dim] - write_dspark_context(kv, positions, mask, pool[:, 0, 0], block_tables, capacities) + win_s = kv_windows[:, s] # [N, win, head_dim] view onto the base buffer + # Read-modify-write so masked-out (g, m) entries keep their current + # value — a graph-safe masked scatter (no dynamic-shape compaction). + cur = win_s[rows, cols] # [G, M, head_dim] + win_s[rows, cols] = torch.where(mask3, kv.to(win_s.dtype), cur) def forward_embed( self, main_hidden: torch.Tensor, bonus_token_ids: torch.Tensor @@ -1373,7 +1946,18 @@ def _forward_stage( **stage._dspark_attn, **self._attn_params, ) - if stage.enable_fused_hc: + if pre_mix is not None: + residual = stage.hc_attn.post_mapping( + x=attn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + ffn_pre, post_mix, comb_mix, layer_input = stage.hc_ffn.pre_mapping_lagged( + residual, attn_pre + ) + layer_input = stage.post_attention_layernorm(layer_input) + elif stage.enable_fused_hc: residual, post_mix, comb_mix, layer_input = stage.hc_ffn.fused_hc( x_prev=attn, residual_prev=residual, @@ -1433,10 +2017,11 @@ def forward( bonus_token_ids: torch.Tensor, start_pos: torch.Tensor, *, - kv_pages: tuple[torch.Tensor, ...], - block_tables: torch.Tensor, - capacities: torch.Tensor, - valid_len: torch.Tensor, + kv_pages: Optional[tuple[torch.Tensor, ...]] = None, + kv_windows: Optional[torch.Tensor] = None, + block_tables: Optional[torch.Tensor] = None, + capacities: Optional[torch.Tensor] = None, + valid_len: Optional[torch.Tensor] = None, temperature: float = 0.0, return_confidence: bool = False, return_logits: bool = False, @@ -1451,10 +2036,22 @@ def forward( zero for synthetic rows. Returns draft tokens and proposal lengths, optionally including the per-position logits. """ + if kv_pages is None: + return self._forward_legacy( + main_hidden, + bonus_token_ids, + start_pos, + kv_windows=kv_windows, + temperature=temperature, + return_confidence=return_confidence, + return_logits=return_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + if block_tables is None or capacities is None or valid_len is None: + raise ValueError("Managed DSpark pages require tables, capacities and valid lengths") if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: raise RuntimeError( - "DSpark attention weights not cached; call " - "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." + "DSpark attention projections are not bound; call post_load_weights() after loading." ) x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) main_x = main_x.unsqueeze(1) # [G, 1, hidden] for the MQA K/V projection @@ -1490,6 +2087,82 @@ def forward( pre_mix=pre_mix, ) + def forward_batched( + self, + main_hidden: torch.Tensor, + bonus_token_ids: torch.Tensor, + start_pos: torch.Tensor, + *, + kv_windows: torch.Tensor, + slots: torch.Tensor, + valid_len: Optional[torch.Tensor] = None, + temperature: float = 0.0, + return_confidence: bool = False, + return_logits: bool = False, + all_rank_num_tokens: Optional[List[int]] = None, + ) -> tuple: + """CUDA-graph-safe batched block-draft forward (all gen requests at once). + + Same computation as :meth:`forward`, but every host-int / data-dependent + operation is tensorized so the whole path can be captured into the target's + CUDA graph (DSpark is a one-engine drafter — its worker runs inside the + graph). ``start_pos`` is a ``[G]`` tensor (one absolute decode position per + gen request); the rolling captured-context windows are written/read through + ``slots`` into the worker-owned ``kv_windows`` buffer; RoPE phases are + gathered from a fixed table. ``forward_head`` always proposes the full + block; ``return_confidence`` only adds a fixed-shape scoring matmul, so + both settings stay capturable. + + Args: + main_hidden: ``[G, num_capture * hidden]`` captured target context. + bonus_token_ids: ``[G]`` last accepted token per gen request. + start_pos: ``[G]`` int tensor of absolute decode positions (> 0). + kv_windows: ``[N, num_stages, window_size, head_dim]`` persistent rolling + windows; written in place through ``slots``. + slots: ``[G]`` int tensor mapping each request to its ``kv_windows`` row. + valid_len: ``[G]`` count of actually written rolling-window entries. + Returns: + ``(draft_tokens [G, block], confidence [G, block] or None)``. + """ + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + raise RuntimeError( + "DSpark attention weights not cached; call " + "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." + ) + x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) + main_x = main_x.unsqueeze(1) # [G, 1, hidden] for the MQA K/V projection + freqs_cis = self._dspark_freqs_table(x.device) + moe_input_ids = draft_ids.reshape(-1) + + h = x + pre_mix = self._initial_pre_mix(x) + for s, stage in enumerate(self.mtp_layers): + stage_window = kv_windows[:, s] # [N, window_size, head_dim] + h = self._forward_stage_legacy( + stage, + h, + main_x, + start_pos, + freqs_cis, + moe_input_ids, + stage_window, + slots, + valid_len, + all_rank_num_tokens=all_rank_num_tokens, + pre_mix=pre_mix, + ) + if self.uses_lagged_pre: + h, pre_mix = h + + return self.forward_head( + h, + bonus_token_ids, + temperature=temperature, + return_confidence=return_confidence, + return_logits=return_logits, + pre_mix=pre_mix, + ) + def _initial_pre_mix(self, hidden_states: torch.Tensor) -> Optional[torch.Tensor]: """V4.1 seeds the lagged input mix with a one-hot on residual stream 0.""" if not self.uses_lagged_pre: @@ -1541,114 +2214,524 @@ def run_moe_lockstep_noop( do_finalize=True, ) - def forward_head( + def forward_head( + self, + block_hidden: torch.Tensor, + bonus_token_ids: torch.Tensor, + *, + temperature: float = 0.0, + return_confidence: bool = False, + return_logits: bool = False, + pre_mix: Optional[torch.Tensor] = None, + ) -> tuple: + """Block-draft head: hc_head + norm + lm_head -> markov refine + confidence. + + ``block_hidden`` is the last stage's mHC residual ``[*, block, hc_mult, hidden]``. + Returns (draft_tokens [*, block], confidence [*, block] or None); with ``return_logits`` + also returns the per-position draft logits [*, block, vocab]. + """ + last = self.mtp_layers[-1] + if self.uses_lagged_pre: + if pre_mix is None: + raise ValueError("V4.1 DSpark head requires the final FFN pre-mix") + h = last.hc_ffn.collapse(block_hidden, pre_mix) + else: + h = last.hc_head(block_hidden) + confidence_hidden = h if return_confidence and last.confidence_head is not None else None + normalized = last.norm(h) + base_logits = self.lm_head(normalized) + proposal_hidden = normalized if self.confidence_uses_normalized_hidden else h + return dspark_propose( + base_logits, + bonus_token_ids=bonus_token_ids, + block_hidden=proposal_hidden, + markov_head=last.markov_head, + confidence_head=last.confidence_head, + block_size=self.block_size, + temperature=temperature, + return_confidence=return_confidence, + return_logits=return_logits, + confidence_hidden=confidence_hidden, + ) + + def _bind_bf16_attn_projections(self, weights: Dict) -> None: + """Bind all five captured-context attention projections as BF16 tensors.""" + for stage_id, stage in enumerate(self.mtp_layers): + prefix = f"mtp.{stage_id}.attn." + device = stage.input_layernorm.weight.device + + def dequant(name: str) -> torch.Tensor: + weight = weights[f"{prefix}{name}.weight"].to(device) + scale = weights[f"{prefix}{name}.scale"].to(device) + return self._block_dequant(weight, scale) + + stage._dspark_attn = dict( + wq_a=dequant("wq_a"), + q_norm_w=weights[f"{prefix}q_norm.weight"].to(device=device, dtype=torch.bfloat16), + wq_b=dequant("wq_b"), + wkv=dequant("wkv"), + kv_norm_w=weights[f"{prefix}kv_norm.weight"].to( + device=device, dtype=torch.bfloat16 + ), + wo_a=dequant("wo_a"), + wo_b=dequant("wo_b"), + attn_sink=weights[f"{prefix}attn_sink"].to(device=device, dtype=torch.float32), + qkv_proj=None, + o_a_proj=None, + o_a_scale=None, + o_a_dequant=None, + use_fused_output_projection=False, + use_fused_qkv_input_projection=False, + use_fused_q_input_projection=False, + ) + self._context_wkv_fp8 = None + + def _bind_attn_projections(self) -> None: + """Bind loaded MLA modules to the DSpark attention functional path. + + DeepSeek-V4 loads q_a+wkv as one block-scaled FP8 ``Linear`` and keeps + o_a as a grouped FP8 weight. Reusing those objects avoids a second BF16 + copy of every attention projection and preserves the Rubin CuTe DSL + weight/scale layouts produced by ``Linear.transform_weights``. + """ + expected_heads = self._attn_params["n_heads"] + expected_groups = self._attn_params["n_groups"] + head_dim = self._attn_params["head_dim"] + for stage in self.mtp_layers: + attn = stage.self_attn + if attn.num_heads_tp != expected_heads or attn.n_local_groups != expected_groups: + raise ValueError( + "DSpark FP8 attention requires unsharded attention projections; " + f"got {attn.num_heads_tp} heads and {attn.n_local_groups} groups, " + f"expected {expected_heads} heads and {expected_groups} groups. " + "Enable attention DP (or use TP=1)." + ) + qkv_in_features = attn.kv_a_proj_with_mqa.in_features + q_lora_rank = attn.q_a_layernorm.weight.shape[0] + use_fused_qkv_input_projection = _can_use_fused_fp8_input_projection( + attn.kv_a_proj_with_mqa, qkv_in_features + ) + use_fused_q_input_projection = _can_use_fused_fp8_input_projection( + attn.q_b_proj, q_lora_rank + ) + if use_fused_qkv_input_projection or use_fused_q_input_projection: + from ..custom_ops.dspark_rmsnorm_rope_custom_op import ( + precompile_dspark_fp8_input_fusions, + ) + + precompile_dspark_fp8_input_fusions( + qkv_in_features if use_fused_qkv_input_projection else None, + q_lora_rank if use_fused_q_input_projection else None, + self._attn_params["eps"], + ) + stage._dspark_attn = dict( + wq_a=None, + q_norm_w=attn.q_a_layernorm.weight, + wq_b=attn.q_b_proj, + wkv=None, + kv_norm_w=attn.kv_a_layernorm.weight, + wo_a=None, + wo_b=attn.o_b_proj, + attn_sink=attn.mqa.attn_sink, + qkv_proj=attn.kv_a_proj_with_mqa, + o_a_proj=attn.o_a_proj, + o_a_scale=attn.o_a_proj_scale, + o_a_dequant=attn.o_a_proj_dequant, + use_fused_output_projection=_can_use_fused_fp8_output_projection( + attn.o_b_proj, + attn.o_a_proj, + attn.n_local_groups, + attn.o_a_proj.shape[1], + ), + use_fused_qkv_input_projection=use_fused_qkv_input_projection, + use_fused_q_input_projection=use_fused_q_input_projection, + ) + # Generation consumes q_a+wkv, while context-window seeding needs only + # the trailing wkv rows. Stack the three transformed WKV weights once so + # CTX runs one N=(stages*head_dim) GEMM instead of three small-N GEMMs. + # Scale layout conversion and the weight copy never enter serving. + self._context_wkv_fp8 = ( + _stack_fp8_projection_output_rows( + [stage.self_attn.kv_a_proj_with_mqa for stage in self.mtp_layers], + head_dim, + ) + if self.mtp_layers + and all( + stage._dspark_attn["use_fused_qkv_input_projection"] for stage in self.mtp_layers + ) + else None + ) + + def configure_generation_autotuner(self, max_batch: int) -> None: + """Give DSpark projection runners the worker's real generation capacity.""" + if self._use_bf16_linear: + return + if max_batch < 1: + raise ValueError(f"DSpark max_batch must be positive, got {max_batch}") + + from ..custom_ops.dspark_attention_custom_op import ( + _get_dspark_generation_mxfp8_tuning_config, + ) + + projection_rows = ( + ("qkv_proj", self.block_size + 1), + ("wq_b", self.block_size), + ("wo_b", self.block_size), + ) + for stage in self.mtp_layers: + for name, rows_per_request in projection_rows: + projection = stage._dspark_attn[name] + projection._dspark_generation_tuning_config = ( + _get_dspark_generation_mxfp8_tuning_config( + rows_per_request, + max_batch, + ) + ) + if name == "wo_b": + projection._dspark_generation_max_batch = max_batch + projection._dspark_generation_rows_per_request = rows_per_request + + def _project_context_wkv(self, main_x: torch.Tensor) -> Optional[torch.Tensor]: + """Quantize once and project every stage's WKV with one stacked GEMM.""" + if self._context_wkv_fp8 is None: + return None + rows = main_x.reshape(-1, main_x.shape[-1]) + input_fp8, input_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(rows) + weight, weight_scale = self._context_wkv_fp8 + output = torch.ops.trtllm.cute_dsl_dspark_context_wkv_gemm_rubin( + input_fp8, weight, input_scale, weight_scale + ) + head_dim = self._attn_params["head_dim"] + return output.view(*main_x.shape[:-1], len(self.mtp_layers), head_dim) + + def write_context_pages( + self, + main_hidden: torch.Tensor, + positions: torch.Tensor, + mask: torch.Tensor, + pools: tuple[torch.Tensor, ...], + block_tables: torch.Tensor, + capacities: torch.Tensor, + ) -> None: + """Project captured context and store only valid tokens in managed pages. + + Inputs have fixed [batch, tokens] leading dimensions. Positions are RoPE + frames (logical token position + 1), including for chunked prefill. + """ + if main_hidden.numel() == 0 or getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + return + rd = int(self._attn_params["rope_head_dim"]) + eps = float(self._attn_params["eps"]) + safe_positions = torch.where(mask, positions, torch.zeros_like(positions)) + freqs = self._dspark_freqs_table(main_hidden.device)[safe_positions] + stage0 = self.mtp_layers[0] + main_x = stage0.main_norm(stage0.main_proj(main_hidden)) + head_dim = int(self._attn_params["head_dim"]) + projected_wkv = self._project_context_wkv(main_x) + for s, (stage, pool) in enumerate(zip(self.mtp_layers, pools)): + a = stage._dspark_attn + kv = ( + projected_wkv[..., s, :] + if projected_wkv is not None + else ( + a["qkv_proj"](main_x)[..., -head_dim:] + if a.get("qkv_proj") is not None + else F.linear(main_x, a["wkv"]) + ) + ) + kv = _rmsnorm(kv, a["kv_norm_w"], eps) # [G, M, head_dim] + kv = _rope_last_dims_batched(kv, rd, freqs) # [G, M, head_dim] + write_dspark_context(kv, positions, mask, pool[:, 0, 0], block_tables, capacities) + + +class DSv4DSparkForCausalLM(nn.Module): + """One-engine draft wrapper for DSpark (mirrors ``DFlashForCausalLM``). + + Wraps :class:`DSv4DSparkDraftModel` (the ``n_mtp_layers``-stage ``mtp.*`` backbone) + for the single-engine external-drafter flow: created by ``get_draft_model``, + appended to the target's epilogue, and driven by ``DSv4DSparkWorker``. + + ``embed_tokens`` / ``lm_head`` are shared with the target model + (:meth:`load_weights_from_target_model`). The draft weights live in the SAME + checkpoint under ``mtp.*``; :meth:`load_weights` remaps them + (``remap_dspark_draft_keys``), loads via ``DeepseekV4WeightLoader``, runs the + fp8 ``post_load_weights`` transforms, and caches the bf16 captured-context + attention weights from the in-memory state dict. + """ + + draft_model_cls = DSv4DSparkDraftModel + + def __init__( + self, + draft_config, + aux_stream_dict=None, + num_stages=None, + block_size=None, + draft_moe_backend: Optional[str] = None, + ): + super().__init__() + self.dspark_model = self.draft_model_cls( + draft_config, + aux_stream_dict, + num_stages=num_stages, + block_size=block_size, + draft_moe_backend=draft_moe_backend, + ) + # Generic handles expected by the loader / weight mappers. + self.model = self.dspark_model + self.model_config = draft_config + self.config = draft_config.pretrained_config + # Worker-facing interface (the worker receives this wrapper as + # ``draft_model`` and calls forward()/reads these properties and scalars). + self.num_stages = self.dspark_model.num_stages + self._attn_params = self.dspark_model._attn_params + self.lm_head = None # shared from the target (load_weights_from_target_model) + self.logits_processor = None # set by the caller after construction + + @property + def block_size(self): + return self.dspark_model.block_size + + @property + def embed_tokens(self): + return self.dspark_model.embed_tokens + + def forward(self, main_hidden, bonus_token_ids, start_pos, **kwargs): + return self.dspark_model.forward(main_hidden, bonus_token_ids, start_pos, **kwargs) + + def forward_batched(self, main_hidden, bonus_token_ids, start_pos, **kwargs): + """CUDA-graph-safe batched draft forward (delegates to the draft model).""" + return self.dspark_model.forward_batched(main_hidden, bonus_token_ids, start_pos, **kwargs) + + def run_moe_lockstep_noop(self, all_rank_num_tokens, device): + """Empty-batch MoE barrier lockstep (delegates to the draft model).""" + return self.dspark_model.run_moe_lockstep_noop(all_rank_num_tokens, device) + + def write_context_windows(self, main_hidden, positions, stage_windows): + """Seed / back-fill the rolling KV windows (delegates to the draft model).""" + return self.dspark_model.write_context_windows(main_hidden, positions, stage_windows) + + def write_context_windows_batched(self, main_hidden, positions, slots, mask, kv_windows): + """Batched + masked window back-fill (delegates to the draft model).""" + return self.dspark_model.write_context_windows_batched( + main_hidden, positions, slots, mask, kv_windows + ) + + def _forward_stage_legacy( + self, + stage: "DSv4DSparkBlock", + h: torch.Tensor, + main_x: torch.Tensor, + start_pos, + freqs_cis: torch.Tensor, + moe_input_ids: torch.Tensor, + stage_window: Optional[torch.Tensor] = None, + slots: Optional[torch.Tensor] = None, + valid_len: Optional[torch.Tensor] = None, + all_rank_num_tokens: Optional[List[int]] = None, + pre_mix: Optional[torch.Tensor] = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """One DSpark stage = reference ``Block.forward`` with captured-context attn. + + ``h`` is the mHC residual stream ``[T, block, hc_mult, hidden]``. The mHC + ``pre_mapping``/``post_mapping`` preserve the leading ``[T, block]`` dims; + the captured-context attention and MoE run on the collapsed token axis. + Mirrors the reference (unfused) mHC boundaries exactly: + ``hc_pre → attn_norm → DSparkAttention → hc_post`` then + ``hc_pre → ffn_norm → MoE → hc_post``. + + ``stage_window`` is this stage's persistent rolling captured-context KV + window ``[T, window_size, head_dim]`` owned by the worker across decode + steps; the attention writes the current ``main_kv`` into it in place. When + ``None`` (golden / single-shot) a zero window is allocated per call. + + When ``slots`` (a ``[G]`` int tensor) is given, the CUDA-graph-safe batched + attention (:func:`dspark_attention_forward_batched`) is used instead: it + takes ``start_pos`` as a ``[G]`` tensor and writes/reads ``stage_window`` + (then shaped ``[N, window_size, head_dim]``) through the ``slots`` index. + """ + T, block, _, hidden = h.shape + + # --- attention sub-block (captured-context, not paged-KV MLA) --- + residual = h + if pre_mix is None: + post_mix, comb_mix, layer_input = stage.hc_attn.pre_mapping(residual) + else: + attn_pre, post_mix, comb_mix, layer_input = stage.hc_attn.pre_mapping_lagged( + residual, pre_mix + ) + layer_input = stage.input_layernorm(layer_input) # [T, block, hidden] + # Rolling-window cache: persist through the worker-owned ``stage_window`` + # for cross-step decode, else a fresh zero window for a single block. + persist = stage_window is not None + kv_cache = ( + stage_window + if persist + else torch.zeros( + T, + self._attn_params["window_size"], + self._attn_params["head_dim"], + dtype=torch.bfloat16, + device=h.device, + ) + ) + if slots is not None: + # Batched, CUDA-graph-safe path (start_pos is a [G] tensor; window is + # written/read through ``slots``). + attn = dspark_attention_forward_batched( + layer_input, + main_x, + start_pos, + kv_cache, + slots, + valid_len=valid_len, + freqs_cis=freqs_cis, + persist=True, + **stage._dspark_attn, + **self._attn_params, + ) + else: + attn = _dspark_attention_forward_legacy( + layer_input, + main_x, + start_pos, + kv_cache, + freqs_cis=freqs_cis, + persist=persist, + **stage._dspark_attn, + **self._attn_params, + ) + if pre_mix is not None: + residual = stage.hc_attn.post_mapping( + x=attn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + ffn_pre, post_mix, comb_mix, layer_input = stage.hc_ffn.pre_mapping_lagged( + residual, attn_pre + ) + layer_input = stage.post_attention_layernorm(layer_input) + elif stage.enable_fused_hc: + residual, post_mix, comb_mix, layer_input = stage.hc_ffn.fused_hc( + x_prev=attn, + residual_prev=residual, + post_mix_prev=post_mix, + comb_mix_prev=comb_mix, + norm_weight=stage.post_attention_layernorm.weight, + norm_eps=stage.post_attention_layernorm.variance_epsilon, + ) + else: + residual = stage.hc_attn.post_mapping( + x=attn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + post_mix, comb_mix, layer_input = stage.hc_ffn.pre_mapping(residual) + layer_input = stage.post_attention_layernorm(layer_input) + num_tokens = T * block + # FUSED_COMM MoE backends (DeepGEMM MegaMoE) size their in-kernel + # NVLink-barrier chunk loop from ``max(all_rank_num_tokens)`` and index + # the local slice by ``moe_ep_rank``, so every EP rank must pass the same + # globally-gathered per-rank list (here: gen tokens = num_gens * block per + # rank). Passing only the local ``[num_tokens]`` desyncs the phase-flip + # barrier across ranks (hang / "unspecified launch failure"). Fall back to + # the local count for single-rank / non-ADP runs where no list is threaded. + moe_all_rank_num_tokens = ( + all_rank_num_tokens if all_rank_num_tokens is not None else [num_tokens] + ) + # The draft captured-context MoE must mirror the target DeepseekV4MoE's TP + # reduction policy. Under attention DP each rank is data-parallel (owns a + # distinct set of requests) and needs no cross-rank reduction; but under + # plain tensor parallelism (attention_dp off, tp_size > 1) the expert-sharded + # MoE output must be all-reduced across ranks -- exactly what the target MoE + # does (enable_allreduce = not (POST_MOE_FUSION or tp_size == 1)). Previously + # this was hard-coded to False, which dropped the reduction on the non-ADP + # path, corrupting the draft block proposals and roughly halving DSpark + # acceptance length whenever attention_dp was disabled. + moe_enable_allreduce = ( + not self.model_config.mapping.enable_attention_dp + and self.model_config.mapping.tp_size > 1 + ) + moe_out = stage.mlp( + layer_input.reshape(num_tokens, hidden), + input_ids=moe_input_ids, + all_rank_num_tokens=moe_all_rank_num_tokens, + final_all_reduce_params=AllReduceParams(enable_allreduce=moe_enable_allreduce), + do_finalize=True, + ).reshape(T, block, hidden) + h = stage.hc_ffn.post_mapping( + x=moe_out, residual=residual, post_layer_mix=post_mix, comb_res_mix=comb_mix + ) + return (h, ffn_pre) if pre_mix is not None else h + + def _forward_legacy( self, - block_hidden: torch.Tensor, + main_hidden: torch.Tensor, bonus_token_ids: torch.Tensor, + start_pos: int, *, + kv_windows: Optional[torch.Tensor] = None, temperature: float = 0.0, return_confidence: bool = False, return_logits: bool = False, - pre_mix: Optional[torch.Tensor] = None, + all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: - """Block-draft head: hc_head + norm + lm_head -> markov refine + confidence. + """Full block-draft forward: chain the ``num_stages`` DSpark stages. - ``block_hidden`` is the last stage's mHC residual ``[*, block, hc_mult, hidden]``. - Returns (draft_tokens [*, block], confidence [*, block] or None); with ``return_logits`` - also returns the per-position draft logits [*, block, vocab]. + Mirrors the reference ``Transformer.forward_spec`` (generation path, + ``start_pos > 0``): ``forward_embed`` builds the block input + captured + context, each stage runs the captured-context backbone, and + ``forward_head`` emits the block draft tokens + per-position confidence. + + Args: + main_hidden: ``[T, num_capture * hidden]`` captured target context. + bonus_token_ids: ``[T]`` last accepted token per request. + start_pos: absolute decode position (must be > 0). + kv_windows: optional persistent per-stage rolling captured-context KV + windows ``[T, num_stages, window_size, head_dim]`` owned by the + worker; updated in place each call. ``None`` allocates fresh zero + windows (single-shot golden / test path). + Returns: + ``(draft_tokens [T, block], confidence [T, block] or None)``. """ - last = self.mtp_layers[-1] - if self.uses_lagged_pre: - if pre_mix is None: - raise ValueError("V4.1 DSpark head requires the final FFN pre-mix") - h = last.hc_ffn.collapse(block_hidden, pre_mix) - else: - h = last.hc_head(block_hidden) - confidence_hidden = h if return_confidence and last.confidence_head is not None else None - normalized = last.norm(h) - base_logits = self.lm_head(normalized) - proposal_hidden = normalized if self.confidence_uses_normalized_hidden else h - return dspark_propose( - base_logits, - bonus_token_ids=bonus_token_ids, - block_hidden=proposal_hidden, - markov_head=last.markov_head, - confidence_head=last.confidence_head, - block_size=self.block_size, + assert start_pos > 0, "DSpark draft runs at generation (start_pos > 0)" + if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: + raise RuntimeError( + "DSpark attention weights not cached; call " + "cache_attn_weights_from_checkpoint(ckpt_dir, weight_map) after loading." + ) + x, main_x, draft_ids = self.forward_embed(main_hidden, bonus_token_ids) + main_x = main_x.unsqueeze(1) # [T, 1, hidden] for the MQA K/V projection + freqs_cis = self._dspark_freqs_table(x.device) + moe_input_ids = draft_ids.reshape(-1) + + h = x + pre_mix = self._initial_pre_mix(x) + for s, stage in enumerate(self.mtp_layers): + stage_window = kv_windows[:, s] if kv_windows is not None else None + h = self._forward_stage_legacy( + stage, + h, + main_x, + start_pos, + freqs_cis, + moe_input_ids, + stage_window, + all_rank_num_tokens=all_rank_num_tokens, + pre_mix=pre_mix, + ) + if self.uses_lagged_pre: + h, pre_mix = h + + return self.forward_head( + h, + bonus_token_ids, temperature=temperature, return_confidence=return_confidence, return_logits=return_logits, - confidence_hidden=confidence_hidden, - ) - - -class DSv4DSparkForCausalLM(nn.Module): - """One-engine draft wrapper for DSpark (mirrors ``DFlashForCausalLM``). - - Wraps :class:`DSv4DSparkDraftModel` (the ``n_mtp_layers``-stage ``mtp.*`` backbone) - for the single-engine external-drafter flow: created by ``get_draft_model``, - appended to the target's epilogue, and driven by ``DSv4DSparkWorker``. - - ``embed_tokens`` / ``lm_head`` are shared with the target model - (:meth:`load_weights_from_target_model`). The draft weights live in the SAME - checkpoint under ``mtp.*``; :meth:`load_weights` remaps them - (``remap_dspark_draft_keys``), loads via ``DeepseekV4WeightLoader``, runs the - fp8 ``post_load_weights`` transforms, and caches the bf16 captured-context - attention weights from the in-memory state dict. - """ - - draft_model_cls = DSv4DSparkDraftModel - - def __init__( - self, - draft_config, - aux_stream_dict=None, - num_stages=None, - block_size=None, - draft_moe_backend: Optional[str] = None, - ): - super().__init__() - self.dspark_model = self.draft_model_cls( - draft_config, - aux_stream_dict, - num_stages=num_stages, - block_size=block_size, - draft_moe_backend=draft_moe_backend, - ) - # Generic handles expected by the loader / weight mappers. - self.model = self.dspark_model - self.model_config = draft_config - self.config = draft_config.pretrained_config - # Worker-facing interface (the worker receives this wrapper as - # ``draft_model`` and calls forward()/reads these properties and scalars). - self.num_stages = self.dspark_model.num_stages - self._attn_params = self.dspark_model._attn_params - self.lm_head = None # shared from the target (load_weights_from_target_model) - self.logits_processor = None # set by the caller after construction - - @property - def block_size(self): - return self.dspark_model.block_size - - @property - def embed_tokens(self): - return self.dspark_model.embed_tokens - - def forward(self, main_hidden, bonus_token_ids, start_pos, **kwargs): - return self.dspark_model.forward(main_hidden, bonus_token_ids, start_pos, **kwargs) - - def run_moe_lockstep_noop(self, all_rank_num_tokens, device): - """Empty-batch MoE barrier lockstep (delegates to the draft model).""" - return self.dspark_model.run_moe_lockstep_noop(all_rank_num_tokens, device) - - def configure_generation_autotuner(self, max_batch): - """Delegate generation autotuner capacity to the DSpark model.""" - self.dspark_model.configure_generation_autotuner(max_batch) - - def write_context_pages(self, main_hidden, positions, mask, pools, block_tables, capacities): - """Project captured context directly into the managed draft pools.""" - return self.dspark_model.write_context_pages( - main_hidden, positions, mask, pools, block_tables, capacities + pre_mix=pre_mix, ) def _prepare_confidence_head_for_load(self, weights: Dict) -> None: @@ -1673,11 +2756,11 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): f"[DSpark] loading {len(remapped)} draft params across {self.num_stages} stages" ) DeepseekV4WeightLoader(self.dspark_model).load_weights(remapped) - self.dspark_model.post_load_weights() - # bf16 captured-context attention path: dequantize the raw mtp.{s}.attn.* - # tensors for ``dspark_attention_forward``. - self.dspark_model.cache_attn_weights_from_state_dict(weights) - logger.info("[DSpark] draft weight load complete") + self.dspark_model.post_load_weights(weights) + logger.info( + f"[DSpark] draft weight load complete " + f"(linear_backend={'bf16' if self.dspark_model._use_bf16_linear else 'fp8'})" + ) def load_weights_from_target_model(self, target_model): """Share the target's embed_tokens / lm_head (DSpark has neither).""" @@ -1687,6 +2770,16 @@ def load_weights_from_target_model(self, target_model): self.lm_head = target_model.lm_head self.dspark_model.lm_head = target_model.lm_head + def configure_generation_autotuner(self, max_batch): + """Delegate generation autotuner capacity to the DSpark model.""" + self.dspark_model.configure_generation_autotuner(max_batch) + + def write_context_pages(self, main_hidden, positions, mask, pools, block_tables, capacities): + """Project captured context directly into the managed draft pools.""" + return self.dspark_model.write_context_pages( + main_hidden, positions, mask, pools, block_tables, capacities + ) + # ---------------------------------------------------------------------------- # MLA-shaped standalone drafter backbone (Inferact/Kimi-K3-DSpark). @@ -2702,8 +3795,14 @@ def _mla_paged_attention(self, query, blk_kv, layer_cache, fixup, block_size, ma if self._mla_block_decode_variant(paged=True) == "cute_dsl": return self._cute_dsl_block_decode(query, layer_cache, fixup, block_size) + out_lower = torch.empty( - batch, block_size, num_heads, self.kv_lora_rank, dtype=query.dtype, device=device + batch, + block_size, + num_heads, + self.kv_lora_rank, + dtype=query.dtype, + device=device, ) lse_lower = torch.empty(batch * block_size, num_heads, dtype=torch.float32, device=device) decode( @@ -2970,9 +4069,8 @@ def dflash_forward( q = attn.q_b_proj(attn.q_a_layernorm(attn.q_a_proj(hs_normed))) q = q.view(batch, block_size, num_heads, self.qk_nope_head_dim + self.qk_rope_head_dim) - # One buffer rather than cat([absorbed, rope]): q's rope half is a 64-slice - # of a 192-wide row, so cat fell off CatArrayBatchedCopy onto the generic - # kernel -- 17.3 us vs 1.6 us, five launches/iter (measured job 2986420). + # Assemble the query in one buffer to avoid concatenating a + # non-contiguous RoPE slice. query = torch.empty( (batch, block_size, num_heads, self.kv_lora_rank + self.qk_rope_head_dim), dtype=q.dtype, @@ -2999,9 +4097,12 @@ def dflash_forward( ) query[..., self.kv_lora_rank :] = q[..., self.qk_nope_head_dim :] q_nope = q[..., : self.qk_nope_head_dim] - # Absorb the nope half into latent space so the block attends the stored - # latent directly (MQA) instead of expanding per head. view(), never - # reshape(): both are strided views and reshape would stage a silent copy. + # Absorb the nope half into latent space so the block attends the + # stored latent directly (MQA), instead of expanding it per head. + # [h, b*s, nope] x [h, nope, kv_lora_rank] straight into the latent + # half. view(), never reshape(): both of these are strided views and + # reshape would silently stage a copy that bmm_out then writes into + # and throws away, whereas view() raises. torch.ops.trtllm.bmm_out( q_nope.view(batch * block_size, num_heads, self.qk_nope_head_dim).transpose(0, 1), attn._k_b_proj.to(q.dtype), @@ -3160,7 +4261,10 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): num_stages = count_dspark_stages(model_config.spec_config.speculative_model) validate_dspark_eplb_layer_base(model_config, draft_config) draft_cls = DSv4DSparkForCausalLM - if draft_config.pretrained_config.model_type in ("deepseek_v41", "deepseek_v41_text"): + if draft_config.pretrained_config.model_type in ( + "deepseek_v41", + "deepseek_v41_text", + ): from .modeling_dspark_v41 import DSv41DSparkForCausalLM draft_cls = DSv41DSparkForCausalLM @@ -3208,11 +4312,356 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): "validate_dspark_eplb_layer_base", "validate_dspark_eplb_stage_layers", # Captured-context attention primitives. + "get_dspark_topk_idxs", + "get_dspark_topk_idxs_batched", + "dspark_sparse_attn", "precompute_dspark_freqs_cis", "apply_dspark_rotary", "apply_dspark_rotary_batched", "dspark_attention_forward", + "dspark_attention_forward_batched", # Block draft I/O. "build_draft_input_ids", "dspark_propose", ] + + +def _linear_projection( + x: torch.Tensor, + weight_or_projection: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], +) -> torch.Tensor: + """Apply either a reference tensor weight or a loaded ``Linear`` module.""" + if callable(weight_or_projection): + return weight_or_projection(x) + return F.linear(x, weight_or_projection) + + +def _prequantized_fp8_projection( + x: tuple[torch.Tensor, torch.Tensor], + projection: Callable[[torch.Tensor], torch.Tensor], +) -> torch.Tensor: + """Run a bind-time-validated Rubin FP8 projection without Linear dispatch.""" + activation, activation_scale = x + tuning_config = getattr(projection, "_dspark_generation_tuning_config", None) + return _run_cute_dsl_dspark_mxfp8_gemm_rubin_prevalidated( + activation, + projection.weight, + activation_scale, + projection.weight_scale, + materialize_launch=True, + tuning_config=tuning_config, + ) + + +def _can_use_fused_fp8_input_projection( + projection: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + in_features: int, +) -> bool: + """Return whether a loaded Linear accepts Rubin R128c4 FP8 input.""" + if ( + get_sm_version() != 107 + or not IS_CUTLASS_DSL_RUBIN_AVAILABLE + or not callable(projection) + or in_features % 128 != 0 + or getattr(projection, "in_features", None) != in_features + or getattr(projection, "has_fp8_block_scales", False) is not True + or not ( + getattr(projection, "use_cute_dsl_blockscaling_mm", False) is True + or getattr(projection, "disable_deep_gemm", False) is True + ) + or getattr(projection, "bias", None) is not None + or getattr(projection, "lora", None) is not None + or getattr(projection, "_ugpu_weight_shards", None) is not None + ): + return False + + weight = getattr(projection, "weight", None) + weight_scale = getattr(projection, "weight_scale", None) + return ( + isinstance(weight, torch.Tensor) + and weight.dtype == torch.float8_e4m3fn + and weight.ndim == 2 + and weight.shape[1] == in_features + and weight.is_contiguous() + and isinstance(weight_scale, torch.Tensor) + and weight_scale.dtype == torch.uint8 + and weight_scale.is_contiguous() + ) + + +def _can_use_fused_fp8_output_projection( + wo_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + o_a_proj: torch.Tensor, + n_groups: int, + o_lora_rank: int, +) -> bool: + """Return whether ``wo_b`` consumes the Rubin R128c4 FP8 layout.""" + if ( + get_sm_version() != 107 + or not IS_CUTLASS_DSL_RUBIN_AVAILABLE + or not callable(wo_b) + or getattr(wo_b, "has_fp8_block_scales", False) is not True + or not ( + getattr(wo_b, "use_cute_dsl_blockscaling_mm", False) is True + or getattr(wo_b, "disable_deep_gemm", False) is True + ) + ): + return False + + weight = getattr(wo_b, "weight", None) + weight_scale = getattr(wo_b, "weight_scale", None) + flat_n = n_groups * o_lora_rank + return ( + o_a_proj.dtype == torch.float8_e4m3fn + and o_lora_rank % 128 == 0 + and tuple(o_a_proj.shape[:2]) == (n_groups, o_lora_rank) + and getattr(wo_b, "in_features", None) == flat_n + and isinstance(weight, torch.Tensor) + and weight.dtype == torch.float8_e4m3fn + and isinstance(weight_scale, torch.Tensor) + and weight_scale.dtype == torch.uint8 + ) + + +def _qkv_projections( + x: torch.Tensor, + main_x: torch.Tensor, + *, + wq_a: Optional[torch.Tensor], + wkv: Optional[torch.Tensor], + q_norm_w: torch.Tensor, + qkv_proj: Optional[Callable[[torch.Tensor], torch.Tensor]], + use_fused_fp8_input: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Project Q-latent and MQA K/V, fusing q_a+wkv when a module is supplied.""" + if qkv_proj is None: + if wq_a is None or wkv is None: + raise ValueError("DSpark attention requires tensor weights or a fused qkv projection") + return F.linear(x, wq_a), F.linear(x, wkv), F.linear(main_x, wkv) + + q_lora_rank = q_norm_w.shape[0] + # Group main and draft rows separately so the last-dimension Q/KV splits + # retain a constant row stride across batches. The views are metadata-only. + # Rubin folds their concatenation into the projection FP8 input producer; + # other paths keep the original concatenation. + main_rows = main_x.view(-1, main_x.shape[-1]) + block_rows = x.view(-1, x.shape[-1]) + if use_fused_fp8_input: + qkv_input = cute_dsl_dspark_concat_fp8_quant(main_rows, block_rows) + else: + qkv_input = torch.cat((main_rows, block_rows), dim=0) + qkv_rows = ( + _prequantized_fp8_projection(qkv_input, qkv_proj) + if use_fused_fp8_input + else qkv_proj(qkv_input) + ) + main_qkv, qkv = qkv_rows.split((main_rows.shape[0], block_rows.shape[0]), dim=0) + main_qkv = main_qkv.view(*main_x.shape[:-1], qkv_rows.shape[-1]) + qkv = qkv.view(*x.shape[:-1], qkv_rows.shape[-1]) + q_latent, kv = qkv.split((q_lora_rank, qkv.shape[-1] - q_lora_rank), dim=-1) + main_kv = main_qkv[..., q_lora_rank:] + return q_latent, kv, main_kv + + +def _q_projection( + q_latent: torch.Tensor, + q_norm_w: torch.Tensor, + eps: float, + wq_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + use_fused_fp8_input: bool, + batched_freqs_cis: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Apply q-latent RMSNorm and q_b, optionally sharing one FP8 producer.""" + if use_fused_fp8_input: + q_input = cute_dsl_dspark_rmsnorm_fp8_quant(q_latent, q_norm_w, eps) + q = _prequantized_fp8_projection(q_input, wq_b) + return q.view(*q_latent.shape[:-1], q.shape[-1]) + + if batched_freqs_cis is not None: + return _linear_projection( + _rmsnorm_rope_batched(q_latent, q_norm_w, eps, 0, batched_freqs_cis), + wq_b, + ) + q = _rmsnorm(q_latent, q_norm_w, eps) + return _linear_projection(q, wq_b) + + +def _output_projection( + o: torch.Tensor, + *, + wo_a: Optional[torch.Tensor], + wo_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + n_groups: int, + o_lora_rank: int, + o_a_proj: Optional[torch.Tensor], + o_a_scale: Optional[torch.Tensor], + o_a_dequant: Optional[torch.Tensor], + use_fused_output_projection: Optional[bool] = None, +) -> torch.Tensor: + """Apply the grouped o_a projection followed by o_b.""" + batch, block = o.shape[:2] + if o_a_proj is None: + if wo_a is None: + raise ValueError("DSpark attention requires a tensor wo_a or a loaded FP8 o_a") + o = o.reshape(batch, block, n_groups, -1) + wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) + o_lora = torch.einsum("bsgd,grd->bsgr", o, wo_a_v).flatten(2) + return _linear_projection(o_lora, wo_b) + + if o_a_proj.dtype != torch.float8_e4m3fn: + grouped_o = o.reshape(batch * block, n_groups, -1).transpose(0, 1) + o_lora = torch.bmm(grouped_o, o_a_proj.transpose(1, 2)) + o_lora = o_lora.transpose(0, 1).reshape(batch, block, n_groups * o_lora_rank) + return _linear_projection(o_lora, wo_b) + + if o_a_scale is None: + raise ValueError("The DSpark FP8 o_a projection requires block scales") + + grouped_o = o.reshape(batch * block, n_groups, -1) + if use_fused_output_projection is None: + use_fused_output_projection = _can_use_fused_fp8_output_projection( + wo_b, o_a_proj, n_groups, o_lora_rank + ) + if use_fused_output_projection: + grouped_o_fp8, grouped_o_scale = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102( + grouped_o + ) + return _fused_fp8_output_projection_from_quantized( + grouped_o_fp8, + grouped_o_scale, + batch=batch, + block=block, + wo_b=wo_b, + n_groups=n_groups, + o_lora_rank=o_lora_rank, + o_a_proj=o_a_proj, + o_a_scale=o_a_scale, + ) + + from ..attention.mla import fp8_block_scaling_bmm_out + + o_lora = o.new_empty((n_groups, batch * block, o_lora_rank)) + fp8_block_scaling_bmm_out( + grouped_o, + o_a_proj, + o_a_scale, + o_lora, + mat2_dequant=o_a_dequant, + use_cute_dsl_blockscaling_bmm=True, + ) + o_lora = o_lora.transpose(0, 1).reshape(batch, block, n_groups * o_lora_rank) + return _linear_projection(o_lora, wo_b) + + +def _fused_fp8_output_projection_from_quantized( + grouped_o_fp8: torch.Tensor, + grouped_o_scale: torch.Tensor, + *, + batch: int, + block: int, + wo_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + n_groups: int, + o_lora_rank: int, + o_a_proj: torch.Tensor, + o_a_scale: torch.Tensor, +) -> torch.Tensor: + """Consume group-major FP8 attention output in the fused OA/OB path.""" + rows = batch * block + flat_n = n_groups * o_lora_rank + fp8_output = grouped_o_fp8.new_empty((rows, flat_n)) + fp8_batched = fp8_output.view(rows, n_groups, o_lora_rank).transpose(0, 1) + scale_m = (rows + 127) // 128 * 128 + scale_k = ((flat_n + 31) // 32 + 3) // 4 * 4 + packed_scale = torch.empty( + (scale_m * scale_k,), + dtype=torch.uint8, + device=grouped_o_fp8.device, + ) + max_batch = getattr(wo_b, "_dspark_generation_max_batch", 0) + rows_per_request = getattr(wo_b, "_dspark_generation_rows_per_request", 0) + torch.ops.trtllm.cute_dsl_dspark_fp8_bmm_quantize_rubin_out( + grouped_o_fp8, + o_a_proj, + grouped_o_scale, + o_a_scale, + fp8_batched, + packed_scale, + True, + max_batch, + rows_per_request, + ) + projected = _prequantized_fp8_projection((fp8_output, packed_scale), wo_b) + return projected.view(batch, block, projected.shape[-1]) + + +def _use_dspark_bf16_linear(mapping) -> bool: + """Select BF16 for an explicit override or sharded attention projections.""" + backend = os.environ.get(_DSPARK_LINEAR_BACKEND_ENV, "fp8").lower() + if backend not in ("fp8", "bf16"): + raise ValueError(f"{_DSPARK_LINEAR_BACKEND_ENV} must be 'fp8' or 'bf16', got {backend!r}") + return backend == "bf16" or (mapping.tp_size > 1 and not mapping.enable_attention_dp) + + +def _stack_fp8_projection_output_rows( + projections: List[Linear], output_rows: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Stack trailing output rows and rebuild their shared R128c4 scales.""" + if output_rows <= 0 or output_rows % 128 != 0: + raise ValueError( + f"DSpark stacked FP8 output_rows must be a positive multiple of 128, got {output_rows}" + ) + + weights = [] + linear_scales = [] + for projection in projections: + weight = projection.weight + if ( + weight.ndim != 2 + or weight.shape[0] % 128 != 0 + or weight.shape[1] % 128 != 0 + or output_rows > weight.shape[0] + ): + raise ValueError( + "DSpark stacked FP8 projections require 2-D 128-aligned weights " + f"with at least {output_rows} rows, got shape {tuple(weight.shape)}" + ) + linear_scale = unswizzle_sf( + projection.weight_scale, + weight.shape[0], + weight.shape[1], + scaling_vector_size=32, + ) + weights.append(weight[-output_rows:]) + linear_scales.append(linear_scale[-output_rows:]) + + stacked_weight = torch.cat(weights, dim=0).contiguous() + stacked_scale = swizzle_sf( + torch.cat(linear_scales, dim=0), + len(projections) * output_rows, + stacked_weight.shape[1], + scaling_vector_size=32, + ) + return stacked_weight, stacked_scale + + +def _yarn_find_correction_dim( + num_rotations: float, dim: int, base: float, max_position_embeddings: int +) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + +def _yarn_find_correction_range( + low_rot: float, high_rot: float, dim: int, base: float, max_position_embeddings: int +) -> tuple[int, int]: + low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)) + high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)) + return max(low, 0), min(high, dim - 1) + + +def _yarn_linear_ramp_mask(minimum: float, maximum: float, dim: int) -> torch.Tensor: + if minimum == maximum: + maximum += 0.001 # Prevent singularity. + linear_func = (torch.arange(dim, dtype=torch.float32) - minimum) / (maximum - minimum) + return torch.clamp(linear_func, 0, 1) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 81443f6c18e6..0fc2fece431c 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -110,9 +110,16 @@ from ...logger import logger from ...mapping import Mapping from ...models.modeling_utils import QuantAlgo, QuantConfig -from ..attention_backend import AttentionMetadata -from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams -from ..kimi_k3_cache_policy import KIMI_K3_BF16_KV_LAYERS_ENV, get_kimi_k3_bf16_kv_layer_ids +from ..attention.attention import _helix_cp_pad +from ..attention.backends import AttentionMetadata +from ..distributed import ( + AllReduce, + AllReduceFusionOp, + AllReduceParams, + AllReduceStrategy, + allgather, +) +from ..distributed.symm_mem_allgather import SymmetricMemoryAllGather from ..model_config import ModelConfig from ..modules.gated_mlp import GatedMLP from ..modules.kimi_kda import KimiKDALinearAttention @@ -1818,9 +1825,10 @@ def _shared_output(): if self._defer_shared_allreduce and not overlap_shared_allreduce: shared_out = self.shared_experts.down_proj.all_reduce(shared_out) if self._reduce_routed_output: - # RMSNorm folded into the routed all-reduce (kARRMSNorm): the + # RMSNorm folded into the routed all-reduce (RMS_NORM fusion): the # MNNVL kernel normalizes straight out of the lamport buffer, so # the reduced row never round-trips through HBM for a norm kernel. + # Strategies without the fusion fall back inside AllReduce. routed_latent = moe_all_reduce( routed_out, all_reduce_params=AllReduceParams( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py deleted file mode 100644 index 110457fa1641..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py +++ /dev/null @@ -1,289 +0,0 @@ -# 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 dataclasses import replace -from typing import Optional, Tuple, Union - -import torch - -from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.models.modeling_utils import QuantAlgo - -from ...utils import ActivationType, Fp4QuantizedTensor -from .fused_moe_cute_dsl import CuteDslFusedMoE -from .fused_moe_cutlass import CutlassFusedMoE -from .impl_contract import ( - MoEDeployment, - MoEEligibility, - MoEProblem, - MoERejectReason, - MoERunContext, - MoEStaticCapability, - require_comm_plan, -) -from .impl_environment import MoEDep -from .interface import _reject - -# Shared MoE output buffer pool, keyed by (max_num_tokens, hidden_size, dtype, -# device). ``B12xMoEWrapper.__init__`` allocates a private -# ``(max_num_tokens, hidden_size)`` output tensor per instance; with one -# wrapper per MoE layer that is ``num_layers * max_num_tokens * hidden_size`` -# bytes of GPU memory holding identical-shape buffers that are written -# sequentially. We fold them into a single shared buffer because MoE layers -# run sequentially on the same CUDA stream, and the wrapper consumes its -# previous output before the next layer is dispatched. -_SHARED_MOE_OUTPUT_BUF: dict = {} - -# ActivationType -> b12x activation string. b12x currently exposes "relu2" -# (Nemotron-style x * relu(x)) and "silu" (SwiGLU-style x * silu(gate)). -_ACTIVATION_MAP = { - ActivationType.Relu2: "relu2", - ActivationType.Swiglu: "silu", -} - - -class CuteDslB12xFusedMoE(CuteDslFusedMoE): - """B12x NVFP4 fused-MoE backend for SM120 / SM121. - - Large prefill chunks use CUTLASS; decode uses FlashInfer's b12x kernel. - """ - - # Restated rather than inherited: the LoRA gate this replaces compared the - # exact class and answered False here, while the DWDP gate used isinstance - # and answered True through CuteDslFusedMoE. - capabilities = MoEStaticCapability(supports_moe_lora=False, supports_dwdp=True) - - # SM versions on which the FlashInfer b12x NVFP4 MoE kernel is available. - # SM120 = desktop Blackwell (RTX 5090 / GB202); SM121 = GB10 / DGX Spark. - _SUPPORTED_SM_VERSIONS = frozenset({120, 121}) - - # Prefill chunks (``x.shape[0] >= threshold``) route via CUTLASS NVFP4 - # GroupGEMM; decode (``x.shape[0] < threshold``) uses b12x. 64 cleanly - # separates conc=1 prefill (m=2048 with ``max_num_tokens=2048``) from - # decode (m=1) and stays robust against future chunked-prefill splits - # that might shrink prefill chunk size. - _PREFILL_VIA_CUTLASS_THRESHOLD = 64 - - @classmethod - def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: - sm_version = d.env.sm - if sm_version not in cls._SUPPORTED_SM_VERSIONS: - sm_list = "/".join(f"SM{v}" for v in sorted(cls._SUPPORTED_SM_VERSIONS)) - return _reject( - MoERejectReason.SM_UNSUPPORTED, - f"CuteDslB12xFusedMoE requires {sm_list}, got SM{sm_version}", - ) - if p.quant_algo not in {QuantAlgo.NVFP4, QuantAlgo.W4A16_NVFP4}: - return _reject( - MoERejectReason.QUANT_UNSUPPORTED, - f"CuteDslB12xFusedMoE only supports NVFP4 or W4A16_NVFP4 quantization " - f"(got quant_algo={p.quant_algo})", - ) - if p.dtype_act not in {torch.float16, torch.bfloat16}: - return _reject( - MoERejectReason.DTYPE_UNSUPPORTED, - f"CuteDslB12xFusedMoE NVFP4 requires float16 or bfloat16 " - f"activation dtype (got {p.dtype_act})", - ) - if p.swiglu_gptoss_style: - return _reject( - MoERejectReason.ACTIVATION_UNSUPPORTED, - "CuteDslB12xFusedMoE does not support swiglu_gptoss_style", - ) - if p.activation_type not in _ACTIVATION_MAP: - supported = ", ".join(a.name for a in _ACTIVATION_MAP) - return _reject( - MoERejectReason.ACTIVATION_UNSUPPORTED, - f"CuteDslB12xFusedMoE does not support activation " - f"{p.activation}; supported: {supported}", - ) - # The decode kernel ships in the FlashInfer wheel. - if not d.env.has_dep(MoEDep.FLASHINFER): - return _reject( - MoERejectReason.DEP_MISSING, - "CuteDslB12xFusedMoE requires the flashinfer package", - ) - # No expert-parallel dispatch/combine kernel: EP must stay at 1. - if d.ep_size != 1: - return _reject( - MoERejectReason.TOPOLOGY_UNSUPPORTED, - f"CuteDslB12xFusedMoE requires ep_size == 1 (got {d.ep_size})", - ) - # Attention-DP is a separate axis from EP: with moe_tp == tp the layer - # can have ep_size == 1 and still sit behind a DP allgather / - # reducescatter that the b12x wrapper has never been exercised under. - # ``use_dp and parallel_size > 1`` is exactly ``mapping.dp_size > 1`` - # (``Mapping.dp_size`` is ``tp_size`` when attention-DP is on). - if d.use_dp and d.parallel_size > 1: - return _reject( - MoERejectReason.TOPOLOGY_UNSUPPORTED, - f"CuteDslB12xFusedMoE does not support attention-DP " - f"(parallel_size={d.parallel_size})", - ) - return MoEEligibility.ok() - - def __init__(self, *args, **kwargs): - # ``ModelConfig`` is consumed by the inherited ``__init__`` for cache - # / mapping setup but isn't kept on ``self``. b12x's wrapper needs the - # ``use_cuda_graph`` flag at construction time, so capture it here - # before delegating. - model_config = kwargs.get("model_config", None) - self._b12x_use_cuda_graph = bool(getattr(model_config, "use_cuda_graph", False)) - - super().__init__(*args, **kwargs) - - # Eligibility (SM / quant / activation / EP) is owned by - # ``can_implement``. ``enable_alltoall`` is a run-time flag not yet on - # ``MoEDeployment``, so keep the construction guard here. - if self.enable_alltoall: - raise ValueError("CuteDslB12xFusedMoE does not support MoE alltoall communication.") - - self._b12x_weights: Optional[dict] = None - self.b12x_wrapper = None - - def _get_quant_method(self): - # Route NVFP4 to the b12x-aware quant method so weight prep - # (SF un-normalization, ``convert_sf_to_mma_layout``, - # ``B12xMoEWrapper`` instantiation) lives next to the rest of the - # NVFP4 quant-method family, while every other quant algo (and the - # unquantized fallback) continues to resolve via the parent. - if ( - self.quant_config is not None - and self.quant_config.layer_quant_mode.has_any_quant(exclude_kv_cache=True) - and self.quant_config.layer_quant_mode.has_nvfp4() - ): - from .quantization import NVFP4CuteDslB12xFusedMoEMethod - - return NVFP4CuteDslB12xFusedMoEMethod() - return super()._get_quant_method() - - def _route_to_cutlass(self, x) -> bool: - """Return ``True`` iff this call should fall back to the inherited - CUTLASS path (NVFP4 prefill chunk). ``Fp4QuantizedTensor`` inputs - always stay on the b12x path (which rejects them) so the existing - error message is preserved.""" - quant_config = getattr(self, "quant_config", None) - if quant_config is not None and quant_config.quant_algo == QuantAlgo.W4A16_NVFP4: - return False - return isinstance(x, torch.Tensor) and x.shape[0] >= self._PREFILL_VIA_CUTLASS_THRESHOLD - - # ``post_load_weights`` is inherited from ``CutlassFusedMoE`` and - # dispatches to ``self.quant_method.transform_weights(self)`` — for this - # backend ``self.quant_method`` is ``NVFP4CuteDslB12xFusedMoEMethod`` - # (see ``_get_quant_method`` override), which performs the SF un-normalization, - # ``convert_sf_to_mma_layout`` reshape, ``B12xMoEWrapper`` instantiation, - # and the cross-layer shared output buffer dance. The wrapper and the - # bundled weight dict are attached to this module as ``self.b12x_wrapper`` - # / ``self._b12x_weights``, which the decode path below consumes. - - @nvtx_range("[b12x] quantize_input") - def quantize_input( - self, - x: Union[torch.Tensor, Fp4QuantizedTensor], - post_quant_comm: bool = True, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """Hybrid dispatch entrypoint for activation handling. - - NVFP4 prefill chunks take the inherited - :meth:`CutlassFusedMoE.quantize_input` path so the downstream - ``run_moe`` can call CUTLASS NVFP4 GroupGEMM. Decode chunks and - W4A16_NVFP4 chunks pass through unchanged because b12x quantizes - activations internally (consumes a bf16 / fp16 ``x`` and produces its - own scale factors). - """ - if self._route_to_cutlass(x): - return CutlassFusedMoE.quantize_input( - self, x, post_quant_comm=post_quant_comm, **kwargs - ) - if isinstance(x, Fp4QuantizedTensor): - raise ValueError( - "CuteDslB12xFusedMoE does not accept Fp4QuantizedTensor input " - "on the b12x decode path; b12x performs its own input quantization." - ) - return x, None - - @nvtx_range("[b12x] run_moe") - def run_moe( - self, - ctx: MoERunContext, - *, - workspace: Optional[dict] = None, - ) -> torch.Tensor: - plan = require_comm_plan(self, ctx) - x = ctx.x - # Join the async input quantization: the scheduler forks the inherited - # ``quantize_input_async`` (which runs the polymorphic - # ``self.quantize_input`` on the MoE aux stream). On the - # prefill-via-CUTLASS route that is a real FP4 quant kernel and - # ``CutlassFusedMoE.run_moe`` below never joins the event; on the b12x - # decode route ``quantize_input`` is a pass-through, so the wait is a - # no-op. - if ctx.x_ready_event is not None: - ctx.x_ready_event.wait() - if self._route_to_cutlass(x): - # ``CutlassFusedMoE.run_moe`` forwards ``output_dtype`` straight - # into the C++ ``trtllm::fused_moe`` op, which requires a concrete - # high-precision ``ScalarType`` (uint8 / FP4-packed activations are - # rejected at the kernel epilogue with "Invalid output type Byte"). - # ``ConfigurableMoE.forward`` always fills ``output_dtype``, so this - # only narrows the type for anything driving ``run_moe`` without it. - _HIGH_PRECISION = {torch.float16, torch.bfloat16, torch.float32} - cutlass_output_dtype = ctx.output_dtype - if cutlass_output_dtype is None: - cutlass_output_dtype = ( - x.dtype - if isinstance(x, torch.Tensor) and x.dtype in _HIGH_PRECISION - else torch.bfloat16 - ) - return CutlassFusedMoE.run_moe( - self, - replace(ctx, output_dtype=cutlass_output_dtype), - workspace=workspace, - ) - token_selected_experts = ctx.token_selected_experts - token_final_scales = ctx.token_final_scales - x_sf = ctx.x_sf - moe_output = plan.moe_output - if self.b12x_wrapper is None or self._b12x_weights is None: - raise RuntimeError( - "CuteDslB12xFusedMoE.run_moe called before process_weights_after_loading completed." - ) - if x_sf is not None: - raise ValueError( - "CuteDslB12xFusedMoE expects unquantized input (x_sf=None) " - "on the b12x decode path; got a precomputed scale factor." - ) - - # Annotate the kwargs spread + wrapper entry separately so we can - # attribute the per-layer Python dispatch cost vs. the kernel cost. - with nvtx_range("[b12x] wrapper.run"): - out = self.b12x_wrapper.run( - x=x, - token_selected_experts=token_selected_experts, - token_final_scales=token_final_scales, - **self._b12x_weights, - ) - - # B12xMoEWrapper allocates its own output buffer for CUDA-graph - # compatibility. If the caller provided ``moe_output`` (e.g. an - # alltoall workspace tensor), copy into it; CuteDslB12xFusedMoE - # currently rejects alltoall in __init__, so this is a defensive - # path for future workspace-driven uses. - if moe_output is not None: - with nvtx_range("[b12x] out_copy"): - moe_output.copy_(out) - return moe_output - return out diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py deleted file mode 100644 index b0f75ee6f637..000000000000 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_fc12.py +++ /dev/null @@ -1,264 +0,0 @@ -# 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. -"""CuteDslFc12FusedMoE: FC1+FC2-fused CuteDSL NVFP4 MoE backend for Rubin (SM107). - -Drives ``cute_dsl_nvfp4_fc12_fused_rubin``: one persistent kernel does gather + -FC1 GEMM + gated epilogue + requant + FC2 GEMM + finalize. The epilogue is a -trace-time specialization of SwiGLU or SiTU (Kimi K3); Relu2 is not gated and -is not supported. uGPU is not enabled for this backend. - -``quantize_input``, ``run_moe``, the weight lifecycle, capabilities, and EPLB -are inherited from :class:`CuteDslFusedMoE`. The SiTU constants live on the -parent (``situ_beta`` / ``situ_linear_beta``) and are forwarded into the fused -op; they are part of the compiled-kernel cache key. - -See ``CuteDslFc12FusedMoE-integration-plan.txt`` for the original integration -plan. -""" - -from typing import Optional - -import torch - -from tensorrt_llm.models.modeling_utils import QuantAlgo - -from ...autotuner import AutoTuner -from ...cute_dsl_utils import IS_CUTLASS_DSL_INTERNAL_AVAILABLE -from ...utils import ActivationType -from .fused_moe_cute_dsl import (CuteDslFusedMoE, CuteDslFusedMoENvfp4Runner, - NvFp4WeightView) -from .impl_contract import (MoEDeployment, MoEEligibility, MoEProblem, - MoERejectReason) -from .interface import _reject - - -class CuteDslFc12FusedMoENvfp4Runner(CuteDslFusedMoENvfp4Runner): - """Outer autotune runner for the fused FC12 backend. - - Identical to the parent except the routing-tile candidate set: the fused - kernel supports the 128-wide tile (1-CTA) and the 256-wide tile (2-CTA, - cluster (2,1)), so restrict ``_tile_sizes`` to ``[128, 256]`` (the parent - also offers 512, which the fused v1 kernel does not support). The inner - fused runner derives mma_tiler_m == tile_size and cluster M == tile_size // - 128 from the selected tile, mirroring the CuteDSL grouped-GEMM runners. - """ - - @staticmethod - def _tile_sizes(): - return [128, 256] - - -class CuteDslFc12FusedMoE(CuteDslFusedMoE): - """FC1+FC2-fused CuteDSL NVFP4 MoE backend (Rubin/SM107). - - SwiGLU and SiTU share the gated FC1 geometry; the fused kernel specializes - the epilogue at trace time. uGPU is not enabled for this backend. - """ - - @classmethod - def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: - """FC12 fused CuteDSL: NVFP4 on Rubin (SM107) only. - - Pure gate -- reads only ``p`` and ``d`` (no ``get_sm_version()``, no - import probes, no ``os.environ``); the frozen environment lives in - ``d.env``. Narrower than the parent (SM107 + NVFP4 only) because the FC12 - fused kernel targets Rubin NVFP4; broaden when other shapes are ported. - """ - sm_version = d.env.sm - - # Output is hardcoded to bfloat16 (inherited), so activation must match. - if p.dtype_act != torch.bfloat16: - return _reject( - MoERejectReason.DTYPE_UNSUPPORTED, - "CuteDslFc12FusedMoE only supports bfloat16 activation, " - f"got {p.dtype_act}") - - if p.activation_type not in (ActivationType.Swiglu, ActivationType.SiTu): - return _reject( - MoERejectReason.ACTIVATION_UNSUPPORTED, - "CuteDslFc12FusedMoE supports only SwiGLU and SiTU, " - f"got {p.activation_type.name}") - - if p.swiglu_gptoss_style: - return _reject( - MoERejectReason.ACTIVATION_UNSUPPORTED, - "CuteDslFc12FusedMoE does not support swiglu_gptoss_style") - - if p.quant_algo != QuantAlgo.NVFP4: - return _reject( - MoERejectReason.QUANT_UNSUPPORTED, - "CuteDslFc12FusedMoE only supports NVFP4, " - f"got quant_algo={p.quant_algo}") - - if sm_version != 107: - return _reject( - MoERejectReason.SM_UNSUPPORTED, - "CuteDslFc12FusedMoE targets Rubin (SM107), " - f"got SM{sm_version}") - - if not IS_CUTLASS_DSL_INTERNAL_AVAILABLE: - return _reject( - MoERejectReason.DEP_MISSING, - "CuteDslFc12FusedMoE (SM107 NVFP4) requires CuTE DSL internal") - - return MoEEligibility.ok() - - # ------------------------------------------------------------------ - # Fused-kernel dispatch. Mirrors the parent's non-uGPU NVFP4 path but - # (1) uses a distinct autotuner key + outer runner so FC12 tactics do - # not collide with the parent CuteDSL backend, and (2) drives the fused - # single-op path in ``run_moe_nvfp4_impl``. uGPU is never used here. - # ------------------------------------------------------------------ - def run_moe_nvfp4( - self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - x_sf: Optional[torch.Tensor] = None, - moe_output: Optional[torch.Tensor] = None, - enable_alltoall: bool = False, - weight_view: Optional[NvFp4WeightView] = None, - x_ready_event: Optional[torch.cuda.Event] = None, - ) -> torch.Tensor: - assert self.has_nvfp4 - assert weight_view is not None - output_dtype = torch.bfloat16 - - if moe_output is None: - moe_output = torch.empty( - (token_final_scales.size(0), self.hidden_size), - dtype=output_dtype, - device=x.device) - else: - assert moe_output.size() == (token_final_scales.size(0), - self.hidden_size) - assert moe_output.dtype == output_dtype - - # Empty micro-batches: skip autotuning (synthetic grouped-GEMM inputs - # require at least one output row). - if token_selected_experts.size(0) == 0: - return moe_output - - # Deferred join for ``quantize_input_async`` (see the parent class): - # the fused kernel is the first consumer of ``x``/``x_sf``, so - # ``run_moe_nvfp4_impl`` joins right after ``moe_sort``. - self._pending_x_ready_event = x_ready_event - - effective_top_k = token_selected_experts.size(-1) - tuner = AutoTuner.get() - runner = CuteDslFc12FusedMoENvfp4Runner( - forward_impl=self.run_moe_nvfp4_impl, - num_experts=self.num_slots, - top_k=effective_top_k, - num_local_experts=weight_view.expert_size_per_partition, - local_expert_offset=weight_view.slot_start, - enable_finalize_fusion=self.use_fused_finalize, - enable_alltoall=enable_alltoall, - # Inner kernel_cache already keys on these; the outer AutoTuner - # unique_id must too, or SwiGLU/SiTU share a tile-size tactic. - workload_identity=( - int(self.activation_type), - self.situ_beta, - self.situ_linear_beta, - ), - ) - inputs = [ - x, token_selected_experts, token_final_scales, x_sf, moe_output, - weight_view - ] - _, best_tactic = tuner.choose_one( - "CuteDslFc12FusedMoE::run_moe_nvfp4", - [runner], - runner.get_tuning_config(), - inputs, - ) - return runner(inputs, tactic=best_tactic) - - def run_moe_nvfp4_impl( - self, - x: torch.Tensor, - token_selected_experts: torch.Tensor, - token_final_scales: Optional[torch.Tensor], - x_sf: torch.Tensor, - moe_output: torch.Tensor, - weight_view: NvFp4WeightView, - enable_alltoall: bool = False, - tile_size: int = 128, - ) -> torch.Tensor: - """Single fused FC1+FC2 op (replaces the parent's two-op sequence).""" - effective_top_k = token_selected_experts.size(1) - esp = weight_view.expert_size_per_partition - slot_start = weight_view.slot_start - - (tile_idx_to_expert_idx, tile_idx_to_mn_limit, - expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, - total_num_padded_tokens, - num_non_exiting_tiles) = torch.ops.trtllm.moe_sort( - token_selected_experts=token_selected_experts, - token_final_scales=token_final_scales, - num_experts=self.num_slots, - top_k=effective_top_k, - local_expert_offset=slot_start, - local_num_experts=esp, - tile_tokens_dim=tile_size, - ) - - # Join the async input quantization (``quantize_input_async``): the - # fused op below is the first consumer of ``x``/``x_sf``; ``moe_sort`` - # above only reads routing tensors, so it overlaps the aux-stream - # quantization. - self._join_pending_x_ready() - - # One fused op: gather + FC1 GEMM + gated act (SwiGLU or SiTU) + requant - # + FC2 GEMM + finalize (scatter-add into moe_output). - # fc1_alpha/fc2_alpha map 1:1 to the two-op path's per-expert global - # scales (the kernel takes split alphas). The three atomic counters are - # allocated + memset inside the op runner. The moe_output zeroing memset - # is issued inside the op too (right before the fused kernel, so it is - # the kernel's immediate stream predecessor and the PDL prologue can - # overlap it); expanded_idx_to_permuted_idx / ep_size / enable_alltoall - # are passed through for that memset. - situ = ActivationType(self.activation_type) == ActivationType.SiTu - torch.ops.trtllm.cute_dsl_nvfp4_fc12_fused_rubin( - input=x.view(torch.float4_e2m1fn_x2), - fc1_weight=weight_view.w3_w1_weight.view(torch.float4_e2m1fn_x2), - input_scale=x_sf.view(torch.uint8), - fc1_weight_scale=weight_view.fc1_weight_scale.view(torch.uint8), - fc1_alpha=weight_view.fc1_global_scale, - tile_idx_to_group_idx=tile_idx_to_expert_idx, - tile_idx_to_mn_limit=tile_idx_to_mn_limit, - permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, - num_non_exiting_tiles=num_non_exiting_tiles, - global_sf=self.fc2_input_scale, - fc2_weight=weight_view.w2_weight.view(torch.float4_e2m1fn_x2), - fc2_weight_scale=weight_view.fc2_weight_scale.view(torch.uint8), - fc2_alpha=weight_view.fc2_global_scale, - output=moe_output, - token_final_scales=token_final_scales, - expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, - num_experts=self.num_slots, - top_k=effective_top_k, - num_local_experts=esp, - local_expert_offset=slot_start, - tile_size=tile_size, - swiglu_limit=self.swiglu_limit_scalar, - ep_size=self.mapping.moe_ep_size, - enable_alltoall=enable_alltoall, - scaling_vector_size=16, - activation_type=int(self.activation_type), - situ_beta=self.situ_beta if situ else -1.0, - situ_linear_beta=self.situ_linear_beta if situ else -1.0, - ) - return moe_output diff --git a/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py b/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py index 97869b5c1a6b..ad39a7429cb3 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py +++ b/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py @@ -285,6 +285,7 @@ def run_moe(self, ctx: MoERunContext, *, workspace: dict | None = None) -> torch moe_output=plan.moe_output, enable_alltoall=plan.enable_alltoall, weight_view=self._build_local_weight_view(), + x_ready_event=ctx.x_ready_event, ) # Fused-kernel dispatch, with its own autotuner key and outer runner so @@ -298,6 +299,7 @@ def run_moe_nvfp4( moe_output: torch.Tensor | None = None, enable_alltoall: bool = False, weight_view: NvFp4WeightView | None = None, + x_ready_event: torch.cuda.Event | None = None, ) -> torch.Tensor: assert self.has_nvfp4 assert weight_view is not None @@ -316,6 +318,8 @@ def run_moe_nvfp4( if token_selected_experts.size(0) == 0: return moe_output + self._pending_x_ready_event = x_ready_event + effective_top_k = token_selected_experts.size(-1) tuner = AutoTuner.get() runner = CuteDslFc12FusedMoENvfp4Runner( @@ -400,6 +404,8 @@ def run_moe_nvfp4_impl( tile_tokens_dim=tile_size, ) + self._join_pending_x_ready() + # One fused op: gather + FC1 GEMM + gated act (SwiGLU or SiTU) + requant # + FC2 GEMM + finalize (scatter-add into moe_output = a2a combine # workspace). diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index 659a8d8b3881..be224347063a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -418,6 +418,8 @@ def run_moe( *, workspace: dict | None = None, ) -> torch.Tensor: + if ctx.x_ready_event is not None: + ctx.x_ready_event.wait() if self._route_to_cutlass(ctx.x): # The inherited ``run_moe`` forwards ``output_dtype`` straight into # the C++ ``trtllm::fused_moe`` op, which requires a concrete diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 994b5dd2c427..6714ba1428c2 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -27,7 +27,6 @@ str_dtype_to_binding, torch_dtype_to_binding, torch_dtype_to_str) from tensorrt_llm.bindings import DataType -from tensorrt_llm.bindings.executor import DecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams # isort: off @@ -1272,6 +1271,7 @@ def _get_draft_cache_cost( kv_cache_config: KvCacheConfig, *, use_separate_draft_kv_cache: bool, + include_request_cost: bool = False, ) -> Optional[CacheCost]: """Return the draft manager's standalone cache cost, if it has one. @@ -1287,8 +1287,7 @@ def _get_draft_cache_cost( helix_cp_scale = self._mapping.cp_size def scaled(cost: CacheCost) -> CacheCost: - return CacheCost(slope=cost.slope * helix_cp_scale, - intercept=cost.intercept) + return dataclasses.replace(cost, slope=cost.slope * helix_cp_scale) if self._draft_model_engine is not None: draft_model_config = self._draft_model_engine.model.model_config @@ -1328,30 +1327,27 @@ def scaled(cost: CacheCost) -> CacheCost: draft_kv_cache_config) if self._speculative_config.spec_dec_mode.is_external_drafter(): # External drafter: layers start from 0, normal PP distribution - # Resolve draft manager class from draft config — may differ - # from target (e.g. hybrid target + plain transformer draft). - draft_kv_cache_manager_cls = get_kv_cache_manager_cls( - effective_draft_config, - draft_kv_cache_config, - is_disagg=self._is_disagg) - draft_cost = self._per_manager_cache_cost( - draft_kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - mapping=draft_mapping, - tokens_per_block=draft_tokens_per_block) - return scaled(draft_cost) + return scaled( + self._per_manager_cache_cost( + draft_kv_cache_manager_cls, + effective_draft_config, + draft_kv_cache_config, + mapping=draft_mapping, + tokens_per_block=draft_tokens_per_block, + is_draft=True, + include_request_cost=include_request_cost)) elif self._mapping.is_last_pp_rank(): # EAGLE3/MTP: draft layers only on last PP rank - draft_cost = self._per_manager_cache_cost( - self._kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - mapping=draft_mapping, - tokens_per_block=draft_tokens_per_block, - num_layers=self._get_num_draft_layers(), - is_draft=True) - return scaled(draft_cost) + return scaled( + self._per_manager_cache_cost( + draft_kv_cache_manager_cls, + effective_draft_config, + draft_kv_cache_config, + mapping=draft_mapping, + tokens_per_block=draft_tokens_per_block, + num_layers=self._get_num_draft_layers(), + is_draft=True, + include_request_cost=include_request_cost)) return None def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -2206,11 +2202,6 @@ def _unified_draft_cache_unsupported_reason(self) -> Optional[str]: return "Unified DSpark KV cache does not yet support MLA drafters." return None - def _get_draft_mapping(self) -> Mapping: - """Use the drafter's dense TP mapping when target attention uses Helix.""" - return (self._mapping.repurpose_helix_cp_to_tp() - if self._mapping.has_cp_helix() else self._mapping) - def _get_draft_scratch_tokens(self) -> int: return 0 if self._is_embedded_dspark( ) else self._speculative_config.max_draft_len + 1 @@ -2279,20 +2270,15 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False if self._uses_unified_standalone_draft_cache(): return False - is_external_drafter = ( - self._speculative_config.spec_dec_mode.is_external_drafter()) - if self._mapping.enable_attention_dp and not is_external_drafter: - # Attention-DP uses the target KV cache manager for MTP draft layers - # so draft cache state follows the DP-sharded target requests. That - # works only because MTP draft layers are target-shaped and can be - # appended to the target pool (get_pp_layers). An external drafter - # has its own architecture and layer count, so there is nothing to - # append and the bail would strand it on the private context arena - # -- dense in max_seq_len, and with KV heads unsharded under - # attention DP that is the allocation that OOMs at long ISL. Its - # own manager follows the DP sharding regardless: same mapping (so - # `tp_size = 1 if enable_attention_dp` applies to it too) and block - # tables copied per rank from attn_metadata.request_ids. + # Narrower than is_external_drafter(): PARD and DRAFT_TARGET_ONE_MODEL + # never reach the arena this carve-out exists for. + spec_dec_mode = self._speculative_config.spec_dec_mode + is_standalone_drafter = (spec_dec_mode.is_dflash() + or spec_dec_mode.is_dspark()) + if self._mapping.enable_attention_dp and not is_standalone_drafter: + # This bail suits MTP, whose draft layers are target-shaped and + # appendable to the target pool. A standalone drafter has nothing to + # append, so it would be stranded on the private arena instead. logger.info( "Attention DP is enabled, separate draft KV cache is not supported." ) @@ -2535,12 +2521,6 @@ def _create_one_model_draft_kv_cache_manager( # the sparse_attention_config. Get it from effective_draft_config which # falls back to the target model's config for MTP mode. sparse_attn_config = effective_draft_config.sparse_attention_config - # The standalone drafter is a plain dense model built against the - # repurposed mapping under helix (CP ranks become TP ranks; each rank - # keeps its full drafter KV). Its paged KV manager must therefore use - # the repurposed, CP-free mapping — helix round-robin ledger semantics - # apply only to the TARGET KV (and KVCacheManagerV2 rejects - # is_draft x helix outright). draft_mapping = self._get_draft_mapping() draft_tokens_per_block = self._get_draft_tokens_per_block() if draft_tokens_per_block != self._tokens_per_block: @@ -3444,6 +3424,9 @@ def _create_kv_cache_manager( head_dim: Optional[Union[int, List[int]]] = None, kv_cache_type=None, is_disagg: bool = False, + disable_overlap_scheduler: bool = False, + cold_page_codec_provider: Optional[object] = None, + kv_events_config: Optional[KVEventsConfig] = None, draft_num_layers: int = 0, draft_num_kv_heads: int = 0, draft_head_dim: int = 0, @@ -3454,11 +3437,18 @@ def _create_kv_cache_manager( draft_window_size: Optional[int] = None, draft_attention_backend: Optional[str] = None, draft_max_position_embeddings: Optional[int] = None, - joint_kv_cache_reuse: bool = False) -> KVCacheManager: + joint_kv_cache_reuse: bool = False, + max_cuda_graph_batch_size: Optional[int] = None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config """ + if cold_page_codec_provider is not None and not issubclass( + kv_cache_manager_cls, KVCacheManagerV2): + raise ValueError( + "Cold-page quantization requires the resolved KV cache manager " + f"to be KVCacheManagerV2; selected {kv_cache_manager_cls.__name__}") + if draft_num_layers and not issubclass(kv_cache_manager_cls, KVCacheManagerV2): raise ValueError("Unified draft cache layers require KVCacheManagerV2.") @@ -3602,6 +3592,22 @@ def _create_kv_cache_manager( "cold_page_codec_provider"] = cold_page_codec_provider manager_extra_kwargs["kv_events_config"] = kv_events_config manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse + manager_extra_kwargs[ + "disable_overlap_scheduler"] = disable_overlap_scheduler + # Vocab size also enables multimodal event decoding and its per-block + # digest scan. Leave it unset for text-only engines. Without an engine + # (e.g. separate one-model draft caches), preserve config resolution: + # a text sub-config alone cannot identify a multimodal deployment. + needs_multimodal_keys = model_engine is None or is_multimodal( + model_engine.model, model_engine.input_processor) + manager_extra_kwargs["vocab_size"] = (resolve_vocab_size(config) if + needs_multimodal_keys else None) + if (needs_multimodal_keys and manager_extra_kwargs["vocab_size"] is None + and kv_cache_config.enable_block_reuse): + logger.warning( + "Could not resolve vocab_size from the model config; " + "multimodal requests will fail when block reuse is on. " + "Disable block reuse to serve them.") manager_extra_kwargs.update( draft_num_layers=draft_num_layers, draft_num_kv_heads=draft_num_kv_heads, @@ -3614,6 +3620,11 @@ def _create_kv_cache_manager( draft_attention_backend=draft_attention_backend, draft_max_position_embeddings=draft_max_position_embeddings, ) + elif kv_events_config is not None and kv_events_config.enable_kv_cache_events: + logger.warning( + "kv_cache_config.kv_events_config is set but streaming KV event " + "publishing requires KV cache manager V2; events will not be " + f"published for {kv_cache_manager_cls.__name__}.") if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): manager_extra_kwargs["is_disagg"] = is_disagg diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 7dbcebefe67c..97c62517cb3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -18,6 +18,8 @@ ) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.distributed import Distributed +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( MoeLoadBalancer, @@ -649,6 +651,22 @@ def _forward_decoder( is_dummy=is_dummy, gather_context_logits=gather_context_logits, ) + spec_worker = None + if enable_spec_decode: + speculative_model = self.model + if isinstance(speculative_model, MultimodalModelMixin): + speculative_model = speculative_model.language_model + if not isinstance(speculative_model, SpecDecOneEngineForCausalLM): + speculative_model = getattr(speculative_model, "llm", None) + if isinstance(speculative_model, SpecDecOneEngineForCausalLM): + spec_worker = speculative_model.spec_worker + if spec_worker is not None: + spec_worker.prepare_managed_draft_cache( + speculative_model.draft_model, + spec_metadata, + attn_metadata, + resource_manager, + ) inputs = prepared.kwargs runtime_draft_len = prepared.runtime_draft_len if execution_promoted_context_ids: @@ -698,6 +716,18 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None: if captures_model_body: return outputs + if ( + spec_worker is not None + and not is_dummy + and not getattr(self._graphs.cuda_graph_runner, "is_warmup_only", False) + ): + history_update = spec_worker.snapshot_managed_draft_history() + if history_update is not None: + # Graph dictionaries are reused; each overlapped step owns + # its own readback until sampling commits the history. + outputs = dict(outputs) + outputs["draft_history_update"] = history_update + if self.forward_pass_callable is not None: self.forward_pass_callable() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a963738f9ff8..77b9c0dce52d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -78,6 +78,7 @@ BatchDesc, BeamIndex, BufferConfig, + CacheDomain, CacheLevel, CacheTier, CacheTierConfig, @@ -137,6 +138,7 @@ scan_kv_buffer_for_invalid_values, ) from ..scheduler import ScheduledRequests +from ..standalone_draft_cache import StandaloneDraftHistory if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata @@ -680,6 +682,7 @@ def _estimate_cache_size_components( scratch: bool, generation_capacity_headroom: int, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + standalone_draft_reserve: int = 0, ) -> tuple[int, int, int]: """Return context/generation bytes per token and generation bytes per request. @@ -708,7 +711,7 @@ def _estimate_cache_size_components( return ( full_attn_size + context_swa_size, full_attn_size + generation_swa_size, - generation_swa_per_request, + generation_swa_per_request + standalone_draft_reserve * sum(layer_sizes), ) @@ -1227,6 +1230,8 @@ def _update_kv_cache_draft_token_location( run_kv_cache_relocation = True if not run_kv_cache_relocation: return + if getattr(cache_manager, "_standalone_draft_layer_ids", ()): + raise ValueError("Unified standalone draft KV does not support tree-token relocation") requests = scheduled_batch.all_requests() ( accepted_draft_token_offsets, @@ -1772,6 +1777,19 @@ class KVCacheManagerV2(BaseResourceManager): # _init_mixed_precision assigns fresh per-instance values. low_precision_dtype: Optional[DataType] = None high_layers: dict[int, int] = {} + + _draft_layer_sizes: tuple[int, ...] = () + + cp_as_tp_layer_ids: frozenset[int] = frozenset() + + draft_window_size: int | None = None + + _standalone_draft_layer_ids: tuple[int, ...] = () + + _standalone_draft_reserve: int = 0 + + _draft_registered_as_model_layers: bool = False + # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True @@ -1831,8 +1849,39 @@ def __init__( joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + draft_num_layers: int = 0, + draft_num_kv_heads: int = 0, + draft_head_dim: int = 0, + draft_dtype: DataType = DataType.BF16, + draft_kv_factor: int = 2, + draft_total_num_kv_heads: int | None = None, + draft_scratch_tokens: int = 0, + draft_window_size: int | None = None, + draft_attention_backend: str | None = None, + draft_max_position_embeddings: int | None = None, **kwargs, ) -> None: + self._draft_layer_sizes = _get_draft_layer_sizes( + draft_num_layers=draft_num_layers, + draft_num_kv_heads=draft_num_kv_heads, + draft_head_dim=draft_head_dim, + draft_dtype=draft_dtype, + draft_kv_factor=draft_kv_factor, + ) + if draft_total_num_kv_heads is not None and draft_total_num_kv_heads < draft_num_kv_heads: + raise ValueError("Total draft KV heads must cover the rank-local heads") + self._draft_total_num_kv_heads = draft_total_num_kv_heads + if draft_scratch_tokens < 0: + raise ValueError("Standalone draft scratch capacity must be nonnegative") + if draft_window_size is not None and draft_window_size <= 0: + raise ValueError("Draft history window must be positive") + self.draft_window_size = draft_window_size + self.draft_attention_backend = draft_attention_backend + self.draft_max_position_embeddings = draft_max_position_embeddings + self._standalone_draft_layer_ids: tuple[int, ...] = () + self.draft_history: dict[int, StandaloneDraftHistory] = {} + self._draft_dummy_request_ids: set[int] = set() + self._standalone_draft_reserve = draft_scratch_tokens self.mapping = mapping self.dtype = dtype self._declared_target_layer_count = num_layers @@ -1853,6 +1902,8 @@ def __init__( self._init_mixed_precision( kv_cache_config, kwargs.get("pretrained_config"), max_seq_len, dtype ) + self._declared_target_layer_count = num_layers + self.layer_specs = tuple(layer_specs) if layer_specs is not None else None self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager @@ -2028,6 +2079,7 @@ def __init__( self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( _get_generation_kv_capacity(spec_config, is_draft=self.is_draft) ) + self._generation_kv_capacity_headroom += self._standalone_draft_reserve self.event_buffer_max_size = kv_cache_config.event_buffer_max_size self.enable_stats = enable_stats @@ -2178,6 +2230,9 @@ def append_to_kv_heads_per_layer( f"unique values={set(self.head_dim_per_layer)}" ) + self.dtype_per_layer = [self.dtype] * self.num_local_layers + self.kv_factor_per_layer = [self.kv_factor] * self.num_local_layers + self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) max_util_for_resume = kv_cache_config.max_util_for_resume @@ -2317,6 +2372,14 @@ def append_to_kv_heads_per_layer( cache_tiers=cache_tiers, ) config = self._build_cache_config(config) + config = self._append_standalone_draft_layers( + config, + num_kv_heads=draft_num_kv_heads, + head_dim=draft_head_dim, + dtype=draft_dtype, + kv_factor=draft_kv_factor, + ) + config = self._configure_unified_draft_pool_ratio(config) config = self._remove_zero_size_buffers(config) config = self._add_mixed_precision_layers(config) has_host_cache_tier = any( @@ -2585,7 +2648,10 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) if self._sparse_layer_group_ids: self.sparse_metadata_batch = Batch( - self.impl, index_mapper_capacity, self.max_blocks_per_seq, self.max_beam_width + self.impl, + index_mapper_capacity, + self.max_blocks_per_seq, + self.max_beam_width, ) self._log_kv_cache_pool_lifecycle_mapping() @@ -2884,10 +2950,13 @@ def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: layer_id = self._pool_layer_ids_by_role[(pool_id, Role.KEY)] roles = self._get_layer_roles_by_local_id(int(layer_id)) return Role.KEY, Role.VALUE if Role.VALUE in roles else None - role_b = None if self.kv_cache_type == CacheTypeCpp.SELFKONLY else Role.VALUE + layer_id = int(self.impl.layer_grouping[pool_id][0]) + role_b = Role.VALUE if self._get_layer_kv_factor_by_local_id(layer_id) == 2 else None return Role.KEY, role_b def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: + if self._is_standalone_draft_layer(local_layer_id): + return self.kv_factor_per_layer[local_layer_id] if self.layer_specs is not None: return self.layer_specs[local_layer_id].kv_factor return self.kv_factor @@ -2900,6 +2969,8 @@ def _get_layer_page_table_roles( return Role.KEY, (Role.VALUE if self.layer_specs[local_layer_id].kv_factor == 2 else None) def _get_layer_dtype_by_local_id(self, local_layer_id: int) -> DataType: + if self._is_standalone_draft_layer(local_layer_id): + return self.dtype_per_layer[local_layer_id] if self.layer_specs is not None: return self.layer_specs[local_layer_id].dtype return self.dtype @@ -3194,7 +3265,17 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Opti layer_sizes.append( self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=Role.ALL) ) - attention_windows.append(self.max_attention_window_vec[local_layer_idx]) + if self._is_standalone_draft_layer(local_layer_idx): + attention_windows.append(self.draft_retention_window_size) + else: + attention_windows.append(self.max_attention_window_vec[local_layer_idx]) + # Quota sizing precedes specialized target-layer construction. Include + # draft storage then; after append it is already in the per-layer arrays. + if self._draft_layer_sizes and not self._draft_registered_as_model_layers: + layer_sizes.extend(self._draft_layer_sizes) + attention_windows.extend( + [self.draft_retention_window_size] * len(self._draft_layer_sizes) + ) return layer_sizes, attention_windows def _get_generation_request_capacity(self) -> int: @@ -3213,19 +3294,24 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: return tokens def _get_max_tokens_from_quota_impl(self, quota: int) -> float: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - ( - context_size_per_token, - generation_size_per_token, - generation_swa_size_per_request, - ) = _estimate_cache_size_components( - layer_sizes, - attention_windows, - self.tokens_per_block, - scratch=self.enable_swa_scratch_reuse, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - layer_specs=self.layer_specs, - ) + if self._draft_layer_sizes: + (context_size_per_token, generation_size_per_token, generation_swa_size_per_request) = ( + self._get_runtime_cache_size_components() + ) + else: + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + ( + context_size_per_token, + generation_size_per_token, + generation_swa_size_per_request, + ) = _estimate_cache_size_components( + layer_sizes, + attention_windows, + self.tokens_per_block, + scratch=self.enable_swa_scratch_reuse, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, + ) size_per_batch = self._get_generation_request_capacity() * generation_swa_size_per_request if quota < size_per_batch: return 0 @@ -3250,19 +3336,24 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: return self._get_quota_from_max_tokens_impl(max_tokens) def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - ( - context_size_per_token, - generation_size_per_token, - generation_swa_size_per_request, - ) = _estimate_cache_size_components( - layer_sizes, - attention_windows, - self.tokens_per_block, - scratch=self.enable_swa_scratch_reuse, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - layer_specs=self.layer_specs, - ) + if self._draft_layer_sizes: + (context_size_per_token, generation_size_per_token, generation_swa_size_per_request) = ( + self._get_runtime_cache_size_components() + ) + else: + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + ( + context_size_per_token, + generation_size_per_token, + generation_swa_size_per_request, + ) = _estimate_cache_size_components( + layer_sizes, + attention_windows, + self.tokens_per_block, + scratch=self.enable_swa_scratch_reuse, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, + ) context_tokens = min(max_tokens, self.max_num_tokens) generation_tokens = max_tokens - context_tokens return int( @@ -3428,6 +3519,12 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: ] for role_idx, role in enumerate(roles): converter = self.impl.get_page_index_converter(layer_id, role) + if converter.expansion != 1 and self._is_standalone_draft_layer(layer_id): + pool_ids[local_layer_idx, role_idx] = pool_id + scales[local_layer_idx, role_idx] = 1 + layer_offsets[local_layer_idx, role_idx] = 0 + scratch_pages[local_layer_idx, role_idx] = 0 + continue if converter.expansion != 1: raise NotImplementedError( "SWA scratch block-table conversion does not support " @@ -3612,6 +3709,8 @@ def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: """Return the primary cache roles allocated for one local layer.""" if self.layer_specs is not None: return list(self.layer_specs[local_layer_idx].buffer_roles) + if self._is_standalone_draft_layer(local_layer_idx): + return [Role.KEY, Role.VALUE][: self.kv_factor_per_layer[local_layer_idx]] roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: roles.append(Role.VALUE) @@ -3641,9 +3740,11 @@ def _build_base_config( """ scratch_reuse_config = None if self.enable_swa_scratch_reuse: - # Context requests allocate num_extra_kv_tokens for spec decoding. - # They should not count toward the scratch range. - scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens) + # Unwritten speculative and draft reserves must stay outside the + # scratch range when a context chunk advances its history. + scratch_reuse_config = SwaScratchReuseConfig( + max_rewind_len=self.num_extra_kv_tokens + self._standalone_draft_reserve + ) # Beam search only replicates blocks from the prompt tail onward, so # KVCacheDesc splits each request into a shared prefix (prompt_length) @@ -4238,6 +4339,8 @@ def target_layer_ids(self) -> tuple[int, ...]: @property def draft_layer_ids(self) -> tuple[int, ...]: + if getattr(self, "_standalone_draft_layer_ids", ()): + return self._standalone_draft_layer_ids return tuple( layer_id for layer_id in self.pp_layers @@ -4246,7 +4349,7 @@ def draft_layer_ids(self) -> tuple[int, ...]: def get_draft_kv_cache_view(self) -> Optional["KVCacheManagerV2"]: """Return this manager when explicit draft layers share its request timeline.""" - if self.layer_specs is None or not self.draft_layer_ids: + if not self.draft_layer_ids: return None return self @@ -4534,7 +4637,11 @@ def get_index_k_buffer( return full_view[:, 0] def get_num_available_tokens( - self, *, token_num_upper_bound: int, batch_size: int = 1, max_num_draft_tokens: int = 0 + self, + *, + token_num_upper_bound: int, + batch_size: int = 1, + max_num_draft_tokens: int = 0, ) -> int: """Clamp ``token_num_upper_bound`` to the allocatable token capacity. @@ -4547,7 +4654,9 @@ def get_num_available_tokens( replicates all tokens on every rank, so both bounds constrain the same request-length variable. """ - extra_tokens = self.num_extra_kv_tokens + max_num_draft_tokens + extra_tokens = ( + self.num_extra_kv_tokens + max_num_draft_tokens + self._standalone_draft_reserve + ) # Token num upper bound is the maximum number of tokens that can be allocated in the kv cache manager. # We need to add extra tokens to the token num upper bound to account for the extra tokens. clamped = ( @@ -4589,10 +4698,10 @@ def get_num_free_blocks(self) -> int: ) # Legacy managers estimate primary-pool capacity from page indices. - max_num_pages = max( + max_num_blocks = max( [ - self.impl.get_page_index_upper_bound(layer_id, Role.KEY) - for layer_id in typed_range(LayerId(self.num_local_layers)) + self._get_attention_pool_num_blocks(local_layer_idx) + for local_layer_idx in range(self.num_local_layers) ] ) # Subtract the pages the reserved sequences hold, not the count of @@ -4600,7 +4709,7 @@ def get_num_free_blocks(self) -> int: # page (the guard is today). Summing ``num_blocks`` keeps this a page # count if a reservation ever spans more than one. reserved_pages = sum(int(self.kv_cache_map[req_id].num_blocks) for req_id in reserved) - return max_num_pages // self.kv_factor - reserved_pages + return max_num_blocks - reserved_pages def commit_scheduled_kv_cache_stats(self, scheduled_batch: ScheduledRequests) -> None: if self.is_draft or (not self.enable_stats and not self._request_stats_enabled_ids): @@ -4904,6 +5013,8 @@ def _context_reuse_tokens( """Radix-tree evidence for one context prefix. Raw prompt, so block hashes stay comparable with KV cache events, the cache-aware router and disagg peers. """ + if self._draft_context_exceeds_limit(req): + return [] all_tokens = self._reuse_token_source(req) # The target must recompute the last prompt token, but a one-model draft # also reads D prompt tokens ahead. Include those tokens as match evidence; @@ -5045,7 +5156,7 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) return None kv_cache.cuda_stream = self._cache_stream(req.py_request_id) - if not self.enable_block_reuse: + if not self.enable_block_reuse or self._draft_context_exceeds_limit(req): kv_cache.stop_committing() else: self._record_branch_snapshot_point(req, kv_cache, num_lookup_tokens) @@ -5074,6 +5185,19 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) kv_cache.enable_swa_scratch_reuse = False if not self._resume_context_cache(req, kv_cache): return None + if ( + self._standalone_draft_layer_ids + and not self._draft_history_uses_kv_history() + and req.py_request_id not in self.draft_history + ): + # The reuse match covers every cache domain. Initialize once, + # after resume succeeds, so overlap/chunk retries cannot rewind + # draft history that has already advanced on the device. + reused = kv_cache.num_committed_tokens + valid_length = reused + if self.draft_window_size is not None: + valid_length = min(valid_length, self.draft_window_size) + self.set_draft_history(req.py_request_id, valid_length, reused) return kv_cache.num_committed_tokens # Subsequent chunk: cache must exist from first chunk. It may be @@ -5148,7 +5272,12 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: if kv_cache is None: return False - target = req.context_current_position + num_tokens + self.num_extra_kv_tokens + target = ( + req.context_current_position + + num_tokens + + self.num_extra_kv_tokens + + self._standalone_draft_reserve + ) capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -5225,7 +5354,12 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: else (req.total_input_len_cp if self._has_cp_helix else req.prompt_len) ) prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens + target = ( + prompt_len + + get_draft_token_length(req) + + self.num_extra_kv_tokens + + self._standalone_draft_reserve + ) capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -6908,6 +7042,7 @@ def try_commit_blocks(self, request: LlmRequest) -> None: self.enable_block_reuse and self._can_publish_block_reuse and not request.is_dummy_request + and not self._draft_context_exceeds_limit(request) ) if not should_block_reuse: return @@ -7038,6 +7173,10 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._disagg_receive_ready.pop(request.py_request_id, None) self._on_allocation_stream.discard(request.py_request_id) self._allocation_ready.pop(request.py_request_id, None) + if self._standalone_draft_layer_ids: + self.draft_history.pop(request.py_request_id, None) + self._draft_dummy_request_ids.discard(request.py_request_id) + self._pending_overlap_slack.pop(request.py_request_id, None) kv_cache = self.kv_cache_map.pop(request.py_request_id, None) if kv_cache is None: self.impl.clear_stats_excluded(request.py_request_id) @@ -7218,9 +7357,9 @@ def get_cache_bytes_per_token(self) -> int: data_roles.append(Role.VALUE_BLOCK_SCALE) return sum( - self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=data_role) - for local_layer_idx in range(self.num_local_layers) - for data_role in data_roles + self.get_layer_bytes_per_token(local_id, Role.ALL) + * (self._helix_cp_size if local_id in self.cp_as_tp_layer_ids else 1) + for local_id in range(len(self.dtype_per_layer)) ) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): @@ -7232,24 +7371,37 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): if self._is_mixed_precision_layer(local_layer_idx): return self._mixed_precision_layer_bytes_per_token(local_layer_idx, data_role) - if self.dtype not in ( + if self.layer_specs is not None: + layer = self.layer_specs[local_layer_idx] + if data_role == Role.ALL: + return layer.bytes_per_token() + return layer.bytes_per_token(data_role) + + if data_role not in (Role.ALL, *self._get_buffer_roles_for_layer(local_layer_idx)): + return 0 + if self._get_layer_dtype_by_local_id(local_layer_idx) not in ( DataType.FP8, DataType.HALF, DataType.BF16, DataType.FLOAT, DataType.NVFP4, ): - raise ValueError(f"Cannot support {self.dtype} KV cache.") + raise ValueError(f"Cannot support {self.dtype_per_layer[local_layer_idx]} KV cache.") if data_role == Role.ALL: - kv_factor = self.kv_factor - elif data_role in [Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: + kv_factor = self.kv_factor_per_layer[local_layer_idx] + elif data_role in [ + Role.KEY, + Role.VALUE, + Role.KEY_BLOCK_SCALE, + Role.VALUE_BLOCK_SCALE, + ]: if data_role in [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: - assert self.dtype == DataType.NVFP4, ( + assert self.dtype_per_layer[local_layer_idx] == DataType.NVFP4, ( "NVFP4 is the only supported dtype for block quant data roles" ) if data_role == Role.VALUE: - assert self.kv_cache_type != CacheTypeCpp.SELFKONLY, ( + assert self.kv_factor_per_layer[local_layer_idx] == 2, ( "VALUE data role is not supported for SELFKONLY cache type" ) kv_factor = 1 @@ -7262,14 +7414,16 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): * self.head_dim_per_layer[local_layer_idx] ) - cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, self.dtype) + cache_size_bytes_per_token = get_size_in_bytes( + cache_size_per_token, self.dtype_per_layer[local_layer_idx] + ) if data_role in [Role.KEY, Role.VALUE]: return cache_size_bytes_per_token quant_size_per_token = 0 - if self.dtype == DataType.NVFP4: + if self.dtype_per_layer[local_layer_idx] == DataType.NVFP4: quant_size_per_token = self.calculate_scaling_factor_size_bytes( cache_size_per_token, quant_vector_size=16, @@ -7362,6 +7516,9 @@ def shutdown(self): for kv_cache in self.kv_cache_map.values(): kv_cache.close() self.kv_cache_map.clear() + if self._standalone_draft_layer_ids: + self.draft_history.clear() + self._draft_dummy_request_ids.clear() self._disagg_receive_ready.clear() self._on_allocation_stream.clear() self._allocation_ready.clear() @@ -7414,6 +7571,8 @@ def get_cache_size_per_token( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, + draft_scratch_tokens: int = 0, + draft_window_size: int | None = None, **kwargs, ): if layer_specs is not None: @@ -7448,6 +7607,32 @@ def get_cache_size_per_token( backoff, max_seq_len, ) + draft_layer_sizes = _get_draft_layer_sizes(**kwargs) + if draft_layer_sizes: + _, headroom = _get_generation_kv_capacity(spec_config, is_draft=is_draft) + headroom += draft_scratch_tokens + context, generation, fixed = _estimate_cache_size_components( + layer_sizes, + attention_windows, + tokens_per_block, + scratch=bool(kv_cache_config and kv_cache_config.enable_swa_scratch_reuse), + generation_capacity_headroom=headroom, + standalone_draft_reserve=draft_scratch_tokens, + ) + draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( + draft_layer_sizes, + tokens_per_block, + scratch_tokens=draft_scratch_tokens, + window_size=draft_window_size, + generation_capacity_headroom=headroom, + helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, + ) + return ( + generation + draft_generation, + (fixed + draft_fixed) * int(max_batch_size or 0) + + (context + draft_context - generation - draft_generation) + * int(max_num_tokens or 0), + ) _, generation_capacity_headroom = _get_generation_kv_capacity( spec_config, is_draft=is_draft ) @@ -7821,6 +8006,7 @@ def update_context_resources(self, scheduled_batch: ScheduledRequests): self.enable_block_reuse and self._can_publish_block_reuse and not req.is_dummy_request + and not self._draft_context_exceeds_limit(req) ) is_all_reusable = self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE # A chunk that ends at a branch point commits there. History has @@ -7915,6 +8101,16 @@ def update_resources( if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) else kv_cache.capacity - rewind_len ) + embedded_draft = KVCacheManagerV2._draft_history_uses_kv_history(self) + if self._standalone_draft_layer_ids and not embedded_draft and new_capacity is not None: + draft_history = self.get_draft_history(req.py_request_id) + if draft_history is not None: + # Rejected target verification slots do not revoke draft + # feature history. Retain it and its next-forward scratch. + new_capacity = max( + new_capacity, + draft_history.position + self._standalone_draft_reserve, + ) history_length = ( None # Reuse (history's consumer) is disabled under helix, and @@ -7922,6 +8118,14 @@ def update_resources( if self.kv_compression_manages_history or self._has_cp_helix else req.max_beam_num_tokens - 1 ) + if embedded_draft and self._has_cp_helix: + # HELIX stores rank-local prompt tokens but generated tokens + # are replicated. The cache ledger uses global positions. + history_length = ( + req.total_input_len_cp + req.max_beam_num_tokens - req.py_prompt_len - 1 + ) + if embedded_draft and new_capacity is not None and history_length is not None: + new_capacity = max(new_capacity, history_length) success = kv_cache.resize(new_capacity, history_length) if not success: raise ValueError( @@ -8092,6 +8296,8 @@ def _create_kv_cache( kv_cache.stop_committing() self.impl.mark_stats_excluded(request_id) kv_cache.discard_pending_stats() + if self._standalone_draft_layer_ids: + self._draft_dummy_request_ids.add(request_id) index = self.index_mapper.add_new_sequence(request_id) if self.sparse_metadata_batch is not None: self.sparse_metadata_batch.add(kv_cache, index) @@ -8193,3 +8399,499 @@ def reset_reuse_state(self): self.conversation_manager.clear() if self._sharing is not None: self._sharing._on_reset() + + def _get_runtime_cache_size_components(self) -> tuple[int, int, int]: + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + num_target_layers = len(layer_sizes) - (len(self._draft_layer_sizes)) + context_size, generation_size, fixed_size = _estimate_cache_size_components( + layer_sizes[:num_target_layers], + attention_windows[:num_target_layers], + self.tokens_per_block, + scratch=self.enable_swa_scratch_reuse, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + standalone_draft_reserve=self._standalone_draft_reserve, + ) + if self._draft_layer_sizes: + draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( + self._draft_layer_sizes, + self.tokens_per_block, + scratch_tokens=self._standalone_draft_reserve, + window_size=self.draft_window_size, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + helix_cp_size=self._helix_cp_size, + ) + context_size += draft_context + generation_size += draft_generation + fixed_size += draft_fixed + return context_size, generation_size, fixed_size + + def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: + """Draft bytes within the unified manager's usable token capacity.""" + assert self._draft_layer_sizes + if self._has_cp_helix: + # Each target ledger block retains all of its draft subpages. + max_tokens = ( + math.ceil(max_tokens / self._ledger_tokens_per_block) + * self._ledger_tokens_per_block + ) + context_size, generation_size, fixed_size = _estimate_draft_cache_size_components( + self._draft_layer_sizes, + self.tokens_per_block, + scratch_tokens=self._standalone_draft_reserve, + window_size=self.draft_window_size, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + ) + context_tokens = ( + max_tokens if self.max_num_tokens is None else min(max_tokens, self.max_num_tokens) + ) + return math.ceil( + context_tokens * context_size + + (max_tokens - context_tokens) * generation_size + + self._get_generation_request_capacity() * fixed_size + ) + + def _append_standalone_draft_layers( + self, + config: KVCacheManagerConfigPy, + *, + num_kv_heads: int, + head_dim: int, + dtype: DataType, + kv_factor: int, + register_model_layers: bool = True, + ) -> KVCacheManagerConfigPy: + if not self._draft_layer_sizes: + return config + first_global_id = max(self.pp_layers, default=-1) + 1 + first_local_id = len(config.layers) + # FP4 MLA adds native HP-tail layers outside the model-layer arrays. + register_model_layers = register_model_layers and first_local_id == len(self.pp_layers) + self._draft_registered_as_model_layers = register_model_layers + # Specialized managers can replace model layers with multiple native layers. + padding = first_local_id - len(self.dtype_per_layer) + assert padding >= 0 + num_draft_layers = len(self._draft_layer_sizes) + self.dtype_per_layer.extend([self.dtype] * padding + [dtype] * num_draft_layers) + self.kv_factor_per_layer.extend([self.kv_factor] * padding + [kv_factor] * num_draft_layers) + self.num_kv_heads_per_layer.extend([0] * padding + [num_kv_heads] * num_draft_layers) + self.head_dim_per_layer.extend([0] * padding + [head_dim] * num_draft_layers) + self.cp_as_tp_layer_ids = frozenset( + range(first_local_id, first_local_id + num_draft_layers) + ) + self._standalone_draft_layer_ids = tuple( + range(first_global_id, first_global_id + len(self._draft_layer_sizes)) + ) + layers = list(config.layers) + for offset, global_id in enumerate(self._standalone_draft_layer_ids): + local_id = first_local_id + offset + layers.append( + AttentionLayerConfig( + layer_id=LayerId(local_id), + buffers=[ + BufferConfig( + role=role, + size=KVCacheManagerV2.get_layer_bytes_per_token(self, local_id, role) + * self.tokens_per_block, + tokens_per_block_override=( + self.tokens_per_block if self._helix_cp_size > 1 else None + ), + ) + for role in self._get_buffer_roles_for_layer(local_id) + ], + sliding_window_size=self.draft_retention_window_size, + cache_domain=CacheDomain.STANDALONE_DRAFT, + ) + ) + self.layer_offsets[global_id] = local_id + if register_model_layers: + self.pp_layers.append(global_id) + self.total_num_kv_heads_per_layer.append( + self._draft_total_num_kv_heads or num_kv_heads + ) + if register_model_layers: + self.num_local_layers = len(self.pp_layers) + self.num_layers += len(self._draft_layer_sizes) + + def reserve_scratch(batch: BatchDesc) -> BatchDesc: + return BatchDesc( + [ + KVCacheDesc( + capacity=desc.capacity + + (self._standalone_draft_reserve if desc.capacity else 0), + history_length=desc.history_length, + ) + for desc in batch.kv_caches + ], + system_prompt_length=batch.system_prompt_length, + ) + + return replace( + config, + layers=layers, + constraints=[reserve_scratch(batch) for batch in config.constraints], + typical_step=reserve_scratch(config.typical_step) + if config.typical_step is not None + else None, + ) + + def _configure_unified_draft_pool_ratio( + self, config: KVCacheManagerConfigPy + ) -> KVCacheManagerConfigPy: + """Preserve target ratios within their share of the unified byte quota.""" + if not self._standalone_draft_layer_ids: + return config + target_ratios = config.initial_pool_ratio + if target_ratios is None: + return config + + # Native lifecycle equality includes window, rounded sink blocks, sparsity, + # and cache domain. All SSM layers share one lifecycle, represented by None. + target_life_cycles: list[AttnLifeCycle | None] = [] + for layer in config.layers: + if isinstance(layer, AttentionLayerConfig): + if layer.cache_domain == CacheDomain.STANDALONE_DRAFT: + continue + life_cycle = AttnLifeCycle.make( + layer.sliding_window_size, + layer.num_sink_tokens, + config.tokens_per_block, + cache_domain=layer.cache_domain, + ) + else: + life_cycle = None + if not any( + type(existing) is type(life_cycle) and existing == life_cycle + for existing in target_life_cycles + ): + target_life_cycles.append(life_cycle) + + if len(target_ratios) != len(target_life_cycles): + raise ValueError( + "kv_cache_config.pool_ratio must have one entry per target layer group " + f"({len(target_life_cycles)}), got {len(target_ratios)}" + ) + # Use the same model-specific capacity accounting as quota selection; + # specialized targets can include compressed KV or recurrent snapshots. + usable_quota = int(config.cache_tiers[GPU_LEVEL].quota * config.max_util_for_resume) + max_tokens = self._get_max_tokens_from_quota(usable_quota) + if math.isinf(max_tokens): + max_tokens = self.max_num_tokens or self.max_seq_len + if usable_quota <= 0 or self._get_quota_from_max_tokens(max_tokens) > usable_quota: + raise ValueError("Unified KV cache quota cannot hold the target and draft reserves") + draft_ratio = self._get_draft_quota_from_max_tokens(max_tokens) / usable_quota + if not 0 < draft_ratio < 1: + raise ValueError("Unified KV cache quota must leave room for target and draft pools") + + # Appended draft layers share one lifecycle, after all target groups. + return replace( + config, + initial_pool_ratio=[ratio * (1 - draft_ratio) for ratio in target_ratios] + + [draft_ratio], + ) + + def _get_layer_buffer_view(self, local_layer_idx: int, kv_layout: str) -> torch.Tensor: + return self._get_kv_buffer_view( + local_layer_idx, + dtype=self.dtype_per_layer[local_layer_idx], + kv_factor=self.kv_factor_per_layer[local_layer_idx], + num_kv_heads=self.num_kv_heads_per_layer[local_layer_idx], + head_dim=self.head_dim_per_layer[local_layer_idx], + kv_layout=kv_layout, + page_expansion=self._helix_cp_size if local_layer_idx in self.cp_as_tp_layer_ids else 1, + ) + + def _get_kv_buffer_view( + self, + layer_id: int, + *, + dtype: DataType | torch.dtype, + kv_factor: int, + num_kv_heads: int, + head_dim: int, + kv_layout: str, + page_expansion: int = 1, + ) -> torch.Tensor: + """View shared pool pages for a manager-local layer and its KV geometry.""" + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported KV layout: {kv_layout}") + addr_key = self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) + if kv_factor == 2: + addr_value = self.impl.get_mem_pool_base_address( + layer_id, Role.VALUE, PageIndexMode.SHARED + ) + page_size_key = self.impl.get_page_stride(layer_id, Role.KEY) + page_size_value = self.impl.get_page_stride(layer_id, Role.VALUE) + if ( + addr_key + page_size_key * page_expansion != addr_value + or page_size_key != page_size_value + ): + raise ValueError("K/V buffers must have adjacent equal-sized pages") + + element_per_container = 1 + if dtype == DataType.NVFP4: + element_per_container = 2 + dtype = torch.int8 + + dimensions = ( + [self.tokens_per_block, num_kv_heads] + if kv_layout == "NHD" + else [num_kv_heads, self.tokens_per_block] + ) + head_dim //= element_per_container + page_upper = self.impl.get_page_index_upper_bound(layer_id, Role.KEY) + num_pages = page_upper // kv_factor + strides = None + if page_expansion > 1: + # Each ledger slot contains all of its K subpages, then all V + # subpages. Expanded KEY indices skip the V spans between slots. + num_pages = page_upper - (kv_factor - 1) * page_expansion + page_elements = self.tokens_per_block * num_kv_heads * head_dim + strides = [ + page_elements, + page_expansion * page_elements, + dimensions[1] * head_dim, + head_dim, + 1, + ] + return convert_to_torch_tensor( + TensorWrapper( + addr_key, + dtype, + [ + num_pages, + kv_factor, + *dimensions, + head_dim, + ], + strides=strides, + ) + ) + + def _is_standalone_draft_layer(self, local_layer_idx: int) -> bool: + return local_layer_idx in self.cp_as_tp_layer_ids + + def get_draft_buffers(self, local_layer_idx: int, kv_layout: str = "HND") -> torch.Tensor: + """View one draft layer using the same storage contract as every other layer.""" + if not 0 <= local_layer_idx < len(self._standalone_draft_layer_ids): + raise ValueError("No standalone draft cache layer at this index") + layer_idx = self._standalone_draft_layer_ids[local_layer_idx] + return self._get_layer_buffer_view(self.layer_offsets[layer_idx], kv_layout) + + @property + def draft_max_blocks_per_seq(self) -> int: + """Maximum physical draft pages, including every Helix ledger subpage.""" + return self.max_blocks_per_seq * self._helix_cp_size + + def get_draft_num_blocks(self, request_id: int) -> int: + """Allocated draft pages covering the request's global token positions.""" + return self.kv_cache_map[request_id].num_blocks * self._helix_cp_size + + @property + def draft_retention_window_size(self) -> int | None: + return self.draft_window_size + 1 if self.draft_window_size is not None else None + + def get_draft_transfer_identity(self) -> dict: + local_layer_idx = self.layer_offsets[self._standalone_draft_layer_ids[0]] + return { + "num_layers": len(self._standalone_draft_layer_ids), + "num_kv_heads": self._draft_total_num_kv_heads + or self.num_kv_heads_per_layer[local_layer_idx], + "head_dim": self.head_dim_per_layer[local_layer_idx], + "dtype": str(binding_to_torch_dtype(self.dtype_per_layer[local_layer_idx])), + "attention_backend": ( + "VANILLA" + if self.draft_attention_backend == "TRTLLM" + else self.draft_attention_backend + ), + "kv_factor": self.kv_factor_per_layer[local_layer_idx], + "window_size": self.draft_window_size, + } + + def _draft_history_uses_kv_history(self) -> bool: + """Embedded DSpark writes the complete live suffix of target history.""" + return getattr(self, "draft_attention_backend", None) == "DSv4" + + def get_draft_block_indices( + self, + request_ids: List[int], + histories: Optional[Sequence[StandaloneDraftHistory]] = None, + ) -> List[List[int]]: + """Validated rank-local mappings for allocated draft pages, without padding.""" + for request_id in request_ids: + cache = self.kv_cache_map.get(request_id) + if cache is None or not cache.is_active: + raise ValueError(f"Standalone draft request {request_id} has no active cache") + expanded = self._helix_cp_size > 1 + layer_idx = self._standalone_draft_layer_ids[0] + batch_indices = self.get_batch_cache_indices(request_ids, layer_idx, raw_indices=expanded) + if expanded: + converter = self.impl.get_page_index_converter(self.layer_offsets[layer_idx], Role.KEY) + batch_indices = [converter(indices, PageIndexMode.SHARED) for indices in batch_indices] + for row, indices in enumerate(batch_indices): + cache = self.kv_cache_map[request_ids[row]] + if ( + len(indices) != self.get_draft_num_blocks(request_ids[row]) + or len(indices) > self.draft_max_blocks_per_seq + ): + raise ValueError("Standalone draft cache has an incomplete or oversized page table") + first_required_block = 0 + if self.draft_window_size is not None: + history = None + if not self._draft_history_uses_kv_history(): + history = ( + histories[row] + if histories is not None + else self.get_draft_history(request_ids[row]) + ) + history_start = ( + history.position - history.valid_length + if history is not None + else max(0, cache.history_length - self.draft_window_size) + ) + first_required_block = history_start // self.tokens_per_block + if any(index == BAD_PAGE_INDEX for index in indices[first_required_block:]): + raise ValueError( + "Draft cache contains missing pages in its required history or scratch" + ) + return batch_indices + + def get_draft_block_table( + self, + request_ids: List[int], + histories: Optional[Sequence[StandaloneDraftHistory]] = None, + ) -> torch.Tensor: + """Current rank-local page mappings; unused tail entries point to page zero.""" + batch_indices = self.get_draft_block_indices(request_ids, histories) + table = torch.zeros((len(request_ids), self.draft_max_blocks_per_seq), dtype=torch.int32) + for row, indices in enumerate(batch_indices): + table[row, : len(indices)] = torch.tensor(indices, dtype=torch.int32) + return table + + def get_draft_history(self, request_id: int) -> Optional[StandaloneDraftHistory]: + return self.draft_history.get(request_id) + + def set_draft_history(self, request_id: int, valid_length: int, position: int) -> None: + cache = self.kv_cache_map.get(request_id) + if cache is None or not cache.is_active: + raise ValueError( + f"Standalone draft request {request_id} has no active cache allocation" + ) + history = StandaloneDraftHistory(valid_length, position) + if history.position > cache.capacity: + raise ValueError("Standalone draft history exceeds allocated capacity") + if self.draft_window_size is not None and history.valid_length > self.draft_window_size: + raise ValueError("Draft history exceeds its retention window") + self.draft_history[request_id] = history + + def export_draft_history(self, request_id: int) -> Optional[dict]: + if not self._standalone_draft_layer_ids: + return None + if self._draft_history_uses_kv_history(): + cache = self.kv_cache_map[request_id] + if not cache.is_active: + raise ValueError("Embedded DSpark transfer requires an active cache") + return { + "valid_length": min(cache.history_length, self.draft_window_size), + "position": cache.history_length, + "layout": self.get_draft_transfer_identity(), + } + history = self.get_draft_history(request_id) + if history is None: + raise ValueError( + f"Standalone draft request {request_id} has no valid history to transfer" + ) + return { + "valid_length": history.valid_length, + "position": history.position, + "layout": self.get_draft_transfer_identity(), + } + + def restore_draft_history(self, request_id: int, metadata: dict) -> None: + if ( + not self._standalone_draft_layer_ids + or metadata.get("layout") != self.get_draft_transfer_identity() + ): + raise ValueError("Standalone draft transfer layout does not match the receiving worker") + valid_length = metadata.get("valid_length") + position = metadata.get("position") + if self._draft_history_uses_kv_history(): + cache = self.kv_cache_map[request_id] + # Receive admission establishes the global prompt watermark; the + # transfer completion path validates that the received KV covers it. + if ( + type(position) is not int + or type(valid_length) is not int + or position != cache.history_length + or valid_length != min(position, self.draft_window_size) + ): + raise ValueError("Embedded DSpark transfer does not match the admitted KV history") + self.get_draft_block_table([request_id]) + return + history = StandaloneDraftHistory(valid_length, position) + # Validate receiver-local allocation before publishing history. + self.get_draft_block_table([request_id], [history]) + self.set_draft_history(request_id, valid_length, position) + + def _get_attention_pool_num_blocks(self, local_layer_idx: int) -> int: + """Count ledger blocks without counting expanded physical subpages.""" + layer_id = LayerId(local_layer_idx) + expansion = self.impl.get_page_index_converter(layer_id, Role.KEY).expansion + return self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // ( + self.get_layer_kv_factor(self.pp_layers[local_layer_idx]) * expansion + ) + + def _draft_context_exceeds_limit(self, req: LlmRequest) -> bool: + """Whether draft prefill skips this prompt, preventing unified block reuse.""" + if ( + not self._standalone_draft_layer_ids + or self.draft_window_size is not None + or self.draft_max_position_embeddings is None + ): + return False + prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len + return prompt_len > min(self.max_seq_len, self.draft_max_position_embeddings) + + +def _get_draft_layer_sizes( + *, + draft_num_layers: int = 0, + draft_num_kv_heads: int = 0, + draft_head_dim: int = 0, + draft_dtype: DataType = DataType.BF16, + draft_kv_factor: int = 2, + **kwargs, +) -> tuple[int, ...]: + if draft_num_layers < 0: + raise ValueError("Draft layer count must be nonnegative") + if draft_num_layers: + if min(draft_num_kv_heads, draft_head_dim) <= 0: + raise ValueError("Standalone draft cache dimensions must be positive") + if draft_dtype not in (DataType.HALF, DataType.BF16, DataType.FP8): + raise ValueError("Standalone draft KV supports FP16, BF16 and FP8 storage") + if draft_kv_factor not in (1, 2): + raise ValueError("KV storage requires one or two planes") + size = get_size_in_bytes(draft_num_kv_heads * draft_head_dim * draft_kv_factor, draft_dtype) + return (size,) * draft_num_layers + + +def _estimate_draft_cache_size_components( + layer_sizes: Sequence[int], + tokens_per_block: int, + *, + generation_capacity_headroom: int, + scratch_tokens: int = 0, + window_size: int | None = None, + helix_cp_size: int = 1, +) -> tuple[int, int, int]: + """Draft storage cost, with history and scratch requirements supplied separately.""" + retention_window = window_size + 1 if window_size is not None else None + context, generation, per_request = _estimate_cache_size_components( + layer_sizes, + [retention_window] * len(layer_sizes), + tokens_per_block, + scratch=False, + generation_capacity_headroom=generation_capacity_headroom, + standalone_draft_reserve=scratch_tokens, + ) + extra_bytes_per_token = (helix_cp_size - 1) * sum(layer_sizes) + return context + extra_bytes_per_token, generation + extra_bytes_per_token, per_request diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py index b88c7f3cfd53..daa57ecf26c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -24,6 +24,11 @@ import torch +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + _estimate_draft_cache_size_components, + _get_draft_layer_sizes, +) + if sys.version_info[:2] >= (3, 12): from typing import override else: @@ -583,6 +588,7 @@ def _estimate_mamba_hybrid_cache_cost( include_explicit_snapshots: bool, cap_partial_attention_snapshots: bool, extra_state_bytes_per_rank: int = 0, + draft_scratch_tokens: int = 0, is_draft: bool = False, use_separate_draft_kv_cache: bool = False, **kwargs, @@ -641,6 +647,22 @@ def _estimate_mamba_hybrid_cache_cost( regular_slope = math.ceil(state_bytes_per_rank / interval) if has_unaligned_periodic_snapshot and not cap_partial_attention_snapshots: regular_slope += math.ceil(attention_block_bytes / interval) + draft_layer_sizes = _get_draft_layer_sizes(**kwargs) + if draft_layer_sizes: + # Hybrid target attention is full attention. Reserve capture/noise + # capacity only in its attention pools and the distinct draft pools; + # recurrent state retains the fixed/snapshot accounting above. + # Hybrid profiling bounds draft storage by full-history attention. + _, draft_slope, draft_fixed = _estimate_draft_cache_size_components( + draft_layer_sizes, + tokens_per_block, + generation_capacity_headroom=1, + scratch_tokens=draft_scratch_tokens, + helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, + ) + draft_fixed += draft_scratch_tokens * attention_slope + attention_slope += draft_slope + intercept += max_batch_size * draft_fixed return attention_slope + regular_slope, intercept diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 3d4dd7b8b952..dea8fc6afe78 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -424,11 +424,16 @@ def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConf ] def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: - return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] + if self._is_standalone_draft_layer(local_layer_idx): + return False + layer_id = self.pp_layers[local_layer_idx] + return layer_id < len(self._mamba_layer_mask) and self._mamba_layer_mask[layer_id] @override def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_standalone_draft_layer(layer_id): + return KVCacheManagerV2._get_pool_roles(self, pool_id) if self._is_local_mamba_layer(layer_id): return MambaRole.SSM_STATE, None return super()._get_pool_roles(pool_id) @@ -574,6 +579,18 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: ) return attention_quota + state_quota + extra_attention_quota + def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: + draft_quota = super()._get_draft_quota_from_max_tokens(max_tokens) + snapshot_slots = self._num_ssm_snapshots_for_capacity(max_tokens, self.kv_cache_config) + if snapshot_slots > 0: + # The snapshot plan's partial-page reserve also covers draft KV. + draft_quota += ( + self._max_resident_sequences() + * sum(self._draft_layer_sizes) + * self.tokens_per_block + ) + return draft_quota + @override def _get_max_tokens_from_quota(self, quota: int) -> float: if self._get_quota_from_max_tokens(0) > quota: @@ -847,9 +864,7 @@ def get_num_free_blocks(self) -> int: // self._ssm_page_index_scale ) else: - attention_pages.append( - self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // self.kv_factor - ) + attention_pages.append(self._get_attention_pool_num_blocks(local_layer_idx)) if attention_pages: return max(attention_pages) return max(ssm_pages) if ssm_pages else 0 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py deleted file mode 100644 index 234128fca0dc..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ /dev/null @@ -1,5326 +0,0 @@ -# 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 hashlib -import math -import os -import sys -from collections import OrderedDict, defaultdict, deque -from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union - -import numpy as np -import torch -from strenum import StrEnum - -from tensorrt_llm._torch.disaggregation.resource.page import MapperKind -from tensorrt_llm._torch.distributed.communicator import Distributed, ReduceOp -from tensorrt_llm._utils import ( - TensorWrapper, - convert_to_torch_tensor, - get_size_in_bytes, - prefer_pinned, -) -from tensorrt_llm.bindings.internal.batch_manager import KvCacheIterationStats, KvCacheStats -from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( - IndexMapper, - copy_batch_block_offsets_to_device, -) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig -from tensorrt_llm.runtime.kv_cache_hash import get_effective_kv_cache_event_hash_algo -from tensorrt_llm.runtime.kv_cache_manager_v2 import _BACKEND as KV_CACHE_MANAGER_V2_BACKEND -from tensorrt_llm.runtime.kv_cache_manager_v2 import ( - _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, - BAD_PAGE_INDEX, - CACHE_LEVEL1, - DEFAULT_BEAM_INDEX, - GPU_LEVEL, - AttentionLayerConfig, - AttnLifeCycle, - BatchDesc, - BufferConfig, - CacheLevel, - CacheTierConfig, - CuError, - DataRole, - DiskCacheTierConfig, - GpuCacheTierConfig, - HostCacheTierConfig, - KVCacheDesc, - KVCacheEventManager, - KVCacheIterationStatsDelta, - LayerId, - PageIndexMode, - PlannedDropHandle, - PoolGroupPeakBlockStats, - ReuseScope, - SwaScratchReuseConfig, - TokenIdExt, - _cpp_introspection, - _introspection, - _KVCache, - exact_div, - gen_multimodal_cache_key_tokens, - typed_range, -) -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as KVCacheManagerPy -from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy -from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfMemoryError as KVCacheOutOfMemoryError -from tensorrt_llm.sampling_params import SamplingParams - -from ..._utils import binding_to_torch_dtype, mpi_rank, nvtx_range, str_dtype_to_torch -from ...logger import logger -from ...mapping import CpType, Mapping -from ..utils import maybe_compile -from .config_utils import uses_vswa_kv_cache_layout -from .connectors.kv_cache_connector import KvCacheConnectorManager -from .kv_cache_stats import ( - KVCacheV2IterationStatsReport, - KVCacheV2LifeCycleIterationStats, - KVCacheV2PoolGroupIterationStats, - KVCacheV2SsmLifeCycleIterationStats, - KVCacheV2SsmSnapshotIterationStats, -) -from .llm_request import LlmRequest, LlmRequestState, SamplingConfig, get_draft_token_length -from .resource_manager import ( - BaseResourceManager, - CacheTypeCpp, - DataType, - KVCacheManager, - ModelConfigCpp, - ModelConfigPython, - _populate_dummy_mrope_config, - get_pp_layers, - request_context, -) -from .scheduler import ScheduledRequests -from .standalone_draft_cache import StandaloneDraftHistory - -if KV_CACHE_MANAGER_V2_BACKEND == "cpp": - from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheDomain - -if TYPE_CHECKING: - from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata - -BASE_GENERATION_TOKEN_COUNT = 1 - -KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS -KV_CACHE_ITERATION_STATS_REUSE_FIELDS = ( - "iter_reused_blocks", - "iter_full_reused_blocks", - "iter_partial_reused_blocks", - "iter_missed_blocks", -) -KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS = tuple( - field_name - for field_name in KV_CACHE_ITERATION_STATS_DELTA_FIELDS - if field_name not in KV_CACHE_ITERATION_STATS_REUSE_FIELDS -) - - -def _create_gpu_cache_tier_config(quota: int, enable_ugpu: bool) -> GpuCacheTierConfig: - if KV_CACHE_MANAGER_V2_BACKEND == "cpp": - if enable_ugpu: - logger.warning_once( - "uGPU KV-cache localization is not supported by the C++ KV cache manager V2; " - "disabling it.", - key="disable_ugpu_for_cpp_kv_cache_manager_v2", - ) - return GpuCacheTierConfig(quota=quota) - return GpuCacheTierConfig(quota=quota, enable_ugpu=enable_ugpu) - - -def _get_num_ugpus(impl) -> int: - """Return the effective uGPU count supported by the selected backend.""" - if KV_CACHE_MANAGER_V2_BACKEND == "cpp": - return 1 - return impl.num_ugpus - - -def _get_kv_cache_ugpu_id(kv_cache) -> int | None: - """Return the cache-local uGPU id supported by the selected backend.""" - if KV_CACHE_MANAGER_V2_BACKEND == "cpp": - return None - return kv_cache.ugpu_id - - -def _create_backend_kv_cache( - impl, - lora_task_id: int | None, - salt: int | None, - input_tokens: Sequence[TokenIdExt] | None, - request_id: int, - expected_prompt_length: int | None, - ugpu_id: int | None, -): - if KV_CACHE_MANAGER_V2_BACKEND == "cpp": - assert ugpu_id is None, "The C++ KV cache manager V2 does not support uGPU localization" - return impl.create_kv_cache( - ReuseScope(lora_id=lora_task_id, salt=salt), - input_tokens, - id=request_id, - expected_prompt_length=expected_prompt_length, - ) - return impl.create_kv_cache( - ReuseScope(lora_id=lora_task_id, salt=salt, ugpu_id=ugpu_id), - input_tokens, - id=request_id, - expected_prompt_length=expected_prompt_length, - ugpu_id=ugpu_id, - ) - - -class Role: - KEY = DataRole("key") - VALUE = DataRole("value") - KEY_BLOCK_SCALE = DataRole("key_block_scale") - VALUE_BLOCK_SCALE = DataRole("value_block_scale") - # Sparse-attention per-layer index-K cache (MiniMax-M3 and similar - # sparse-block-selection backends). Registered as a native V2 - # BufferConfig on sparse layers via the extra_buffers_per_layer hook on - # _build_base_config, so allocation, free, slot reuse, and prefix reuse - # share the lifecycle of the main K/V buffers for the same layer. - INDEX_KEY = DataRole("index_key") - MLA_V_SCALE = DataRole("mla_v_scale") - MLA_V_PACKED = DataRole("mla_v_packed") - MLA_HP_TAIL = DataRole("mla_hp_tail") - ALL = DataRole("all") - - -class BlockReusePolicy(StrEnum): - ALL_REUSABLE = "all_reusable" - PER_REQUEST = "per_request" - PER_CONVERSATION = "per_conversation" - - -def _request_conversation_id(request: LlmRequest) -> Optional[str]: - if request.is_dummy_request: - return None - conversation_params = request.py_conversation_params - if conversation_params is None: - return None - conversation_id = conversation_params.conversation_id.strip() - return conversation_id or None - - -@dataclass(slots=True) -class _ConversationState: - current_request_id: Optional[int] = None - planned_drop_handles: deque[PlannedDropHandle] = field(default_factory=deque) - - -class ConversationManager: - """Track the current request and drop plan for each conversation.""" - - def __init__(self, max_num_turns: int) -> None: - self._max_num_turns = max_num_turns - self._conversation_states: Dict[str, _ConversationState] = {} - - def save_drop_plan(self, request: LlmRequest, kv_cache: _KVCache) -> None: - """Save a completed context's drop plan and apply the preceding plan on success.""" - request_id = request.py_request_id - conversation_id = _request_conversation_id(request) - if conversation_id is None: - return - - state = self._conversation_states[conversation_id] - if state.current_request_id != request_id: - return - - drop_handle = kv_cache.plan_committed_block_drop() - if drop_handle is None: - logger.warning( - f"Committed blocks for request {request_id} in conversation " - f"{conversation_id} have been dropped." - ) - else: - state.planned_drop_handles.append(drop_handle) - if len(state.planned_drop_handles) > self._max_num_turns: - state.planned_drop_handles.popleft().drop() - - self.finish_request(request) - - def prepare_request(self, request: LlmRequest) -> None: - """Register a context request unless its conversation has another active one.""" - conversation_id = _request_conversation_id(request) - if conversation_id is None: - return - request_id = request.py_request_id - state = self._conversation_states.setdefault(conversation_id, _ConversationState()) - current_request_id = state.current_request_id - if current_request_id is not None and current_request_id != request_id: - logger.warning( - f"Conversation {conversation_id} already has current request " - f"{current_request_id}. Request {request_id} will ignore " - "conversation params." - ) - return - - state.current_request_id = request_id - - def finish_request(self, request: LlmRequest) -> None: - """Clear a request as active while preserving any saved drop plan.""" - conversation_id = _request_conversation_id(request) - if conversation_id is None: - return - state = self._conversation_states.get(conversation_id) - if state is None or state.current_request_id != request.py_request_id: - return - - state.current_request_id = None - if not state.planned_drop_handles: - self._conversation_states.pop(conversation_id) - - def clear(self) -> None: - """Clear state after reusable KV-cache blocks have been cleared.""" - self._conversation_states.clear() - - -def _estimate_full_attn_size_per_token( - layer_sizes: Sequence[int], attention_windows: Sequence[Optional[int]] -) -> int: - return sum( - layer_size - for layer_size, window_size in zip(layer_sizes, attention_windows) - if window_size is None or window_size <= 0 - ) - - -def _compute_auto_host_tier_quota( - quota: int, - local_ranks: int, - mem_available: float, - memlock_limit: float, -) -> int: - """Compute the auto-provisioned host cache tier quota for a single rank. - - The host tier backs the MAX_UTILIZATION scheduler's suspend/resume path, - so it must be positive. It defaults to the device quota, capped by the - per-node available-memory budget shared with co-located ranks and by the - pinnable-memory (RLIMIT_MEMLOCK) limit. - - Args: - quota: Device (GPU) KV cache quota in bytes; must be positive. - local_ranks: Number of ranks co-located on this physical node. - mem_available: Available host memory in bytes, or ``float("inf")`` - if unknown. - memlock_limit: RLIMIT_MEMLOCK soft limit in bytes, or - ``float("inf")`` if unlimited or unknown. - - Returns: - Host cache tier quota in bytes (always positive). - """ - candidates = [quota] - if mem_available != float("inf"): - candidates.append(int(mem_available / local_ranks * 0.5)) - if memlock_limit != float("inf"): - candidates.append(int(memlock_limit * 0.8)) - host_quota = min(candidates) - if host_quota <= 0: - logger.warning( - f"KV cache manager v2 auto host tier sizing computed a " - f"non-positive quota ({host_quota}); falling back to the " - f"device quota {quota / (1 << 30):.2f}GiB" - ) - host_quota = quota - return host_quota - - -def _sync_host_tier_quota(host_quota: int, mapping: Mapping) -> int: - """Reduce the auto-provisioned host cache tier quota to the fleet minimum. - - ``_compute_auto_host_tier_quota`` reads rank-local host state - (``SC_AVPHYS_PAGES``, ``RLIMIT_MEMLOCK``), so co-scheduled ranks can arrive - at divergent host quotas (observed up to 10x). Divergent host-tier - retention makes per-rank MAX_UTILIZATION schedulers disagree about which - suspended requests can resume, which wedges collectives on - non-attention-DP TP. Reducing to the fleet minimum makes the - most-constrained rank set the value for everyone, mirroring the device - quota sync. A single-rank job needs no collective and is returned as-is. - - Args: - host_quota: This rank's locally-computed host tier quota in bytes. - mapping: Parallelism mapping; ``world_size`` selects whether a - collective is needed. - - Returns: - The globally-agreed host tier quota in bytes (the fleet ``MIN``). - """ - if mapping.world_size > 1: - host_quota = Distributed.get(mapping).allreduce(host_quota, op=ReduceOp.MIN) - return host_quota - - -def _estimate_swa_cache_size( - layer_sizes: Sequence[int], - attention_windows: Sequence[Optional[int]], - tokens_per_block: int, - *, - context: bool, - scratch: bool, - generation_capacity_headroom: Optional[int] = None, -) -> tuple[int, int]: - tokens_per_block = int(tokens_per_block) - size_per_token = 0 - size_per_request = 0 - scratch_keys = set() - for layer_size, window_size in zip(layer_sizes, attention_windows): - if window_size is not None and window_size > 0: - if generation_capacity_headroom is None: - window_tokens = math.ceil(window_size / tokens_per_block) * tokens_per_block - else: - window_blocks = ( - math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) - + 1 - ) - window_tokens = window_blocks * tokens_per_block - if not context: - size_per_request += window_tokens * layer_size - elif not scratch: - size_per_token += layer_size - else: - scratch_key = (int(window_size), layer_size) - if scratch_key in scratch_keys: - size_per_request += window_tokens * layer_size - else: - scratch_keys.add(scratch_key) - size_per_token += layer_size - return size_per_token, size_per_request - - -def _get_generation_kv_capacity(spec_config, *, is_draft: bool) -> tuple[int, int]: - """Return draft-token reserve and total capacity headroom over history.""" - from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens - - if spec_config is None: - return 0, BASE_GENERATION_TOKEN_COUNT - reserve = spec_config.max_total_draft_tokens - if ( - is_draft - and getattr(spec_config, "use_dynamic_tree", False) - and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0 - ): - draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len - reserve = max(reserve, draft_loop_tokens) - dynamic_reserve = reserve - spec_config.max_total_draft_tokens - headroom = ( - get_num_extra_kv_tokens(spec_config) - + spec_config.tokens_per_gen_step - + dynamic_reserve - + reserve # K3's overlap scheduler can be one verification round ahead. - ) - return reserve, headroom - - -def _estimate_cache_size_components( - layer_sizes: Sequence[int], - attention_windows: Sequence[int | None], - tokens_per_block: int, - *, - scratch: bool, - generation_capacity_headroom: int, - standalone_draft_reserve: int = 0, -) -> tuple[int, int, int]: - """Return context/generation bytes per token and generation bytes per request. - - Static profiling and runtime quota conversion must charge the same SWA - retention pages and context scratch space. The standalone draft reserve - is a shared allocation envelope, including windowed attention groups; - it does not represent committed history. Resume-watermark normalization - is separate from these usable-capacity costs. Callers estimate target and - draft layers separately so only target layers can share context scratch. - """ - full_attn_size = _estimate_full_attn_size_per_token(layer_sizes, attention_windows) - context_swa_size, _ = _estimate_swa_cache_size( - layer_sizes, - attention_windows, - tokens_per_block, - context=True, - scratch=scratch, - ) - generation_swa_size, generation_swa_per_request = _estimate_swa_cache_size( - layer_sizes, - attention_windows, - tokens_per_block, - context=False, - scratch=False, - generation_capacity_headroom=generation_capacity_headroom, - ) - return ( - full_attn_size + context_swa_size, - full_attn_size + generation_swa_size, - generation_swa_per_request + standalone_draft_reserve * sum(layer_sizes), - ) - - -def _get_draft_layer_sizes( - *, - draft_num_layers: int = 0, - draft_num_kv_heads: int = 0, - draft_head_dim: int = 0, - draft_dtype: DataType = DataType.BF16, - draft_kv_factor: int = 2, - **kwargs, -) -> tuple[int, ...]: - if draft_num_layers < 0: - raise ValueError("Draft layer count must be nonnegative") - if draft_num_layers: - if min(draft_num_kv_heads, draft_head_dim) <= 0: - raise ValueError("Standalone draft cache dimensions must be positive") - if draft_dtype not in (DataType.HALF, DataType.BF16, DataType.FP8): - raise ValueError("Standalone draft KV supports FP16, BF16 and FP8 storage") - if draft_kv_factor not in (1, 2): - raise ValueError("KV storage requires one or two planes") - size = get_size_in_bytes(draft_num_kv_heads * draft_head_dim * draft_kv_factor, draft_dtype) - return (size,) * draft_num_layers - - -def _estimate_draft_cache_size_components( - layer_sizes: Sequence[int], - tokens_per_block: int, - *, - generation_capacity_headroom: int, - scratch_tokens: int = 0, - window_size: int | None = None, - helix_cp_size: int = 1, -) -> tuple[int, int, int]: - """Draft storage cost, with history and scratch requirements supplied separately.""" - retention_window = window_size + 1 if window_size is not None else None - context, generation, per_request = _estimate_cache_size_components( - layer_sizes, - [retention_window] * len(layer_sizes), - tokens_per_block, - scratch=False, - generation_capacity_headroom=generation_capacity_headroom, - standalone_draft_reserve=scratch_tokens, - ) - extra_bytes_per_token = (helix_cp_size - 1) * sum(layer_sizes) - return context + extra_bytes_per_token, generation + extra_bytes_per_token, per_request - - -def _extend_swa_windows_for_reuse( - attention_windows: Sequence[Optional[int]], - reuse_match_backoff: int, - max_seq_len: Optional[int], -) -> List[Optional[int]]: - """Extend physical SWA retention to cover a backed-off reuse match.""" - if reuse_match_backoff <= 0: - return list(attention_windows) - - retention_windows: List[Optional[int]] = [] - for window_size in attention_windows: - if window_size is None or window_size <= 0: - retention_windows.append(window_size) - continue - - retention_window = int(window_size) + reuse_match_backoff - retention_windows.append( - None - if max_seq_len is not None and retention_window >= int(max_seq_len) - else retention_window - ) - return retention_windows - - -def _get_static_cache_size_layer_components( - model_config: ModelConfigPython, - mapping: Mapping, - num_layers: Optional[int] = None, - **kwargs, -) -> tuple[List[int], List[Optional[int]]]: - config = model_config.pretrained_config - max_seq_len = kwargs.get("max_seq_len") - kv_cache_config = kwargs.get("kv_cache_config") - - num_key_value_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) - if isinstance(num_key_value_heads, Iterable): - num_key_value_heads = sum(num_key_value_heads) / len(num_key_value_heads) - - mla = hasattr(config, "kv_lora_rank") and config.kv_lora_rank is not None - if mla: - head_dim = config.kv_lora_rank + config.qk_rope_head_dim - kv_factor = 1 - else: - tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size - head_dim = getattr(config, "head_dim", None) - if not isinstance(head_dim, int): - head_dim = config.hidden_size // config.num_attention_heads - head_dim = head_dim * num_key_value_heads // tp_size - kv_factor = 2 - - cache_size_per_token = kv_factor * head_dim - quant_config = model_config.quant_config - if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): - layer_size = cache_size_per_token - elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache(): - layer_size = math.ceil(cache_size_per_token / 2) + math.ceil(cache_size_per_token / 16) - else: - assert quant_config is None or (not quant_config.quant_mode.has_kv_cache_quant()), ( - "Quantized kv cache is not expected" - ) - layer_size = cache_size_per_token * 2 - - num_attention_layers = KVCacheManager._resolve_num_attention_layers( - model_config, mapping, num_layers - ) - layer_sizes = [layer_size] * num_attention_layers - window_pattern = kv_cache_config.max_attention_window if kv_cache_config is not None else None - - def get_window_size(layer_idx: int) -> Optional[int]: - if window_pattern is None or not isinstance(window_pattern, (list, tuple)): - return None - window_size = window_pattern[layer_idx % len(window_pattern)] - if window_size is None or window_size <= 0: - return None - window_size = int(window_size) - if max_seq_len is not None and window_size == int(max_seq_len): - return None - return window_size - - attention_windows = [get_window_size(layer_idx) for layer_idx in range(num_attention_layers)] - return layer_sizes, attention_windows - - -class _MmRunMetadata(NamedTuple): - """CPU tensors for exact multimodal run lookup. - - `num_runs` is the total number of flat multimodal runs across all - multimodal items, i.e. `item_run_cu_offsets[-1]`. All members are 1-D CPU - int64 tensors with shape `[num_runs]`. - """ - - # Shape [num_runs]; full-prompt start index for each flat run. - run_positions: torch.Tensor - # Shape [num_runs]; one-past-last full-prompt index for each flat run. - run_ends: torch.Tensor - # Shape [num_runs]; logical multimodal item index for each flat run. - run_item_indices: torch.Tensor - # Shape [num_runs]; item-local token offset where each flat run begins. - run_item_offsets: torch.Tensor - - -def _hash_to_digest(hash_ints: Sequence[int]) -> bytes: - # Convert 8 x int32 hash chunks to the 32-byte digest used by C++ block - # keys. The byte order matches getNthByte(), which extracts MSB first. - try: - hash_len = len(hash_ints) - except TypeError as exc: - raise ValueError("Expected 8 int32 hash values, got non-sized input") from exc - if hash_len != 8: - raise ValueError(f"Expected 8 int32 hash values, got {hash_len}") - if not all(isinstance(value, int) for value in hash_ints): - raise ValueError("Expected multimodal hash values to be integers") - return b"".join(v.to_bytes(4, "big", signed=True) for v in hash_ints) - - -def _ensure_int64_cpu_tensor(values: Sequence[int] | torch.Tensor) -> torch.Tensor: - # Block-reuse augmentation is Python-side index math. The metadata is - # produced by host mm preprocessing and carried in - # MultimodalInput metadata. A non-CPU tensor here means the upstream - # contract drifted and would introduce an unexpected sync in the - # block-reuse path, so fail loudly. - if isinstance(values, torch.Tensor): - if values.device.type != "cpu": - raise ValueError( - f"multimodal block-reuse metadata must be CPU-resident, got {values.device}" - ) - return values.to(dtype=torch.int64) - return torch.as_tensor(values, dtype=torch.int64) - - -_MM_RUN_METADATA_UNSET = object() - - -def _resolve_multimodal_run_metadata(req: LlmRequest) -> Optional[_MmRunMetadata]: - # Memoize on the request. Block reuse resolves this on the first context - # chunk, again on every commit during chunked prefill, and on prefetch, yet - # the result is a pure function of the request's immutable multimodal run - # layout. Caching on the request bounds the cache to the request lifetime, - # so no separate teardown is needed. - cached = getattr(req, "py_mm_run_metadata", _MM_RUN_METADATA_UNSET) - if cached is not _MM_RUN_METADATA_UNSET: - return cached - - # Worked example for one logical multimodal item split by text: - # - # prompt index: 0 1 2 3 4 5 - # prompt token: T0 MM0:0 MM0:1 T1 MM0:2 MM0:3 - # flat run: run0 run0 run1 run1 - # - # Inputs encode the prompt layout as: - # - # item_run_cu_offsets = [0, 2] # item 0 owns flat runs [0, 2) - # run_positions = [1, 4] # run0 starts at prompt 1, run1 at 4 - # run_lengths = [2, 2] # token counts for run0 and run1 - # - # The derived arrays below let the augmentation loop map any overlapping - # flat run back to the item digest and to the item-local token offset: - # - # item_run_counts = [2] # item 0 has two runs - # cumulative_run_lengths = [0,2,4] # tokens before run0, run1, and end - # item_starts = [0] # item 0 starts at flat run 0 - # run_item_indices = [0,0] # run0 and run1 both belong to item 0 - # run_item_offsets = [0,2] # run1 begins at MM0 token offset 2 - # - # So a chunk slice covering prompt [4, 6) uses MM0's digest with token - # offset 2; a slice covering prompt [5, 6) uses token offset 3. - - # Shape [num_items + 1]; item i owns flat runs [offsets[i], offsets[i + 1]). - item_run_cu_offsets = req.multimodal_item_run_cu_offsets - # Shape [num_runs]; full-prompt start index for each flat run. - run_positions = req.multimodal_run_positions - # Shape [num_runs]; token count for each flat run. - run_lengths = req.multimodal_run_lengths - - if all(field is None for field in (item_run_cu_offsets, run_positions, run_lengths)): - req.py_mm_run_metadata = None - return None - - if item_run_cu_offsets is None or run_positions is None or run_lengths is None: - raise ValueError( - "multimodal run metadata must be validated before block reuse and provided together" - ) - - item_run_cu_offsets = _ensure_int64_cpu_tensor(item_run_cu_offsets) - run_positions = _ensure_int64_cpu_tensor(run_positions) - run_lengths = _ensure_int64_cpu_tensor(run_lengths) - - # Shape [num_items]; number of flat runs owned by each logical item. - item_run_counts = item_run_cu_offsets[1:] - item_run_cu_offsets[:-1] - - # Shape [num_runs + 1]; item-token count before each flat run, plus end. - cumulative_run_lengths = torch.cat( - (torch.zeros(1, dtype=torch.int64), torch.cumsum(run_lengths, dim=0)) - ) - - # Shape [num_items]; first flat-run index for each logical item. - item_starts = item_run_cu_offsets[:-1] - - # Shape [num_runs]; logical multimodal item index for each flat run. - run_item_indices = torch.repeat_interleave( - torch.arange(item_run_counts.numel(), dtype=torch.int64), item_run_counts - ) - - # Shape [num_runs]; item-local token offset where each flat run begins. - run_item_offsets = ( - cumulative_run_lengths[:-1] - cumulative_run_lengths[item_starts][run_item_indices] - ) - # Shape [num_runs]; one-past-last full-prompt index for each flat run. - run_ends = run_positions + run_lengths - - metadata = _MmRunMetadata( - run_positions=run_positions, - run_ends=run_ends, - run_item_indices=run_item_indices, - run_item_offsets=run_item_offsets, - ) - req.py_mm_run_metadata = metadata - return metadata - - -def _augment_tokens_with_mm_run_metadata( - vocab_size: int, - result: list[TokenIdExt], - multimodal_hashes: Sequence[Sequence[int]], - metadata: _MmRunMetadata, - chunk_start: int, - chunk_end: int, -) -> list[TokenIdExt]: - # Only rewrite multimodal runs that overlap the materialized prompt slice. - overlap_mask = (metadata.run_ends > chunk_start) & (metadata.run_positions < chunk_end) - # Shape [num_overlapping_runs]; flat-run indices that touch this chunk. - overlap_run_indices = torch.nonzero(overlap_mask).flatten() - if overlap_run_indices.numel() == 0: - return result - - # `result` is indexed relative to this chunk, but multimodal cache-key - # tokens are indexed relative to the logical multimodal item. The rewrite - # below therefore computes, for each selected run segment: - # - chunk_result_offset = prompt position - chunk_start - # - item_token_offset = item-local run start + intra-run prompt offset - # Shape [num_overlapping_runs]; item index for each selected flat run. - overlap_run_item_indices = metadata.run_item_indices[overlap_run_indices] - # Materialize the selected run metadata once before the Python loop. These - # are CPU tensors by construction, and the loop below is request hot-path - # bookkeeping. - overlap_run_positions = metadata.run_positions[overlap_run_indices] - prompt_overlap_starts = torch.clamp(overlap_run_positions, min=chunk_start) - prompt_overlap_ends = torch.clamp(metadata.run_ends[overlap_run_indices], max=chunk_end) - item_token_offsets = ( - metadata.run_item_offsets[overlap_run_indices] - + prompt_overlap_starts - - overlap_run_positions - ) - chunk_result_offsets = prompt_overlap_starts - chunk_start - lengths = prompt_overlap_ends - prompt_overlap_starts - - current_item_idx: Optional[int] = None - digest = b"" - for item_idx, chunk_result_offset, item_token_offset, length in zip( - overlap_run_item_indices.tolist(), - chunk_result_offsets.tolist(), - item_token_offsets.tolist(), - lengths.tolist(), - strict=True, - ): - if item_idx != current_item_idx: - current_item_idx = item_idx - digest = _hash_to_digest(multimodal_hashes[item_idx]) - # Feed the coarse item property (content digest) and granular run - # properties (item-local offset and span length) into the key - # generator, so cache keys reflect the actual multimodal tokens being - # rewritten. - result[chunk_result_offset : chunk_result_offset + length] = ( - gen_multimodal_cache_key_tokens( - vocab_size, digest, length, token_offset=item_token_offset - ) - ) - - return result - - -def _augment_tokens_with_contiguous_mm_metadata( - vocab_size: int, - result: list[TokenIdExt], - multimodal_hashes: Sequence[Sequence[int]], - multimodal_positions: Sequence[int] | torch.Tensor, - multimodal_lengths: Sequence[int] | torch.Tensor, - chunk_start: int, - chunk_end: int, -) -> list[TokenIdExt]: - # Legacy metadata assumes every item is one contiguous prompt span. Exact - # run metadata is preferred for video layouts that interleave text tokens. - positions = _ensure_int64_cpu_tensor(multimodal_positions) - lengths = _ensure_int64_cpu_tensor(multimodal_lengths) - - item_ends = positions + lengths - overlap_mask = (item_ends > chunk_start) & (positions < chunk_end) - overlap_item_indices = torch.nonzero(overlap_mask).flatten() - for item_idx_tensor in overlap_item_indices: - item_idx = int(item_idx_tensor) - pos = positions[item_idx].item() - length = lengths[item_idx].item() - overlap_start = max(pos, chunk_start) - overlap_end = min(pos + length, chunk_end) - source_offset = overlap_start - pos - result_offset = overlap_start - chunk_start - overlap_length = overlap_end - overlap_start - result[result_offset : result_offset + overlap_length] = gen_multimodal_cache_key_tokens( - vocab_size, - _hash_to_digest(multimodal_hashes[item_idx]), - overlap_length, - token_offset=source_offset, - ) - - return result - - -def _locate_accepted_draft_tokens(requests: List[LlmRequest]): - num_accepted_draft_tokens = [] - accepted_draft_tokens_indices = [] - rewind_draft_token_separate_adjustments = [] - # for context requests, the py_num_accepted_draft_tokens = 0, and py_num_accepted_draft_tokens_indices = [] - for seq in requests: - num_accepted_draft_tokens.append(seq.py_num_accepted_draft_tokens) - rewind_draft_token_separate_adjustments.append( - seq.py_rewind_draft_token_separate_adjustment - ) - accepted_draft_tokens_indices.extend(seq.py_num_accepted_draft_tokens_indices) - batch_size = len(requests) - num_accepted_draft_tokens_offset = torch.zeros(batch_size + 1, dtype=torch.int32, device="cuda") - num_accepted_draft_tokens_offset[1:] = torch.cumsum( - torch.tensor(num_accepted_draft_tokens, dtype=torch.int32), dim=0 - ) - accepted_draft_tokens_indices = torch.tensor( - accepted_draft_tokens_indices, dtype=torch.int32, device="cuda" - ) - rewind_draft_token_separate_adjustments = torch.tensor( - rewind_draft_token_separate_adjustments, dtype=torch.int32, device="cuda" - ) - return ( - num_accepted_draft_tokens_offset, - accepted_draft_tokens_indices, - rewind_draft_token_separate_adjustments, - ) - - -def _update_kv_cache_draft_token_location( - cache_manager, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata", - kv_cache_dtype_byte_size: float, -): - run_kv_cache_relocation = False - for request in scheduled_batch.generation_requests: - if request.state != LlmRequestState.GENERATION_COMPLETE: - if ( - request.py_num_accepted_draft_tokens > 0 - and len(request.py_num_accepted_draft_tokens_indices) > 0 - ): - run_kv_cache_relocation = True - if not run_kv_cache_relocation: - return - if getattr(cache_manager, "draft_layer_ids", ()): - raise ValueError("Unified standalone draft KV does not support tree-token relocation") - requests = scheduled_batch.all_requests() - ( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - rewind_draft_token_separate_adjustments, - ) = _locate_accepted_draft_tokens(requests) - past_key_value_lengths = attn_metadata.kv_lens_cuda[: len(requests)] - if ( - attn_metadata.kv_cache_block_offsets is not None - and attn_metadata.host_kv_cache_pool_pointers is not None - and attn_metadata.host_kv_cache_pool_mapping is not None - ): - use_paged_kv_cache = True - else: - use_paged_kv_cache = False - assert use_paged_kv_cache, "Only paged kv cache is supported" - assert len(cache_manager.max_attention_window_vec) == 1, ( - "Currently, only one max attention window size is supported." - ) - - if use_paged_kv_cache: - assert len(set(cache_manager.num_kv_heads_per_layer)) == 1, ( - "update_kv_cache_draft_token_location requires uniform num_kv_heads across all layers, " - f"but got {cache_manager.num_kv_heads_per_layer}" - ) - torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( - accepted_draft_token_offsets, - packed_accepted_draft_tokens_indices, - past_key_value_lengths, - True, - cache_manager.num_layers, - # Use TP-sharded num_kv_heads (per-rank) instead of the unsharded - # total so the C++ kernel computes correct strides and grid dims. - cache_manager.num_kv_heads_per_layer[0], - int(cache_manager.head_dim * kv_cache_dtype_byte_size), - cache_manager.max_total_draft_tokens, - cache_manager.max_attention_window_vec[0], - rewind_draft_token_separate_adjustments, - None, - cache_manager.kv_cache_pool_pointers, - attn_metadata.kv_cache_block_offsets, - cache_manager.max_blocks_per_seq, - cache_manager.tokens_per_block, - None, - ) - - -@maybe_compile(options={"max-autotune": True}) -def _copy_swa_block_offsets_with_scratch_compiled( - block_offsets: torch.Tensor, - copy_idx: torch.Tensor, - pool_ids: torch.Tensor, - scales: torch.Tensor, - layer_offsets: torch.Tensor, - scratch_pages: torch.Tensor, - block_positions: torch.Tensor, - scratch_begs: torch.Tensor, - scratch_ends: torch.Tensor, - scratch_slots: torch.Tensor, - num_contexts: torch.Tensor, - output: torch.Tensor, -) -> None: - base = block_offsets[pool_ids[:, :, None], copy_idx[None, None, :], 0, :] - converted = torch.where( - base == BAD_PAGE_INDEX, - BAD_PAGE_INDEX, - base * scales[:, :, None, None] + layer_offsets[:, :, None, None], - ) - - context_positions = torch.arange( - scratch_begs.shape[1], - dtype=torch.int32, - device=scratch_begs.device, - ) - active_context = context_positions < num_contexts - scratch_mask_by_pool = ( - (block_positions >= scratch_begs[:, :, None]) - & (block_positions < scratch_ends[:, :, None]) - & active_context[None, :, None] - ) - range_index = torch.where( - scratch_mask_by_pool, - block_positions - scratch_begs[:, :, None], - 0, - ) - total_offset = range_index[pool_ids] * scratch_pages[:, :, None, None] - slot_idx = (total_offset // scales[:, :, None, None]).clamp(max=scratch_slots.shape[-1] - 1) - slot_id = scratch_slots[pool_ids].gather(-1, slot_idx.long()) - offset = total_offset % scales[:, :, None, None] - scratch_index = ( - slot_id * scales[:, :, None, None] - + (offset + layer_offsets[:, :, None, None]) % scales[:, :, None, None] - ) - scratch_mask = scratch_mask_by_pool[pool_ids] - converted = torch.where(scratch_mask, scratch_index, converted) - - output.copy_(converted.permute(0, 2, 1, 3)) - - -def _rewind_context_cursor(req: LlmRequest, reuse: int) -> None: - """Pull a context cursor back to a shallower prefix. ``setPrepopulatedPromptLen`` - only moves forward, so leave the forward case to it -- it re-floors chunk size too. - """ - if reuse < req.context_current_position: - req.context_current_position = reuse - - -class KVCacheManagerV2(BaseResourceManager): - _draft_layer_sizes: tuple[int, ...] = () - cp_as_tp_layer_ids: frozenset[int] = frozenset() - draft_window_size: int | None = None - draft_layer_ids: tuple[int, ...] = () - _standalone_draft_reserve: int = 0 - _draft_registered_as_model_layers: bool = False - # Read by KvCacheCreator: a subclass that overrides the context - # commit/history protocol opts out of generic reuse-match backoff. - _supports_reuse_match_backoff = True - - def __init__( - self, - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: Union[int, List[int]], - tokens_per_block: int, - # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. - # It's derived from the model's BuildConfig for consistency with the C++ backend. - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - dtype: DataType = DataType.HALF, - spec_config=None, - layer_mask: Optional[List[bool]] = None, - vocab_size: int = None, - max_num_tokens: int = 8192, - model_config: Optional[ModelConfigCpp] = None, - max_beam_width: int = 1, - is_draft: bool = False, - kv_connector_manager: Optional[KvCacheConnectorManager] = None, - execution_stream: Optional[torch.cuda.Stream] = None, - is_disagg: bool = False, - enable_stats: bool = False, - num_reserved_index_slots: int = 1, - is_estimating_kv_cache: bool = False, - enable_ugpu: bool = False, - joint_kv_cache_reuse: bool = False, - draft_num_layers: int = 0, - draft_num_kv_heads: int = 0, - draft_head_dim: int = 0, - draft_dtype: DataType = DataType.BF16, - draft_kv_factor: int = 2, - draft_total_num_kv_heads: int | None = None, - draft_scratch_tokens: int = 0, - draft_window_size: int | None = None, - draft_attention_backend: str | None = None, - draft_max_position_embeddings: int | None = None, - **kwargs, - ) -> None: - self._draft_layer_sizes = _get_draft_layer_sizes( - draft_num_layers=draft_num_layers, - draft_num_kv_heads=draft_num_kv_heads, - draft_head_dim=draft_head_dim, - draft_dtype=draft_dtype, - draft_kv_factor=draft_kv_factor, - ) - if self._draft_layer_sizes and KV_CACHE_MANAGER_V2_BACKEND != "cpp": - raise ValueError("Unified standalone draft KV requires the C++ V2 cache backend") - if draft_total_num_kv_heads is not None and draft_total_num_kv_heads < draft_num_kv_heads: - raise ValueError("Total draft KV heads must cover the rank-local heads") - self._draft_total_num_kv_heads = draft_total_num_kv_heads - if draft_scratch_tokens < 0: - raise ValueError("Standalone draft scratch capacity must be nonnegative") - if draft_window_size is not None and draft_window_size <= 0: - raise ValueError("Draft history window must be positive") - self.draft_window_size = draft_window_size - self.draft_attention_backend = draft_attention_backend - self.draft_max_position_embeddings = draft_max_position_embeddings - self.draft_layer_ids: tuple[int, ...] = () - self.draft_history: dict[int, StandaloneDraftHistory] = {} - self._draft_dummy_request_ids: set[int] = set() - self._standalone_draft_reserve = draft_scratch_tokens - self.mapping = mapping - self.dtype = dtype - self.is_disagg = is_disagg - - assert kv_connector_manager is None, ( - "kv_connector_manager is not supported for KVCacheManagerV2" - ) - assert max_beam_width == 1, "max_beam_width must be 1 for KVCacheManagerV2" - assert not (mapping.cp_config.get("cp_type") == CpType.STAR), ( - "Star attention is not supported for KVCacheManagerV2" - ) - - self.kv_cache_type = kv_cache_type - self.pp_layers, self.num_layers = get_pp_layers( - num_layers, - mapping, - spec_config=spec_config, - layer_mask=layer_mask, - ) - self.is_draft = is_draft - - # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) - # can distinguish the throwaway estimation-phase managers from the - # final ones: the estimation cache is sized with no headroom for - # retained dummy requests. - self.is_estimating_kv_cache = is_estimating_kv_cache - - # Set True by a compression manager; generation-step resize then leaves history untouched. - self.kv_compression_manages_history: bool = False - self.enable_swa_scratch_reuse = ( - kv_cache_config.enable_swa_scratch_reuse and not self.is_draft - ) - block_reuse_config = kv_cache_config.block_reuse_config - self.block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) - self.num_local_layers = len(self.pp_layers) - self.layer_offsets = {idx: offset for offset, idx in enumerate(self.pp_layers)} - self.max_beam_width = max_beam_width - - tp_size = mapping.tp_size - if mapping.enable_attention_dp: - tp_size = 1 - - self.num_kv_heads = num_kv_heads - self.head_dim = head_dim - self.tokens_per_block = tokens_per_block - # Helix super-block ledger: bookkeeping runs in GLOBAL tokens, one - # ledger block = cp_size physical pages (one per rank), while - # self.tokens_per_block stays the PHYSICAL kernel page size. - self._has_cp_helix = mapping.has_cp_helix() - self._helix_cp_rank = mapping.cp_rank if self._has_cp_helix else 0 - self._helix_cp_size = mapping.cp_size if self._has_cp_helix else 1 - self._ledger_tokens_per_block = tokens_per_block * self._helix_cp_size - if self._has_cp_helix: - if kv_cache_config.enable_block_reuse: - raise ValueError( - "KVCacheManagerV2 does not support block reuse with " - "helix context parallelism: a ledger block's bytes span " - "all CP ranks, so publishing it requires a group-wide " - "protocol. Set kv_cache_config.enable_block_reuse=False." - ) - if is_draft: - raise ValueError( - "KVCacheManagerV2 does not support a draft cache " - "manager with helix context parallelism." - ) - if mapping.enable_attention_dp: - raise ValueError( - "KVCacheManagerV2 does not support attention-DP with " - "helix context parallelism: disagg transfer-completion " - "consensus is skipped under attention-DP, so the " - "scheduler's request view would not be rank-invariant " - "across the CP group." - ) - self.max_seq_len = max_seq_len - self.max_batch_size = max_batch_size - self.max_num_tokens = max_num_tokens - self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 - from ..speculative import draft_prompt_lookahead, get_num_extra_kv_tokens - - self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) - # Mirror V1: expose max_draft_len so the native disagg AuxBuffer - # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. - self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 - self.max_total_draft_tokens = ( - spec_config.max_total_draft_tokens if spec_config is not None else 0 - ) - # KvCacheCreator owns the pairing decision; never re-derive it here, or - # the two predicates drift apart and so do the two pools. - self.enable_joint_kv_cache_reuse = joint_kv_cache_reuse - # A draft manager only publishes to its radix tree when it is paired. - self._can_publish_block_reuse = joint_kv_cache_reuse or not self.is_draft - # Unsupported adapters keep their main-like reuse endpoint. - self.reuse_match_backoff = draft_prompt_lookahead(spec_config) or 0 - if not self._supports_reuse_match_backoff: - self.reuse_match_backoff = 0 - # Mirror V1's KV reserve sizing (see V1 __init__ for rationale). - self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( - _get_generation_kv_capacity(spec_config, is_draft=self.is_draft) - ) - self._generation_kv_capacity_headroom += self._standalone_draft_reserve - - self.event_buffer_max_size = kv_cache_config.event_buffer_max_size - self.enable_stats = enable_stats - kv_cache_event_hash_algo = get_effective_kv_cache_event_hash_algo( - kv_cache_config.kv_cache_event_hash_algo, - use_kv_cache_manager_v2=True, - ) - - self._stream = ( - execution_stream if execution_stream is not None else torch.cuda.current_stream() - ) - logger.info(f"[KVCacheManager] execution_stream: {self._stream}") - - # Determine max_attention_window_vec - if kv_cache_config.max_attention_window is not None: - self.max_attention_window_vec = ( - kv_cache_config.max_attention_window.copy() - ) # Make a copy to avoid modifying original - # Clamp all window sizes to max_seq_len before calculating the - # number of KV cache blocks. This prevents the KV cache pool from - # being skewed by the largest window values. - self.max_attention_window_vec = [ - min(max_seq_len, w) for w in self.max_attention_window_vec - ] - - self.max_attention_window_vec = [ - None if w == max_seq_len else w for w in self.max_attention_window_vec - ] - - else: - self.max_attention_window_vec = [None] - - # A D-token match backoff needs D more physical SWA history at the - # committed endpoint so commit_min_snapshot still covers the rewind. - if kv_cache_config.enable_block_reuse and self.reuse_match_backoff: - self.max_attention_window_vec = _extend_swa_windows_for_reuse( - self.max_attention_window_vec, - self.reuse_match_backoff, - max_seq_len, - ) - - event_window_size = max( - self.max_seq_len if window_size is None else int(window_size) - for window_size in self.max_attention_window_vec - ) - self.event_manager: Optional[KVCacheEventManager] = None - if self.event_buffer_max_size > 0: - if mapping.enable_attention_dp: - self.event_manager = KVCacheEventManager( - self.event_buffer_max_size, - window_size=event_window_size, - attention_dp_rank=mapping.rank, - attention_dp_gather=Distributed.get(mapping).allgather, - hash_algo=kv_cache_event_hash_algo, - ) - elif mpi_rank() == 0: - self.event_manager = KVCacheEventManager( - self.event_buffer_max_size, - window_size=event_window_size, - hash_algo=kv_cache_event_hash_algo, - ) - - if isinstance(num_kv_heads, int): - self.num_kv_heads_per_layer = [ - (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_local_layers) - ] - self.total_num_kv_heads_per_layer = [ - (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_layers) - ] - else: - assert len(num_kv_heads) == self.num_layers - - def append_to_kv_heads_per_layer( - num_kv_heads_per_layer: List[int], kv_head: Optional[int] - ): - if kv_head is not None: - num_kv_heads_per_layer.append((kv_head + tp_size - 1) // tp_size) - else: - num_kv_heads_per_layer.append(0) - - self.num_kv_heads_per_layer = [] - if self.num_local_layers > 0: - for i in self.pp_layers: - kv_head = num_kv_heads[i] - append_to_kv_heads_per_layer(self.num_kv_heads_per_layer, kv_head) - - self.total_num_kv_heads_per_layer = [] - for i in range(self.num_layers): - kv_head = num_kv_heads[i] - append_to_kv_heads_per_layer(self.total_num_kv_heads_per_layer, kv_head) - - # Build per-layer head_dim (similar to num_kv_heads_per_layer) - if isinstance(head_dim, int): - self.head_dim_per_layer = [head_dim for _ in range(self.num_local_layers)] - else: - assert len(head_dim) == self.num_layers, ( - f"head_dim list length ({len(head_dim)}) must match num_layers ({self.num_layers})" - ) - self.head_dim_per_layer = [] - if self.num_local_layers > 0: - for i in self.pp_layers: - self.head_dim_per_layer.append(head_dim[i]) - if len(set(self.head_dim_per_layer)) > 1: - logger.info( - f"Per-layer head_dim: {len(self.head_dim_per_layer)} layers, " - f"unique values={set(self.head_dim_per_layer)}" - ) - - self.dtype_per_layer = [self.dtype] * self.num_local_layers - self.kv_factor_per_layer = [self.kv_factor] * self.num_local_layers - - self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) - - max_util_for_resume = kv_cache_config.max_util_for_resume - quota = sys.maxsize - if ( - kv_cache_config.max_gpu_total_bytes is not None - and kv_cache_config.max_gpu_total_bytes > 0 - ): - quota = int(kv_cache_config.max_gpu_total_bytes) - logger.info(f"max_gpu_total_bytes is provided. New quota is {quota / (1 << 30)}GiB") - if kv_cache_config.max_tokens is not None: - quota_from_max_tokens = int( - math.ceil( - self._get_quota_from_max_tokens(kv_cache_config.max_tokens) - / max_util_for_resume - ) - ) - quota = min(quota, quota_from_max_tokens) - logger.info( - f"max_tokens {kv_cache_config.max_tokens} is provided. " - f"Allowed quota from max_tokens is {quota_from_max_tokens / (1 << 30)}GiB. " - f"New quota is {quota / (1 << 30)}GiB" - ) - - assert quota < sys.maxsize, ( - "Quota not set. Check kv_cache_config.max_tokens or kv_cache_config.max_gpu_total_bytes" - ) - - # Sync resumable token capacity across ranks so the scheduler produces - # identical batches. Normalize to tokens because cache costs vary - # across PP ranks, including fixed per-rank costs. - if mapping.world_size > 1: - dist = Distributed.get(mapping) - resumable_quota = int(quota * max_util_for_resume) - max_tokens = self._get_max_tokens_from_quota(resumable_quota) - max_tokens = dist.allreduce(max_tokens, op=ReduceOp.MIN) - # inf max_tokens means all layers are SWA and every rank quota can - # fit all SWA fixed cache. - if not math.isinf(max_tokens): - # allreduce(MIN) must never increase the local quota. The - # token↔quota round-trip is not identity when SWA layers - # dominate (full_attn_size_per_token==0), so clamp to guard - # against a bogus inflation (nvbugs/6418103). - synced_quota = int( - math.ceil(self._get_quota_from_max_tokens(max_tokens) / max_util_for_resume) - ) - quota = min(quota, synced_quota) - - logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") - - cache_tiers: List[CacheTierConfig] = [ - _create_gpu_cache_tier_config(int(quota), enable_ugpu) - ] - if kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: - host_quota = kv_cache_config.host_cache_size - else: - # The V2 MAX_UTILIZATION scheduler relies on suspend/resume to - # evict and later restore KV cache pages. Without a secondary - # tier, suspended held pages cannot migrate out of GPU, so - # suspension cannot free capacity and scheduling can deadlock. - # - # Automatically provision a host tier matching the GPU quota so - # suspend/resume works out of the box. Cap at available host - # memory and pinnable memory limit to avoid allocation failures. - import resource - - # Rank-aware host tier sizing: divide the per-node memory budget - # by the number of ranks sharing this physical node. Each rank - # independently provisions a host tier; without this, N - # co-located ranks each reserving a device-quota-sized block can - # OOM the host (observed on GB300 NVL72 with 4 ranks/node and - # ~170GiB device quota each on a 975GiB node). - local_ranks = max(1, Distributed.get(mapping).local_world_size) - - try: - mem_available = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_AVPHYS_PAGES") - except (ValueError, OSError): - mem_available = float("inf") - try: - _soft, _hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) - memlock_limit = _soft if _soft != resource.RLIM_INFINITY else float("inf") - except (ValueError, OSError): - memlock_limit = float("inf") - host_quota = _compute_auto_host_tier_quota( - quota, local_ranks, mem_available, memlock_limit - ) - logger.info( - f"KV cache manager v2 auto host tier sizing: " - f"{local_ranks} co-located rank(s) on this node, " - f"available host memory {mem_available / (1 << 30):.2f}GiB" - ) - # Reduce the rank-local auto quota to the fleet minimum; see - # _sync_host_tier_quota for why divergence wedges collectives. - host_quota = _sync_host_tier_quota(host_quota, mapping) - if host_quota > 0: - cache_tiers.append(HostCacheTierConfig(quota=int(host_quota))) - logger.info( - f"KV cache manager v2 host cache quota set to {host_quota / (1 << 30):.2f}GiB" - ) - disk_cache_size = kv_cache_config.disk_cache_size - if disk_cache_size is not None and disk_cache_size > 0: - disk_cache_path = kv_cache_config.disk_cache_path - assert disk_cache_path is not None - cache_tiers.append( - DiskCacheTierConfig(quota=int(disk_cache_size), path=disk_cache_path) - ) - logger.info( - f"KV cache manager v2 disk cache quota set to {disk_cache_size / (1 << 30):.2f}GiB at {disk_cache_path}" - ) - - self.vocab_size = vocab_size - - config = self._build_base_config( - kv_cache_config, - tokens_per_block=tokens_per_block, - cache_tiers=cache_tiers, - ) - config = self._build_cache_config(config) - config = self._append_standalone_draft_layers( - config, - num_kv_heads=draft_num_kv_heads, - head_dim=draft_head_dim, - dtype=draft_dtype, - kv_factor=draft_kv_factor, - ) - config = self._configure_unified_draft_pool_ratio(config) - has_host_cache_tier = any( - isinstance(tier, HostCacheTierConfig) for tier in config.cache_tiers - ) - - self.kv_cache_manager_py_config = config - - try: - self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) - except (CuError, KVCacheOutOfMemoryError): - if has_host_cache_tier: - logger.warning( - "Failed to initialize KV cache manager with host cache " - "tier (cuMemHostRegister may have failed). " - "Retrying without host cache tier." - ) - cache_tiers_without_host = [ - tier for tier in config.cache_tiers if not isinstance(tier, HostCacheTierConfig) - ] - config = replace(config, cache_tiers=cache_tiers_without_host) - self.kv_cache_manager_py_config = config - self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) - else: - raise - self.can_evict = len(config.cache_tiers) > 1 - if self.event_manager is not None: - self.event_manager.set_layer_group_window_sizes( - self._get_event_window_sizes_by_layer_group() - ) - self.event_manager.add_created_event( - self._get_event_num_blocks_per_cache_level(config.cache_tiers, tokens_per_block), - self._get_event_layer_group_ids(), - ) - - # Both backends build layer_grouping on demand, and the layer order - # within a group is not part of its contract. Cache a stable physical- - # layout representative for each role from the public pool descriptors. - - # Cache the fork-join attention flag at construction time so the - # manager's per-uGPU metadata invariants stay stable for its lifetime. - self._fork_join_attn = self.num_ugpus > 1 and enable_ugpu - self.num_pools = len(self.impl.layer_grouping) - self._pool_layer_ids_by_role: Dict[Tuple[int, DataRole], LayerId] = {} - for pool_group in self.impl.pool_group_descs: - for variant in pool_group.slot_desc.variants: - pool_id = int(variant.layer_group_id) - for coalesced in variant.coalesced_buffers: - for buffer_id in coalesced.buffer_ids: - self._pool_layer_ids_by_role.setdefault( - (pool_id, buffer_id.role), buffer_id.layer_id - ) - # num_pools is the logical layer-group count. With SWA scratch reuse, - # scratch slot IDs are only valid with per-layer page indices, so the - # attention op sees one virtual pool per local layer while the - # underlying manager can still group layers. - if self.enable_swa_scratch_reuse: - self.num_attention_op_pools = self.num_local_layers - else: - self.num_attention_op_pools = self.num_pools - - num_layers = len(config.layers) - self.layer_to_pool_mapping_dict: dict[int, int] = { - layer_id: self.impl.get_layer_group_id(layer_id) - for layer_id in typed_range(LayerId(num_layers)) - } - - self.kv_cache_map: dict[int, _KVCache] = {} - - # Tracks the draft length allocated by try_allocate_generation per - # request. Used by extend_capacity_for_tokens to compute the exact - # padding delta instead of blindly extending, which would cause - # unbounded capacity growth. - self._allocated_draft_lens: dict[int, int] = {} - - # Overlap-scheduler slack (the extra draft_len in - # _required_gen_capacity's 1 + 2 * draft_len growth) granted to each - # request and not yet reclaimed. The sampler's py_rewind_len only - # covers the rejected part of the recurring 1 + draft_len growth, so - # update_resources must additionally trim this slack — otherwise it - # compounds, leaking draft_len tokens of capacity per generation - # iteration until the request's block count overruns - # max_blocks_per_seq ("User-provided base page indices is too short"). - self._pending_overlap_slack: dict[int, int] = {} - - # Defensive cap for get_num_available_tokens: when host cache is - # enabled, clamp_max_seq_len_for_mem may return a value that spans - # both GPU and host tiers. Storing the explicit max_tokens (if set) - # lets us cap the result to GPU-only capacity so callers like CUDA - # graph warmup don't over-allocate beyond the GPU pool. - # None when max_tokens is not explicitly configured — other config - # paths (max_gpu_total_bytes, free_gpu_memory_fraction) are already - # bounded by the GPU quota passed to GpuCacheTierConfig. - self._gpu_max_tokens = kv_cache_config.max_tokens - - max_num_tokens = self.get_num_available_tokens(token_num_upper_bound=max_seq_len) - - if max_seq_len > max_num_tokens: - logger.warning( - f"max_seq_len {max_seq_len} is greater than max_num_tokens {max_num_tokens} " - "that can be allocated in kv cache manager, setting " - f"max_seq_len to {max_num_tokens}" - ) - # max_num_tokens is a float from clamp_max_seq_len_for_mem; cast - # so downstream int-only consumers (torch.randint size, range) - # stay int. - self.max_seq_len = int(max_num_tokens) - - # Pad max_blocks_per_seq to next multiple of 4 (copy_block_offsets kernel). - # Account for max single-sequence capacity = seq_len + extra KV tokens + - # _kv_reserve_draft_tokens (see __init__) + 1 base decode token. - max_seq_capacity = ( - self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1 - ) - self.max_blocks_per_seq = ( - max_seq_capacity + self._ledger_tokens_per_block - 1 - ) // self._ledger_tokens_per_block - if self.max_blocks_per_seq % 4 != 0: - self.max_blocks_per_seq = ((self.max_blocks_per_seq + 3) // 4) * 4 - - self.enable_block_reuse = kv_cache_config.enable_block_reuse - self.enable_partial_reuse = kv_cache_config.enable_partial_reuse - self.disk_prefetch_num_reqs = kv_cache_config.disk_prefetch_num_reqs - enable_conversation_manager = ( - self.enable_block_reuse - and self.block_reuse_policy == BlockReusePolicy.PER_CONVERSATION - and self._can_publish_block_reuse - ) - self.conversation_manager = ( - ConversationManager(block_reuse_config.max_num_turns) - if enable_conversation_manager - else None - ) - - # With pipeline parallelism, multiple microbatches can be in-flight - # simultaneously, so we need slots for all concurrent sequences. - # Reserve stable slots for persistent request IDs such as CUDA-graph - # padding requests. The default preserves the main-branch allocation. - # In disaggregated mode, use a coefficient of 2: at any moment up to - # `max_num_sequences` requests can be actively generating while another - # up to `max_num_sequences` requests are still in KV transfer - # (TRANS_IN_PROGRESS) and continue to hold their index slots. The 2x - # capacity lets the next batch of active requests acquire slots without - # waiting for the previous batch's transfers to finish. - max_num_sequences = max_batch_size * mapping.pp_size - assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" - index_mapper_capacity = ( - max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots - ) - logger.info( - f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " - f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " - f"num_reserved_index_slots={num_reserved_index_slots}, " - f"max_beam_width={max_beam_width})" - ) - self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) - self._early_freed_index_requests: set[int] = set() - self._prepare_page_table_tensor(index_mapper_capacity) - - self._log_kv_cache_pool_lifecycle_mapping() - - def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: - """Return the roles represented by the two page-table index lanes. - - When present, role B must be addressable from role A using a constant - page-index offset. - """ - layer_id = int(self.impl.layer_grouping[pool_id][0]) - return Role.KEY, Role.VALUE if self.kv_factor_per_layer[layer_id] == 2 else None - - def _get_block_scale_role(self, role_a: DataRole, layer_id: int) -> Optional[DataRole]: - """Select block scales for a role and local layer in the pool mapping.""" - if self.dtype_per_layer[layer_id] != DataType.NVFP4 or role_a != Role.KEY: - return None - return Role.KEY_BLOCK_SCALE - - def _build_pool_mapping_tensors(self): - """Build the (kv_cache_pool_pointers, kv_cache_pool_mapping) tensors. - - An overridable hook for subclasses whose pools coalesce extra - per-layer buffers alongside K/V. - """ - has_block_scales = DataType.NVFP4 in self.dtype_per_layer - kv_cache_pool_pointers_list = [] - kv_cache_pool_mapping_list = [] - block_scale_pool_pointers_list = [] - if self.enable_swa_scratch_reuse: - for layer_id in typed_range(LayerId(self.num_local_layers)): - pool_id = self.impl.get_layer_group_id(layer_id) - role_a, _ = self._get_pool_roles(pool_id) - kv_cache_pool_pointers_list.append( - [ - self.impl.get_mem_pool_base_address( - layer_id, role_a, PageIndexMode.PER_LAYER - ), - 0, - ] - ) - if has_block_scales: - block_scale_role = self._get_block_scale_role(role_a, layer_id) - block_scale_pool_pointers_list.append( - [ - self.impl.get_mem_pool_base_address( - layer_id, block_scale_role, PageIndexMode.PER_LAYER - ) - if block_scale_role is not None - else 0, - 0, - ] - ) - kv_cache_pool_mapping_list.append([int(layer_id), 0]) - else: - for pool_id in range(self.num_pools): - role_a, _ = self._get_pool_roles(pool_id) - layer_id = self._pool_layer_ids_by_role[(pool_id, role_a)] - key_base_addr = self.impl.get_mem_pool_base_address( - layer_id, role_a, PageIndexMode.SHARED - ) - kv_cache_pool_pointers_list.append([key_base_addr, 0]) - if has_block_scales: - # The KEY/scale pointers are a (rep-layer, 0-offset) origin - # against which each layer's kv_cache_pool_mapping offset is - # resolved. The block-scale origin must reproduce the SAME - # per-layer offset() as KEY, so mirror the KEY base for the - # same representative layer and shift it back by that layer's - # offset. For the base manager offset(rep) == 0, so this is - # just the rep layer's scale base; for address-ranked - # subclasses (MiniMax-M3) offset(rep) may be non-zero, and the - # shift lands the origin on the pool's slot-0 scale address. - # This keeps block_scale_offset == offset without depending on - # the non-contractual layer_grouping order. - block_scale_role = self._get_block_scale_role(role_a, layer_id) - if block_scale_role is not None: - rep_offset = self._kv_pool_mapping_offset(layer_id, pool_id, key_base_addr) - scale_stride = ( - self.get_layer_bytes_per_token(layer_id, block_scale_role) - * self.kv_factor_per_layer[layer_id] - * self.tokens_per_block - ) - scale_base_addr = ( - self.impl.get_mem_pool_base_address( - layer_id, block_scale_role, PageIndexMode.SHARED - ) - - rep_offset * scale_stride - ) - else: - scale_base_addr = 0 - block_scale_pool_pointers_list.append([scale_base_addr, 0]) - - for layer_id in typed_range(LayerId(self.num_local_layers)): - layer_group_id = self.impl.get_layer_group_id(layer_id) - role_a, role_b = self._get_pool_roles(layer_group_id) - index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - if role_a == Role.KEY: - offset = self._kv_pool_mapping_offset(layer_id, layer_group_id, index_base_addr) - else: - addr_offset = ( - self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED) - - index_base_addr - ) - offset_divisor = self.impl.get_page_stride(layer_id, role_a) - if role_b is not None: - offset_divisor *= self.kv_factor_per_layer[layer_id] - offset = exact_div( - addr_offset, - offset_divisor, - ) - - if self.dtype_per_layer[layer_id] != DataType.NVFP4 or role_a != Role.KEY: - block_scale_offset = None - else: - block_scale_role = self._get_block_scale_role(role_a, layer_id) - if block_scale_role is None: - block_scale_offset = None - else: - block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] - block_scale_addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, block_scale_role, PageIndexMode.SHARED - ) - - block_scale_base_addr - ) - block_scale_offset = exact_div( - block_scale_addr_offset, - self.get_layer_bytes_per_token(layer_id, block_scale_role) - * self.kv_factor_per_layer[layer_id] - * self.tokens_per_block, - ) - - if block_scale_offset is not None: - assert block_scale_offset == offset, ( - "Block scale offset and offset should be the same" - ) - - kv_cache_pool_mapping_list.append([layer_group_id, offset]) - - if has_block_scales: - for pool_id, block_scale_pool_pointers in enumerate(block_scale_pool_pointers_list): - pool_pointers = kv_cache_pool_pointers_list[pool_id] - kv_cache_pool_pointers_list[pool_id] = [ - [pool_pointers[0], block_scale_pool_pointers[0]], - [pool_pointers[1], block_scale_pool_pointers[1]], - ] - - kv_cache_pool_pointers = torch.tensor( - kv_cache_pool_pointers_list, - dtype=torch.int64, - device="cpu", - pin_memory=prefer_pinned(), - ) - kv_cache_pool_mapping = torch.tensor( - kv_cache_pool_mapping_list, - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), - ) - return kv_cache_pool_pointers, kv_cache_pool_mapping - - def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: - self.kv_cache_pool_pointers, self.kv_cache_pool_mapping = self._build_pool_mapping_tensors() - self.index_scales = torch.empty( - self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu" - ) - self.kv_offset = torch.empty( - self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu" - ) - for pool_id in range(self.num_pools): - role_a, role_b = self._get_pool_roles(pool_id) - layer_id = self._pool_layer_ids_by_role[(pool_id, role_a)] - self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, role_a) - if role_b is not None: - self.kv_offset[pool_id] = exact_div( - self.impl.get_mem_pool_base_address(layer_id, role_b, PageIndexMode.SHARED) - - self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), - self.impl.get_page_stride(layer_id, role_a), - ) - else: - self.kv_offset[pool_id] = 0 - # Plain-int mirror of index_scales so the per-step block-table build - # does not index a tensor per request (see get_batch_cache_indices*). - self._index_scale_ints: List[int] = self.index_scales.tolist() - - # Keep unused block offsets as safe block index 0. - self.host_kv_cache_block_offsets = torch.zeros( - self.num_pools, - index_mapper_capacity * self.max_beam_width, - 2, # key and value - self.max_blocks_per_seq, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - if self.enable_swa_scratch_reuse: - self._prepare_swa_scratch_copy_tensors(index_mapper_capacity) - - def _kv_pool_mapping_offset( - self, layer_id: LayerId, layer_group_id: int, key_base_addr: int - ) -> int: - """Per-layer offset recorded in ``kv_cache_pool_mapping``. - - The default derives the layer's position within its pool from the K - base address, assuming every layer contributes exactly K(+V) to the - pool slot so the layer stride is uniform. Managers whose pool slots - may interleave extra per-layer buffers between layers (non-uniform - layer strides — e.g. MiniMax M3 when the index-K buffer coalesces - into the K/V pool) must override this with a positional formula. - """ - addr_offset = ( - self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) - - key_base_addr - ) - return exact_div( - addr_offset, - self.get_layer_bytes_per_token(layer_id, Role.KEY) - * self.get_layer_kv_factor(self.pp_layers[layer_id]) - * self.tokens_per_block, - ) - - def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: - layer_sizes = [] - attention_windows = [] - pattern_len = len(self.max_attention_window_vec) - for local_layer_idx in range(self.num_local_layers): - layer_sizes.append( - self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=Role.ALL) - ) - if self._is_standalone_draft_layer(local_layer_idx): - attention_windows.append(self.draft_retention_window_size) - else: - attention_windows.append( - self.max_attention_window_vec[self.pp_layers[local_layer_idx] % pattern_len] - ) - # Quota sizing precedes specialized target-layer construction. Include - # draft storage then; after append it is already in the per-layer arrays. - if self._draft_layer_sizes and not self._draft_registered_as_model_layers: - layer_sizes.extend(self._draft_layer_sizes) - attention_windows.extend( - [self.draft_retention_window_size] * len(self._draft_layer_sizes) - ) - return layer_sizes, attention_windows - - def _get_generation_request_capacity(self) -> int: - """Return the resident generation requests used for SWA sizing.""" - return self.max_batch_size - - def _get_runtime_cache_size_components(self) -> tuple[int, int, int]: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - num_target_layers = len(layer_sizes) - (len(self._draft_layer_sizes)) - context_size, generation_size, fixed_size = _estimate_cache_size_components( - layer_sizes[:num_target_layers], - attention_windows[:num_target_layers], - self.tokens_per_block, - scratch=self.enable_swa_scratch_reuse, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - standalone_draft_reserve=self._standalone_draft_reserve, - ) - if self._draft_layer_sizes: - draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( - self._draft_layer_sizes, - self.tokens_per_block, - scratch_tokens=self._standalone_draft_reserve, - window_size=self.draft_window_size, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - helix_cp_size=self._helix_cp_size, - ) - context_size += draft_context - generation_size += draft_generation - fixed_size += draft_fixed - return context_size, generation_size, fixed_size - - def _get_max_tokens_from_quota(self, quota: int) -> float: - """Rank-local byte quota -> token capacity (GLOBAL tokens under helix).""" - tokens = self._get_max_tokens_from_quota_impl(quota) - if self._has_cp_helix and not math.isinf(tokens): - # Floor to whole physical pages before scaling: a ledger block - # allocates one full page on every CP rank, so a partial - # trailing page in the rank-local budget is never usable. - tokens = int(tokens) // self.tokens_per_block * self.tokens_per_block - tokens *= self._helix_cp_size - return tokens - - def _get_max_tokens_from_quota_impl(self, quota: int) -> float: - if not self._draft_layer_sizes: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - full_attn_size_per_token = _estimate_full_attn_size_per_token( - layer_sizes, attention_windows - ) - context_swa_size_per_token, _ = _estimate_swa_cache_size( - layer_sizes, - attention_windows, - self.tokens_per_block, - context=True, - scratch=self.enable_swa_scratch_reuse, - ) - ( - generation_swa_size_per_token, - generation_swa_size_per_request, - ) = _estimate_swa_cache_size( - layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False - ) - size_per_batch = ( - self._get_generation_request_capacity() * generation_swa_size_per_request - ) - if quota < size_per_batch: - return 0 - context_size_per_token = full_attn_size_per_token + context_swa_size_per_token - context_limit_quota = self.max_num_tokens * context_size_per_token + size_per_batch - if quota <= context_limit_quota: - if context_size_per_token <= 0: - return float("inf") - return (quota - size_per_batch) / context_size_per_token - - generation_size_per_token = full_attn_size_per_token + generation_swa_size_per_token - if generation_size_per_token <= 0: - return float("inf") - return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token - - ( - context_size_per_token, - generation_size_per_token, - generation_size_per_request, - ) = self._get_runtime_cache_size_components() - size_per_batch = self._get_generation_request_capacity() * generation_size_per_request - if quota < size_per_batch: - return 0 - context_limit_quota = self.max_num_tokens * context_size_per_token + size_per_batch - if quota <= context_limit_quota: - if context_size_per_token <= 0: - return float("inf") - return (quota - size_per_batch) / context_size_per_token - - if generation_size_per_token <= 0: - return float("inf") - return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token - - def _get_quota_from_max_tokens(self, max_tokens: int) -> int: - """Token capacity (GLOBAL tokens under helix) -> rank-local byte quota.""" - if self._has_cp_helix: - # Round up to whole ledger blocks first: allocation is page- - # granular on every rank, so a request of N global tokens costs - # ceil(N / ledger_tpb) full physical pages per rank. - blocks = -(-int(max_tokens) // self._ledger_tokens_per_block) - max_tokens = blocks * self.tokens_per_block - return self._get_quota_from_max_tokens_impl(max_tokens) - - def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: - if not self._draft_layer_sizes: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - full_attn_size_per_token = _estimate_full_attn_size_per_token( - layer_sizes, attention_windows - ) - ( - context_swa_size_per_token, - _, - ) = _estimate_swa_cache_size( - layer_sizes, - attention_windows, - self.tokens_per_block, - context=True, - scratch=self.enable_swa_scratch_reuse, - ) - ( - generation_swa_size_per_token, - generation_swa_size_per_request, - ) = _estimate_swa_cache_size( - layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False - ) - context_tokens = min(max_tokens, self.max_num_tokens) - generation_tokens = max_tokens - context_tokens - generation_quota = ( - max_tokens * full_attn_size_per_token - + generation_tokens * generation_swa_size_per_token - + self._get_generation_request_capacity() * generation_swa_size_per_request - ) - context_extra_quota = context_tokens * context_swa_size_per_token - return int(generation_quota + context_extra_quota) - - ( - context_size_per_token, - generation_size_per_token, - generation_size_per_request, - ) = self._get_runtime_cache_size_components() - context_tokens = min(max_tokens, self.max_num_tokens) - generation_tokens = max_tokens - context_tokens - return int( - context_tokens * context_size_per_token - + generation_tokens * generation_size_per_token - + self._get_generation_request_capacity() * generation_size_per_request - ) - - def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: - """Draft bytes within the unified manager's usable token capacity.""" - assert self._draft_layer_sizes - if self._has_cp_helix: - # Each target ledger block retains all of its draft subpages. - max_tokens = ( - math.ceil(max_tokens / self._ledger_tokens_per_block) - * self._ledger_tokens_per_block - ) - context_size, generation_size, fixed_size = _estimate_draft_cache_size_components( - self._draft_layer_sizes, - self.tokens_per_block, - scratch_tokens=self._standalone_draft_reserve, - window_size=self.draft_window_size, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - ) - context_tokens = ( - max_tokens if self.max_num_tokens is None else min(max_tokens, self.max_num_tokens) - ) - return math.ceil( - context_tokens * context_size - + (max_tokens - context_tokens) * generation_size - + self._get_generation_request_capacity() * fixed_size - ) - - def _get_event_num_blocks_per_cache_level( - self, - cache_tiers: List[CacheTierConfig], - tokens_per_block: int, - ) -> List[int]: - bytes_per_block = self.get_cache_bytes_per_token() * tokens_per_block - if bytes_per_block <= 0: - return [] - return [int(tier.quota // bytes_per_block) for tier in cache_tiers] - - def _get_event_layer_group_ids(self) -> List[int]: - return [int(layer_group_id) for layer_group_id in range(len(self.impl.layer_grouping))] - - def _get_event_window_sizes_by_layer_group(self) -> Dict[int, int]: - # Assumes every layer in a group shares the same sliding_window_size, - # which is how `impl.layer_grouping` partitions layers today. Only the - # first layer's window is read; if the grouping policy ever permits - # mixed windows in one group, this needs to fan out per-layer. - - def get_event_window_size(layer_id: int) -> int: - layer_config = self.kv_cache_manager_py_config.layers[layer_id] - window_size = getattr(layer_config, "sliding_window_size", None) - return self.max_seq_len if window_size is None else int(window_size) - - return { - int(layer_group_id): get_event_window_size(int(layer_ids[0])) - for layer_group_id, layer_ids in enumerate(self.impl.layer_grouping) - } - - def _format_kv_cache_pool_lifecycle_entry(self, layer_id: LayerId, role: DataRole) -> str: - for pool_group in self.impl.pool_group_descs: - for variant in pool_group.slot_desc.variants: - for coalesced in variant.coalesced_buffers: - for buffer_id in coalesced.buffer_ids: - if int(buffer_id.layer_id) == int(layer_id) and str(buffer_id.role) == str( - role - ): - return ( - f"role={role!s}, " - f"pool_group_id={int(pool_group.pool_group_index)}, " - f"layer_group_id={int(variant.layer_group_id)}" - ) - return f"role={role!s}, pool_group_id=?, layer_group_id=?" - - def _format_kv_cache_pool_group_sizes(self) -> list[str]: - """Describe how many slots each pool group actually got. - - The aggregate KV utilization reported per iteration cannot answer - "which pool group is the binding constraint", because it sums every - pool group: a run can sit at 13% aggregate while one small pool group - is at its admission ceiling and refusing to activate requests. The - per-group slot count is what makes that case legible, and comparing it - against ``max_batch_size`` is what makes it actionable -- a pool group - sized at roughly one slot per declared request has no room for the - per-request slot count to vary, which windowed life cycles do as their - window slides across block boundaries. - - Read through ``_introspection`` so the numbers are the same under - either the Python or the C++ backend. - """ - try: - stats = _introspection.storage_statistics(self.impl, GPU_LEVEL) - except (AttributeError, RuntimeError) as exc: - # Introspection is a debugging aid, never a correctness - # dependency: a manager that cannot report its own layout must - # still start. - logger.debug(f"KV cache pool group sizes unavailable: {exc}") - return [] - - def flat_sizes(value) -> list[int]: - """Per-pool byte sizes, whatever shape the backend reported them in. - - `slot_size` is a scalar under one backend and a list under another, - and wrapping a list in another list made `sum()` raise - ``TypeError: unsupported operand type(s) for +: 'int' and 'list'`` - -- which is how a debugging aid took the startup log line with it. - """ - if isinstance(value, int): - return [value] - sizes: list[int] = [] - for item in value: - sizes.extend(flat_sizes(item)) - return sizes - - entries = [] - for pool_group_id, stat in enumerate(stats): - reported = getattr(stat, "slot_sizes", None) - if reported is None: - reported = stat.slot_size - slot_sizes = flat_sizes(reported) - entries.append( - f"pool_group_id={pool_group_id}, " - f"num_slots={stat.total}, " - f"slot_size={slot_sizes}, " - f"bytes={stat.total * sum(slot_sizes)}" - ) - return entries - - def _log_kv_cache_pool_lifecycle_mapping(self) -> None: - entries = OrderedDict() - for layer in self.kv_cache_manager_py_config.layers: - for buffer in layer.buffers: - entries.setdefault( - self._format_kv_cache_pool_lifecycle_entry(layer.layer_id, buffer.role), None - ) - - if not entries: - return - - logger.info(f"{type(self).__name__} role-to-pool/lifecycle mapping:") - for entry in entries: - logger.info(entry) - - for entry in self._format_kv_cache_pool_group_sizes(): - logger.info(entry) - - def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: - pool_ids = torch.empty( - self.num_attention_op_pools, - 2, - dtype=torch.long, - device="cpu", - ) - scales = torch.empty( - self.num_attention_op_pools, - 2, - dtype=torch.int32, - device="cpu", - ) - layer_offsets = torch.empty_like(scales) - scratch_pages = torch.empty_like(scales) - for local_layer_idx in range(self.num_local_layers): - layer_id = LayerId(local_layer_idx) - pool_id = self.layer_to_pool_mapping_dict[layer_id] - role_a, role_b = self._get_pool_roles(pool_id) - roles = [ - role_a, - role_a if role_b is None else role_b, - ] - for role_idx, role in enumerate(roles): - converter = self.impl.get_page_index_converter(layer_id, role) - if converter.expansion != 1 and self._is_standalone_draft_layer(layer_id): - pool_ids[local_layer_idx, role_idx] = pool_id - scales[local_layer_idx, role_idx] = 1 - layer_offsets[local_layer_idx, role_idx] = 0 - scratch_pages[local_layer_idx, role_idx] = 0 - continue - if converter.expansion != 1: - raise NotImplementedError( - "SWA scratch block-table conversion does not support " - f"expanded page indices yet: layer={layer_id}, role={role}, " - f"expansion={converter.expansion}" - ) - pool_ids[local_layer_idx, role_idx] = pool_id - scales[local_layer_idx, role_idx] = int(converter.scale) - layer_offsets[local_layer_idx, role_idx] = int(converter.layer_offset) - scratch_pages[local_layer_idx, role_idx] = int(converter.scratch_pages_per_block) - - staging_capacity = index_mapper_capacity * self.max_beam_width - device = torch.device("cuda", torch.cuda.current_device()) - self._device_kv_cache_block_offsets_input = torch.empty_like( - self.host_kv_cache_block_offsets, - device=device, - ) - self._device_attention_op_block_offsets_staging = torch.empty( - self.num_attention_op_pools, - staging_capacity, - 2, - self.max_blocks_per_seq, - dtype=torch.int32, - device=device, - ) - self._device_copy_idx_staging = torch.zeros( - staging_capacity, - dtype=torch.long, - device=device, - ) - self._device_num_contexts = torch.empty((), dtype=torch.int32, device=device) - self._device_attention_op_pool_ids = pool_ids.to(device=device) - self._device_attention_op_scales = scales.to(device=device) - self._device_attention_op_layer_offsets = layer_offsets.to(device=device) - self._device_attention_op_scratch_pages = scratch_pages.to(device=device) - self._device_block_positions = torch.arange( - self.max_blocks_per_seq, - dtype=torch.int32, - device=device, - ) - - min_scale = int(scales.min().item()) if scales.numel() > 0 else 1 - max_scratch_pages = int(scratch_pages.max().item()) if scratch_pages.numel() > 0 else 1 - self._max_scratch_slots = max( - 1, - (self.max_blocks_per_seq * max_scratch_pages + min_scale - 1) // min_scale, - ) - scratch_slots_shape = ( - self.num_pools, - staging_capacity, - self._max_scratch_slots, - ) - self._host_scratch_begs_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - self._host_scratch_ends_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - self._host_scratch_slots_staging = torch.zeros( - scratch_slots_shape, - dtype=torch.int32, - pin_memory=prefer_pinned(), - device="cpu", - ) - self._device_scratch_begs_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - device=device, - ) - self._device_scratch_ends_staging = torch.zeros( - self.num_pools, - staging_capacity, - dtype=torch.int32, - device=device, - ) - self._device_scratch_slots_staging = torch.zeros( - scratch_slots_shape, - dtype=torch.int32, - device=device, - ) - - def _copy_idx_to_device(self, copy_idx: torch.Tensor) -> torch.Tensor: - self._device_copy_idx_staging[: copy_idx.size(0)].copy_(copy_idx, non_blocking=True) - return self._device_copy_idx_staging - - def _copy_scratch_metadata_to_device( - self, - request_ids: List[int], - num_contexts: int, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - host_begs = self._host_scratch_begs_staging[:, :num_contexts] - host_ends = self._host_scratch_ends_staging[:, :num_contexts] - host_slots = self._host_scratch_slots_staging[:, :num_contexts, :] - host_begs.zero_() - host_ends.zero_() - host_slots.zero_() - - for pool_id in range(self.num_pools): - for context_idx, request_id in enumerate(request_ids[:num_contexts]): - desc = self.kv_cache_map[request_id].get_scratch_desc(pool_id) - if desc is None: - continue - slot_ids = desc.slot_ids - if len(slot_ids) > self._max_scratch_slots: - raise RuntimeError( - f"Scratch slot count {len(slot_ids)} exceeds staging capacity " - f"{self._max_scratch_slots}" - ) - host_begs[pool_id, context_idx] = int(desc.range.beg) - host_ends[pool_id, context_idx] = int(desc.range.end) - for slot_idx, slot_id in enumerate(slot_ids): - host_slots[pool_id, context_idx, slot_idx] = int(slot_id) - - self._device_scratch_begs_staging.copy_( - self._host_scratch_begs_staging, - non_blocking=True, - ) - self._device_scratch_ends_staging.copy_( - self._host_scratch_ends_staging, - non_blocking=True, - ) - self._device_scratch_slots_staging.copy_( - self._host_scratch_slots_staging, - non_blocking=True, - ) - return ( - self._device_scratch_begs_staging, - self._device_scratch_ends_staging, - self._device_scratch_slots_staging, - ) - - def _copy_batch_block_offsets_per_layer( - self, - dst_tensor: torch.Tensor, - request_ids: List[int], - copy_idx: torch.Tensor, - num_contexts: int, - num_seqs: int, - ) -> None: - device_copy_idx = self._copy_idx_to_device(copy_idx) - self._device_kv_cache_block_offsets_input.copy_( - self.host_kv_cache_block_offsets, - non_blocking=True, - ) - scratch_begs, scratch_ends, scratch_slots = self._copy_scratch_metadata_to_device( - request_ids, - num_contexts, - ) - self._device_num_contexts.fill_(num_contexts) - _copy_swa_block_offsets_with_scratch_compiled( - self._device_kv_cache_block_offsets_input, - device_copy_idx, - self._device_attention_op_pool_ids, - self._device_attention_op_scales, - self._device_attention_op_layer_offsets, - self._device_attention_op_scratch_pages, - self._device_block_positions, - scratch_begs, - scratch_ends, - scratch_slots, - self._device_num_contexts, - self._device_attention_op_block_offsets_staging, - ) - dst_tensor[: self.num_attention_op_pools, :num_seqs].copy_( - self._device_attention_op_block_offsets_staging[ - : self.num_attention_op_pools, - :num_seqs, - ], - non_blocking=True, - ) - - def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: - """Return the primary cache roles allocated for one local layer.""" - roles = [Role.KEY, Role.VALUE][: self.kv_factor_per_layer[local_layer_idx]] - if self.dtype_per_layer[local_layer_idx] == DataType.NVFP4: - assert self.head_dim_per_layer[local_layer_idx] % 2 == 0, ( - "NVFP4 cache head_dim must be divisible by two" - ) - roles += [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE][ - : self.kv_factor_per_layer[local_layer_idx] - ] - return roles - - def _build_base_config( - self, - kv_cache_config: KvCacheConfig, - *, - tokens_per_block: int, - cache_tiers: List[CacheTierConfig], - ) -> KVCacheManagerConfigPy: - """Build the general cache configuration used by most models. - - Models that need a custom configuration should subclass - :class:`KVCacheManagerV2` and override :meth:`_build_cache_config` to - update the necessary fields. - """ - scratch_reuse_config = None - if self.enable_swa_scratch_reuse: - # Unwritten speculative and draft reserves must stay outside the - # scratch range when a context chunk advances its history. - scratch_reuse_config = SwaScratchReuseConfig( - max_rewind_len=self.num_extra_kv_tokens + self._standalone_draft_reserve - ) - - typical_step = None - constraints = [] - if kv_cache_config.pool_ratio is None: - typical_seq_len = self._get_typical_seq_len(kv_cache_config) - if typical_seq_len is not None and typical_seq_len > self.max_seq_len: - raise ValueError( - f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or " - f"equal to max_seq_len ({self.max_seq_len})" - ) - - if typical_seq_len is not None: - # Model one context request and enough generation requests to fill - # the resident request capacity without over-provisioning windowed - # cache pools. - generation_request_capacity = self._get_generation_request_capacity() - context_capacity = ( - self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len - ) + self.num_extra_kv_tokens - generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) - typical_step = BatchDesc( - [KVCacheDesc(capacity=context_capacity, history_length=0)] - + [ - KVCacheDesc( - capacity=typical_seq_len, - history_length=generation_history_length, - ) - ] - * (generation_request_capacity - 1) - ) - - # CUDA graph generation warmup uses one request at max_seq_len and - # enough minimal decode requests to fill the resident capacity. - min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens - constraints.append( - BatchDesc( - [ - KVCacheDesc( - capacity=self.max_seq_len, - history_length=self.max_seq_len - 1, - ) - ] - + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)] - * (generation_request_capacity - 1) - ) - ) - # General and chunked-prefill warmup uses one fresh context request - # at the per-iteration token budget. - if self.max_num_tokens is not None: - constraints.append( - BatchDesc( - [ - KVCacheDesc( - capacity=self.max_num_tokens + self.num_extra_kv_tokens, - history_length=0, - ) - ] - ) - ) - - if kv_cache_config.dtype == "nvfp4": - for layer_idx, hd in enumerate(self.head_dim_per_layer): - buffer_roles = self._get_buffer_roles_for_layer(layer_idx) - if Role.KEY_BLOCK_SCALE not in buffer_roles: - continue - assert hd % 2 == 0, ( - f"head_dim must be divisible by 2 for nvfp4 kv cache, but layer {layer_idx} has head_dim={hd}" - ) - - # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional - # per-layer BufferConfig entries — for example a sparse index-K - # buffer — without overriding the K/V/NVFP4 scale wiring above. - # The dict maps local layer id -> list of extra BufferConfig. Each - # extra buffer's role must be unique within the layer (asserted by - # AttentionLayerConfig.__post_init__) and its size must be in bytes - # per block (= bytes_per_token * tokens_per_block). - extra_buffers_per_layer = ( - self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} - ) - - layer_configs: List[AttentionLayerConfig] = [] - for layer_id in typed_range(LayerId(self.num_local_layers)): - buffer_roles = self._get_buffer_roles_for_layer(layer_id) - buffers = [ - BufferConfig( - role=role, - size=self.get_layer_bytes_per_token(local_layer_idx=layer_id, data_role=role) - * tokens_per_block, - ) - for role in buffer_roles - ] - for extra in extra_buffers_per_layer.get(int(layer_id), ()): - assert extra.role not in buffer_roles, ( - f"extra buffer role {extra.role!r} for layer " - f"{int(layer_id)} duplicates a standard K/V/scale role" - ) - buffers.append(extra) - layer_configs.append( - AttentionLayerConfig( - layer_id=layer_id, - buffers=buffers, - sliding_window_size=self.max_attention_window_vec[ - self.pp_layers[layer_id] % len(self.max_attention_window_vec) - ], - num_sink_tokens=None, - ) - ) - - return KVCacheManagerConfigPy( - # Used by the backend only for token<->block arithmetic and - # radix hashing; BufferConfig.size above stays the physical page. - tokens_per_block=self._ledger_tokens_per_block, - cache_tiers=cache_tiers, - layers=layer_configs, - typical_step=typical_step, - constraints=constraints, - max_util_for_resume=kv_cache_config.max_util_for_resume, - enable_partial_reuse=kv_cache_config.enable_partial_reuse, - # Keep the lookahead evidence and its backoff in the same tree match. - # A paired scheduler caps the claim by retaining D evidence past the - # common usable depth, so this still trims exactly once. - reuse_match_backoff=self.reuse_match_backoff, - enable_stats=self.enable_stats, - swa_scratch_reuse=scratch_reuse_config, - commit_min_snapshot=( - kv_cache_config.enable_block_reuse - and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE - ), - initial_pool_ratio=kv_cache_config.pool_ratio, - ) - - def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: - """Customize the general cache config for a specialized cache manager.""" - return config - - def _append_standalone_draft_layers( - self, - config: KVCacheManagerConfigPy, - *, - num_kv_heads: int, - head_dim: int, - dtype: DataType, - kv_factor: int, - register_model_layers: bool = True, - ) -> KVCacheManagerConfigPy: - if not self._draft_layer_sizes: - return config - first_global_id = max(self.pp_layers, default=-1) + 1 - first_local_id = len(config.layers) - # FP4 MLA adds native HP-tail layers outside the model-layer arrays. - register_model_layers = register_model_layers and first_local_id == len(self.pp_layers) - self._draft_registered_as_model_layers = register_model_layers - # Specialized managers can replace model layers with multiple native layers. - padding = first_local_id - len(self.dtype_per_layer) - assert padding >= 0 - num_draft_layers = len(self._draft_layer_sizes) - self.dtype_per_layer.extend([self.dtype] * padding + [dtype] * num_draft_layers) - self.kv_factor_per_layer.extend([self.kv_factor] * padding + [kv_factor] * num_draft_layers) - self.num_kv_heads_per_layer.extend([0] * padding + [num_kv_heads] * num_draft_layers) - self.head_dim_per_layer.extend([0] * padding + [head_dim] * num_draft_layers) - self.cp_as_tp_layer_ids = frozenset( - range(first_local_id, first_local_id + num_draft_layers) - ) - self.draft_layer_ids = tuple( - range(first_global_id, first_global_id + len(self._draft_layer_sizes)) - ) - layers = list(config.layers) - for offset, global_id in enumerate(self.draft_layer_ids): - local_id = first_local_id + offset - layers.append( - AttentionLayerConfig( - layer_id=LayerId(local_id), - buffers=[ - BufferConfig( - role=role, - size=KVCacheManagerV2.get_layer_bytes_per_token(self, local_id, role) - * self.tokens_per_block, - tokens_per_block_override=( - self.tokens_per_block if self._helix_cp_size > 1 else None - ), - ) - for role in self._get_buffer_roles_for_layer(local_id) - ], - sliding_window_size=self.draft_retention_window_size, - cache_domain=CacheDomain.STANDALONE_DRAFT, - ) - ) - self.layer_offsets[global_id] = local_id - if register_model_layers: - self.pp_layers.append(global_id) - self.total_num_kv_heads_per_layer.append( - self._draft_total_num_kv_heads or num_kv_heads - ) - if register_model_layers: - self.num_local_layers = len(self.pp_layers) - self.num_layers += len(self._draft_layer_sizes) - - def reserve_scratch(batch: BatchDesc) -> BatchDesc: - return BatchDesc( - [ - KVCacheDesc( - capacity=desc.capacity - + (self._standalone_draft_reserve if desc.capacity else 0), - history_length=desc.history_length, - ) - for desc in batch.kv_caches - ], - system_prompt_length=batch.system_prompt_length, - ) - - return replace( - config, - layers=layers, - constraints=[reserve_scratch(batch) for batch in config.constraints], - typical_step=reserve_scratch(config.typical_step) - if config.typical_step is not None - else None, - ) - - def _get_typical_seq_len(self, kv_cache_config: KvCacheConfig) -> int | None: - """Return the configured typical sequence length, if any.""" - return kv_cache_config.avg_seq_len - - def _configure_unified_draft_pool_ratio( - self, config: KVCacheManagerConfigPy - ) -> KVCacheManagerConfigPy: - """Preserve target ratios within their share of the unified byte quota.""" - if not self.draft_layer_ids: - return config - target_ratios = config.initial_pool_ratio - if target_ratios is None: - return config - - # Native lifecycle equality includes window, rounded sink blocks, sparsity, - # and cache domain. All SSM layers share one lifecycle, represented by None. - target_life_cycles: list[AttnLifeCycle | None] = [] - for layer in config.layers: - if isinstance(layer, AttentionLayerConfig): - if layer.cache_domain == CacheDomain.STANDALONE_DRAFT: - continue - life_cycle = AttnLifeCycle.make( - layer.sliding_window_size, - layer.num_sink_tokens, - config.tokens_per_block, - cache_domain=layer.cache_domain, - ) - else: - life_cycle = None - if not any( - type(existing) is type(life_cycle) and existing == life_cycle - for existing in target_life_cycles - ): - target_life_cycles.append(life_cycle) - - if len(target_ratios) != len(target_life_cycles): - raise ValueError( - "kv_cache_config.pool_ratio must have one entry per target layer group " - f"({len(target_life_cycles)}), got {len(target_ratios)}" - ) - # Use the same model-specific capacity accounting as quota selection; - # specialized targets can include compressed KV or recurrent snapshots. - usable_quota = int(config.cache_tiers[GPU_LEVEL].quota * config.max_util_for_resume) - max_tokens = self._get_max_tokens_from_quota(usable_quota) - if math.isinf(max_tokens): - max_tokens = self.max_num_tokens or self.max_seq_len - if usable_quota <= 0 or self._get_quota_from_max_tokens(max_tokens) > usable_quota: - raise ValueError("Unified KV cache quota cannot hold the target and draft reserves") - draft_ratio = self._get_draft_quota_from_max_tokens(max_tokens) / usable_quota - if not 0 < draft_ratio < 1: - raise ValueError("Unified KV cache quota must leave room for target and draft pools") - - # Appended draft layers share one lifecycle, after all target groups. - return replace( - config, - initial_pool_ratio=[ratio * (1 - draft_ratio) for ratio in target_ratios] - + [draft_ratio], - ) - - def _extra_buffers_per_layer( - self, *, tokens_per_block: int - ) -> Optional[dict[int, List[BufferConfig]]]: - """Return per-local-layer extra BufferConfig entries to register - alongside the standard K/V/NVFP4 scale buffers. - - Default implementation returns ``None``. Subclasses override this - to register additional buffers — for example, MiniMax-M3 registers - a sparse index-K buffer for each sparse local layer. Each - ``BufferConfig.size`` is interpreted as bytes per block (i.e., - ``bytes_per_token * tokens_per_block``), matching the standard - buffers built in :meth:`_build_base_config`. The block storage - groups buffers by lifecycle and size with an opaque role key, so - new roles do not require C++ changes. - """ - return None - - def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: - """Map native cache roles to disaggregation mapper kinds. - - ``Role.ALL`` is the required fallback for roles without an explicit - entry. The default is the head-major (HND) ``INDEXED`` layout written - by the TRTLLM attention kernels — correct for V1 and standard V2 - managers. ``Role.INDEX_KEY`` defaults to ``REPLICATED``: every - index-key side cache shipped so far (DSA indexer-K on V1, MiniMax M3 - on V2) computes its projection replicated across TP ranks, so the - cache bytes are identical per rank; the entry is inert unless a - subclass actually registers ``INDEX_KEY`` buffers via - ``_extra_buffers_per_layer``. Model-specific managers may declare - logical layouts without requiring the shared extractor to inspect - private attributes or role names. MiniMax M3, for example, maps - ordinary K/V to ``NHD`` and keeps index-key ``REPLICATED``. - - This declaration does not influence storage pooling: V2 storage - coalesces buffers purely by ``(life_cycle, buffer size)``, so roles - with different transfer semantics may share a pool slot when their - per-block sizes coincide (e.g. MiniMax M3 at TP degrees where - K == V == INDEX_KEY bytes per block). The disagg page-table builder - splits each physical pool into one logical view per mapper kind, so - transfer correctness never depends on the coalescing outcome. - - Pool memory is layout-agnostic; this declaration describes what the - manager's paired attention backend actually writes. A static mapping - is valid only for a fixed manager/backend pair. A manager whose backend - selects the layout at runtime must derive the mapping from that - backend's configuration. - """ - return {Role.ALL: MapperKind.INDEXED, Role.INDEX_KEY: MapperKind.REPLICATED} - - def get_disagg_transfer_roles(self) -> Optional[frozenset[DataRole]]: - """Return the native roles included in disaggregated transfer. - - ``None`` preserves the default policy and transfers every registered - role. Specialized managers may omit deterministic derived roles that - are rebuilt on the receiver after canonical cache data arrives. - """ - return None - - @property - def blocks_in_primary_pool(self) -> int: - """ - Get the number of blocks in the primary pool. - """ - return self.impl.get_page_index_upper_bound(0, Role.KEY) - - def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch.Tensor]: - return self._get_layer_buffer_view(self.layer_offsets[layer_idx], kv_layout) - - def _get_layer_buffer_view(self, local_layer_idx: int, kv_layout: str) -> torch.Tensor: - return self._get_kv_buffer_view( - local_layer_idx, - dtype=self.dtype_per_layer[local_layer_idx], - kv_factor=self.kv_factor_per_layer[local_layer_idx], - num_kv_heads=self.num_kv_heads_per_layer[local_layer_idx], - head_dim=self.head_dim_per_layer[local_layer_idx], - kv_layout=kv_layout, - page_expansion=self._helix_cp_size if local_layer_idx in self.cp_as_tp_layer_ids else 1, - ) - - def _get_kv_buffer_view( - self, - layer_id: int, - *, - dtype: DataType | torch.dtype, - kv_factor: int, - num_kv_heads: int, - head_dim: int, - kv_layout: str, - page_expansion: int = 1, - ) -> torch.Tensor: - """View shared pool pages for a manager-local layer and its KV geometry.""" - if kv_layout not in ("NHD", "HND"): - raise ValueError(f"Unsupported KV layout: {kv_layout}") - addr_key = self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) - if kv_factor == 2: - addr_value = self.impl.get_mem_pool_base_address( - layer_id, Role.VALUE, PageIndexMode.SHARED - ) - page_size_key = self.impl.get_page_stride(layer_id, Role.KEY) - page_size_value = self.impl.get_page_stride(layer_id, Role.VALUE) - if ( - addr_key + page_size_key * page_expansion != addr_value - or page_size_key != page_size_value - ): - raise ValueError("K/V buffers must have adjacent equal-sized pages") - - element_per_container = 1 - if dtype == DataType.NVFP4: - element_per_container = 2 - dtype = torch.int8 - - dimensions = ( - [self.tokens_per_block, num_kv_heads] - if kv_layout == "NHD" - else [num_kv_heads, self.tokens_per_block] - ) - head_dim //= element_per_container - page_upper = self.impl.get_page_index_upper_bound(layer_id, Role.KEY) - num_pages = page_upper // kv_factor - strides = None - if page_expansion > 1: - # Each ledger slot contains all of its K subpages, then all V - # subpages. Expanded KEY indices skip the V spans between slots. - num_pages = page_upper - (kv_factor - 1) * page_expansion - page_elements = self.tokens_per_block * num_kv_heads * head_dim - strides = [ - page_elements, - page_expansion * page_elements, - dimensions[1] * head_dim, - head_dim, - 1, - ] - return convert_to_torch_tensor( - TensorWrapper( - addr_key, - dtype, - [ - num_pages, - kv_factor, - *dimensions, - head_dim, - ], - strides=strides, - ) - ) - - def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: - """Return the storage dtype for a layer, or the manager default.""" - return ( - self.dtype if layer_idx is None else self.dtype_per_layer[self.layer_offsets[layer_idx]] - ) - - def _is_standalone_draft_layer(self, local_layer_idx: int) -> bool: - return local_layer_idx in self.cp_as_tp_layer_ids - - def get_layer_kv_factor(self, layer_idx: int) -> int: - return self.kv_factor_per_layer[self.layer_offsets[layer_idx]] - - def get_draft_buffers(self, local_layer_idx: int, kv_layout: str = "HND") -> torch.Tensor: - """View one draft layer using the same storage contract as every other layer.""" - if not 0 <= local_layer_idx < len(self.draft_layer_ids): - raise ValueError("No standalone draft cache layer at this index") - layer_idx = self.draft_layer_ids[local_layer_idx] - return self._get_layer_buffer_view(self.layer_offsets[layer_idx], kv_layout) - - @property - def draft_max_blocks_per_seq(self) -> int: - """Maximum physical draft pages, including every Helix ledger subpage.""" - return self.max_blocks_per_seq * self._helix_cp_size - - def get_draft_num_blocks(self, request_id: int) -> int: - """Allocated draft pages covering the request's global token positions.""" - return self.kv_cache_map[request_id].num_blocks * self._helix_cp_size - - @property - def draft_retention_window_size(self) -> int | None: - return self.draft_window_size + 1 if self.draft_window_size is not None else None - - def get_draft_transfer_identity(self) -> dict: - local_layer_idx = self.layer_offsets[self.draft_layer_ids[0]] - return { - "num_layers": len(self.draft_layer_ids), - "num_kv_heads": self._draft_total_num_kv_heads - or self.num_kv_heads_per_layer[local_layer_idx], - "head_dim": self.head_dim_per_layer[local_layer_idx], - "dtype": str(binding_to_torch_dtype(self.dtype_per_layer[local_layer_idx])), - "attention_backend": ( - "VANILLA" - if self.draft_attention_backend == "TRTLLM" - else self.draft_attention_backend - ), - "kv_factor": self.kv_factor_per_layer[local_layer_idx], - "window_size": self.draft_window_size, - } - - def _draft_history_uses_kv_history(self) -> bool: - """Embedded DSpark writes the complete live suffix of target history.""" - return getattr(self, "draft_attention_backend", None) == "DSv4" - - def get_draft_block_indices( - self, - request_ids: List[int], - histories: Optional[Sequence[StandaloneDraftHistory]] = None, - ) -> List[List[int]]: - """Validated rank-local mappings for allocated draft pages, without padding.""" - for request_id in request_ids: - cache = self.kv_cache_map.get(request_id) - if cache is None or not cache.is_active: - raise ValueError(f"Standalone draft request {request_id} has no active cache") - expanded = self._helix_cp_size > 1 - layer_idx = self.draft_layer_ids[0] - batch_indices = self.get_batch_cache_indices(request_ids, layer_idx, raw_indices=expanded) - if expanded: - converter = self.impl.get_page_index_converter(self.layer_offsets[layer_idx], Role.KEY) - batch_indices = [converter(indices, PageIndexMode.SHARED) for indices in batch_indices] - for row, indices in enumerate(batch_indices): - cache = self.kv_cache_map[request_ids[row]] - if ( - len(indices) != self.get_draft_num_blocks(request_ids[row]) - or len(indices) > self.draft_max_blocks_per_seq - ): - raise ValueError("Standalone draft cache has an incomplete or oversized page table") - first_required_block = 0 - if self.draft_window_size is not None: - history = None - if not self._draft_history_uses_kv_history(): - history = ( - histories[row] - if histories is not None - else self.get_draft_history(request_ids[row]) - ) - history_start = ( - history.position - history.valid_length - if history is not None - else max(0, cache.history_length - self.draft_window_size) - ) - first_required_block = history_start // self.tokens_per_block - if any(index == BAD_PAGE_INDEX for index in indices[first_required_block:]): - raise ValueError( - "Draft cache contains missing pages in its required history or scratch" - ) - return batch_indices - - def get_draft_block_table( - self, - request_ids: List[int], - histories: Optional[Sequence[StandaloneDraftHistory]] = None, - ) -> torch.Tensor: - """Current rank-local page mappings; unused tail entries point to page zero.""" - batch_indices = self.get_draft_block_indices(request_ids, histories) - table = torch.zeros((len(request_ids), self.draft_max_blocks_per_seq), dtype=torch.int32) - for row, indices in enumerate(batch_indices): - table[row, : len(indices)] = torch.tensor(indices, dtype=torch.int32) - return table - - def get_draft_history(self, request_id: int) -> Optional[StandaloneDraftHistory]: - return self.draft_history.get(request_id) - - def set_draft_history(self, request_id: int, valid_length: int, position: int) -> None: - cache = self.kv_cache_map.get(request_id) - if cache is None or not cache.is_active: - raise ValueError( - f"Standalone draft request {request_id} has no active cache allocation" - ) - history = StandaloneDraftHistory(valid_length, position) - if history.position > cache.capacity: - raise ValueError("Standalone draft history exceeds allocated capacity") - if self.draft_window_size is not None and history.valid_length > self.draft_window_size: - raise ValueError("Draft history exceeds its retention window") - self.draft_history[request_id] = history - - def export_draft_history(self, request_id: int) -> Optional[dict]: - if not self.draft_layer_ids: - return None - if self._draft_history_uses_kv_history(): - cache = self.kv_cache_map[request_id] - if not cache.is_active: - raise ValueError("Embedded DSpark transfer requires an active cache") - return { - "valid_length": min(cache.history_length, self.draft_window_size), - "position": cache.history_length, - "layout": self.get_draft_transfer_identity(), - } - history = self.get_draft_history(request_id) - if history is None: - raise ValueError( - f"Standalone draft request {request_id} has no valid history to transfer" - ) - return { - "valid_length": history.valid_length, - "position": history.position, - "layout": self.get_draft_transfer_identity(), - } - - def restore_draft_history(self, request_id: int, metadata: dict) -> None: - if not self.draft_layer_ids or metadata.get("layout") != self.get_draft_transfer_identity(): - raise ValueError("Standalone draft transfer layout does not match the receiving worker") - valid_length = metadata.get("valid_length") - position = metadata.get("position") - if self._draft_history_uses_kv_history(): - cache = self.kv_cache_map[request_id] - # Receive admission establishes the global prompt watermark; the - # transfer completion path validates that the received KV covers it. - if ( - type(position) is not int - or type(valid_length) is not int - or position != cache.history_length - or valid_length != min(position, self.draft_window_size) - ): - raise ValueError("Embedded DSpark transfer does not match the admitted KV history") - self.get_draft_block_table([request_id]) - return - history = StandaloneDraftHistory(valid_length, position) - # Validate receiver-local allocation before publishing history. - self.get_draft_block_table([request_id], [history]) - self.set_draft_history(request_id, valid_length, position) - - def get_index_k_buffer( - self, - layer_idx: int, - *, - num_heads: int = 1, - head_dim: int, - dtype: Union[torch.dtype, "DataType", str] = torch.bfloat16, - kv_layout: str = "NHD", - ) -> Optional[torch.Tensor]: - """Return a torch view over the V2-managed paged ``Role.INDEX_KEY`` - buffer for ``layer_idx``, or ``None`` when the layer has no - INDEX_KEY buffer registered (e.g. dense layers in a sparse model, - or non-local layers on the current PP rank). - - For ``kv_layout="NHD"``, the view has shape ``[num_pages, - tokens_per_block, num_heads, head_dim]``. For ``kv_layout="HND"``, - it has shape ``[num_pages, num_heads, tokens_per_block, head_dim]``. - Sparse modeling code decomposes the per-token slot id used by the - main paged K/V cache into ``(page, within_page)`` and indexes the - token axis selected by ``kv_layout``. - - Because :class:`BufferConfig` only carries an opaque byte ``size`` - per block, the dtype and head shape are caller-side contracts. - The caller must pass the same ``num_heads``, ``head_dim``, and - ``dtype`` it used to compute the registered - ``size = num_heads * head_dim * dtype_bytes * tokens_per_block`` - in ``_extra_buffers_per_layer``. The accessor validates - ``num_heads * head_dim * dtype_bytes * tokens_per_block == - page_stride`` so a wiring mismatch fails loudly instead of - returning a view with silently wrong stride. - - Follows the same ``TensorWrapper`` / ``convert_to_torch_tensor`` - pattern as :meth:`get_buffers`. The returned tensor is a - zero-copy view over V2-managed pool memory and stays valid for - the lifetime of the cache manager — writes through the view - propagate to the pool, and successive calls return views over - the same backing storage. - """ - if kv_layout not in ("NHD", "HND"): - raise ValueError(f"Unsupported kv_layout: {kv_layout}") - if layer_idx not in self.layer_offsets: - return None - layer_offset = self.layer_offsets[layer_idx] - try: - addr = self.impl.get_mem_pool_base_address(layer_offset, Role.INDEX_KEY) - except (KeyError, IndexError): - # INDEX_KEY not registered for this layer (default V2 manager - # registers only K/V/scale; sparse subclasses register - # INDEX_KEY only on sparse layers via - # ``_extra_buffers_per_layer``). - # - # The python backend raises ``KeyError`` from the missing dict - # lookup; the C++ backend's ``getBufferAttr`` throws - # ``std::out_of_range`` for an unknown buffer id, which nanobind - # maps to ``IndexError``. Catch both so the "role not registered - # -> None" contract holds on either backend. - return None - # INDEX_KEY is registered; the remaining lookups must succeed. Keep them - # outside the try so genuine failures surface instead of returning None. - page_stride = self.impl.get_page_stride(layer_offset, Role.INDEX_KEY) - page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.INDEX_KEY) - converter = self.impl.get_page_index_converter(layer_offset, Role.INDEX_KEY) - - if isinstance(dtype, DataType): - torch_dtype = binding_to_torch_dtype(dtype) - elif isinstance(dtype, str): - torch_dtype = str_dtype_to_torch(dtype) - else: - torch_dtype = dtype - - elem_bytes = torch.tensor([], dtype=torch_dtype).element_size() - expected_stride = num_heads * head_dim * elem_bytes * self.tokens_per_block - assert page_stride == expected_stride, ( - f"INDEX_KEY page stride mismatch for layer {layer_idx}: " - f"V2 reports page_stride={page_stride}, but the caller " - f"supplied num_heads={num_heads} * head_dim={head_dim} * " - f"elem_bytes={elem_bytes} * tokens_per_block=" - f"{self.tokens_per_block} = {expected_stride}. Re-check " - f"the BufferConfig.size used to register Role.INDEX_KEY " - f"in _extra_buffers_per_layer." - ) - - # Multi-layer coalescing: when INDEX_KEY shares the V2 storage - # pool with K/V (production MiniMax-M3 at TP=8 makes K, V, and - # INDEX_KEY all 256 bytes/token so V2 coalesces them by - # ``(life_cycle_id, single_buffer_size)``), ``page_upper`` is - # not the per-layer page count — it is the max page index from - # this buffer's base in the coalesced pool. Decompose it into - # ``num_slots = (page_upper + layer_offset_pages) // scale`` so - # the per-layer view has the right shape. ``scale`` here is the - # buffers-per-slot count; for a non-coalesced INDEX_KEY pool it - # equals 1 and ``num_slots == page_upper``, matching the legacy - # ``shape[0] == page_upper`` behavior. - scale = int(converter.scale) - layer_offset_pages = int(converter.layer_offset) - num_slots_total = page_upper + layer_offset_pages - assert num_slots_total % scale == 0, ( - f"V2 storage inconsistency for INDEX_KEY of layer " - f"{layer_idx}: page_upper + layer_offset_pages = " - f"{num_slots_total} is not divisible by scale = {scale}." - ) - num_slots = num_slots_total // scale - if kv_layout == "NHD": - page_shape = [self.tokens_per_block, num_heads, head_dim] - else: - page_shape = [num_heads, self.tokens_per_block, head_dim] - - if scale == 1: - # Non-coalesced INDEX_KEY pool: the per-buffer stride is - # the entire page, so either layout is a contiguous view. - shape = [page_upper, *page_shape] - return convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, shape)) - - # Coalesced pool: include the buffers-per-slot dimension, then - # slice ``[:, 0]`` to extract this layer's INDEX_KEY data while - # preserving dim-0 stride = ``scale * page_stride`` bytes. - full_slot_shape = [num_slots, scale, *page_shape] - full_view = convert_to_torch_tensor(TensorWrapper(addr, torch_dtype, full_slot_shape)) - return full_view[:, 0] - - def get_num_available_tokens( - self, *, token_num_upper_bound: int, batch_size: int = 1, max_num_draft_tokens: int = 0 - ) -> int: - """Clamp ``token_num_upper_bound`` to the allocatable token capacity. - - Unit note: under helix the backend runs on the ledger - ``tokens_per_block``, so the returned capacity (like - ``token_num_upper_bound`` and ``max_seq_len``) is in GLOBAL ledger - tokens - the coordinate request lengths are expressed in. Callers - that additionally bound the result by per-forward budgets (e.g. - ``max_num_tokens``) stay consistent because a helix context forward - replicates all tokens on every rank, so both bounds constrain the - same request-length variable. - """ - extra_tokens = ( - self.num_extra_kv_tokens + max_num_draft_tokens + self._standalone_draft_reserve - ) - # Token num upper bound is the maximum number of tokens that can be allocated in the kv cache manager. - # We need to add extra tokens to the token num upper bound to account for the extra tokens. - clamped = ( - self.impl.clamp_max_seq_len_for_mem(batch_size, token_num_upper_bound + extra_tokens) - - extra_tokens - ) - # clamp_max_seq_len_for_mem considers all tiers (GPU + host). When - # max_tokens is explicitly set, cap by GPU-only capacity so callers - # (e.g. CUDA graph warmup) don't exceed the GPU pool. - if self._gpu_max_tokens is not None: - clamped = min(clamped, self._gpu_max_tokens - extra_tokens) - return clamped - - def _get_attention_pool_num_blocks(self, local_layer_idx: int) -> int: - """Count ledger blocks without counting expanded physical subpages.""" - layer_id = LayerId(local_layer_idx) - expansion = self.impl.get_page_index_converter(layer_id, Role.KEY).expansion - return self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // ( - self.get_layer_kv_factor(self.pp_layers[local_layer_idx]) * expansion - ) - - def get_num_free_blocks(self) -> int: - # NOTE This method is used to get the number of blocks in the primary pool not the FREE blocks. - # However, since we only use this function when the kv cache manager is empty, so it is safe to do so. - assert len(self.kv_cache_map) == 0, ( - "get_num_free_blocks is only used when the kv cache manager is empty" - ) - max_num_blocks = max( - [ - self._get_attention_pool_num_blocks(local_layer_idx) - for local_layer_idx in range(self.num_local_layers) - ] - ) - return max_num_blocks - - def _get_buffer_views_for_invalid_value_check( - self, layer_idx: int, kv_layout: str = "NHD" - ) -> list[torch.Tensor]: - """Return contiguous tensor views covering the mapped KV pages to scan.""" - if self.num_ugpus <= 1: - return [self.get_buffers(layer_idx, kv_layout=kv_layout).flatten(0, 1)] - - layer_offset = self.layer_offsets[layer_idx] - storage = self.impl._storage - lc_id = storage._layer_to_life_cycle_ids[LayerId(layer_offset)] - attr = storage.get_buffer_attr(LayerId(layer_offset), Role.KEY) - pg_idx = storage.get_pool_group_index(lc_id) - gpu_storage = storage._levels[GPU_LEVEL].storage - pool_group = gpu_storage._pool_groups[pg_idx] - slot_size = pool_group.slot_size[attr.pool_index] - element_per_container = 1 - dtype = self.dtype - if dtype == DataType.NVFP4: - element_per_container = 2 - dtype = torch.int8 - - views: list[torch.Tensor] = [] - for ugpu_id in range(self.num_ugpus): - local_page_upper_bound = ( - exact_div(slot_size, attr.size) * pool_group.num_slots(ugpu_id) - - exact_div(attr.offset, attr.size) - ) * attr.expansion - if kv_layout == "NHD": - shape = [ - local_page_upper_bound // self.kv_factor, - self.kv_factor, - self.tokens_per_block, - self.num_kv_heads_per_layer[layer_offset], - self.head_dim // element_per_container, - ] - else: - shape = [ - local_page_upper_bound // self.kv_factor, - self.kv_factor, - self.num_kv_heads_per_layer[layer_offset], - self.tokens_per_block, - self.head_dim // element_per_container, - ] - addr_key = ( - int(gpu_storage.slot_address(pg_idx, attr.pool_index, 0, ugpu_id)) + attr.offset - ) - views.append( - convert_to_torch_tensor(TensorWrapper(addr_key, dtype, shape)).flatten(0, 1) - ) - return views - - @property - def num_ugpus(self) -> int: - """Number of uGPUs. Returns 1 for non-localized configurations.""" - return _get_num_ugpus(self.impl) - - def get_per_ugpu_free_slots(self) -> list[int]: - """Free slot count per uGPU, summed across all pool groups.""" - return self.impl.get_per_ugpu_free_slots() - - def pick_ugpu(self, request_id: int) -> int | None: - """Default uGPU placement for newly-created requests.""" - if self.num_ugpus <= 1: - return None - return request_id % self.num_ugpus - - def _request_ugpu_id(self, req: LlmRequest) -> int | None: - """Return the request-owned uGPU id after validating the invariant.""" - ugpu_id = getattr(req, "py_ugpu_id", None) - if self.num_ugpus <= 1: - assert ugpu_id is None or ugpu_id == 0, ( - "Non-localized requests must not use a nonzero ugpu_id" - ) - return None - - assert ugpu_id is not None, ( - f"Request {req.py_request_id} must carry py_ugpu_id before localized KV-cache placement" - ) - assert 0 <= ugpu_id < self.num_ugpus, ( - f"ugpu_id must be in [0, {self.num_ugpus}), got {ugpu_id}" - ) - return ugpu_id - - def get_ugpu(self, request_id: int) -> int: - """Return the uGPU index assigned to an allocated request.""" - if self.num_ugpus <= 1: - return 0 - kv_cache = self.kv_cache_map.get(request_id) - assert kv_cache is not None, ( - f"get_ugpu: request {request_id} has no allocated KV cache " - f"(num_ugpus={self.num_ugpus}, " - f"kv_cache_map size={len(self.kv_cache_map)})" - ) - assert _get_kv_cache_ugpu_id(kv_cache) is not None, ( - f"get_ugpu: request {request_id} has no concrete ugpu_id " - f"(num_ugpus={self.num_ugpus}, " - f"kv_cache_map size={len(self.kv_cache_map)})" - ) - return _get_kv_cache_ugpu_id(kv_cache) - - @property - def fork_join_attn(self) -> bool: - """Whether to drive single-forward fork-join attention.""" - return self._fork_join_attn - - def commit_scheduled_kv_cache_stats(self, scheduled_batch: ScheduledRequests) -> None: - if self.is_draft or not self.enable_stats: - return - dirty_req_ids = self.impl.get_dirty_stats_kv_cache_ids() - for req in scheduled_batch.all_requests(): - if req.py_request_id in dirty_req_ids: - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - continue - request_stats = kv_cache.commit_pending_stats() - if not req.is_dummy and not request_stats.empty: - req.update_kv_cache_perf_metrics( - request_stats.alloc_total_blocks, - request_stats.alloc_new_blocks, - request_stats.reused_blocks, - request_stats.missed_blocks, - ) - - # ---- Scheduling API (called by KVCacheV2Scheduler) ---- - - def is_request_active(self, request_id: int) -> bool: - """Return True if *request_id* has a live, non-suspended KV cache.""" - kv_cache = self.kv_cache_map.get(request_id) - return kv_cache is not None and kv_cache.is_active - - def _effective_draft_len(self, req: LlmRequest) -> int: - """Draft token length to use for next-step KV capacity calculation. - - During the context-to-generation transition, ``py_draft_tokens`` is - still empty when the scheduler asks for capacity. Prefer transferred - context draft tokens; if there are none, reserve the generation - worker's configured draft capacity before its first forward pass. - """ - draft_len = get_draft_token_length(req) - if ( - draft_len == 0 - and req.is_disagg_generation_transmission_complete - and req.context_phase_params is not None - ): - ctx_draft_tokens = req.context_phase_params.draft_tokens - if ctx_draft_tokens is not None: - draft_len = len(ctx_draft_tokens) - if draft_len == 0 and not self.is_draft and not req.py_disable_speculative_decoding: - draft_len = self.max_total_draft_tokens - return draft_len - - def _helix_local_len(self, global_len: int) -> int: - """Tokens of the first ``global_len`` owned by this CP rank - (continuation round-robin: page b lives on rank b %% cp).""" - phys = self.tokens_per_block - full, rem = divmod(global_len, self._ledger_tokens_per_block) - return full * phys + min(max(rem - self._helix_cp_rank * phys, 0), phys) - - def _set_helix_rank_fields(self, req: LlmRequest) -> None: - """Derive the per-rank helix fields from the global position. - - The decode-step index is manager-owned (committed in - ``try_allocate_generation`` on successful resize) rather than derived - from ``py_decoding_iter``: the sampler advances that counter after - scheduling under the overlap loop, so a schedule-time read is one - step behind and would repeat the first decode position, overwriting - the first generated token's KV. Multi-token verify groups advance - py_helix_decode_group_index per committed group, so this formula - stays exact under DSpark speculation. - """ - step = req.py_helix_decode_group_index + 1 - pos = req.total_input_len_cp + step - 1 - owner = (pos // self.tokens_per_block) % self._helix_cp_size - active = owner == self._helix_cp_rank - req.py_helix_is_inactive_rank = not active - # Convention shared with model_engine: the active rank's seqlen - # includes the in-flight token (past_seen = seqlen - 1 there). - req.seqlen_this_rank_cp = self._helix_local_len(pos) + (1 if active else 0) - - def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int: - """Compute generation KV cache capacity for a request. - - Grows *current_capacity* by 1 + draft tokens, plus another draft_len - of slack for one-model speculative decoding: under the overlap - scheduler the device-side KV position runs one verification round - ahead of host bookkeeping (up to draft_len accepted-but-uncommitted - tokens), and the MTP draft layers then append draft KV beyond that - position. Without the slack, the draft KV append can address a block - ordinal the host has not allocated yet whenever the overrun crosses a - tokens_per_block boundary — an illegal memory access in the MLA rope - generation kernel (observed with DSV4 DEP+MTP3 at 128k, where the - first window-slide boundary crossing after prefill faults). - - The slack must stay a *constant* offset over the request's lifetime: - py_rewind_len only rewinds the rejected part of the recurring - 1 + draft_len growth, so every call site records the extra draft_len - in _pending_overlap_slack and update_resources trims it back. - """ - return current_capacity + 1 + 2 * self._generation_draft_slots(req) - - def _generation_draft_slots(self, req: LlmRequest) -> int: - """Physical draft width reserved for one iteration. Dynamic-tree draft pools - reserve wider than target pools; each reclaims its own in ``update_resources``. - """ - slots = self._effective_draft_len(req) - if self.is_draft: - slots = max(slots, self._kv_reserve_draft_tokens) - return slots - - def try_allocate_generation(self, req: LlmRequest) -> bool: - """Try to allocate one additional KV cache slot for a generation request. - - Resumes from suspended state if needed, then resizes capacity by 1 (+ - draft tokens). Returns True on success, False if allocation failed. - """ - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - return False - - if not kv_cache.is_active: - if not kv_cache.resume(self._stream.cuda_stream): - return False - self._restore_page_index_bufs(req.py_request_id, kv_cache) - - request_id = req.py_request_id - draft_slots = self._generation_draft_slots(req) - self._allocated_draft_lens.pop(request_id, None) - is_helix_req = self._has_cp_helix and not req.is_dummy_request - if is_helix_req: - self._set_helix_rank_fields(req) - if not kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)): - return False - if is_helix_req: - # Commit only on success so a same-pass retry recomputes the - # same step instead of skipping one. - req.py_helix_decode_group_index += 1 - self._allocated_draft_lens[request_id] = draft_slots - if draft_slots > 0: - self._pending_overlap_slack[request_id] = ( - self._pending_overlap_slack.get(request_id, 0) + draft_slots - ) - return True - - def revert_allocate_generation(self, req: LlmRequest) -> None: - """Undo the capacity growth from try_allocate_generation. - - When attention DP causes can_queue=False after scheduling, the - forward pass is skipped but the scheduler already grew each - generation request's KV cache capacity by 1 (+draft tokens). - This method shrinks capacity back to undo that spurious growth - so it does not accumulate across iterations and overflow the - host page-index buffer. - - Keyed on ``_allocated_draft_lens``: no entry means no growth to give back. - Recording the width keeps the paired target/draft reverts in step. - """ - request_id = req.py_request_id - draft_slots = self._allocated_draft_lens.pop(request_id, None) - if draft_slots is None: - return - kv_cache = self.kv_cache_map.get(request_id) - if kv_cache is None or not kv_cache.is_active: - return - if self._has_cp_helix and not req.is_dummy_request and req.py_helix_decode_group_index > 0: - # The forward pass for this step is skipped; give the step back. - req.py_helix_decode_group_index -= 1 - # Mirror the 1 + 2 * draft_slots growth in _required_gen_capacity. - reverted_cap = kv_cache.capacity - 1 - 2 * draft_slots - if reverted_cap < 0: - return - # The reverted growth included draft_slots of overlap slack; deduct it - # so update_resources does not trim slack that no longer exists. - if draft_slots > 0: - remaining = self._pending_overlap_slack.get(request_id, 0) - draft_slots - if remaining > 0: - self._pending_overlap_slack[request_id] = remaining - else: - self._pending_overlap_slack.pop(request_id, None) - if not kv_cache.resize(reverted_cap): - raise RuntimeError( - f"Failed to revert KV cache capacity for request " - f"{req.py_request_id} from {kv_cache.capacity} to " - f"{reverted_cap}" - ) - - def revert_allocate_context(self, req: LlmRequest) -> bool: - """Undo this iteration's context resize. False means the cache was dropped, - not shrunk (history outran pre-resize capacity); the caller drops any draft pool. - """ - pre_cap = getattr(req, "py_ctx_pre_resize_cap", None) - if pre_cap is None: - return True - req.py_ctx_pre_resize_cap = None - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None or not kv_cache.is_active: - return True - if pre_cap >= kv_cache.capacity: - return True - if kv_cache.history_length > pre_cap: - self.free_resources(req) - # The setter asserts prepopulated < promptLen; a Helix CP rank that - # holds none of the prompt has promptLen 0 and nothing to reset. - if req.prompt_len > 0 and req.prepopulated_prompt_len > 0: - req.set_prepopulated_prompt_len(0, self.tokens_per_block) - # setPrepopulatedPromptLen only moves forward, so rewind the same - # fields LlmRequest::pause() resets for a recompute pause. - req.context_current_position = 0 - req.context_chunk_size = req.prompt_len - req.estimated_reusable_tokens = 0 - return False - history_length = min(kv_cache.history_length, pre_cap) - if not kv_cache.resize(pre_cap, history_length): - raise RuntimeError( - f"Failed to revert KV cache capacity for context " - f"request {req.py_request_id} from {kv_cache.capacity} " - f"to {pre_cap}" - ) - if pre_cap > 0: - kv_cache.suspend() - return True - - def _restore_page_index_bufs(self, request_id: int, kv_cache) -> None: - """Re-connect host page-index buffers after resume(). - - suspend() clears the base_page_index_buf pointers (sets them to - None) so the KV cache stops writing page indices to the host - buffer. After resume(), the KV cache has re-locked pages but - copy_batch_block_offsets still reads from the host buffer, so we - must re-connect the buffers to avoid stale/zero page indices that - would cause illegal memory accesses during the forward pass. - """ - index = self.index_mapper.get_index(request_id) - for i in range(self.max_beam_width): - for pool_idx in range(self.num_pools): - buffer: torch.Tensor = self.host_kv_cache_block_offsets[ - pool_idx, index * self.max_beam_width + i, 0 - ] - kv_cache.set_base_page_index_buf(i, pool_idx, memoryview(buffer.numpy())) - - def _resume_and_restore(self, req_id: int, kv_cache) -> bool: - """Resume a suspended KV cache and restore its page index buffers. - - Returns True if the cache is (or becomes) active, False on failure. - """ - if kv_cache.is_active: - return True - if not kv_cache.resume(self._stream.cuda_stream): - return False - self._restore_page_index_bufs(req_id, kv_cache) - return True - - def _draft_context_exceeds_limit(self, req: LlmRequest) -> bool: - """Whether draft prefill skips this prompt, preventing unified block reuse.""" - if ( - not self.draft_layer_ids - or self.draft_window_size is not None - or self.draft_max_position_embeddings is None - ): - return False - prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - return prompt_len > min(self.max_seq_len, self.draft_max_position_embeddings) - - def _context_reuse_tokens( - self, req: LlmRequest, reuse_limit: int | None = None - ) -> Sequence[TokenIdExt]: - """Radix-tree evidence for one context prefix. Raw prompt, so block hashes - stay comparable with KV cache events, the cache-aware router and disagg peers. - """ - if self._draft_context_exceeds_limit(req): - return [] - all_tokens = self._reuse_token_source(req) - # The target must recompute the last prompt token, but a one-model draft - # also reads D prompt tokens ahead. Include those tokens as match evidence; - # the core drops D before claiming in both reuse protocols. - reuse_end = min( - len(all_tokens), - max(len(all_tokens) - 1, 0) + self.reuse_match_backoff, - ) - if reuse_limit is not None: - # ``reuse_limit`` caps the claimed depth, not the evidence. Retain - # the D following tokens that prove the draft KV at that depth. - evidence_limit = 0 if reuse_limit <= 0 else reuse_limit + self.reuse_match_backoff - reuse_end = min(reuse_end, evidence_limit) - return self._augment_tokens_for_block_reuse(all_tokens, req, end=reuse_end) - - def _record_branch_snapshot_point( - self, req: LlmRequest, kv_cache: _KVCache, num_lookup_tokens: Optional[int] - ) -> None: - """Hook for recurrent-state managers; attention-only caches have no snapshots.""" - - def probe_context_reuse(self, req: LlmRequest) -> int | None: - """Reusable prefix depth for this pool, without claiming pages. - - The lookup includes this pool's prompt lookahead as evidence, and the core - applies ``reuse_match_backoff`` in the same match. A paired caller takes - the common usable depth. ``None``: nothing to probe -- reuse is off, or the - request already holds a settled cache claim. - """ - if not self.enable_block_reuse or req.py_request_id in self.kv_cache_map: - return None - return self.impl.probe_reuse( - ReuseScope(lora_id=req.lora_task_id, salt=self._derive_reuse_salt(req.cache_salt)), - self._context_reuse_tokens(req), - ) - - def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) -> int | None: - """Create/resume a context cache without touching request cursors. ``reuse_limit`` - caps the claimed depth while retaining its lookahead evidence; the returned - match lets the scheduler check both pools agree. - """ - if req.is_first_context_chunk: - if self.conversation_manager is not None: - self.conversation_manager.prepare_request(req) - kv_cache = self.kv_cache_map.get(req.py_request_id) - # Stays None when the cache was resumed rather than freshly - # matched, which is also the case where no branch point applies. - num_lookup_tokens = None - if kv_cache is None: - if self.enable_block_reuse: - # Empty (not None) for a zero limit: lookup is off, but - # newly computed blocks may still be committed. - tokens = self._context_reuse_tokens(req, reuse_limit) - else: - tokens = None - ugpu_id = self._request_ugpu_id(req) - # Taken from the augmented sequence rather than derived as - # prompt_len - 1, because augmentation can change the length - # for multimodal requests. - num_lookup_tokens = len(tokens) if tokens is not None else None - kv_cache = self._create_kv_cache( - req.py_request_id, - req.lora_task_id, - tokens, - cache_salt=req.cache_salt, - is_dummy=req.is_dummy, - expected_prompt_length=( - req.total_input_len_cp if self._has_cp_helix else req.prompt_len - ) - - 1, - ugpu_id=ugpu_id, - ) - if kv_cache is None: - return None - assert _get_kv_cache_ugpu_id(kv_cache) == ugpu_id - kv_cache.cuda_stream = self._stream.cuda_stream - - if not self.enable_block_reuse or self._draft_context_exceeds_limit(req): - kv_cache.stop_committing() - else: - self._record_branch_snapshot_point(req, kv_cache, num_lookup_tokens) - - if req.is_disagg_generation_init_state: - # Disagg generation receives prompt KV from the context worker; - # scratch blocks are only valid for local prefill chunks. - kv_cache.enable_swa_scratch_reuse = False - if not self._resume_and_restore(req.py_request_id, kv_cache): - return None - if ( - self.draft_layer_ids - and not self._draft_history_uses_kv_history() - and req.py_request_id not in self.draft_history - ): - # The reuse match covers every cache domain. Initialize once, - # after resume succeeds, so overlap/chunk retries cannot rewind - # draft history that has already advanced on the device. - reused = kv_cache.num_committed_tokens - valid_length = reused - if self.draft_window_size is not None: - valid_length = min(valid_length, self.draft_window_size) - self.set_draft_history(req.py_request_id, valid_length, reused) - return kv_cache.num_committed_tokens - - # Subsequent chunk: cache must exist from first chunk. It may be - # suspended (e.g. evicted between chunks), so resume it here. - kv_cache = self.kv_cache_map.get(req.py_request_id) - assert kv_cache is not None, ( - f"KV cache missing for non-first context chunk, request {req.py_request_id}" - ) - if not self._resume_and_restore(req.py_request_id, kv_cache): - return None - return kv_cache.num_committed_tokens - - def prepare_context(self, req: LlmRequest) -> bool: - """Create _KVCache, handle block reuse, and resume. Does NOT resize.""" - assert not req.is_disagg_generation_init_state, ( - f"req {req.py_request_id}: use prepare_disagg_gen_init" - ) - reused = self.prepare_context_cache(req) - if reused is None: - return False - # First chunk only: num_committed_tokens holds at the initial prefix - # until context end, so reapplying later would rewind the cursor. - if req.is_first_context_chunk and self.enable_block_reuse: - _rewind_context_cursor(req, reused) - req.set_prepopulated_prompt_len(reused, self.tokens_per_block) - return True - - def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: - """Resize KV cache to cover context_current_position + num_tokens. - - Returns True on success, False if resize failed (first chunk is - suspended on failure). - """ - assert not req.is_disagg_generation_init_state, ( - f"req {req.py_request_id}: use prepare_disagg_gen_init" - ) - if self._has_cp_helix and not req.is_dummy_request: - raise ValueError( - "resize_context is not helix-aware: its rank-local chunk " - "target would under-size the global ledger. Helix requests " - "are disagg-generation-only and must never take the context " - "path." - ) - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - return False - - target = ( - req.context_current_position - + num_tokens - + self.num_extra_kv_tokens - + self._standalone_draft_reserve - ) - capacity = max(kv_cache.capacity, target) - pre_cap = kv_cache.capacity - - success = kv_cache.resize(capacity) - if not success: - if req.is_first_context_chunk: - kv_cache.suspend() - return False - req.py_ctx_pre_resize_cap = pre_cap if capacity > pre_cap else None - return True - - def try_allocate_draft_context(self, req: LlmRequest, num_tokens: int) -> bool: - """Resize a draft context cache during admission. Joint reuse shares one - cursor; the draft manager owns only its radix tree and capacity. - """ - assert self.is_draft - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None or not self._resume_and_restore(req.py_request_id, kv_cache): - return False - - target = req.context_current_position + num_tokens + self.num_extra_kv_tokens - return kv_cache.resize(max(kv_cache.capacity, target)) - - def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: - """Prepare KV cache for a disagg generation init request. - - Allocates capacity for the full prompt (+ draft) and sets - ``kv_cache.history_length`` to ``prompt_len``. Returns True on - success, False if preparation or resize failed (cache is suspended - on resize failure). - """ - reused = self.prepare_context_cache(req) - if reused is None: - return False - if self.enable_block_reuse: - _rewind_context_cursor(req, reused) - req.set_prepopulated_prompt_len(reused, self.tokens_per_block) - - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - return False - - # prompt_len is the full incoming prompt length, robust to block - # reuse (which may leave a non-zero context_current_position). - # Helix requests carry the rank-local strided slice in prompt_len; - # the global ledger sizes off the full prompt instead. - prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - target = ( - prompt_len - + get_draft_token_length(req) - + self.num_extra_kv_tokens - + self._standalone_draft_reserve - ) - capacity = max(kv_cache.capacity, target) - pre_cap = kv_cache.capacity - - success = kv_cache.resize(capacity, prompt_len) - if not success: - if req.is_first_context_chunk: - kv_cache.suspend() - return False - req.py_ctx_pre_resize_cap = pre_cap if capacity > pre_cap else None - return True - - def get_history_length(self, req: LlmRequest) -> int | None: - """Return the cache's current history_length, or None if no cache. - - Exposes the per-request SWA history watermark so callers - (e.g., the disagg transceiver) can verify scheduler/cache contracts - without reaching into ``kv_cache_map`` directly. - """ - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - return None - return kv_cache.history_length - - def extend_capacity_for_tokens(self, request: LlmRequest) -> None: - """Extend KV cache capacity for the CUDA-graph padding delta. - - ``try_allocate_generation`` allocated capacity for the schedule-reduced - draft length. After padding restores ``py_draft_tokens`` to the static - max, we must extend by exactly the difference so that the subsequent - rewind (which operates on the padded length) does not underflow. - - The delta is computed from ``_allocated_draft_lens`` (recorded by - ``try_allocate_generation``) vs the current draft length (post-padding). - """ - allocated = self._allocated_draft_lens.get(request.py_request_id) - if allocated is None: - return - current_draft_len = get_draft_token_length(request) - # Growth is 1 + 2 * draft_len (see _required_gen_capacity), so the - # padding delta scales by 2 as well. - delta = 2 * (current_draft_len - allocated) - if delta <= 0: - return - # Half of the delta tops up the overlap slack; record it so - # update_resources reclaims the full slack for this iteration. - self._pending_overlap_slack[request.py_request_id] = self._pending_overlap_slack.get( - request.py_request_id, 0 - ) + (current_draft_len - allocated) - kv_cache = self.kv_cache_map[request.py_request_id] - new_capacity = kv_cache.capacity + delta - success = kv_cache.resize(new_capacity) - if not success: - raise ValueError( - f"Failed to extend capacity of KV cache for request " - f"{request.py_request_id} by {delta} tokens " - f"(target capacity {new_capacity})" - ) - self._allocated_draft_lens[request.py_request_id] = current_draft_len - - def suspend_request(self, req: LlmRequest) -> None: - """Suspend a request's KV cache, allowing pages to migrate to a secondary tier.""" - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is not None and kv_cache.is_active: - kv_cache.suspend() - - def resume_request(self, req: LlmRequest) -> bool: - """Resume a previously-suspended KV cache for *req*. - - Returns True if the cache is (or becomes) active on GPU, False if - resume was refused (e.g. GPU pressure above max_util_for_resume) - or no cache exists for the request. - """ - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - return False - return self._resume_and_restore(req.py_request_id, kv_cache) - - # ---- prepare_resources ---- - - @nvtx_range("prepare_resources_kv_cache_manager_v2") - def prepare_resources(self, scheduled_batch: ScheduledRequests): - if self.is_draft: - # Mirror the main manager. Under one-model spec decoding the - # scheduler may already have created the context cache. - self._prepare_draft_resources(scheduled_batch) - return - - def _mirror_draft_kv_cache(self, req: LlmRequest): - """The draft manager's entry for ``req``, created on first sight. - - Mirrors what the target manager does in ``_prepare_context_impl``: - create when missing rather than fail, passing ``is_dummy`` through. The - draft mirror never looks up block reuse (``tokens=None``) and never - commits, so this cannot share the target's implementation. - - Reached from both the context and the generation loop, because not - every generation request passed through context on this worker: - attention DP injects its idle placeholder straight as a generation - request when a rank has no real work (ATTENTION_DP_DUMMY_REQUEST_ID, - see py_executor._adp_dummy_is_gen), and a disaggregated generation - worker receives prompt KV rather than prefilling it. - - Returns None when the IndexMapper is saturated. The caller must treat - that as fatal: the request stays in the scheduled batch, so skipping it - only defers the failure to copy_batch_block_offsets(), which asserts in - C++ on the unmapped request ID. Slots are released early on a - context-only worker (release_index_slot), so saturation here means the - draft mirror is out of sync with the target, not merely busy. - """ - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is not None: - return kv_cache - ugpu_id = self._request_ugpu_id(req) - kv_cache = self._create_kv_cache( - req.py_request_id, - req.lora_task_id, - None, - cache_salt=req.cache_salt, - is_dummy=req.is_dummy, - ugpu_id=ugpu_id, - ) - if kv_cache is None: - return None - assert _get_kv_cache_ugpu_id(kv_cache) == ugpu_id - kv_cache.stop_committing() - return kv_cache - - def _draft_pool_diagnostic(self) -> str: - """Draft-pool occupancy, for the resize-failure messages below. - - The draft manager mirrors the target's tokens but is sized from its own - byte budget, and the capacity scheduler admits on the TARGET pool alone - (`scheduler_v2` touches `draft_kv_cache_manager` only to suspend/free). - A draft pool smaller in tokens than the target therefore cannot - backpressure -- it can only raise, and the raise kills every rank. When - that happens the first question is always "how big was the draft pool - and how full was it", so answer it in the message rather than leaving - it to post-hoc arithmetic over the budget-split log line. - - Best-effort: never let a diagnostic mask the failure it describes. - """ - try: - live = sum(c.capacity for c in self.kv_cache_map.values()) - return ( - f" [draft pool: {len(self.kv_cache_map)} live caches holding " - f"{live} tokens, gpu_max_tokens={self._gpu_max_tokens}]" - ) - except Exception: # noqa: BLE001 - diagnostic only - return "" - - def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): - """Create/resize KV caches in the draft V2 manager for scheduled requests. - - The main V2 scheduler manages capacity in the primary KV cache - manager. The draft manager mirrors scheduled requests so that its - IndexMapper contains the correct request IDs. With block reuse, the - scheduler pre-creates context caches to coordinate the reusable prefix. - """ - # Joint V2 reuse deliberately shares the target request progress. A - # separate draft engine still uses the legacy draft request view. - with request_context(not self.enable_joint_kv_cache_reuse, scheduled_batch): - for req in scheduled_batch.context_requests: - kv_cache = self._mirror_draft_kv_cache(req) - if kv_cache is None: - raise RuntimeError( - f"Draft KV cache mirror exhausted its IndexMapper on " - f"context request {req.py_request_id}" - ) - if not self._resume_and_restore(req.py_request_id, kv_cache): - raise RuntimeError( - f"Failed to resume draft KV cache for request {req.py_request_id}" - ) - draft_len = get_draft_token_length(req) if req.is_last_context_chunk else 0 - capacity = ( - req.context_current_position - + req.context_chunk_size - + draft_len - + self.num_extra_kv_tokens - ) - if not kv_cache.resize(capacity): - raise RuntimeError( - f"Draft KV cache context resize failed for request " - f"{req.py_request_id}: could not resize to {capacity} tokens" - f"{self._draft_pool_diagnostic()}" - ) - - for req in scheduled_batch.generation_requests: - kv_cache = self._mirror_draft_kv_cache(req) - if kv_cache is None: - raise RuntimeError( - f"Draft KV cache mirror exhausted its IndexMapper on " - f"generation request {req.py_request_id}" - ) - if not self._resume_and_restore(req.py_request_id, kv_cache): - raise RuntimeError( - f"Failed to resume draft KV cache for request {req.py_request_id}" - ) - # A paired scheduler already grew both pools atomically; the - # recorded width says so, so do not charge this step twice. - if req.py_request_id in self._allocated_draft_lens: - continue - new_cap = self._required_gen_capacity(req, kv_cache.capacity) - if not kv_cache.resize(new_cap): - raise RuntimeError( - f"Draft KV cache generation resize failed for request " - f"{req.py_request_id}: could not resize to {new_cap} tokens" - f"{self._draft_pool_diagnostic()}" - ) - slack = self._effective_draft_len(req) - if slack > 0: - self._pending_overlap_slack[req.py_request_id] = ( - self._pending_overlap_slack.get(req.py_request_id, 0) + slack - ) - - def _reuse_token_source(self, req: LlmRequest) -> Sequence[int]: - """Beam-0 tokens for block reuse, in the form the active backend consumes. - - The C++ backend ingests a zero-copy int32 view (get_tokens_view) — no per-token - PyLong allocation. The pure-Python backend cannot consume a numpy array, so it gets - a plain list (get_tokens), preserving pre-optimization behavior. ``_cpp_introspection`` - is None exactly when the Python backend is active. - """ - if _cpp_introspection is not None: - return req.get_tokens_view(DEFAULT_BEAM_INDEX) - return req.get_tokens(DEFAULT_BEAM_INDEX) - - def _augment_tokens_for_block_reuse( - self, tokens: Sequence[int], req: LlmRequest, start: int = 0, end: int | None = None - ) -> Sequence[TokenIdExt]: - """Augment token sequence with multimodal content digests for block reuse. - - Multimodal placeholder tokens (e.g. image_token_id) share the same ID - regardless of the underlying content. This method replaces each - multimodal token region with TokenIdExt values produced by - gen_multimodal_cache_key_tokens(), embedding the content digest - (Blake3 hash) into the token sequence so that the radix tree can - distinguish blocks belonging to different images/videos. - - When *start*/*end* are given, they define the chunk bounds; only - `tokens[start:end]` is materialized and returned. This avoids - re-augmenting the full prompt on every chunk during chunked prefill. - - For text-only requests this is a no-op. - """ - if end is None: - end = len(tokens) - chunk_start = start - chunk_end = end - is_sliced = chunk_start != 0 or chunk_end != len(tokens) - - if ( - req.multimodal_hashes is None - or req.multimodal_positions is None - or req.multimodal_lengths is None - ): - return tokens[chunk_start:chunk_end] if is_sliced else tokens - - # Multimodal path: materialize a Python-int list (digest bytes get spliced in below), - # which flows through the per-element binding fallback. tokens may be a zero-copy numpy - # int32 view (get_tokens_view) — use tolist() so elements are Python ints, not np.int32. - chunk = tokens[chunk_start:chunk_end] - result: list[TokenIdExt] = chunk.tolist() if hasattr(chunk, "tolist") else list(chunk) - run_metadata = _resolve_multimodal_run_metadata(req) - if run_metadata is not None: - return _augment_tokens_with_mm_run_metadata( - self.vocab_size, result, req.multimodal_hashes, run_metadata, chunk_start, chunk_end - ) - - return _augment_tokens_with_contiguous_mm_metadata( - self.vocab_size, - result, - req.multimodal_hashes, - req.multimodal_positions, - req.multimodal_lengths, - chunk_start, - chunk_end, - ) - - def _stats_window_size(self, window_size: Optional[int]) -> int: - return self.max_seq_len if window_size is None else int(window_size) - - def _stats_life_cycle_window_size(self, life_cycle) -> Optional[int]: - if not isinstance(life_cycle, AttnLifeCycle): - return None - return self._stats_window_size(life_cycle.window_size) - - def _get_storage_statistics(self, cache_level: CacheLevel): - return _introspection.storage_statistics(self.impl, cache_level) - - def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str]]: - # life cycle (== layer group) -> pool group is static structure exposed by - # the public pool_group_descs API; no introspection needed. - pool_groups_by_life_cycle = { - int(variant.layer_group_id): int(pool_group.pool_group_index) - for pool_group in self.impl.pool_group_descs - for variant in pool_group.slot_desc.variants - } - - metadata: dict[int, tuple[int, Optional[int], str]] = {} - for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): - if not layer_ids: - continue - layer = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] - is_attention = isinstance(layer, AttentionLayerConfig) - metadata[life_cycle_id] = ( - int(pool_groups_by_life_cycle[life_cycle_id]), - self._stats_window_size(layer.sliding_window_size) if is_attention else None, - "attention" if is_attention else "ssm", - ) - return metadata - - def _storage_pool_groups_by_window(self) -> dict[int, set[int]]: - pool_groups_by_window: dict[int, set[int]] = defaultdict(set) - for pool_group_id, window_size, _ in self._stats_life_cycle_metadata().values(): - if window_size is not None: - pool_groups_by_window[window_size].add(pool_group_id) - return pool_groups_by_window - - def _get_and_reset_iteration_peak_block_stats(self, cache_level: CacheLevel): - get_peak_stats = getattr(self.impl, "get_and_reset_iteration_peak_block_stats", None) - if get_peak_stats is not None: - return get_peak_stats(cache_level) - return [ - PoolGroupPeakBlockStats( - available=stats.available, - unavailable=stats.total - stats.available, - evictable=stats.evictable, - ) - for stats in self._get_storage_statistics(cache_level) - ] - - @staticmethod - def _windows_by_pool_group( - pool_groups_by_window: dict[int, set[int]], - ) -> dict[int, tuple[int, ...]]: - windows_by_pool_group: dict[int, set[int]] = defaultdict(set) - for window_size, pool_group_ids in pool_groups_by_window.items(): - for pool_group_id in pool_group_ids: - windows_by_pool_group[pool_group_id].add(window_size) - return { - pool_group_id: tuple(sorted(window_sizes)) - for pool_group_id, window_sizes in windows_by_pool_group.items() - } - - @staticmethod - def _filter_iteration_stats_delta(delta, field_names) -> KVCacheIterationStatsDelta: - filtered = KVCacheIterationStatsDelta() - for field_name in field_names: - setattr(filtered, field_name, getattr(delta, field_name)) - return filtered - - @staticmethod - def _add_iteration_stats_delta( - bucket: dict[int, KVCacheIterationStatsDelta], key: int, delta: KVCacheIterationStatsDelta - ) -> None: - if delta.empty: - return - if key not in bucket: - bucket[key] = delta.copy() - return - bucket[key].add(delta) - - @staticmethod - def _iteration_cache_hit_rate(stats) -> float: - total = stats.iter_reused_blocks + stats.iter_missed_blocks - if stats.iter_reused_blocks == 0 or total == 0: - return 0.0 - return stats.iter_reused_blocks / total - - @staticmethod - def _apply_iteration_stats_delta( - stats, delta, field_names=KV_CACHE_ITERATION_STATS_DELTA_FIELDS - ) -> None: - if delta is None: - return - for field_name in field_names: - setattr(stats, field_name, getattr(delta, field_name)) - stats.iter_cache_hit_rate = KVCacheManagerV2._iteration_cache_hit_rate(stats) - - def _build_iteration_stats( - self, - pool_group_ids: Iterable[int], - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - delta, - field_names=KV_CACHE_ITERATION_STATS_DELTA_FIELDS, - ): - pool_group_ids = tuple(pool_group_ids) - stats = KvCacheIterationStats() - stats.primary_max_num_blocks = sum( - primary_stats[pool_group_id].total for pool_group_id in pool_group_ids - ) - stats.primary_free_num_blocks = sum( - primary_stats[pool_group_id].available for pool_group_id in pool_group_ids - ) - stats.primary_used_num_blocks = stats.primary_max_num_blocks - stats.primary_free_num_blocks - stats.primary_evictable_num_blocks = sum( - primary_stats[pool_group_id].evictable for pool_group_id in pool_group_ids - ) - stats.primary_peak_free_num_blocks = sum( - primary_peak_stats[pool_group_id].available for pool_group_id in pool_group_ids - ) - stats.primary_peak_used_num_blocks = sum( - primary_peak_stats[pool_group_id].unavailable for pool_group_id in pool_group_ids - ) - stats.primary_peak_evictable_num_blocks = sum( - primary_peak_stats[pool_group_id].evictable for pool_group_id in pool_group_ids - ) - stats.secondary_max_num_blocks = sum( - level_stats[pool_group_id].total - for level_stats in secondary_stats_by_level - for pool_group_id in pool_group_ids - ) - stats.secondary_free_num_blocks = sum( - level_stats[pool_group_id].available - for level_stats in secondary_stats_by_level - for pool_group_id in pool_group_ids - ) - stats.secondary_used_num_blocks = ( - stats.secondary_max_num_blocks - stats.secondary_free_num_blocks - ) - stats.secondary_evictable_num_blocks = sum( - level_stats[pool_group_id].evictable - for level_stats in secondary_stats_by_level - for pool_group_id in pool_group_ids - ) - stats.secondary_peak_free_num_blocks = sum( - peak_stats[pool_group_id].available - for peak_stats in secondary_peak_stats_by_level - for pool_group_id in pool_group_ids - ) - stats.secondary_peak_used_num_blocks = sum( - peak_stats[pool_group_id].unavailable - for peak_stats in secondary_peak_stats_by_level - for pool_group_id in pool_group_ids - ) - stats.secondary_peak_evictable_num_blocks = sum( - peak_stats[pool_group_id].evictable - for peak_stats in secondary_peak_stats_by_level - for pool_group_id in pool_group_ids - ) - self._apply_iteration_stats_delta(stats, delta, field_names) - return stats - - def _collect_iteration_stats_deltas( - self, raw_iteration_stats, life_cycle_metadata - ) -> tuple[dict, dict, dict, dict]: - reuse_deltas_by_window: dict[int, KVCacheIterationStatsDelta] = {} - reuse_deltas_by_life_cycle: dict[int, KVCacheIterationStatsDelta] = {} - pool_group_deltas_by_window: dict[int, KVCacheIterationStatsDelta] = {} - pool_group_deltas: dict[int, KVCacheIterationStatsDelta] = {} - - for life_cycle_id, delta in raw_iteration_stats.items(): - pool_group_id, window_size, _ = life_cycle_metadata[int(life_cycle_id)] - - pool_group_delta = self._filter_iteration_stats_delta( - delta, KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS - ) - self._add_iteration_stats_delta(pool_group_deltas, pool_group_id, pool_group_delta) - if window_size is not None: - self._add_iteration_stats_delta( - pool_group_deltas_by_window, window_size, pool_group_delta - ) - - reuse_delta = self._filter_iteration_stats_delta( - delta, KV_CACHE_ITERATION_STATS_REUSE_FIELDS - ) - if reuse_delta.empty: - continue - reuse_deltas_by_life_cycle[int(life_cycle_id)] = reuse_delta.copy() - if window_size is not None: - self._add_iteration_stats_delta(reuse_deltas_by_window, window_size, reuse_delta) - - return ( - reuse_deltas_by_window, - reuse_deltas_by_life_cycle, - pool_group_deltas_by_window, - pool_group_deltas, - ) - - def _build_window_iteration_stats( - self, - window_size: int, - pool_groups_by_window: dict[int, set[int]], - windows_by_pool_group: dict[int, tuple[int, ...]], - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - pool_group_delta, - reuse_delta, - ): - pool_group_ids = tuple( - pool_group_id - for pool_group_id in pool_groups_by_window.get(window_size, set()) - if windows_by_pool_group.get(pool_group_id) == (window_size,) - ) - stats = self._build_iteration_stats( - pool_group_ids, - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - pool_group_delta, - KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS, - ) - self._apply_iteration_stats_delta(stats, reuse_delta, KV_CACHE_ITERATION_STATS_REUSE_FIELDS) - return stats - - def _build_pool_group_iteration_stats( - self, - pool_group_id: int, - windows_by_pool_group: dict[int, tuple[int, ...]], - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - pool_group_delta, - ) -> KVCacheV2PoolGroupIterationStats: - primary_pool_group_stats = primary_stats[pool_group_id] - slot_size = ( - primary_pool_group_stats.slot_sizes - if hasattr(primary_pool_group_stats, "slot_sizes") - else primary_pool_group_stats.slot_size - ) - return KVCacheV2PoolGroupIterationStats( - pool_group_id=pool_group_id, - slot_size=tuple(slot_size), - window_sizes=windows_by_pool_group.get(pool_group_id, ()), - stats=self._build_iteration_stats( - (pool_group_id,), - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - pool_group_delta, - KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS, - ), - ) - - def _build_attention_life_cycle_iteration_stats( - self, - life_cycle_id: int, - life_cycle_metadata, - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - reuse_delta, - ) -> KVCacheV2LifeCycleIterationStats: - pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] - assert kind == "attention" - return KVCacheV2LifeCycleIterationStats( - life_cycle_id=life_cycle_id, - pool_group_id=pool_group_id, - window_size=window_size, - kind=kind, - stats=self._build_iteration_stats( - (), - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - reuse_delta, - KV_CACHE_ITERATION_STATS_REUSE_FIELDS, - ), - ) - - @staticmethod - def _build_ssm_life_cycle_iteration_stats( - life_cycle_id: int, - life_cycle_metadata, - snapshot_delta, - ) -> KVCacheV2SsmLifeCycleIterationStats: - pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] - assert kind == "ssm" - assert window_size is None - return KVCacheV2SsmLifeCycleIterationStats( - life_cycle_id=life_cycle_id, - pool_group_id=pool_group_id, - snapshot_stats=KVCacheV2SsmSnapshotIterationStats( - iter_snapshot_lookups=snapshot_delta.iter_snapshot_lookups, - iter_snapshot_hits=snapshot_delta.iter_snapshot_hits, - iter_snapshot_misses=snapshot_delta.iter_snapshot_misses, - iter_reused_tokens=snapshot_delta.iter_reused_tokens, - iter_unreused_tokens=snapshot_delta.iter_unreused_tokens, - iter_aligned_snapshot_hits=snapshot_delta.iter_aligned_snapshot_hits, - iter_unaligned_snapshot_hits=snapshot_delta.iter_unaligned_snapshot_hits, - ), - ) - - def get_kv_cache_stats(self): - kv_cache_stats = KvCacheStats() - pool_group_stats = self._get_storage_statistics(GPU_LEVEL) - max_num_blocks = sum(stat.total for stat in pool_group_stats) - free_num_blocks = sum(stat.available for stat in pool_group_stats) - committed_stats = self.impl.get_committed_stats() - - kv_cache_stats.max_num_blocks = max_num_blocks - kv_cache_stats.free_num_blocks = free_num_blocks - kv_cache_stats.used_num_blocks = max_num_blocks - free_num_blocks - kv_cache_stats.tokens_per_block = self.tokens_per_block - kv_cache_stats.alloc_total_blocks = committed_stats.alloc_total_blocks - kv_cache_stats.alloc_new_blocks = committed_stats.alloc_new_blocks - kv_cache_stats.reused_blocks = committed_stats.reused_blocks - kv_cache_stats.missed_blocks = committed_stats.missed_blocks - total = kv_cache_stats.reused_blocks + kv_cache_stats.missed_blocks - kv_cache_stats.cache_hit_rate = ( - 0.0 - if kv_cache_stats.reused_blocks == 0 or total == 0 - else kv_cache_stats.reused_blocks / total - ) - kv_cache_stats.num_free_blocks_per_window_size = { - window_size: sum( - pool_group_stats[pool_group_id].available for pool_group_id in pool_group_ids - ) - for window_size, pool_group_ids in self._storage_pool_groups_by_window().items() - } - kv_cache_stats.allocated_bytes = self.impl.get_quota(GPU_LEVEL) - - return kv_cache_stats - - def flush_iteration_events(self): - if self.event_manager is not None: - self.event_manager.flush_iteration_events() - - def get_latest_events(self, timeout_ms: Optional[float] = None): - if self.event_manager is None: - return [] - return self.event_manager.get_latest_events(timeout_ms) - - def get_iteration_stats(self): - if not self.enable_stats: - return None - - life_cycle_metadata = self._stats_life_cycle_metadata() - pool_groups_by_window = self._storage_pool_groups_by_window() - windows_by_pool_group = self._windows_by_pool_group(pool_groups_by_window) - raw_iteration_stats = self.impl.get_and_reset_iteration_stats() - raw_ssm_snapshot_iteration_stats = self.impl.get_and_reset_ssm_snapshot_iteration_stats() - primary_peak_stats = self._get_and_reset_iteration_peak_block_stats(GPU_LEVEL) - num_cache_levels = len(self.impl.cache_tier_list) - secondary_peak_stats_by_level = [ - self._get_and_reset_iteration_peak_block_stats(CacheLevel(level)) - for level in range(1, num_cache_levels) - ] - ( - reuse_deltas_by_window, - reuse_deltas_by_life_cycle, - pool_group_deltas_by_window, - pool_group_deltas, - ) = self._collect_iteration_stats_deltas(raw_iteration_stats, life_cycle_metadata) - - windows = set(pool_groups_by_window) - windows.update(reuse_deltas_by_window) - windows.update(pool_group_deltas_by_window) - primary_stats = self._get_storage_statistics(GPU_LEVEL) - secondary_stats_by_level = [ - self._get_storage_statistics(CacheLevel(level)) for level in range(1, num_cache_levels) - ] - - stats_by_window = { - window_size: self._build_window_iteration_stats( - window_size, - pool_groups_by_window, - windows_by_pool_group, - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - pool_group_deltas_by_window.get(window_size), - reuse_deltas_by_window.get(window_size), - ) - for window_size in sorted(windows) - } - - all_pool_group_ids = set(range(len(primary_stats))) - pool_group_ids = sorted( - all_pool_group_ids | set(windows_by_pool_group) | set(pool_group_deltas) - ) - stats_by_pool_group = { - pool_group_id: self._build_pool_group_iteration_stats( - pool_group_id, - windows_by_pool_group, - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - pool_group_deltas.get(pool_group_id), - ) - for pool_group_id in pool_group_ids - } - - stats_by_life_cycle = { - life_cycle_id: self._build_attention_life_cycle_iteration_stats( - life_cycle_id, - life_cycle_metadata, - primary_stats, - secondary_stats_by_level, - primary_peak_stats, - secondary_peak_stats_by_level, - reuse_delta, - ) - for life_cycle_id, reuse_delta in sorted(reuse_deltas_by_life_cycle.items()) - } - for life_cycle_id, snapshot_delta in sorted(raw_ssm_snapshot_iteration_stats.items()): - assert int(life_cycle_id) not in stats_by_life_cycle - stats_by_life_cycle[int(life_cycle_id)] = self._build_ssm_life_cycle_iteration_stats( - int(life_cycle_id), - life_cycle_metadata, - snapshot_delta, - ) - stats_by_life_cycle = dict(sorted(stats_by_life_cycle.items())) - - return KVCacheV2IterationStatsReport( - stats_by_window, stats_by_pool_group, stats_by_life_cycle - ) - - def get_block_ids_per_seq(self, request_ids: List[int]) -> torch.Tensor: - block_ids_per_seq = self.get_batch_cache_indices(request_ids) - block_ids_per_seq_tensors = [ - torch.tensor( - [i // self.num_local_layers if i != BAD_PAGE_INDEX else 0 for i in sublist], - dtype=torch.int, - ) - for sublist in block_ids_per_seq - ] - padded_tensor = torch.nn.utils.rnn.pad_sequence( - block_ids_per_seq_tensors, batch_first=True, padding_value=0 - ) - return padded_tensor - - def add_dummy_requests( - self, - request_ids: List[int], - # Note that token_nums should be past_kv_len + input_len (without - # spec decoding). The draft tokens will be added in this function, - # so we don't need to take care of it in the caller. When preparing - # token_nums, we should not take the draft tokens into account, so - # don't use the kv_cache_manager.max_seq_len, which includes both - # extra tokens and draft tokens. - token_nums: Optional[List[int]] = None, - is_gen: bool = False, - prepare_resource: bool = True, - max_num_draft_tokens: int = 0, - kv_reserve_draft_tokens: Optional[int] = None, - use_mrope: bool = False, - max_beam_width: int = 1, - encoder_output_lens: Optional[List[int]] = None, - num_extra_decoding_steps: int = 0, - draft_kv_cache_manager: Optional["BaseResourceManager"] = None, - ugpu_ids: Optional[List[int | None]] = None, - ): - if ugpu_ids is not None: - assert len(ugpu_ids) == len(request_ids), ( - "ugpu_ids must have the same length as request_ids" - ) - - _kv_draft = ( - kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens - ) - - beam_width = max_beam_width - requests = [] - - def release_resources( - current_request: LlmRequest, free_draft_resources: bool = False - ) -> None: - for req in requests: - self.free_resources(req) - self.free_resources(current_request) - if draft_kv_cache_manager is not None: - for req in requests: - draft_kv_cache_manager.free_resources(req) - if free_draft_resources: - draft_kv_cache_manager.free_resources(current_request) - - for i, req_id in enumerate(request_ids): - # exact choice of n can be ignored for dummy requests - sampling_params = SamplingParams( - n=beam_width, best_of=beam_width, use_beam_search=beam_width > 1 - ) - # Here 1+max_num_draft_tokens is used to extend the prompt length to - # a non-zero number to skip illegal memory access issue in MLA kernel - # during warmup. - token_num = token_nums[i] if token_nums is not None else 1 + max_num_draft_tokens - if self._has_cp_helix: - # Keep the frozen dummy fields self-consistent (the active - # rank's past_seen = seqlen - 1 must stay >= 1). - token_num = max(token_num, 2) - # token_num - 1 is the past history length in generation. - history_hint = max(0, token_num - 1) if is_gen else None - encoder_output_len = encoder_output_lens[i] if encoder_output_lens is not None else None - encoder_input_tokens = ( - [1] * encoder_output_len if encoder_output_len is not None else None - ) - # Using 1 instead of 0 prevents NaN during warmup in e.g. Deepseek - input_tokens = [1 for _ in range(token_num)] - request_ugpu_id = ( - self.pick_ugpu(req_id) if ugpu_ids is None or ugpu_ids[i] is None else ugpu_ids[i] - ) - req = LlmRequest( - request_id=req_id, - max_new_tokens=1, - input_tokens=input_tokens, - sampling_config=SamplingConfig(sampling_params._get_sampling_config()), - is_streaming=False, - encoder_input_tokens=encoder_input_tokens, - ugpu_id=request_ugpu_id, - encoder_output_len=encoder_output_len, - ) - req.is_dummy_request = True - req.paged_kv_block_ids = [] - if prepare_resource: - # Dummy/warmup request. ``stop_committing()`` below blocks all - # writes to the radix tree, so the choice of branch does not - # affect committed state. ``cache_salt`` is left defaulted - # to None to avoid coupling synthetic data to any salted branch. - ugpu_id = self._request_ugpu_id(req) - kv_cache = self._create_kv_cache( - req.py_request_id, - req.lora_task_id, - input_tokens, - is_dummy=req.is_dummy, - ugpu_id=ugpu_id, - ) - # Saturated IndexMapper (e.g. disagg gen trans in progress) - # returns None; retry next iter. - if kv_cache is None: - logger.warning_once( - f"Unable to create KV cache for dummy request " - f"{req.py_request_id}; IndexMapper capacity is " - f"temporarily exhausted", - key="dummy_kv_cache_create_failed", - ) - release_resources(req) - return None - assert _get_kv_cache_ugpu_id(kv_cache) == ugpu_id - assert kv_cache.num_committed_tokens == 0 - success = kv_cache.resume(self._stream.cuda_stream) - if not success: - logger.warning_once( - f"Unable to resume KV cache for dummy request {req.py_request_id}", - key="dummy_kv_cache_resume_failed", - ) - release_resources(req) - return None - kv_cache.stop_committing() - dummy_capacity = ( - token_num - + self.num_extra_kv_tokens - + num_extra_decoding_steps - + self._standalone_draft_reserve - ) - if is_gen: - kv_cache.enable_swa_scratch_reuse = False - # Need to hint the committed history to activate stale-block - # optimization and match the solver's pool budget. - success = kv_cache.resize(dummy_capacity, history_length=history_hint) - if not success: - logger.warning_once( - f"Unable to resize KV cache for dummy request " - f"{req.py_request_id} to capacity {dummy_capacity} " - f"(history_length={history_hint})", - key="dummy_kv_cache_resize_failed", - ) - release_resources(req) - return None - draft_kv_cache = None - if draft_kv_cache_manager is not None: - ugpu_id = self._request_ugpu_id(req) - draft_kv_cache = draft_kv_cache_manager._create_kv_cache( - req.py_request_id, - req.lora_task_id, - input_tokens, - is_dummy=req.is_dummy, - ugpu_id=ugpu_id, - ) - # Dummy path: see comment above, no salt. - if draft_kv_cache is None: - logger.warning_once( - f"Unable to create draft KV cache for dummy " - f"request {req.py_request_id}", - key="dummy_draft_kv_cache_create_failed", - ) - release_resources(req) - return None - assert _get_kv_cache_ugpu_id(draft_kv_cache) == ugpu_id - success = draft_kv_cache.resume(draft_kv_cache_manager._stream.cuda_stream) - if not success: - logger.warning_once( - f"Unable to resume draft KV cache for dummy " - f"request {req.py_request_id}", - key="dummy_draft_kv_cache_resume_failed", - ) - release_resources(req, free_draft_resources=True) - return None - draft_kv_cache.stop_committing() - success = draft_kv_cache.resize(dummy_capacity) - if not success: - logger.warning_once( - f"Unable to resize draft KV cache for dummy " - f"request {req.py_request_id} to capacity " - f"{dummy_capacity}", - key="dummy_draft_kv_cache_resize_failed", - ) - release_resources(req, free_draft_resources=True) - return None - - if is_gen: - req.state = LlmRequestState.GENERATION_IN_PROGRESS - req.prompt_len = token_num - 1 - req.py_prompt_len = req.prompt_len - req.py_draft_tokens = [1] * max_num_draft_tokens - if self._has_cp_helix: - # Frozen fields (V1 parity): dummies never pass the - # per-step derivation — last CP rank active, shared - # synthetic global length. - if self._helix_cp_rank == self._helix_cp_size - 1: - req.py_helix_is_inactive_rank = False - req.prompt_len = token_num - 1 - else: - req.py_helix_is_inactive_rank = True - req.prompt_len = token_num - req.py_prompt_len = req.prompt_len - req.seqlen_this_rank_cp = req.prompt_len - req.total_input_len_cp = token_num * self._helix_cp_size - 1 - req.py_decoding_iter = 1 - if prepare_resource: - new_capacity = kv_cache.capacity + _kv_draft + 1 - success = kv_cache.resize(new_capacity, history_length=history_hint) - if not success: - logger.warning_once( - f"Unable to extend generation dummy request " - f"{req.py_request_id} to capacity {new_capacity} " - f"(history_length={history_hint})", - key="generation_dummy_kv_cache_resize_failed", - ) - release_resources(req, free_draft_resources=draft_kv_cache is not None) - return None - if draft_kv_cache is not None: - success = draft_kv_cache.resize(new_capacity) - if not success: - logger.warning_once( - f"Unable to extend generation draft dummy " - f"request {req.py_request_id} to capacity " - f"{new_capacity}", - key="generation_dummy_draft_kv_cache_resize_failed", - ) - release_resources(req, free_draft_resources=True) - return None - - if use_mrope: - _populate_dummy_mrope_config(req, token_num, is_gen) - requests.append(req) - - return requests - - def try_commit_blocks(self, request: LlmRequest) -> None: - should_block_reuse = ( - self.enable_block_reuse - and self._can_publish_block_reuse - and not request.is_dummy_request - and not self._draft_context_exceeds_limit(request) - ) - if not should_block_reuse: - return - - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is None: - return - - commit_end = request.context_current_position - if commit_end > kv_cache.num_committed_tokens: - tokens = self._augment_tokens_for_block_reuse( - self._reuse_token_source(request), - request, - start=kv_cache.num_committed_tokens, - end=commit_end, - ) - # TODO: On a disaggregated prefill server, pass is_end=True for - # the last context chunk to improve performance. - kv_cache.commit(tokens) - if request.context_remaining_length == 0: - kv_cache.stop_committing() - - def release_index_slot(self, request_id: int) -> None: - """Release IndexMapper slot early while keeping KV cache blocks allocated. - - After prefill completes on a context-only worker, the IndexMapper slot - (used for host_kv_cache_block_offsets during model forward) is no longer - needed. Releasing it early allows new requests to be scheduled while - the KV cache blocks are still being transferred via NIXL/UCX. - """ - kv_cache = self.kv_cache_map.get(request_id) - if self.is_draft and (kv_cache is None or request_id in self._early_freed_index_requests): - # The draft mirror only holds a slot for requests it actually - # mirrored, and the target may release the same request twice. - return - if kv_cache is not None: - for i in range(self.max_beam_width): - for pool_idx in range(self.num_pools): - kv_cache.set_base_page_index_buf(i, pool_idx, None) - self.index_mapper.remove_sequence(request_id) - self._early_freed_index_requests.add(request_id) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - if self.conversation_manager is not None: - self.conversation_manager.finish_request(request) - self._allocated_draft_lens.pop(request.py_request_id, None) - if self.draft_layer_ids: - self.draft_history.pop(request.py_request_id, None) - self._draft_dummy_request_ids.discard(request.py_request_id) - self._pending_overlap_slack.pop(request.py_request_id, None) - kv_cache = self.kv_cache_map.pop(request.py_request_id, None) - if kv_cache is None: - self.impl.clear_stats_excluded(request.py_request_id) - return - kv_cache.discard_pending_stats() - kv_cache.close() - self.impl.clear_stats_excluded(request.py_request_id) - if request.py_request_id in self._early_freed_index_requests: - self._early_freed_index_requests.discard(request.py_request_id) - else: - self.index_mapper.remove_sequence(request.py_request_id) - - def get_layer_page_index_scale(self, layer_idx: int) -> int: - """Page-index scale of this layer's KV buffer. Layers in one pool can - have different scales (e.g. different head_dim), so per-layer callers - must not use the pool-level scale.""" - return int(self.impl.get_page_index_scale(self.layer_offsets[layer_idx], Role.KEY)) - - def get_batch_cache_indices( - self, - request_ids: List[int], - layer_idx: Optional[int] = None, - num_blocks_per_seq: Optional[Sequence[int]] = None, - *, - raw_indices: bool = False, - ) -> List[List[int]]: - """Return cache indices for a layer, or pool 0 when no layer is given. - - Set raw_indices for slot-major views that need base slot IDs without - page-index scaling or KV aggregation. - """ - if layer_idx is None: - pool_id = 0 - index_scale = 1 if raw_indices else None - else: - pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] - index_scale = 1 if raw_indices else self.get_layer_page_index_scale(layer_idx) - return self._get_batch_cache_indices_by_pool_id( - request_ids, - pool_id=pool_id, - is_kv_aggregate=not raw_indices, - num_blocks_per_seq=num_blocks_per_seq, - index_scale=index_scale, - kv_factor=self.get_layer_kv_factor(layer_idx) if layer_idx is not None else None, - ) - - def _get_batch_cache_indices_by_pool_id( - self, - request_ids: List[int], - *, - pool_id: int = 0, - is_kv_aggregate: bool = True, - num_blocks_per_seq: Optional[Sequence[int]] = None, - index_scale: Optional[int] = None, - kv_factor: Optional[int] = None, - ) -> List[List[int]]: - if is_kv_aggregate: - # Div by kv_factor to index kv cache with size - # [num_blocks, kv_factor, tokens_per_block, num_kv_heads, head_dim] - div_factor = ( - self.kv_factor_per_layer[int(self.impl.layer_grouping[pool_id][0])] - if kv_factor is None - else kv_factor - ) - else: - div_factor = 1 - - if index_scale is None: - index_scale = int(self.index_scales[pool_id]) - res = [] - - for req_idx, req_id in enumerate(request_ids): - kv_cache = self.kv_cache_map[req_id] - # Zero-copy page-index buffers are padded to max_blocks_per_seq. - # Only convert blocks owned by this request; attention callers - # discard the padded tail immediately. - num_blocks = kv_cache.num_blocks - if num_blocks_per_seq is not None: - num_blocks = min(num_blocks, num_blocks_per_seq[req_idx]) - base_page_indices = kv_cache.get_base_page_indices(pool_id)[:num_blocks] - res.append( - [ - base_page_index * index_scale // div_factor - if base_page_index != BAD_PAGE_INDEX - else BAD_PAGE_INDEX - for base_page_index in base_page_indices - ] - ) - - return res - - def get_batch_cache_indices_flat( - self, - request_ids: List[int], - num_blocks: List[int], - layer_idx: Optional[int] = None, - ) -> torch.Tensor: - """Concatenated per-request block tables, trimmed to real widths. - - Equivalent to concatenating - ``get_batch_cache_indices(request_ids, layer_idx)[i][:num_blocks[i]]`` - over all requests, but never materializes the padded-to-capacity - per-request lists: the page indices are host data maintained in place - by the KV cache, so only ``num_blocks[i]`` entries per request are - gathered and a single vectorized transform runs over the result. - - Returns a CPU int32 tensor (pinned when supported) ready for an - async H2D copy. - """ - if layer_idx is None: - pool_id = 0 - scale = self._index_scale_ints[pool_id] - else: - pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] - # Layers sharing a pool can still require different page-index - # scales (e.g. Gemma4 sliding/global head_dim). Use the per-layer - # scale so this flat block table matches get_batch_cache_indices() - # and never feeds out-of-range page ids to FlashInfer. - scale = self.get_layer_page_index_scale(layer_idx) - div_factor = ( - self.get_layer_kv_factor(layer_idx) - if layer_idx is not None - else self.kv_factor_per_layer[int(self.impl.layer_grouping[pool_id][0])] - ) - - out_tensor = torch.empty(sum(num_blocks), dtype=torch.int32, pin_memory=prefer_pinned()) - out = out_tensor.numpy() - offset = 0 - for req_id, n in zip(request_ids, num_blocks): - out[offset : offset + n] = np.frombuffer( - self.kv_cache_map[req_id].get_base_page_indices(pool_id), - dtype=np.int32, - count=n, - ) - offset += n - - # One batched transform over the real widths; BAD_PAGE_INDEX entries - # (e.g. evicted out-of-window SWA blocks) stay untouched, matching - # get_batch_cache_indices. - valid = out != BAD_PAGE_INDEX - np.copyto(out, out * scale // div_factor, where=valid) - return out_tensor - - def get_cache_bytes_per_token(self) -> int: - return sum( - self.get_layer_bytes_per_token(local_id, Role.ALL) - * (self._helix_cp_size if local_id in self.cp_as_tp_layer_ids else 1) - for local_id in range(len(self.dtype_per_layer)) - ) - - def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): - if data_role not in (Role.ALL, *self._get_buffer_roles_for_layer(local_layer_idx)): - return 0 - if self.dtype_per_layer[local_layer_idx] not in ( - DataType.FP8, - DataType.HALF, - DataType.BF16, - DataType.FLOAT, - DataType.NVFP4, - ): - raise ValueError(f"Cannot support {self.dtype_per_layer[local_layer_idx]} KV cache.") - - if data_role == Role.ALL: - kv_factor = self.kv_factor_per_layer[local_layer_idx] - elif data_role in [Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: - if data_role in [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: - assert self.dtype_per_layer[local_layer_idx] == DataType.NVFP4, ( - "NVFP4 is the only supported dtype for block quant data roles" - ) - if data_role == Role.VALUE: - assert self.kv_factor_per_layer[local_layer_idx] == 2, ( - "VALUE data role is not supported for SELFKONLY cache type" - ) - kv_factor = 1 - else: - raise ValueError(f"Invalid data role: {data_role}") - - cache_size_per_token = ( - kv_factor - * self.num_kv_heads_per_layer[local_layer_idx] - * self.head_dim_per_layer[local_layer_idx] - ) - - cache_size_bytes_per_token = get_size_in_bytes( - cache_size_per_token, self.dtype_per_layer[local_layer_idx] - ) - - if data_role in [Role.KEY, Role.VALUE]: - return cache_size_bytes_per_token - - quant_size_per_token = 0 - - if self.dtype_per_layer[local_layer_idx] == DataType.NVFP4: - quant_size_per_token = self.calculate_scaling_factor_size_bytes( - cache_size_per_token, - quant_vector_size=16, - scaling_factor_dtype=DataType.FP8, - ) - - if data_role in [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: - return quant_size_per_token - - # Role.ALL combines both - return cache_size_bytes_per_token + quant_size_per_token - - @staticmethod - def calculate_scaling_factor_size_bytes( - cache_size: int, quant_vector_size: int, scaling_factor_dtype: DataType - ) -> int: - assert cache_size % quant_vector_size == 0, ( - "NVFP4 cache size must be divisible by quant vector size" - ) - return get_size_in_bytes(cache_size // quant_vector_size, scaling_factor_dtype) - - def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: - pool_handled = set() - for layer_id, layer_offset in self.layer_offsets.items(): - pool_id = self.layer_to_pool_mapping_dict[layer_offset] - if pool_id in pool_handled: - continue - for buffer in self._get_buffer_views_for_invalid_value_check(layer_id): - yield buffer - pool_handled.add(pool_id) - - @torch.inference_mode() - 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() - ) - - for buffer in self._iter_cache_buffers_for_invalid_check(): - if torch.is_floating_point(buffer) or torch.is_complex(buffer): - # Process page-major physical views in bounded chunks. Packed - # integer cache data cannot contain NaN or Inf. - 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_() - 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) - - def shutdown(self): - for kv_cache in self.kv_cache_map.values(): - kv_cache.close() - self.kv_cache_map.clear() - if self.draft_layer_ids: - self.draft_history.clear() - self._draft_dummy_request_ids.clear() - self.impl.shutdown() - if self.conversation_manager is not None: - self.conversation_manager.clear() - - def get_max_resource_count(self) -> int: - # TODO: implement this - return 1 - - def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - # TODO: implement this - # context_token_count = request.orig_prompt_len - # num_context_blocks = context_token_count // self.tokens_per_block - # remaining_tokens = context_token_count + request.max_new_tokens - num_context_blocks * self.tokens_per_block - # need_blocks = num_context_blocks + math.ceil( - # remaining_tokens / self.tokens_per_block) - # return need_blocks - return 0 - - # TODO: refactor get_cache_size_per_token and get_cache_bytes_per_token to use the same logic - @classmethod - def get_cache_size_per_token( - cls, - model_config: ModelConfigPython, - mapping: Mapping, - num_layers: Optional[int] = None, - draft_scratch_tokens: int = 0, - draft_window_size: int | None = None, - **kwargs, - ): - layer_sizes, attention_windows = _get_static_cache_size_layer_components( - model_config, mapping, num_layers=num_layers, **kwargs - ) - # Upstream takes these as named parameters; here they ride **kwargs. - kv_cache_config = kwargs.get("kv_cache_config") - reuse_backoff_enabled = ( - kv_cache_config is not None - and kv_cache_config.enable_block_reuse - and cls._supports_reuse_match_backoff - ) - if reuse_backoff_enabled: - from ..speculative import draft_prompt_lookahead - - # Mirror the runtime W+D retention window before page rounding. - backoff = draft_prompt_lookahead(kwargs.get("spec_config")) or 0 - attention_windows = _extend_swa_windows_for_reuse( - attention_windows, - backoff, - kwargs.get("max_seq_len"), - ) - draft_layer_sizes = _get_draft_layer_sizes(**kwargs) - if draft_layer_sizes: - _, headroom = _get_generation_kv_capacity( - kwargs.get("spec_config"), is_draft=kwargs.get("is_draft", False) - ) - headroom += draft_scratch_tokens - context, generation, fixed = _estimate_cache_size_components( - layer_sizes, - attention_windows, - kwargs["tokens_per_block"], - scratch=bool(kv_cache_config and kv_cache_config.enable_swa_scratch_reuse), - generation_capacity_headroom=headroom, - standalone_draft_reserve=draft_scratch_tokens, - ) - draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( - draft_layer_sizes, - kwargs["tokens_per_block"], - scratch_tokens=draft_scratch_tokens, - window_size=draft_window_size, - generation_capacity_headroom=headroom, - helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, - ) - return ( - generation + draft_generation, - (fixed + draft_fixed) * int(kwargs.get("max_batch_size") or 0) - + (context + draft_context - generation - draft_generation) - * int(kwargs.get("max_num_tokens") or 0), - ) - full_attn_size_per_token = _estimate_full_attn_size_per_token( - layer_sizes, attention_windows - ) - swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( - layer_sizes, - attention_windows, - kwargs["tokens_per_block"], - context=False, - scratch=False, - ) - max_batch_size = int(kwargs.get("max_batch_size") or 0) - return ( - full_attn_size_per_token + swa_size_per_token, - swa_size_per_request * max_batch_size, - ) - - def update_context_resources(self, scheduled_batch: ScheduledRequests): - """Update KV cache for context requests in the current batch. - - This is separated from update_resources (which handles generation - requests only) because the overlap executor needs context KV cache - updates to happen before next batch scheduling. Otherwise, the scheduler - would under-estimate available KV cache for sliding-window attention - layer. In non-overlap scheduler, you should call it together with - update_resources(). - """ - for req in scheduled_batch.context_requests: - if req.py_request_id not in self.kv_cache_map: - continue - kv_cache = self.kv_cache_map[req.py_request_id] - # In the overlap scheduler, iteration N+1's eviction may - # suspend a ctx request's KV cache while iteration N's - # update still needs to process it. The request will be resumed by - # the scheduler next iteration. - if not kv_cache.is_active: - continue - should_block_reuse = ( - self.enable_block_reuse - and self._can_publish_block_reuse - and not req.is_dummy_request - and not self._draft_context_exceeds_limit(req) - ) - is_all_reusable = self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE - should_resize = not should_block_reuse or not is_all_reusable - should_commit = is_all_reusable or req.context_remaining_length == 0 - - if should_resize: - self._resize_context_history(req, kv_cache, req.context_current_position) - if should_commit: - self.try_commit_blocks(req) - if req.context_remaining_length == 0: - if self.conversation_manager is not None: - self.conversation_manager.save_drop_plan(req, kv_cache) - # Scratch blocks are only for prefill chunks. Disable them at - # the context/generation boundary so generation uses normal KV - # pages before the first generation allocation. - kv_cache.enable_swa_scratch_reuse = False - - def _resize_context_history(self, req: LlmRequest, kv_cache, history_length: int) -> None: - """Advance a context cache's history marker, or raise.""" - if not kv_cache.resize(None, history_length): - raise ValueError( - "Failed to resize history length of KV cache for request " - f"{req.py_request_id} to {history_length} tokens at context update" - ) - - def update_resources( - self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None, - ): - if not self.is_draft: - _update_kv_cache_draft_token_location( - self, scheduled_batch, attn_metadata, kv_cache_dtype_byte_size - ) - # Context request KV cache updates are handled by - # update_context_resources, called separately from the executor loop. - for req in scheduled_batch.generation_requests: - if req.py_request_id not in self.kv_cache_map: - continue - kv_cache = self.kv_cache_map[req.py_request_id] - # In the overlap scheduler, the scheduler for iteration N+1 - # may suspend a gen request's KV cache (via self-eviction or - # victim eviction) while iteration N's update_resources still - # needs to process it. Skip suspended caches — the request - # will be resumed by the scheduler on the next iteration. - if not kv_cache.is_active: - continue - accepted_draft_len = req.py_num_accepted_draft_tokens - runtime_draft_len = req.py_rewind_len + accepted_draft_len - # No entry: this manager did not grow through the scheduler (e.g. - # disagg), so fall back to _generation_draft_slots' width. - allocated_draft_slots = self._allocated_draft_lens.get( - req.py_request_id, - max(runtime_draft_len, self._kv_reserve_draft_tokens) - if self.is_draft - else runtime_draft_len, - ) - # Reclaim every reserved slot that was not occupied by an accepted - # draft token. - rewind_len = max(allocated_draft_slots - accepted_draft_len, 0) - # _required_gen_capacity grows by 1 + 2 * slots; rewind_len gives - # back the unused half of the first slots term, this the second. - overlap_slack = self._pending_overlap_slack.pop(req.py_request_id, 0) - new_capacity = ( - None - if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) - else max( - kv_cache.capacity - rewind_len - overlap_slack, - req.max_beam_num_tokens - 1, - ) - ) - embedded_draft = KVCacheManagerV2._draft_history_uses_kv_history(self) - if self.draft_layer_ids and not embedded_draft and new_capacity is not None: - draft_history = self.get_draft_history(req.py_request_id) - if draft_history is not None: - # Rejected target verification slots do not revoke draft - # feature history. Retain it and its next-forward scratch. - new_capacity = max( - new_capacity, draft_history.position + self._standalone_draft_reserve - ) - history_length = ( - None - # Reuse (history's consumer) is disabled under helix, and - # max_beam_num_tokens mixes rank-local and global counts. - if self.kv_compression_manages_history or self._has_cp_helix - else req.max_beam_num_tokens - 1 - ) - if embedded_draft and self._has_cp_helix: - # HELIX stores rank-local prompt tokens but generated tokens - # are replicated. The cache ledger uses global positions. - history_length = ( - req.total_input_len_cp + req.max_beam_num_tokens - req.py_prompt_len - 1 - ) - if embedded_draft and new_capacity is not None and history_length is not None: - new_capacity = max(new_capacity, history_length) - success = kv_cache.resize(new_capacity, history_length) - if not success: - raise ValueError( - f"Failed to resize KV cache for request {req.py_request_id} " - f"to capacity {new_capacity} and history length " - f"{history_length} tokens at generation update" - ) - self._allocated_draft_lens.pop(req.py_request_id, None) - - def copy_batch_block_offsets( - self, - dst_tensor: torch.Tensor, - request_ids: List[int], - beam_width: int, - num_contexts: int, - num_seqs: int, - max_blocks: Optional[int] = None, - ): - # max_blocks is accepted for signature parity with KVCacheManager; the - # device-side copy op here already scales with allocated blocks only. - assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" - - copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) - assert copy_idx.shape[0] == num_seqs - - if self.enable_swa_scratch_reuse: - self._copy_batch_block_offsets_per_layer( - dst_tensor, request_ids, copy_idx, num_contexts, num_seqs - ) - return - - copy_batch_block_offsets_to_device( - self.host_kv_cache_block_offsets, - dst_tensor, - copy_idx, - self.index_scales, - self.kv_offset, - torch.cuda.current_stream().cuda_stream, - ) - - @staticmethod - def _derive_reuse_salt(cache_salt: str | None) -> int | None: - """Derive ``ReuseScope.salt`` (int|None) from the ``cache_salt`` string. - - Deterministic so the same string yields the same reuse namespace across - processes (matches C++ blockKey hashing on cacheSalt). Shared by cache - creation, prefetch, and the cache-aware router probe so they all hit the - same radix-tree namespace. - """ - if cache_salt is None: - return None - digest = hashlib.sha256(cache_salt.encode("utf-8")).digest() - return int.from_bytes(digest[:8], "little") - - def _create_kv_cache( - self, - request_id: int, - lora_task_id: int | None, - input_tokens: Sequence[TokenIdExt] | None, - *, - cache_salt: str | None = None, - is_dummy: bool = False, - expected_prompt_length: int | None = None, - ugpu_id: int | None = None, - ): - assert request_id not in self.kv_cache_map, ( - f"KV cache for request {request_id} already exists" - ) - if self.num_ugpus > 1: - assert ugpu_id is not None, ( - "Localized KV cache managers require a concrete ugpu_id for KV-cache creation" - ) - assert 0 <= ugpu_id < self.num_ugpus, ( - f"ugpu_id must be in [0, {self.num_ugpus}), got {ugpu_id}" - ) - else: - assert ugpu_id is None, "Non-localized KV cache managers require ugpu_id=None" - if self.index_mapper.num_free_slots() == 0: - logger.warning( - "No free IndexMapper slots for request %s " - "(%d/%d slots in use, likely held by DISAGG_GENERATION_TRANS_IN_PROGRESS requests). " - "Skipping KV cache creation; request will retry next iteration.", - request_id, - self.index_mapper.size(), - self.index_mapper.size(), - ) - return None - salt_int = self._derive_reuse_salt(cache_salt) - kv_cache = _create_backend_kv_cache( - self.impl, - lora_task_id, - salt_int, - input_tokens, - request_id, - expected_prompt_length, - ugpu_id, - ) - self.kv_cache_map[request_id] = kv_cache - if is_dummy: - self.impl.mark_stats_excluded(request_id) - kv_cache.discard_pending_stats() - if self.draft_layer_ids: - self._draft_dummy_request_ids.add(request_id) - index = self.index_mapper.add_new_sequence(request_id) - for i in range(self.max_beam_width): - for pool_idx in range(self.num_pools): - buffer: torch.Tensor = self.host_kv_cache_block_offsets[ - pool_idx, index * self.max_beam_width + i, 0 - ] - kv_cache.set_base_page_index_buf(i, pool_idx, memoryview(buffer.numpy())) - return kv_cache - - def probe_prefix_match_length(self, input_tokens, lora_task_id=None, cache_salt=None): - """Probe the v2 KV cache radix tree for prefix match length. - - Returns the number of prefix tokens already cached on this rank, - without acquiring page ownership. Mirrors the v1 ``KVCacheManager`` - adapter so ``KVCacheAwareADPRouter`` works on both KV cache backends. - - ``cache_salt`` (and ``lora_task_id``) must match the values used by - ``_create_kv_cache``; the salt is derived from the ``cache_salt`` - string the same way, so the probe queries the same reuse namespace. - Otherwise the router would see an incorrect match length. - """ - if not self.enable_block_reuse: - return 0 - if len(input_tokens) == 0: - return 0 - salt_int = self._derive_reuse_salt(cache_salt) - return self.impl.probe_reuse( - ReuseScope(lora_id=lora_task_id, salt=salt_int), - input_tokens, - ) - - def prefetch_for_context_tokens(self, requests: list) -> bool: - """Prefetch radix-tree blocks from disk→host for upcoming context requests. - - Returns True if all prefetches succeeded, False if any failed. - """ - if not self.enable_block_reuse: - return False - # Prefetch via a transient KV cache that holds the reuse-matched blocks, - # prefetches disk->host, then closes. Holding blocks costs no GPU space - # (never resumed) and close() needs no stream sync. The transient cache - # is NOT registered in kv_cache_map / IndexMapper. - success = True - for req in requests: - tokens = self._context_reuse_tokens(req) - # Use the same salt derivation as _create_kv_cache so the transient - # cache hits the same radix-tree blocks. - salt_int = self._derive_reuse_salt(req.cache_salt) - kv_cache = self.impl.create_kv_cache( - ReuseScope(lora_id=req.lora_task_id, salt=salt_int), tokens - ) - # Prefetch to the first tier below GPU (host if present, otherwise - # disk). prefetch() is a best-effort hint either way. - if not kv_cache.prefetch(CACHE_LEVEL1): - logger.warning("prefetch failed for request %s", req.py_request_id) - success = False - kv_cache.close() - return success - - def reset_reuse_state(self): - self.impl.clear_reusable_blocks() - if self.conversation_manager is not None: - self.conversation_manager.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py deleted file mode 100644 index 26139027838b..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ /dev/null @@ -1,4723 +0,0 @@ -# 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 math -import os -from abc import ABC, abstractmethod -from dataclasses import dataclass, replace -from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, - Optional, Protocol, Sequence, Tuple, Union, cast) - -import torch -import triton -import triton.language as tl - -if TYPE_CHECKING: - from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata - from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig - -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( - BlockReusePolicy, KVCacheManagerV2, Role, - _estimate_draft_cache_size_components, _get_draft_layer_sizes) -from tensorrt_llm._torch.pyexecutor.kv_cache_stats import \ - KVCacheV2IterationStatsReport -from tensorrt_llm._torch.pyexecutor.llm_request import ( - ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest) -from tensorrt_llm._torch.pyexecutor.resource_manager import ( - BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, - PoolConfiguration, get_pp_layers) -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, - nvtx_range, prefer_pinned, - torch_dtype_to_binding) -from tensorrt_llm.bindings.internal.batch_manager import ( - LinearAttentionMetadata, LinearCacheType) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, - BatchDesc, BufferConfig, - DataRole, KVCacheDesc) -from tensorrt_llm.runtime.kv_cache_manager_v2 import \ - KVCacheManagerConfig as KVCacheManagerConfigPy -from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, - SsmLayerConfig, - TokenIdExt, _KVCache) - -# Shared with the KV budget estimator so allocator and budgeting can never -# diverge on the sharding rule (config_utils is import-cycle-free). -from .config_utils import mamba_effective_tp_size as _mamba_effective_tp_size - -GB = 1 << 30 - - -class _PrefixReuseDiagnostics(Protocol): - - def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: - ... - - def _get_num_reusable_tokens_before_pruning(self) -> int: - ... - - -# Replay kernels pad the token/window dimension to at least 16 for tensor-core -# tiles, so history sizes below 16 are no faster when not writing and do -# expensive checkpoint writes more often. This minimum tensor-core tile size is -# present in all tensor-core generations, with newer GPUs adding larger tile -# operations, not smaller ones. In the other direction, history sizes above 16 -# currently pad to 32, which is substantially slower with the current kernel -# design. Keep the default floor at 16 while still allowing larger T, which -# degenerates to checkpointing every step instead of rejecting the request. If -# larger T becomes common, we could consider padding larger histories to powers -# of 2 or multiples of 16, but this is untested. For smaller T, we could explore -# combining larger histories with new kernel designs that stay efficient when -# the window is only partly full. -MIN_REPLAY_HISTORY_SIZE = 16 - -_KDA_BETA_CACHE_ALIGNMENT_BYTES = 16 -# KVCM copies 16-byte payloads; the replay kernel reads only the first int32. -_KDA_REPLAY_COUNT_BYTES = 16 - - -def _allocate_kda_beta_cache( - shape: Tuple[int, ...], - device: Optional[torch.device], -) -> torch.Tensor: - """Allocate a logically shaped bf16 beta cache with aligned physical rows.""" - logical_num_heads = shape[-1] - heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // - torch.bfloat16.itemsize) - padded_num_heads = ((logical_num_heads + heads_per_alignment - 1) // - heads_per_alignment) * heads_per_alignment - # CuTe requires every nested cache view to be 16-byte aligned. Keep the - # logical head count while padding the physical stride (six fp32 heads are - # 24 bytes, so an unpadded row makes alternating views only 8-byte aligned). - return torch.zeros( - *shape[:-1], - padded_num_heads, - dtype=torch.bfloat16, - device=device, - )[..., :logical_num_heads] - - -def _get_num_cuda_graph_padding_dummy_slots( - spec_config: Optional["DecodingBaseConfig"], - max_batch_size: int, -) -> int: - """Return the number of persistent CUDA-graph padding dummy IDs. - - This is computed before ``ModelEngine`` exists and covers draft lengths - reachable at every batch size, including the zero-length acceptance-rate - fallback. ``ModelEngine._compute_dynamic_draft_len_mapping`` is created - later and covers only configured CUDA-graph batch sizes, so it cannot size - this persistent ID set. - """ - if spec_config is None: - return 1 - - draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) - spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) - supports_dynamic_draft_len = (spec_dec_mode is not None and hasattr( - spec_dec_mode, "support_dynamic_draft_len") - and spec_dec_mode.support_dynamic_draft_len()) - if draft_len_schedule and supports_dynamic_draft_len: - runtime_draft_lengths = set() - first_uncovered_batch_size = 1 - for batch_size_threshold, draft_len in draft_len_schedule.items(): - if first_uncovered_batch_size > max_batch_size: - break - if batch_size_threshold >= first_uncovered_batch_size: - runtime_draft_lengths.add(draft_len) - first_uncovered_batch_size = batch_size_threshold + 1 - if first_uncovered_batch_size <= max_batch_size: - runtime_draft_lengths.add(0) - else: - max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 - max_total_draft_tokens = (getattr(spec_config, "max_total_draft_tokens", - 0) or 0) - is_linear_tree = getattr( - spec_config, - "is_linear_tree", - max_draft_len == max_total_draft_tokens, - ) - static_draft_len = (max_draft_len - if is_linear_tree else max_total_draft_tokens) - runtime_draft_lengths = {static_draft_len or 0} - - if ((getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and - (getattr(spec_config, "acceptance_rate_threshold", 0) or 0) > 0): - runtime_draft_lengths.add(0) - return len(runtime_draft_lengths) - - -class MambaRole: - """V2 buffer roles owned only by the hybrid Mamba manager.""" - - SSM_STATE = DataRole("ssm_state") - CONV_STATE = DataRole("conv_state") - - -def get_tensor_size_bytes(tensor): - """Calculate tensor size in bytes.""" - if isinstance(tensor, torch.Tensor): - return tensor.element_size() * tensor.nelement() - elif isinstance(tensor, list): - return sum(get_tensor_size_bytes(t) for t in tensor) - return 0 - - -# Mamba SSM stochastic-rounding Philox seed plumbing. -# -# Both the replay kernel and flashinfer's selective_state_update consume a -# `rand_seed` int64 tensor. Historically the non-replay paths created this -# tensor via `torch.randint(..., (1,))` on every decode step, which (a) is -# non-deterministic across runs, (b) allocates a fresh CUDA tensor per call -# and is therefore unfriendly to CUDA-graph capture, and (c) has no notion -# of cache-slot identity (so slot reuse cannot rotate the stream). -# -# The functions below produce per-slot int64 seeds via SplitMix64 finalization -# of (counter, slot, rank). Adjacent inputs yield uncorrelated outputs so -# consecutive cache slots and consecutive request counters do not leave -# structural fingerprints in the Philox input stream. All outputs live in -# (0, 2**62) so they avoid Philox's degenerate seed=0 case while staying -# within int64. -_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 -_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 -_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 -_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D -_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 -_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 - - -def _splitmix64(x: int) -> int: - """SplitMix64 finalizer; pure function, no torch.""" - x = (x + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 30) - x = (x * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 27) - x = (x * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 31) - return x & _MAMBA_SSM_UINT64_MASK - - -def _compute_deterministic_mamba_seed(counter: int, slot: int, - rank_offset: int) -> int: - """Deterministic int64 seed in (0, 2**62) from (counter, slot, rank). - - Pure function (no RNG, no torch.randint). Identical inputs across - process invocations produce identical outputs, which is what the - acceptance criteria require for cross-run reproducibility. - """ - folded = (_MAMBA_SSM_SEED_BASE + counter * _MAMBA_SSM_SEED_MIX_COUNTER + - slot * _MAMBA_SSM_SEED_MIX_SLOT + - rank_offset * _MAMBA_SSM_SEED_MIX_RANK) - folded &= _MAMBA_SSM_UINT64_MASK - value = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK - if value == 0: - value = 1 - return value - - -def _allocate_mamba_seed_buffer(cache_size: int, rank_offset: int, - device: torch.device) -> torch.Tensor: - """Allocate a (cache_size,) int64 CUDA buffer of deterministic seeds. - - counter=0 at allocation time; per-slot reset on fresh request assignment - bumps a host-side counter and rewrites only the freshly-assigned slot. - """ - slot_seeds = [ - _compute_deterministic_mamba_seed(0, i, rank_offset) - for i in range(cache_size) - ] - return torch.tensor(slot_seeds, dtype=torch.int64, device=device) - - -def _host_slots_to_device(slots: List[int], - device: torch.device) -> torch.Tensor: - """Copy host slot ids to ``device`` without a host sync. - - A pageable copy would wait for the queued step. Each call stages through a - fresh pinned tensor, so no later host write can race the copy. - """ - return torch.tensor(slots, - dtype=torch.long, - pin_memory=device.type == "cuda" - and prefer_pinned()).to(device, non_blocking=True) - - -def _copy_from_host_snapshot(dst: torch.Tensor, src: torch.Tensor) -> None: - """Copy reusable host staging ``src`` into ``dst`` without a host sync. - - Disaggregated admission and the execution batch can both prepare in one - executor iteration, refilling ``src`` before the first copy has run. Each - copy therefore reads a fresh pinned snapshot; ``dst`` keeps its address. - """ - snapshot = torch.empty(src.shape, - dtype=src.dtype, - pin_memory=dst.device.type == "cuda" - and prefer_pinned()) - snapshot.copy_(src) - dst.copy_(snapshot, non_blocking=True) - - -def _mamba_rank_offset(mapping: Mapping) -> int: - """Distinct seed offset per (tp_rank, pp_rank) so independent ranks - don't draw identical streams when not coordinated.""" - return (mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + - mapping.rank * 1_009) - - -def use_py_mamba_cache_manager() -> bool: - """Check if PythonMambaCacheManager should be forced (agg mode override). - - Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. - - Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager - with PythonMambaCacheManager inside instead of the configured manager. - Disagg mode is unaffected — its compatibility routing selects Mixed or - Cpp based on the transceiver configuration. - """ - return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' - - -class ReplayStateUpdateMetadata(NamedTuple): - """Shared tensors and fixed sizes for replay state updates.""" - prev_num_accepted_tokens: torch.Tensor - cache_buf_idx: torch.Tensor - replay_step_width: int - replay_history_size: int - - -def _advance_replay_state( - replay_metadata: ReplayStateUpdateMetadata, - state_indices: torch.Tensor, - accepted_tokens: torch.Tensor, - is_dummy_request: Optional[torch.Tensor] = None, -) -> None: - """Advance replay bookkeeping after a speculative generation step.""" - slots = state_indices.long() - accepted_tokens = accepted_tokens.to( - replay_metadata.prev_num_accepted_tokens.dtype) - prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens[slots] - wrote_checkpoint = (prev_num_accepted_tokens + - replay_metadata.replay_step_width - > replay_metadata.replay_history_size) - next_num_accepted_tokens = torch.where( - wrote_checkpoint, - accepted_tokens, - prev_num_accepted_tokens + accepted_tokens, - ) - cache_buf_idx = replay_metadata.cache_buf_idx[slots] - next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx) - if is_dummy_request is not None: - next_num_accepted_tokens = torch.where( - is_dummy_request, - prev_num_accepted_tokens, - next_num_accepted_tokens, - ) - next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, - next_cache_buf_idx) - replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens - replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx - - -class BaseMambaCacheManager(ABC): - """Abstract interface for accessing mamba/recurrent state caches.""" - - @abstractmethod - def get_state_indices(self, *args, **kwargs) -> torch.Tensor: - """Return slot indices of each request. - - Shape: [max_batch_size] - """ - ... - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - """Return replay metadata tensors and fixed replay sizes.""" - return None - - @abstractmethod - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - """Return conv states for specific layer. - - Shape: [slot_size, conv_dim, d_conv - 1] - """ - ... - - @abstractmethod - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - """Return SSM states for specific layer. - - Shape: [slot_size, num_heads, head_dim, d_state] - """ - ... - - @abstractmethod - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - ... - - @abstractmethod - def is_speculative(self) -> bool: - ... - - @abstractmethod - def mamba_layer_cache( - self, layer_idx: int - ) -> Union['PythonMambaCacheManager.State', - 'PythonMambaCacheManager.SpeculativeState', None]: - ... - - -class PythonMambaCacheManager(BaseResourceManager): - """Pure-Python mamba state manager with speculative decoding support. - - Manages only mamba states (conv + SSM) using PyTorch tensors on GPU. - Supports speculative decoding and disaggregated serving. - """ - - @dataclass(frozen=True, kw_only=True) - class State: - """Base state container for Mamba cache.""" - conv: torch.Tensor - # Represents last state, or "two back" if prev_num_accepted_tokens is > 0 - temporal: torch.Tensor - - # Subclasses override to list fields shared across layers (not indexed by layer). - _SHARED_FIELDS = frozenset() - - def at_layer_idx(self, layer: int): - kwargs = {} - for k, v in vars(self).items(): - if v is None: - kwargs[k] = None - elif k in self._SHARED_FIELDS: - kwargs[k] = v - else: - kwargs[k] = v[layer] - return type(self)(**kwargs) - - @dataclass(frozen=True, kw_only=True) - class SpeculativeState(State): - """Speculative state with intermediate states for draft tokens. - - Supports three SSM update paths (only one set of tensors is - allocated): - - Legacy: caches full intermediate SSM states (intermediate_ssm) - - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) - - KDA replay: per-slot draft-token caches consumed by the fused - ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted - drafts and commits states in place (kda_conv_*, kda_*_cache) - """ - _SHARED_FIELDS = frozenset({ - "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" - }) - - # Allocated for the legacy and Mamba2-replay paths; None for the - # KDA replay path (the kernel commits conv windows in place). - intermediate_conv_window: torch.Tensor | None = None - - # Legacy path: full intermediate SSM states at each step - intermediate_ssm: torch.Tensor | None = None - - # KDA replay path (fused multi-token verify, kimi_linear). - # Pool invariant under this path: `temporal` holds the state after - # the LAST GOLDEN token; the accepted drafts recorded in - # prev_num_accepted_tokens are pending in these caches and are - # replayed by the kernel at the start of the next verify round. - # Extended conv caches [layers, slots, dim, (W-1) + num_spec] fp32, - # dim-contiguous (stride(dim) == 1): columns [0, W-1) are the - # committed raw-input window, tail columns the pending drafts' raw - # inputs. - kda_conv_q: torch.Tensor | None = None - kda_conv_k: torch.Tensor | None = None - kda_conv_v: torch.Tensor | None = None - # Post-processed per-draft quantities for replay: k - # [layers, slots, num_spec, H*K] bf16, gate [.., num_spec, H*K] fp32 - # (feeds the cumulative-gate exponents), v [.., num_spec, H*V] bf16, - # beta [.., num_spec, H] bf16. - kda_k_cache: torch.Tensor | None = None - kda_g_cache: torch.Tensor | None = None - kda_v_cache: torch.Tensor | None = None - kda_beta_cache: torch.Tensor | None = None - - # Replay path: compact double-buffered cache - # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). - # 0 means temporal saved state is actually the last state, not two back. - prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers - cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers - # Per-cache-slot Philox seeds. Replay bumps them in-place per launch so - # CUDA graph replay uses fresh SR draws without allocating RNG tensors. - # (cache,) int64 - shared across layers - mamba_ssm_rand_seed: torch.Tensor | None = None - old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) - old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) - # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. - old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 - old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 - - def __init__( - self, - d_state: int, - d_conv: int, - num_heads: int, - n_groups: int, - head_dim: int, - num_layers: int, - max_batch_size: int, - spec_state_size: int, - mapping: Mapping, - dtype: torch.dtype, - ssm_cache_dtype: torch.dtype, - layer_mask: Optional[List[bool]] = None, - speculative_num_draft_tokens: Optional[int] = None, - model_type: str = "nemotron_hybrid", - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - ) -> None: - - self.mamba_ssm_cache_dtype = ssm_cache_dtype - self.speculative_num_draft_tokens = speculative_num_draft_tokens - self.spec_state_size = spec_state_size - self._use_replay_state_update = use_replay_state_update - # KDA replay path (kimi_linear fused multi-token verify). Mutually - # exclusive with use_replay_state_update; requires speculative mode. - self._kda_replay_num_spec = kda_replay_num_spec - self._use_kda_replay_update = kda_replay_num_spec is not None - if self._use_kda_replay_update: - assert not use_replay_state_update, ( - "kda_replay_num_spec and use_replay_state_update are " - "mutually exclusive") - assert speculative_num_draft_tokens is not None, ( - "KDA replay caches require speculative decoding") - assert kda_replay_num_spec == speculative_num_draft_tokens, ( - f"KDA replay cache width ({kda_replay_num_spec}) must match " - f"the draft length ({speculative_num_draft_tokens}): the " - "fused verify kernel is compiled with a static NUM_SPEC") - self.replay_history_size: Optional[int] = None - self.replay_step_width: Optional[int] = None - # When True, allocate the per-slot Philox seed buffer even outside - # the replay path so the non-replay flashinfer SR kernel reads a - # persistent deterministic seed instead of a per-call torch.randint. - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - # Host-side counter bumped per fresh cache-slot assignment. Combined - # with slot index and rank offset to produce reproducible per-slot - # seed values. Starts at 0 so the post-init "reset" stream is - # disjoint from the counter=0 stream used at allocation time. - self._seed_request_counter = 0 - - # get tp size - tp_size = _mamba_effective_tp_size(mapping) - - # derive mamba parameters for conv and ssm states - d_inner = head_dim * num_heads - conv_dim = d_inner + 2 * n_groups * d_state - nheads = num_heads - - # check that can be partitioned - assert nheads % tp_size == 0, "nheads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - - # partition conv_dim and nheads - d_inner_local = d_inner // tp_size - ng_ds_local = n_groups * d_state // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - d_inner = d_inner // tp_size - - # Per-section dims for conv_state. - # Qwen3-Next: [Q | K | V] = [ng*ds, ng*ds, d_inner] - # Nemotron_hybrid: [x | B | C] = [d_inner, ng*ds, ng*ds] - if model_type == "qwen3_next": - self.conv_section_dims = [ng_ds_local, ng_ds_local, d_inner_local] - elif model_type == "nemotron_hybrid": - self.conv_section_dims = [d_inner_local, ng_ds_local, ng_ds_local] - else: - raise ValueError(f"Unsupported model type: {model_type}") - - # conv and ssm states device - device = torch.device("cuda") - - pp_layers, num_layers = get_pp_layers( - num_layers, - mapping, - layer_mask=layer_mask, - ) - num_local_layers = len(pp_layers) - self.mamba_layer_offsets = { - idx: offset - for offset, idx in enumerate(pp_layers) - } - - conv_state_shape = (conv_dim, d_conv - 1) - ssm_state_shape = (nheads, head_dim, d_state) - - # create mamba conv and ssm states - conv_states = torch.zeros( - size=(num_local_layers, max_batch_size) + conv_state_shape, - dtype=dtype, - device=device, - ) - - ssm_states = torch.zeros( - size=(num_local_layers, max_batch_size) + ssm_state_shape, - dtype=self.mamba_ssm_cache_dtype, - device=device, - ) - - # Per-slot Philox seeds. Allocated whenever stochastic rounding can - # fire, even outside the replay path and even when MTP/spec is off, - # so all consumers (replay kernel, MTP non-replay flashinfer, - # non-MTP flashinfer) read a persistent deterministic seed from the - # cache manager instead of calling torch.randint per forward. - self._mamba_ssm_rand_seed: Optional[torch.Tensor] = None - if (self._use_replay_state_update - or self._mamba_ssm_stochastic_rounding): - self._mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - max_batch_size, self._seed_rank_offset, device) - - # create state container - if speculative_num_draft_tokens is not None: - T = speculative_num_draft_tokens + 1 - self.replay_step_width = T - - # Conv intermediate cache — legacy and Mamba2-replay paths only. - # The KDA replay kernel commits conv windows in place, so the - # per-step window scratch is not needed. - intermediate_conv_window_cache = None - if not self._use_kda_replay_update: - intermediate_conv_window_cache = torch.zeros( - size=(num_local_layers, self.spec_state_size, T) + - conv_state_shape, - dtype=dtype, - device=device, - ) - - # SSM speculative cache — path-specific tensors - spec_kwargs = {} - # Share the manager-level seed buffer through SpeculativeState - # so the MTP path can still read it via layer_cache. - if self._mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self._mamba_ssm_rand_seed - if self._use_kda_replay_update: - # KDA replay caches for the fused multi-token verify kernel - # (trtllm::kda_mtp_decode). Per cache slot (persistent - # across rounds — the next round replays from them), unlike - # the batch-row-indexed intermediate buffers. - M = self._kda_replay_num_spec - # conv_dim covers the [q | k | v] short-conv sections; the - # kernel consumes them as three dim-contiguous caches. - assert conv_dim % 3 == 0, ( - "KDA replay caches expect the [q | k | v] conv-state " - "sectioning (3 equal sections)") - section_dim = conv_dim // 3 - # d_conv is KDA's convolution width; the live pool stores - # its W - 1 committed input columns. - w_kernel = d_conv - extended_s = w_kernel - 1 + M - - def _dim_contiguous_conv_cache(): - # raw bf16 projection inputs: bf16 is exact and halves - # the per-draft replay traffic of the verify kernel - return torch.zeros(num_local_layers, - max_batch_size, - extended_s, - section_dim, - dtype=torch.bfloat16, - device=device).transpose(-1, -2) - - spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( - max_batch_size, dtype=torch.int32, device=device) - spec_kwargs['kda_conv_q'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_conv_k'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_conv_v'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_k_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.bfloat16, - device=device) - spec_kwargs['kda_g_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.float32, - device=device) - spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.bfloat16, - device=device) - spec_kwargs['kda_beta_cache'] = _allocate_kda_beta_cache( - (num_local_layers, max_batch_size, M, nheads), device) - ssm_spec_cache = [ - spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], - spec_kwargs['kda_conv_v'], spec_kwargs['kda_k_cache'], - spec_kwargs['kda_g_cache'], spec_kwargs['kda_v_cache'], - spec_kwargs['kda_beta_cache'] - ] - spec_path_label = "kda-replay" - elif self._use_replay_state_update: - assert n_groups % tp_size == 0, \ - "replay state update requires n_groups divisible by tp_size" - n_groups_per_rank = n_groups // tp_size - self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) - - # Compact replay cache. - spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( - max_batch_size, dtype=torch.int32, device=device) - spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size, - dtype=torch.int32, - device=device) - spec_kwargs['old_x'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - self.replay_history_size, - nheads, - head_dim, - dtype=dtype, - device=device) - spec_kwargs['old_B'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - self.replay_history_size, - n_groups_per_rank, - d_state, - dtype=dtype, - device=device) - spec_kwargs['old_dt'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - nheads, - self.replay_history_size, - dtype=torch.float32, - device=device) - spec_kwargs['old_dA_cumsum'] = torch.zeros( - num_local_layers, - max_batch_size, - 2, - nheads, - self.replay_history_size, - dtype=torch.float32, - device=device) - ssm_spec_cache = [ - spec_kwargs['old_x'], spec_kwargs['old_B'], - spec_kwargs['old_dt'], spec_kwargs['old_dA_cumsum'] - ] - spec_path_label = "replay" - else: - # Legacy: full intermediate SSM states at each step - spec_kwargs['intermediate_ssm'] = torch.zeros( - size=(num_local_layers, self.spec_state_size, T) + - ssm_state_shape, - dtype=self.mamba_ssm_cache_dtype, - device=device) - ssm_spec_cache = [spec_kwargs['intermediate_ssm']] - spec_path_label = "legacy" - - self.mamba_cache = self.SpeculativeState( - conv=conv_states, - temporal=ssm_states, - intermediate_conv_window=intermediate_conv_window_cache, - **spec_kwargs, - ) - - logger.info( - f"Mamba Cache ({spec_path_label}) is allocated. " - f"max_mamba_cache_size: {max_batch_size}, " - f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " - f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB, " - f"ssm_spec_cache size: {get_tensor_size_bytes(ssm_spec_cache) / GB:.2f}GB, " - "intermediate_conv_window_cache size: " - f"{get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" - ) - else: - self.mamba_cache = self.State( - conv=conv_states, - temporal=ssm_states, - ) - - logger.info( - f"Mamba Cache is allocated. " - f"max_mamba_cache_size: {max_batch_size}, " - f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " - f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB" - ) - - # mamba cache available blocks - self.mamba_cache_free_blocks = [i for i in range(max_batch_size)] - - # mamba cache index, maps request_id -> state indices - self.mamba_cache_index: Dict[int, int] = {} - self._dummy_request_ids: set[int] = set() - # Batch-order mask aligned with state_indices; duplicate dummy request - # IDs mark every batch row even when they share one cache slot. - self._dummy_request_mask = torch.zeros(max_batch_size, - dtype=torch.bool, - device=device) - self._dummy_request_mask_host = torch.zeros(max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned()) - - # Permanent slot shared by every CUDA-graph padding sentinel id - # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per - # draft length). Pool sizing must include +1 headroom for this; - # see MixedMambaHybridCacheManager. - self._padding_slot: int = self.mamba_cache_free_blocks.pop() - - # Reserved slot for the attention-DP padding dummy. Keep it out of the - # free pool so dummy insertion never consumes real-request capacity. - self._attention_dp_dummy_slot: Optional[int] = ( - self.mamba_cache_free_blocks.pop() - if mapping.enable_attention_dp else None) - - # save intermediate state indices for requests - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device=device) - - # Physical tensor rows include reserved dummy slots. Resource capacity - # is the number of real request slots remaining after those - # reservations. - self._max_batch_size = max_batch_size - self._max_resource_count = len(self.mamba_cache_free_blocks) - - def get_max_resource_count(self) -> int: - """Return the maximum number of real requests that can be cached.""" - return self._max_resource_count - - def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: - """Return the prefix of *context_requests* that fits in the - available Mamba cache blocks. Requests that already have a - cached block do not consume a free block.""" - free = len(self.mamba_cache_free_blocks) - result = [] - for r in context_requests: - if r.py_request_id in self.mamba_cache_index: - result.append(r) - elif free > 0: - result.append(r) - free -= 1 - else: - break - return result - - def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - """For Mamba cache manager, we always need one slot per request.""" - return 1 - - @torch.inference_mode() - def _prepare_mamba_cache_blocks(self, request_ids: List[int]): - for r in request_ids: - if r in self.mamba_cache_index: - continue - if len(self.mamba_cache_free_blocks) == 0: - raise RuntimeError("run out of mamba cache blocks") - block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - # Fresh request: no drafts pending in the replay caches. - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - if self._mamba_ssm_rand_seed is not None: - # Deterministic per-slot rotation on fresh assignment. - # `block` is pulled from mamba_cache_free_blocks, which - # excludes _padding_slot by construction (see __init__), - # so padding sentinels never reach this branch. - self._seed_request_counter += 1 - self._mamba_ssm_rand_seed[block] = ( - _compute_deterministic_mamba_seed( - self._seed_request_counter, block, - self._seed_rank_offset)) - - @torch.inference_mode() - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - """Seed the fused-verify KDA replay conv caches from the conv pool. - - On a disaggregated generation server the ctx->gen transfer populates - only the base ``conv`` / ``temporal`` pools (see - ``disaggregation/resource/page.py``); the per-slot ``kda_conv_*`` - replay caches consumed by ``trtllm::kda_mtp_decode`` are normally - seeded by a *local* prefill/decode via - ``_sync_kda_replay_conv_window`` and would otherwise hold zeros (or a - previous occupant's window) for a transferred request — corrupting - the recurrent state from the first verify step onward. Call this - after the state transfer completes, before the first generation - forward. Mirrors ``_sync_kda_replay_conv_window``: the conv pool row - stores the committed ``W - 1`` raw inputs directly. Pending-draft - scratch is cleared (no drafts are pending for a freshly transferred - request). - """ - if not (self._use_kda_replay_update - and isinstance(self.mamba_cache, self.SpeculativeState)): - return - blocks = [ - self.mamba_cache_index[rid] for rid in request_ids - if rid in self.mamba_cache_index - ] - if not blocks: - return - conv = self.mamba_cache.conv # [L, slots, 3D, W - 1] - idx = torch.tensor(sorted(set(blocks)), - dtype=torch.long, - device=conv.device) - d = conv.shape[2] // 3 - committed = conv.shape[3] - cs = conv.index_select(1, idx) - for cache, section in ( - (self.mamba_cache.kda_conv_q, cs[:, :, :d]), - (self.mamba_cache.kda_conv_k, cs[:, :, d:2 * d]), - (self.mamba_cache.kda_conv_v, cs[:, :, 2 * d:]), - ): - # cache: [L, slots, D, committed + num_spec]; zero the draft - # tail columns and seed the committed window in one copy. - seeded = torch.zeros( - (cache.shape[0], idx.numel()) + cache.shape[2:], - dtype=cache.dtype, - device=cache.device) - seeded[:, :, :, :committed] = section.to(cache.dtype) - cache.index_copy_(1, idx, seeded) - for buf in (self.mamba_cache.kda_k_cache, self.mamba_cache.kda_g_cache, - self.mamba_cache.kda_v_cache, - self.mamba_cache.kda_beta_cache): - buf.index_fill_(1, idx, 0) - self.mamba_cache.prev_num_accepted_tokens[idx] = 0 - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - context_ids = [ - i.py_request_id for i in scheduled_batch.context_requests - ] - generation_ids = [ - i.py_request_id for i in scheduled_batch.generation_requests - ] - request_ids = context_ids + generation_ids - self._prepare_mamba_cache_blocks(request_ids) - self._refresh_dummy_request_mask([req.is_dummy for req in requests]) - - def _is_padding_sentinel(self, request_id: int) -> bool: - # cuda_graph_runner caches one dummy per runtime_draft_len value - # (see _get_padded_batch), so any id in the range of dummy request IDs - # may be live concurrently. - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ - CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl = self.speculative_num_draft_tokens or 0 - return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= - CUDA_GRAPH_DUMMY_REQUEST_ID) - - @torch.inference_mode() - def add_dummy_requests(self, request_ids: List[int], **kwargs): - # Sentinels alias to the shared _padding_slot; non-sentinel - # dummies (warmup, attention-DP idle padding) get their own - # slot and are freed individually. - if not request_ids: - return - self._dummy_request_ids.update(request_ids) - for r in request_ids: - if r in self.mamba_cache_index: - block = self.mamba_cache_index[r] - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - continue - if self._is_padding_sentinel(r): - block = self._padding_slot - elif (r == ATTENTION_DP_DUMMY_REQUEST_ID - and self._attention_dp_dummy_slot is not None): - block = self._attention_dp_dummy_slot - else: - if len(self.mamba_cache_free_blocks) == 0: - raise RuntimeError("run out of mamba cache blocks") - block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - - def free_resources(self, request: LlmRequest): - request_id = request.py_request_id - if request_id not in self.mamba_cache_index: - return - self._dummy_request_ids.discard(request_id) - block = self.mamba_cache_index.pop(request_id) - # Reserved slots must not re-enter the real-request free pool. - if block != self._padding_slot and \ - block != self._attention_dp_dummy_slot: - self.mamba_cache_free_blocks.append(block) - - def get_state_indices(self, request_ids: List[int], - is_padding: List[bool]) -> List[int]: - assert len(request_ids) == len(is_padding) - indices = [self.mamba_cache_index[rid] for rid in request_ids] - is_dummy = [ - rid in self._dummy_request_ids or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.as_tensor(is_dummy, dtype=torch.bool)) - self._dummy_request_mask.copy_(self._dummy_request_mask_host, - non_blocking=True) - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).conv - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).temporal - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if not isinstance(self.mamba_cache, self.SpeculativeState): - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).intermediate_ssm - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - """Get intermediate conv states for speculative decoding.""" - if not isinstance(self.mamba_cache, self.SpeculativeState): - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx( - layer_offset).intermediate_conv_window - - def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_x[layer_offset] - - def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_B[layer_offset] - - def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_dt[layer_offset] - - def get_replay_old_dA_cumsum(self, - layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_dA_cumsum[layer_offset] - - def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - return self.mamba_cache.cache_buf_idx - - def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - return self.mamba_cache.prev_num_accepted_tokens - - def is_speculative(self) -> bool: - return isinstance(self.mamba_cache, self.SpeculativeState) - - def mamba_layer_cache(self, - layer_idx: int) -> Union[State, SpeculativeState]: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.mamba_ssm_cache_dtype - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - """Return the persistent (cache_size,) int64 Philox seed buffer or - None when stochastic rounding is not active for this manager. - - Used by mamba2_mixer non-MTP paths that don't hold a SpeculativeState. - Callers must bump in-place (`.add_(1)` or slice-and-add) and pass a - view that matches the consuming kernel's expected shape. - """ - return self._mamba_ssm_rand_seed - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - if (not self._use_replay_state_update - or not isinstance(self.mamba_cache, self.SpeculativeState) - or self.mamba_cache.prev_num_accepted_tokens is None - or self.mamba_cache.cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=( - self.mamba_cache.prev_num_accepted_tokens), - cache_buf_idx=self.mamba_cache.cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def shutdown(self): - """Release tensor memory.""" - # Clear mamba cache states - empty = torch.tensor([]) - - def _drop(tensor): - return empty if tensor is not None else None - - if isinstance(self.mamba_cache, self.SpeculativeState): - self.mamba_cache = self.SpeculativeState( - conv=empty, - temporal=empty, - intermediate_conv_window=_drop( - self.mamba_cache.intermediate_conv_window), - intermediate_ssm=_drop(self.mamba_cache.intermediate_ssm), - prev_num_accepted_tokens=_drop( - self.mamba_cache.prev_num_accepted_tokens), - cache_buf_idx=_drop(self.mamba_cache.cache_buf_idx), - mamba_ssm_rand_seed=_drop(self.mamba_cache.mamba_ssm_rand_seed), - old_x=_drop(self.mamba_cache.old_x), - old_B=_drop(self.mamba_cache.old_B), - old_dt=_drop(self.mamba_cache.old_dt), - old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), - kda_conv_q=_drop(self.mamba_cache.kda_conv_q), - kda_conv_k=_drop(self.mamba_cache.kda_conv_k), - kda_conv_v=_drop(self.mamba_cache.kda_conv_v), - kda_k_cache=_drop(self.mamba_cache.kda_k_cache), - kda_g_cache=_drop(self.mamba_cache.kda_g_cache), - kda_v_cache=_drop(self.mamba_cache.kda_v_cache), - kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), - ) - else: - self.mamba_cache = self.State(conv=empty, temporal=empty) - - torch.cuda.empty_cache() - - @torch.compile(options={"max-autotune": True}) - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor, - accepted_leaf_positions: Optional[torch.Tensor] = None): - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - num_accepted_draft_tokens = num_accepted_tokens[ - num_contexts:num_contexts + num_gens] - 1 - # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. - accepted_positions = (accepted_leaf_positions if accepted_leaf_positions - is not None else num_accepted_draft_tokens) - state_indices_d = state_indices[num_contexts:num_contexts + num_gens] - src_state_indices = self.intermediate_state_indices[:num_gens] - - if self._use_kda_replay_update: - # KDA replay: the fused verify kernel already committed the SSM - # state and conv windows (in place, after the golden token) and - # cached this round's drafts. All that remains is recording how - # many of those drafts the sampler accepted, so the next round's - # kernel launch replays exactly that prefix. - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - prev = self.mamba_cache.prev_num_accepted_tokens - current = prev[state_indices_d] - accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) - prev[state_indices_d] = torch.where(is_dummy_request, current, - accepted) - return - - if self._use_replay_state_update: - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - replay_metadata = self.get_replay_state_update_metadata() - assert replay_metadata is not None - _advance_replay_state( - replay_metadata, - state_indices_d, - num_accepted_tokens[num_contexts:num_contexts + num_gens], - is_dummy_request, - ) - else: - # Legacy: copy accepted SSM state from intermediate cache. - ssm_states = self.mamba_cache.temporal - intermediate_ssm_cache = self.mamba_cache.intermediate_ssm - accepted_ssm_state = intermediate_ssm_cache[:, src_state_indices, - accepted_positions] - ssm_states[:, state_indices_d, :] = accepted_ssm_state - - # Conv: both paths save all intermediate conv windows, carry over the accepted one. - conv_states = self.mamba_cache.conv - intermediate_conv_window_cache = self.mamba_cache.intermediate_conv_window - accepted_conv_state = intermediate_conv_window_cache[:, - src_state_indices, - accepted_positions] - conv_states[:, state_indices_d, :] = accepted_conv_state - - -class MambaCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Facade for standalone mamba state management (no KV cache). - - Delegates to PythonMambaCacheManager. - """ - - def __init__( - self, - d_state: int, - d_conv: int, - num_heads: int, - n_groups: int, - head_dim: int, - num_layers: int, - max_batch_size: int, - spec_state_size: int, - mapping: Mapping, - dtype: torch.dtype, - ssm_cache_dtype: torch.dtype, - layer_mask: Optional[List[bool]] = None, - stream: Optional[torch.cuda.Stream] = None, - speculative_num_draft_tokens: Optional[int] = None, - model_type: str = "nemotron_hybrid", - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - ) -> None: - max_num_sequences = max_batch_size * mapping.pp_size - - self._impl = PythonMambaCacheManager( - d_state=d_state, - d_conv=d_conv, - num_heads=num_heads, - n_groups=n_groups, - head_dim=head_dim, - num_layers=num_layers, - max_batch_size=max_num_sequences, - spec_state_size=spec_state_size, - mapping=mapping, - dtype=dtype, - ssm_cache_dtype=ssm_cache_dtype, - layer_mask=layer_mask, - speculative_num_draft_tokens=speculative_num_draft_tokens, - model_type=model_type, - use_replay_state_update=use_replay_state_update, - mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - kda_replay_num_spec=kda_replay_num_spec, - ) - - def get_max_resource_count(self) -> int: - return self._impl.get_max_resource_count() - - def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: - return self._impl.filter_ctx_requests_by_capacity(context_requests) - - def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - return self._impl.get_needed_resource_to_completion(request) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - self._impl.prepare_resources(scheduled_batch) - - def free_resources(self, request: LlmRequest): - self._impl.free_resources(request) - - def add_dummy_requests(self, request_ids: List[int], **kwargs): - self._impl.add_dummy_requests(request_ids, **kwargs) - - def get_state_indices( - self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None - ) -> Union[torch.Tensor, List[int]]: - return self._impl.get_state_indices(request_ids, is_padding) - - @property - def mamba_cache_free_blocks(self) -> List[int]: - return self._impl.mamba_cache_free_blocks - - @property - def mamba_cache_index(self) -> Dict[int, int]: - return self._impl.mamba_cache_index - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self._impl.get_conv_states(layer_idx) - - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - self._impl.seed_kda_replay_caches_for_disagg_gen(request_ids) - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self._impl.get_ssm_states(layer_idx) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self._impl.get_mamba_ssm_cache_dtype() - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - """Delegate to the underlying Python manager.""" - return self._impl.get_mamba_ssm_rand_seed() - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - return self._impl.get_replay_state_update_metadata() - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_intermediate_ssm_states(layer_idx) - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_intermediate_conv_states(layer_idx) - - def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_x(layer_idx) - - def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_B(layer_idx) - - def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_dt(layer_idx) - - def get_replay_old_dA_cumsum(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_dA_cumsum(layer_idx) - - def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: - return self._impl.get_replay_cache_buf_idx() - - def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: - return self._impl.get_replay_prev_num_accepted_tokens() - - def is_speculative(self) -> bool: - return self._impl.is_speculative() - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - return self._impl.mamba_layer_cache(layer_idx) - - def shutdown(self): - self._impl.shutdown() - - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor, - accepted_leaf_positions: Optional[torch.Tensor] = None): - # Non-speculative configs don't allocate intermediate state; the - # promotion is a clean no-op. - if not self._impl.is_speculative(): - return - self._impl.update_mamba_states(attn_metadata, num_accepted_tokens, - state_indices, accepted_leaf_positions) - - -class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Shared interface and state plumbing for hybrid Mamba managers. - - Concrete storage, state-index, and resource lifecycles remain owned by the - Cpp and V2 implementations. - """ - - _supports_additional_snapshot_offsets = False - - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size: int) -> None: - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if spec_config is None or self.local_num_mamba_layers == 0: - return - - tokens_per_gen_step = spec_config.tokens_per_gen_step - if not self._use_replay_state_update: - self.intermediate_ssm_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, - tokens_per_gen_step - ] + self.ssm_state_shape, - dtype=self.ssm_state_dtype, - device="cuda", - ) - self.intermediate_conv_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step - ] + self.conv_state_shape, - dtype=self.conv_state_dtype, - device="cuda", - ) - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - - def _allocate_pool_replay_buffers( - self, - spec_config, - cache_size: int, - device: Optional[torch.device], - ) -> bool: - """Allocate replay tensors shared by the Cpp and V2 state pools.""" - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - - if (self.local_num_mamba_layers == 0 - or (not self._use_replay_state_update - and not self._mamba_ssm_stochastic_rounding)): - return False - - assert device is not None - self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - cache_size, self._seed_rank_offset, device) - if spec_config is None or not self._use_replay_state_update: - return False - - history_size = self.replay_history_size - assert history_size is not None - nheads, head_dim, d_state = self.ssm_state_shape - common_shape = [self.local_num_mamba_layers, cache_size, 2] - self.prev_num_accepted_tokens = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.cache_buf_idx = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.old_x = torch.zeros( - common_shape + [history_size, nheads, head_dim], - dtype=self.conv_state_dtype, - device=device, - ) - self.old_B = torch.zeros( - common_shape + [history_size, self._n_groups_per_rank, d_state], - dtype=self.conv_state_dtype, - device=device, - ) - self.old_dt = torch.zeros( - common_shape + [nheads, history_size], - dtype=torch.float32, - device=device, - ) - self.old_dA_cumsum = torch.zeros( - common_shape + [nheads, history_size], - dtype=torch.float32, - device=device, - ) - return True - - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - if self._dummy_request_mask is None: - return - - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.tensor(is_dummy, dtype=torch.bool)) - _copy_from_host_snapshot(self._dummy_request_mask, - self._dummy_request_mask_host) - - def _reset_context_mamba_slots(self, num_contexts: int) -> None: - if num_contexts == 0: - return - - context_slots = self.cuda_state_indices[:num_contexts].long() - if (self._use_replay_state_update - and self.prev_num_accepted_tokens is not None - and self.cache_buf_idx is not None): - self.prev_num_accepted_tokens[context_slots] = 0 - self.cache_buf_idx[context_slots] = 0 - if self.old_x is not None: - self.old_x[:, context_slots] = 0 - if self.old_B is not None: - self.old_B[:, context_slots] = 0 - if self.old_dt is not None: - self.old_dt[:, context_slots] = 0 - if self.old_dA_cumsum is not None: - self.old_dA_cumsum[:, context_slots] = 0 - - if self.mamba_ssm_rand_seed is None: - return - self._seed_request_counter += 1 - counter = self._seed_request_counter - rank_offset = self._seed_rank_offset - host_slots = self._host_state_indices[:num_contexts].tolist() - new_seeds = [ - _compute_deterministic_mamba_seed(counter, slot, rank_offset) - for slot in host_slots - ] - self.mamba_ssm_rand_seed[context_slots] = torch.tensor( - new_seeds, - dtype=torch.int64, - device=self.mamba_ssm_rand_seed.device, - ) - - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - """Set reusable Mamba snapshot boundaries before scheduling.""" - if not self.enable_block_reuse: - for request in requests: - request.expect_snapshot_points = [] - return - - state_config = self.kv_cache_config.mamba_state_config - interval = state_config.periodic_snapshot_interval - for request in requests: - snapshot_points = set() - if interval is not None and interval > 0: - snapshot_points.update( - range(interval, request.prompt_len + 1, interval)) - if self._supports_additional_snapshot_offsets: - for offset in state_config.additional_snapshot_offsets_from_start: - if offset <= request.prompt_len: - snapshot_points.add(offset) - for offset in state_config.additional_snapshot_offsets_from_end: - point = request.prompt_len - offset - if point > 0: - snapshot_points.add(point) - request.expect_snapshot_points = sorted(snapshot_points) - - def is_speculative(self) -> bool: - return self.spec_config is not None - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_ssm_states is None: - return None - return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_conv_states is None: - return None - return self.intermediate_conv_states[ - self.mamba_layer_offsets[layer_idx]] - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - conv = self.get_conv_states(layer_idx) - ssm = self.get_ssm_states(layer_idx) - if self.spec_config is not None: - layer_offset = self.mamba_layer_offsets[layer_idx] - spec_kwargs = {} - if self.mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed - if self._use_replay_state_update: - spec_kwargs['old_x'] = self.old_x[layer_offset] - spec_kwargs['old_B'] = self.old_B[layer_offset] - spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] - spec_kwargs['cache_buf_idx'] = self.cache_buf_idx - spec_kwargs['prev_num_accepted_tokens'] = ( - self.prev_num_accepted_tokens) - else: - spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ - layer_offset] - return PythonMambaCacheManager.SpeculativeState( - conv=conv, - temporal=ssm, - intermediate_conv_window=self. - intermediate_conv_states[layer_offset], - **spec_kwargs, - ) - return PythonMambaCacheManager.State(conv=conv, temporal=ssm) - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', - None) - cache_buf_idx = getattr(self, 'cache_buf_idx', None) - if (not self._use_replay_state_update - or prev_num_accepted_tokens is None or cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=prev_num_accepted_tokens, - cache_buf_idx=cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.ssm_state_dtype - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - return getattr(self, 'mamba_ssm_rand_seed', None) - - -def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: - """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" - pool_size = max_batch_size - # One permanent slot is shared by every CUDA-graph padding sentinel. - pool_size += 1 - if mapping.enable_attention_dp: - # Attention-DP can insert a transient dummy request on an otherwise - # idle rank. Keep this headroom internal so scheduler-visible - # max_batch_size still limits real requests. - pool_size += 1 - return pool_size - - -class MixedMambaHybridCacheManager(KVCacheManager, MambaCacheManager, - MambaHybridCacheManager): - """Hybrid cache manager combining separate KVCacheManager and MambaCacheManager. - - Manages KV cache and mamba states in independent pools, with support of - speculative decoding and disaggregated serving. - Does not support block reuse / prefix caching for mamba states. - """ - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - # kv cache parameters - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - layer_mask: List[bool], - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. - # It's derived from the model's BuildConfig for consistency with the C++ backend. - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - max_num_tokens: int = 8192, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - is_estimating_kv_cache: bool = False, - execution_stream: Optional[torch.cuda.Stream] = None, - model_type: str = "nemotron_hybrid", - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - # Per-pool configurations forwarded to the C++ KVCacheManager ctor. - # Lets a single manager host pools with mixed shapes (e.g. Gemma4 - # hybrid attention). See KVCacheManager.__init__. - pool_configurations: Optional[List[PoolConfiguration]] = None, - ) -> None: - - # mamba hybrid cache requires block reuse to be disabled in KV cache config - assert not kv_cache_config.enable_block_reuse, ( - "mamba hybrid cache requires block reuse to be disabled in KV cache config" - ) - - # Host-drafter spec modes (NGram) have no spec worker to call - # update_mamba_states; this manager promotes accepted states itself - # in update_resources. One-model modes (MTP/Eagle/DFlash) and the - # suffix-automaton worker promote from their spec workers and must - # NOT be promoted twice. - self._promote_states_in_update_resources = ( - spec_config is not None - and getattr(spec_config, "decoding_type", None) == "NGram") - - pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) - - MambaCacheManager.__init__( - self, - mamba_d_state, - mamba_d_conv, - mamba_num_heads, - mamba_n_groups, - mamba_head_dim, - mamba_num_layers, - pool_size, - max_batch_size, - mapping, - mamba_cache_dtype, - mamba_ssm_cache_dtype, - mamba_layer_mask, - execution_stream, - speculative_num_draft_tokens=(spec_config.tokens_per_gen_step - 1 - if spec_config is not None else None), - model_type=model_type, - use_replay_state_update=use_replay_state_update, - mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - kda_replay_num_spec=kda_replay_num_spec, - ) - - # initialize kv cache manager - KVCacheManager.__init__( - self, - kv_cache_config, - kv_cache_type, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - execution_stream=execution_stream, - is_draft=is_draft, - pool_configurations=pool_configurations, - ) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - MambaCacheManager.prepare_resources(self, scheduled_batch) - KVCacheManager.prepare_resources(self, scheduled_batch) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - MambaCacheManager.free_resources(self, request) - KVCacheManager.free_resources(self, request, pin_on_release) - - def add_dummy_requests(self, request_ids: List[int], **kwargs): - MambaCacheManager.add_dummy_requests(self, request_ids) - return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) - - def shutdown(self): - MambaCacheManager.shutdown(self) - KVCacheManager.shutdown(self) - - def update_resources(self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None): - KVCacheManager.update_resources(self, scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - self._maybe_promote_drafter_states(scheduled_batch, attn_metadata) - - def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): - """Commit accepted verification states for host-drafter spec modes. - - One-model spec workers (MTP/Eagle/DFlash) and the suffix-automaton - worker call update_mamba_states themselves right after on-device - acceptance. Host-drafter modes (NGram) have no spec worker: - acceptance lands on the requests as - ``py_num_accepted_draft_tokens`` during sampler update, and this - hook — running right after, alongside the KV rewind — promotes the - accepted step's intermediate state into the live pools (or, on the - KDA replay path, records the accepted-draft count for the next - round's replay). Without it, the pools would keep the - pre-verification state and the next forward would resume from a - stale prefix. - """ - if not self._promote_states_in_update_resources: - return - if not self.is_speculative() or attn_metadata is None: - return - gen_requests = scheduled_batch.generation_requests - if not gen_requests: - return - drafted = [ - r for r in gen_requests - if r.py_draft_tokens is not None and len(r.py_draft_tokens) > 0 - ] - if not drafted: - # Drafter skipped this step (e.g. speculation gated off): the - # forward ran the plain in-place decode path; nothing to promote. - return - assert len(drafted) == len(gen_requests), ( - "mixed drafted/undrafted generation batch is not supported for " - "hybrid state promotion (drafts are padded to the static max)") - device = self._impl.mamba_cache.temporal.device - num_contexts = len(scheduled_batch.context_requests) - num_accepted = torch.tensor( - [0] * num_contexts + - [r.py_num_accepted_draft_tokens + 1 for r in gen_requests], - dtype=torch.int32, - device=device) - # Batch-ordered slots (contexts then gens), matching the ordering - # the forward used for the intermediate scratch buffers. Requests - # that finished this step were already freed by response handling - # (which runs before update_resources) and are gone from the index; - # their rows must stay in place for alignment, so redirect them to - # the reserved padding slot (a harmless scratch write). - slot_index = self.mamba_cache_index - padding_slot = self._impl._padding_slot - state_indices = torch.tensor([ - slot_index.get(r.py_request_id, padding_slot) - for r in scheduled_batch.context_requests + gen_requests - ], - dtype=torch.int32, - device=device) - self.update_mamba_states(attn_metadata, num_accepted, state_indices) - - -@triton.jit -def _promote_mamba_state_kernel( - src_ptr, - dst_ptr, - src_idx_ptr, - accepted_position_source_ptr, - num_accepted_tokens_ptr, - blk_ptr, - replay_pnat_ptr, - replay_cache_buf_idx_ptr, - dummy_request_mask_ptr, - num_gens, - count, - src_s_layer, - src_s_row, - src_s_step, - dst_s_layer, - dst_s_block, - POSITION_SOURCE_IS_TOKEN_COUNT: tl.constexpr, - REPLAY_STEP_WIDTH: tl.constexpr, - REPLAY_HISTORY_SIZE: tl.constexpr, - UPDATE_REPLAY_BOOKKEEPING: tl.constexpr, - BLOCK: tl.constexpr, -): - # One program copies a BLOCK-sized tile of the contiguous inner state for a - # single (layer, gen) pair. grid = (num_layers * num_gens, ceil(count/BLOCK)). - pair = tl.program_id(0) - tile = tl.program_id(1) - layer = (pair // num_gens).to(tl.int64) - g = pair % num_gens - row = tl.load(src_idx_ptr + g).to(tl.int64) - accepted_position = tl.load(accepted_position_source_ptr + g) - if POSITION_SOURCE_IS_TOKEN_COUNT: - accepted_position -= 1 - acc = accepted_position.to(tl.int64) - blk = tl.load(blk_ptr + g).to(tl.int64) - - # The replay metadata update is batch-scoped, so layer 0's first copy tile - # owns it. This folds PNAT and active-buffer maintenance into the existing - # accepted-convolution-state promotion launch. - if UPDATE_REPLAY_BOOKKEEPING: - if (layer == 0) & (tile == 0): - num_accepted_tokens = tl.load(num_accepted_tokens_ptr + g) - is_dummy_request = tl.load(dummy_request_mask_ptr + g) - pnat = tl.load(replay_pnat_ptr + blk) - cache_buf_idx = tl.load(replay_cache_buf_idx_ptr + blk) - wrote_checkpoint = pnat + REPLAY_STEP_WIDTH > REPLAY_HISTORY_SIZE - next_pnat = tl.where(wrote_checkpoint, num_accepted_tokens, - pnat + num_accepted_tokens) - next_cache_buf_idx = tl.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx) - tl.store(replay_pnat_ptr + blk, next_pnat, mask=~is_dummy_request) - tl.store(replay_cache_buf_idx_ptr + blk, - next_cache_buf_idx, - mask=~is_dummy_request) - - # int64 throughout: per-layer strides are O(1e8), so layer*stride overflows - # int32 and would corrupt addresses / fault. - src_base = layer * src_s_layer + row * src_s_row + acc * src_s_step - dst_base = layer * dst_s_layer + blk * dst_s_block - offs = tile * BLOCK + tl.arange(0, BLOCK) - mask = offs < count - v = tl.load(src_ptr + src_base + offs, mask=mask) - tl.store(dst_ptr + dst_base + offs, v, mask=mask) - - -def _promote_mamba_state_triton( - dst: torch.Tensor, - intermediate: torch.Tensor, - src_state_indices: torch.Tensor, - accepted_position_source: torch.Tensor, - dst_state_indices: torch.Tensor, - num_accepted_tokens: Optional[torch.Tensor] = None, - position_source_is_token_count: bool = False, - replay_pnat: Optional[torch.Tensor] = None, - replay_cache_buf_idx: Optional[torch.Tensor] = None, - dummy_request_mask: Optional[torch.Tensor] = None, - replay_step_width: int = 0, - replay_history_size: int = 0, - BLOCK: int = 2048) -> None: - """Scatter each generation request's accepted draft-step recurrent state - from the per-request intermediate buffer into the unified C++ pool view. - - Args: - dst: ``[num_layers, num_blocks, *state_shape]`` view of the C++ pool. - The block dim (dim 1) is strided (the pool interleaves - ``ssm_bytes | conv_bytes`` per block), so ``dst`` is non-contiguous. - The kernel writes through ``dst``'s real strides, touching only this - state's bytes -- the case that defeats torch.compile / aot_autograd - (dtype-view mutation) and Inductor codegen (``XBLOCK`` on the uint8 - pool). - intermediate: ``[num_layers, max_batch_size, T, *state_shape]`` dense. - src_state_indices, accepted_position_source, dst_state_indices: - ``[num_gens]`` int. - num_accepted_tokens: ``[num_gens]`` accepted-token counts. Required - when replay bookkeeping is enabled. - position_source_is_token_count: Whether ``accepted_position_source`` - contains accepted-token counts instead of zero-based intermediate - state positions. When true, the kernel subtracts one while - computing the source position. - replay_pnat, replay_cache_buf_idx, dummy_request_mask: Optional replay - metadata. When supplied together, the kernel advances PNAT and the - active history buffer while promoting the convolution state. - - For each gen ``g``:: - - accepted_position = ( - accepted_position_source[g] - 1 - if position_source_is_token_count - else accepted_position_source[g] - ) - dst[:, dst_state_indices[g]] = intermediate[ - :, src_state_indices[g], accepted_position - ] - - Pure gather->scatter copy; bandwidth-bound (~85% of HBM peak), one launch. - """ - num_layers = dst.shape[0] - num_gens = src_state_indices.shape[0] - if num_gens == 0: - return - update_replay_bookkeeping = replay_pnat is not None - if update_replay_bookkeeping: - assert num_accepted_tokens is not None - assert replay_cache_buf_idx is not None - assert dummy_request_mask is not None - assert replay_step_width > 0 - assert replay_history_size >= replay_step_width - assert replay_pnat.dtype == torch.int32 - assert replay_cache_buf_idx.dtype == torch.int32 - assert dummy_request_mask.dtype == torch.bool - assert dummy_request_mask.numel() >= num_gens - else: - # Compile-time-disabled pointer arguments still require tensors. - num_accepted_tokens = accepted_position_source - assert replay_cache_buf_idx is None - assert dummy_request_mask is None - replay_pnat = dst_state_indices - replay_cache_buf_idx = dst_state_indices - dummy_request_mask = dst_state_indices - count = 1 - for s in dst.shape[2:]: - count *= s - # Grid dim order matters: dim 0 -> CUDA grid.x (limit 2^31-1), dim 1 -> - # grid.y (limit 65535). Put the (layer, gen) pairs in dim 0 since that is - # what grows with batch size (num_layers*num_gens stays far below 2^31 for - # any real config); the tile count in dim 1 is small (~count/BLOCK). Do NOT - # swap them: pairs would hit the 65535 y-limit at num_layers*num_gens>65535. - grid = (num_layers * num_gens, triton.cdiv(count, BLOCK)) - _promote_mamba_state_kernel[grid]( - intermediate, - dst, - src_state_indices, - accepted_position_source, - num_accepted_tokens, - dst_state_indices, - replay_pnat, - replay_cache_buf_idx, - dummy_request_mask, - num_gens, - count, - intermediate.stride(0), - intermediate.stride(1), - intermediate.stride(2), - dst.stride(0), - dst.stride(1), - POSITION_SOURCE_IS_TOKEN_COUNT=position_source_is_token_count, - REPLAY_STEP_WIDTH=replay_step_width, - REPLAY_HISTORY_SIZE=replay_history_size, - UPDATE_REPLAY_BOOKKEEPING=update_replay_bookkeeping, - BLOCK=BLOCK, - ) - - -def _mamba_snapshot_rule_counts( - kv_cache_config: KvCacheConfig, - max_seq_len: Optional[int], - tokens_per_block: int, -) -> Tuple[int, int]: - """Return reachable fixed rules and their partial-block upper bound.""" - if not kv_cache_config.enable_block_reuse: - return 0, 0 - - num_rules = 0 - num_unaligned_rules = 0 - state_config = kv_cache_config.mamba_state_config - for offset in set(state_config.additional_snapshot_offsets_from_start): - if max_seq_len is not None and offset > max_seq_len: - continue - num_rules += 1 - num_unaligned_rules += int(offset % tokens_per_block != 0) - for offset in set(state_config.additional_snapshot_offsets_from_end): - # A from-end offset is reachable iff some valid prompt is longer than - # the offset. Its absolute alignment depends on that prompt. - if max_seq_len is not None and offset >= max_seq_len: - continue - num_rules += 1 - num_unaligned_rules += 1 - return num_rules, num_unaligned_rules - - -def _mamba_regular_snapshot_interval( - kv_cache_config: KvCacheConfig, - max_seq_len: Optional[int], -) -> Optional[int]: - if not kv_cache_config.enable_block_reuse: - return None - interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval - if interval is None or interval <= 0: - return None - if max_seq_len is not None and interval > max_seq_len: - return None - return interval - - -def _get_local_mamba_cache_layout( - model_config, - mapping: Mapping, - *, - spec_config=None, - is_draft: bool = False, - use_separate_draft_kv_cache: bool = False, -): - """Return normalized params and local Mamba/attention layer counts. - - Cache construction and affine sizing must follow the model's PP layout: - partition base layers first, then place appended speculative layers on the - last PP rank. The normalized params retain target masks and the appended - draft-layer count so estimation selects the same combined or per-manager - layout as runtime. - """ - from tensorrt_llm._torch.pyexecutor.config_utils import \ - extract_mamba_kv_cache_params - - params = extract_mamba_kv_cache_params( - model_config.pretrained_config, - spec_config=spec_config, - quant_config=model_config.quant_config, - ) - mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( - is_draft=is_draft, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - combined_layer_mask = [ - is_mamba or is_attention for is_mamba, is_attention in zip( - mamba_layer_mask, full_attention_layer_mask) - ] - local_layer_indices, _ = get_pp_layers( - sum(combined_layer_mask), - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - local_mamba_layers = sum(mamba_layer_mask[layer_idx] - for layer_idx in local_layer_indices) - local_attention_layers = sum(full_attention_layer_mask[layer_idx] - for layer_idx in local_layer_indices) - return params, local_mamba_layers, local_attention_layers - - -def _estimate_mamba_hybrid_cache_cost( - model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - tokens_per_block: int, - max_seq_len: Optional[int], - num_reserved_dummy_slots: int, - include_explicit_snapshots: bool, - cap_partial_attention_snapshots: bool, - is_draft: bool = False, - use_separate_draft_kv_cache: bool = False, - attention_slope_override: Optional[int] = None, - attention_page_bytes_override: int = 0, - draft_scratch_tokens: int = 0, - **kwargs, -) -> Tuple[int, int]: - spec_config = kwargs.get("spec_config") - params, local_mamba_layers, local_attention_layers = ( - _get_local_mamba_cache_layout( - model_config, - mapping, - spec_config=spec_config, - is_draft=is_draft, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - )) - attention_slope = attention_slope_override - if attention_slope is None: - attention_slope = (KVCacheManager.get_cache_size_per_token( - model_config, - mapping, - num_layers=local_attention_layers, - **kwargs, - ) if local_attention_layers > 0 else 0) - state_bytes_per_rank = (local_mamba_layers * - params.get_states_bytes_per_layer(mapping)) - max_resident_sequences = max_batch_size * mapping.pp_size - - if include_explicit_snapshots: - fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( - kv_cache_config, max_seq_len, tokens_per_block) - else: - fixed_rules = 0 - unaligned_fixed_rules = 0 - fixed_state_slots = (max_resident_sequences + num_reserved_dummy_slots + - max_resident_sequences * fixed_rules) - attention_block_bytes = attention_slope * tokens_per_block - - interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) - has_unaligned_periodic_snapshot = (interval is not None - and interval % tokens_per_block != 0) - if cap_partial_attention_snapshots: - # Snapshot alignment is unknown while estimating cache cost. Once a - # snapshot is possible, reserve one retained partial attention page - # per resident lineage. Dummy requests carry no attention capacity. - has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None - partial_attention_slots = (max_resident_sequences - if has_non_live_ssm_capacity else 0) - else: - partial_attention_slots = (max_resident_sequences * - unaligned_fixed_rules) - intercept = (fixed_state_slots * state_bytes_per_rank + - partial_attention_slots * - (attention_block_bytes + attention_page_bytes_override)) - - if interval is None: - regular_slope = 0 - else: - regular_slope = math.ceil(state_bytes_per_rank / interval) - if (has_unaligned_periodic_snapshot - and not cap_partial_attention_snapshots): - regular_slope += math.ceil(attention_block_bytes / interval) - draft_layer_sizes = _get_draft_layer_sizes(**kwargs) - if draft_layer_sizes: - # Hybrid target attention is full attention. Reserve capture/noise - # capacity only in its attention pools and the distinct draft pools; - # recurrent state retains the fixed/snapshot accounting above. - # Hybrid profiling bounds draft storage by full-history attention. - _, draft_slope, draft_fixed = _estimate_draft_cache_size_components( - draft_layer_sizes, - tokens_per_block, - generation_capacity_headroom=1, - scratch_tokens=draft_scratch_tokens, - helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, - ) - draft_fixed += draft_scratch_tokens * attention_slope - attention_slope += draft_slope - intercept += max_batch_size * draft_fixed - return attention_slope + regular_slope, intercept - - -class CppMambaHybridCacheManager(KVCacheManager, MambaHybridCacheManager): - """Hybrid cache manager storing mamba states inside the KVCacheManager pool. - - Both KV cache blocks and recurrent state blocks are managed by the unified - C++ KVCacheManager, enabling block reuse / prefix caching across attention - and mamba layers. This compatibility manager remains available through the - manager preference override and legacy disaggregated routing. - - """ - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. - # It's derived from the model's BuildConfig for consistency with the C++ backend. - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - layer_mask: Optional[ - List[bool]] = None, # this is the full attention layer mask - is_estimating_kv_cache: bool = False, - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - model_type: str = "nemotron_hybrid", - **kwargs, - ) -> None: - # 3 kinds of layers: - # 1) Mamba layers (mamba_layer_mask is True) - # 2) Full attention layers (full_attention_layer_mask is True) - # 3) Not managed layers (both masks are False) - total_layers = len(mamba_layer_mask) - if layer_mask is None: - full_attention_layer_mask = [False] * total_layers - elif len(layer_mask) != total_layers: - raise ValueError( - f"layer_mask length ({len(layer_mask)}) must match " - f"mamba_layer_mask length ({total_layers})") - else: - full_attention_layer_mask = list(layer_mask) - layer_mask = [ - mamba_layer_mask[i] or full_attention_layer_mask[i] - for i in range(total_layers) - ] - # PP sharding is done across all layers. - # This is called again in the super().__init__, but we want it to run first - # to set up mtp states before the C++ backend is initialized. - self.pp_layers, _ = get_pp_layers( - mamba_num_layers + num_layers, - mapping, - spec_config=spec_config, - layer_mask=layer_mask, - ) - self.mamba_pp_layers = [ - layer_idx for layer_idx in self.pp_layers - if mamba_layer_mask[layer_idx] - ] - self.local_num_mamba_layers = len(self.mamba_pp_layers) - self.requests = [] - # Seed externally visible mamba fields before any early return so that - # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work - # on ranks with no local mamba layers. - self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = ( - use_replay_state_update and model_type == "qwen3_next" - and self.local_num_mamba_layers > 0) - self.replay_step_width: Optional[int] = ( - spec_config.tokens_per_gen_step - if spec_config is not None and use_replay_state_update else None) - self.replay_history_size: Optional[int] = ( - max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) - if self.replay_step_width is not None else None) - # Same allocation gate as PythonMambaCacheManager: the rand_seed - # buffer must exist whenever SR can fire, not only on the replay path. - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - # Host-side counter bumped per fresh context-slot assignment; combined - # with the slot index and rank offset to produce reproducible per-slot - # seed values without any torch.randint. - self._seed_request_counter = 0 - self.ssm_state_dtype = mamba_ssm_cache_dtype - # Keep the shared Mamba interface valid on PP ranks that do not own a - # local Mamba layer. - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - - if self.local_num_mamba_layers == 0: - logger.info( - "No local mamba layers for this rank, skipping mamba cache initialization" - ) - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=full_attention_layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - is_draft=is_draft, - **kwargs, - ) - # PP ranks replay the same scheduling decisions, so a rank without - # local Mamba layers must still publish the configured boundaries. - self.kv_cache_config = kv_cache_config - self.linear_attention_metadata = LinearAttentionMetadata() - self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_config.periodic_snapshot_interval - if kv_cache_config.enable_block_reuse else 0) - return - - # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) - tp_size = _mamba_effective_tp_size(mapping) - d_inner = mamba_head_dim * mamba_num_heads - conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state - nheads = mamba_num_heads - assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - if use_replay_state_update: - assert mamba_n_groups % tp_size == 0, \ - "replay state update requires mamba_n_groups divisible by tp_size" - self._n_groups_per_rank = mamba_n_groups // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - self.conv_state_shape = [conv_dim, mamba_d_conv - 1] - self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] - self.conv_state_dtype = mamba_cache_dtype - - # Store GLOBAL (pre-TP-division) mamba params for disagg RnnModelConfig. - # d_inner is computed at the top of __init__ before TP division. - d_inner_global = d_inner # = mamba_head_dim * mamba_num_heads (GLOBAL) - conv_dim_global = d_inner_global + 2 * mamba_n_groups * mamba_d_state - self._rnn_d_state = mamba_d_state - self._rnn_d_conv = mamba_d_conv - self._rnn_num_heads = mamba_num_heads # GLOBAL - self._rnn_n_groups = mamba_n_groups - self._rnn_head_dim = mamba_head_dim - self._rnn_hidden_size = d_inner_global # GLOBAL = head_dim * num_heads - self._rnn_conv_dim_size = conv_dim_global # GLOBAL conv dim - self._rnn_num_layers = mamba_num_layers - - # Conv section layout for section-aware split/concat in TP mismatch. - # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; - # we just need to tell C++ which ordering to use. - self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" - self.ssm_bytes = (math.prod(self.ssm_state_shape) * - self.ssm_state_dtype.itemsize) - self.conv_bytes = (math.prod(self.conv_state_shape) * - self.conv_state_dtype.itemsize) - - total_bytes = self.ssm_bytes + self.conv_bytes - if total_bytes % self.ssm_state_dtype.itemsize != 0: - raise RuntimeError( - f"Total state bytes ({total_bytes}) not divisible by " - f"ssm_state_dtype size ({self.ssm_state_dtype.itemsize})") - if total_bytes % self.conv_state_dtype.itemsize != 0: - raise RuntimeError( - f"Total state bytes ({total_bytes}) not divisible by " - f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") - if self.ssm_bytes % self.conv_state_dtype.itemsize != 0: - raise RuntimeError( - f"SSM state bytes ({self.ssm_bytes}) not divisible by " - f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") - self.linear_attention_metadata = LinearAttentionMetadata() - self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value - self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes - self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_config.periodic_snapshot_interval - if kv_cache_config.enable_block_reuse else 0) - # RNN model params for disagg TP-mismatch split/concat. - conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} - self.linear_attention_metadata.rnn_num_heads = self._rnn_num_heads - self.linear_attention_metadata.rnn_head_dim = self._rnn_head_dim - self.linear_attention_metadata.rnn_d_state = self._rnn_d_state - self.linear_attention_metadata.rnn_d_conv = self._rnn_d_conv - self.linear_attention_metadata.rnn_n_groups = self._rnn_n_groups - self.linear_attention_metadata.rnn_conv_section_layout = conv_section_map.get( - self._rnn_conv_section_layout, 0) - self.linear_attention_metadata.rnn_ssm_bytes = self.ssm_bytes - self.linear_attention_metadata.rnn_ssm_dtype_size = self.ssm_state_dtype.itemsize - self.linear_attention_metadata.rnn_conv_dtype_size = self.conv_state_dtype.itemsize - kv_cache_config = kv_cache_config.model_copy(deep=True) - if kv_cache_config.enable_partial_reuse: - logger.warning( - "Partial reuse is not supported for mamba hybrid models, disabling partial reuse" - ) - kv_cache_config.enable_partial_reuse = False - - kv_cache_config.max_attention_window = [] - for i in range(len(layer_mask)): - if layer_mask[i]: - kv_cache_config.max_attention_window.append( - LinearCacheType.RECURRENT_STATES. - value if mamba_layer_mask[i] else max_seq_len) - - recurrent_states_window = LinearCacheType.RECURRENT_STATES.value - local_windows = { - recurrent_states_window - if mamba_layer_mask[layer_idx] else max_seq_len - for layer_idx in self.pp_layers - } - kwargs["pool_configurations"] = [ - PoolConfiguration( - window_size=window_size, - head_dim=head_dim, - dtype=torch_dtype_to_binding(self.ssm_state_dtype) - if window_size == recurrent_states_window else dtype, - ) for window_size in sorted(local_windows) - ] - - # Normalize num_kv_heads to a per-layer list and zero out mamba - # layer positions: those layers carry SSM/conv state instead of KV - # heads, so the parent KV cache should not allocate KV head storage - # for them. - if isinstance(num_kv_heads, int): - per_layer_kv_heads = [num_kv_heads] * total_layers - else: - if len(num_kv_heads) != total_layers: - raise ValueError( - f"num_kv_heads list length ({len(num_kv_heads)}) does not " - f"match total layers ({total_layers})") - per_layer_kv_heads = list(num_kv_heads) - for i, is_mamba in enumerate(mamba_layer_mask): - if is_mamba: - per_layer_kv_heads[i] = 0 - - self._setup_mtp_intermediate_states(spec_config, max_batch_size) - - # pass remaining arguments to super class - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=mamba_num_layers + num_layers, - num_kv_heads=per_layer_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - is_draft=is_draft, - linear_attention_metadata=self.linear_attention_metadata, - **kwargs, - ) - - assert self.local_num_mamba_layers > 0, "At least one mamba layer is required" - self.mamba_layer_offsets = {} - for idx, layer_id in enumerate(self.mamba_pp_layers): - self.mamba_layer_offsets[layer_id] = idx - - self.host_block_offsets = torch.zeros([ - self.impl.num_pools, self.max_batch_size, 2, self.max_blocks_per_seq - ], - dtype=torch.int32, - device="cpu") - self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ - self.layer_offsets[self.mamba_pp_layers[0]]][0] - - self.cuda_state_indices = torch.zeros([self.max_batch_size], - dtype=torch.int32, - device="cuda") - self._host_state_indices = torch.zeros([self.max_batch_size], - dtype=torch.int32, - pin_memory=prefer_pinned()) - self._row_indices = torch.arange(self.max_batch_size, - dtype=torch.long, - device="cpu") - self._request_id_to_state_index = {} - self._request_id_to_is_dummy = {} - self.kv_cache_config = kv_cache_config - self.is_estimating_kv_cache = is_estimating_kv_cache - - self._setup_states() - self._setup_replay_buffers(spec_config) - if use_replay_state_update and model_type == "qwen3_next": - logger.info_once( - "Configured GDN cached replay commit mode: small-batch fused, " - "large-batch all-layer", - key="gdn_cached_replay_commit_mode_fused", - ) - - @staticmethod - def get_cache_size_per_token( - model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - num_layers: Optional[int] = None, - tokens_per_block: int = 32, - max_seq_len: Optional[int] = None, - **kwargs, - ): - """Affine memory model for the unified hybrid KV pool. - - Returns ``(slope_bytes_per_token, intercept_bytes)``: - - * ``slope`` = attention KV bytes per token plus conservative periodic - Mamba-state and partial-attention-snapshot costs. - * ``intercept`` = live and CUDA-graph dummy Mamba state. - - Memory budget -> max tokens then becomes - ``T = (budget - intercept) // slope`` instead of plain - ``T = budget // bytes_per_token``. - """ - return _estimate_mamba_hybrid_cache_cost( - model_config, - mapping, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - num_reserved_dummy_slots=1, - include_explicit_snapshots=False, - cap_partial_attention_snapshots=False, - **kwargs, - ) - - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return self._use_gdn_cached_replay_all_layer_commit - - def _commit_gdn_cached_replay_history_layers( - self, - attn_metadata: "AttentionMetadata", - num_decodes: int, - ) -> None: - """Synchronously advance every local GDN checkpoint in one launch.""" - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers) - - if (not self._use_gdn_cached_replay_all_layer_commit - or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): - return - if (self.all_ssm_states is None or self.old_x is None - or self.old_B is None or self.old_dt is None - or self.replay_history_size is None): - raise RuntimeError( - "GDN cached replay all-layer commit requires replay state buffers." - ) - - mamba_metadata = attn_metadata.mamba_metadata - if mamba_metadata.replay_num_decodes != num_decodes: - raise RuntimeError( - "GDN replay metadata contains " - f"{mamba_metadata.replay_num_decodes} decode requests, " - f"but state update received {num_decodes}.") - commit_gdn_cached_replay_history_layers( - ssm_states=self.all_ssm_states, - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=mamba_metadata.replay_work_items[:num_decodes], - n_writes=mamba_metadata.replay_n_writes, - history_size=self.replay_history_size, - ) - - def shutdown(self): - # Release tensor views into the pool before the pool memory is freed, - # so their deleters don't see stale pointers. - self.all_ssm_states = None - self.all_conv_states = None - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - super().shutdown() - - def add_dummy_requests( - self, - request_ids: List[int], - # Note that token_nums should be past_kv_len + input_len (without - # spec decoding). The draft tokens will be added in this function, - # so we don't need to take care of it in the caller. When preparing - # token_nums, we should not take the draft tokens into account, so - # don't use the kv_cache_manager.max_seq_len, which includes both - # extra tokens and draft tokens. - token_nums: Optional[List[int]] = None, - is_gen: bool = False, - prepare_resource: bool = True, - max_num_draft_tokens: int = 0, - kv_reserve_draft_tokens: Optional[int] = None, - use_mrope: bool = False, - max_beam_width: int = 1, - encoder_output_lens: Optional[List[int]] = None, - # For capturable drafting loops. During normal inference, the draft model always - # has enough KV cache space to fit all of our draft tokens. During warmup, however, - # we need to make the KV cache manager aware that multiple autoregressive steps will - # occur. - num_extra_decoding_steps: int = 0, - draft_kv_cache_manager: Optional[KVCacheManager] = None, - ) -> List[LlmRequest]: - requests = super().add_dummy_requests( - request_ids=request_ids, - token_nums=token_nums, - is_gen=is_gen, - prepare_resource=prepare_resource, - max_num_draft_tokens=max_num_draft_tokens, - kv_reserve_draft_tokens=kv_reserve_draft_tokens, - use_mrope=use_mrope, - max_beam_width=max_beam_width, - encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, - draft_kv_cache_manager=draft_kv_cache_manager, - ) - if requests: - self.requests.extend(requests) - # Only process the newly added requests, not all of self.requests. - # self.requests may contain stale entries from the previous - # _prepare_resources call (e.g., disagg transfer-pending requests) - # that would exceed max_batch_size and cause out-of-bounds access. - self._setup_state_indices(requests) - return requests - - def update_resources(self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None): - super().update_resources(scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - - @nvtx_range("hybrid_prepare_resources") - def _prepare_resources(self, scheduled_batch: ScheduledRequests): - self.requests = scheduled_batch.context_requests + \ - scheduled_batch.generation_requests - # Issue all async block onboards; defer the syncTransfers() (refresh_blocks) - # until just before forward so the CPU-side prep (attn metadata, input - # tensor assembly, draft model prep, etc.) can overlap with the in-flight - # cudaMemcpyAsync calls instead of being serialized behind them. - # copy_linear_attention_block returns True iff it actually issued a copy; - # we skip refresh_blocks entirely when nothing was scheduled. - self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( - self.requests) - self._setup_state_indices() - self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - super().prepare_resources(scheduled_batch) - if self.local_num_mamba_layers == 0: - return - self._prepare_resources(scheduled_batch) - - @nvtx_range("hybrid_flush_state_transfers") - def flush_state_transfers(self) -> None: - """Complete any deferred state-block onboards scheduled by - prepare_resources(). Must be called before forward() reads recurrent - state blocks. Cheap no-op when nothing was scheduled. - """ - if getattr(self, "_pending_state_transfers", False): - self.impl.refresh_blocks() - self._pending_state_transfers = False - - @nvtx_range("hybrid_update_mamba_states") - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None, - accepted_leaf_positions: Optional[torch.Tensor] = None): - if self.local_num_mamba_layers == 0: - return - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + - num_gens] - # Dynamic tree supplies explicit tree-node leaf positions. Linear MTP - # passes accepted-token counts and lets the promotion kernel subtract - # one, avoiding a separate GPU subtraction/cast. - if accepted_leaf_positions is None: - accepted_position_source = accepted_tokens - position_source_is_token_count = True - else: - accepted_position_source = accepted_leaf_positions.to(torch.int32) - position_source_is_token_count = False - # Match the API of MambaCacheManager.update_mamba_states: callers - # may pass per-request state slot indices explicitly (e.g. MTP via - # attn_metadata.mamba_metadata.state_indices). Fall back to this - # manager's own slot mapping when not provided. - if state_indices is None: - state_indices = self.get_state_indices() - state_indices_d = state_indices[num_contexts:num_contexts + - num_gens].to(torch.int32) - src_state_indices = self.intermediate_state_indices[:num_gens] - - # The accepted SSM/conv promotion is a bandwidth-bound gather->scatter - # into the dtype-reinterpreted, strided C++ pool view. torch.compile - # can't handle the dtype-view mutation (and Inductor chokes on the - # uint8 pool with "XBLOCK too large"), so a dedicated Triton kernel - # writes through the view's real strides (~85% of HBM peak, one launch - # per state). - if self._use_replay_state_update: - # Every GDN layer has finished reading the old checkpoint and - # writing its candidate history. Advance all local checkpoints - # in one launch before PNAT and the active history buffer change. - self._commit_gdn_cached_replay_history_layers( - attn_metadata, num_gens) - assert self._dummy_request_mask is not None - else: - # Legacy: copy the accepted SSM state from the intermediate buffer. - _promote_mamba_state_triton( - self.all_ssm_states, - self.intermediate_ssm_states, - src_state_indices, - accepted_position_source, - state_indices_d, - position_source_is_token_count=position_source_is_token_count) - - # Conv: both paths save all intermediate conv windows, carry over the - # accepted one. The replay path also advances PNAT and the active cache - # buffer in this same launch. - replay_dummy_request_mask = None - if self._use_replay_state_update: - replay_dummy_request_mask = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - _promote_mamba_state_triton( - self.all_conv_states, - self.intermediate_conv_states, - src_state_indices, - accepted_position_source, - state_indices_d, - num_accepted_tokens=accepted_tokens, - position_source_is_token_count=position_source_is_token_count, - replay_pnat=(self.prev_num_accepted_tokens - if self._use_replay_state_update else None), - replay_cache_buf_idx=(self.cache_buf_idx - if self._use_replay_state_update else None), - dummy_request_mask=replay_dummy_request_mask, - replay_step_width=(self.replay_step_width - if self._use_replay_state_update else 0), - replay_history_size=(self.replay_history_size - if self._use_replay_state_update else 0)) - - def get_num_available_tokens(self, - token_num_upper_bound: int, - max_num_draft_tokens: int = 0, - **kwargs) -> int: - # Base bound from attention KV cache pool (the parent's behaviour). - result = super().get_num_available_tokens(token_num_upper_bound, - max_num_draft_tokens, - **kwargs) - # Also bound by the recurrent-state pool capacity: each request needs - # roughly ceil(N / states_snapshot_interval) recurrent-state blocks - # (+1 for the corner case where N is a multiple of tokens_per_block). - # When block reuse is disabled, only one snapshot is needed per - # request, so no additional capping is required here. - interval = (self.linear_attention_metadata.states_snapshot_interval - if self.linear_attention_metadata is not None else 0) - # Attention-only PP ranks keep the interval so every rank publishes - # identical scheduling boundaries, but they have no recurrent-state - # pool whose capacity should constrain their attention KV cache. - if self.local_num_mamba_layers > 0 and interval and interval > 0: - stats = self.impl.get_kv_cache_stats() - rs_free = stats.num_free_blocks_per_window_size.get( - LinearCacheType.RECURRENT_STATES.value, 0) - # Reserve 1 block for the always-allocated last block (corner case - # / final live state) so we don't promise more tokens than the - # pool can actually back at allocation time. - usable_rs_blocks = max(0, rs_free - 1) - rs_token_cap = usable_rs_blocks * interval - result = min(result, rs_token_cap) - return max(result, 0) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - if request in self.requests: - self.requests.remove(request) - self._request_id_to_state_index.pop(request.py_request_id, None) - self._request_id_to_is_dummy.pop(request.py_request_id, None) - super().free_resources(request, pin_on_release) - - def _setup_state_indices(self, requests=None) -> None: - if self.local_num_mamba_layers == 0: - return - if requests is None: - requests = self.requests - block_indices = [] - for req in requests: - if req.is_context_finished: - next_step = self.get_num_tokens(req) - 1 - elif self.kv_cache_config.enable_block_reuse: - next_step = (req.context_current_position - 1 + - req.context_chunk_size) - else: - next_step = req.prompt_len - 1 - block_indices.append(next_step // self.tokens_per_block) - self.impl.copy_batch_block_offsets( - self.host_block_offsets, [req.py_request_id for req in requests], 1, - 0) - max_blocks = self.blocks_per_window[ - LinearCacheType.RECURRENT_STATES.value][0] - n = len(requests) - self._host_state_indices.zero_() - if n > 0: - # Vectorized gather: replace per-element Python loop with a single - # tensor index op. block_indices is a small Python list so - # torch.tensor() conversion here is O(n) at C level, much cheaper - # than n round-trips through Python indexing. - bi = torch.tensor(block_indices, dtype=torch.long) - rows = self._row_indices[:n] - # host_block_offsets: [num_pools, max_batch_size, 2, max_blocks_per_seq] - values = self.host_block_offsets[self.recurrent_states_pool_index, - rows, 0, bi] - invalid_mask = (values < 0) | (values >= max_blocks) - # The C++ recurrent-state manager uses null page-table entries for - # logical blocks that are not snapshot boundaries. Usually a - # context chunk ends exactly at an allocated snapshot (or at the - # final live-state block), but the scheduler may shorten a chunk - # further when KV capacity is tight. In that case the Mamba - # kernel must keep accumulating into the next allocated snapshot - # or final block, matching copyLinearAttentionBlock(), which also - # skips placeholders when it advances the live state. - for bad_i in invalid_mask.nonzero( - as_tuple=False).flatten().tolist(): - req = requests[bad_i] - if req.is_context_finished: - continue - last_prompt_block = (req.prompt_len - - 1) // self.tokens_per_block - row = self.host_block_offsets[self.recurrent_states_pool_index, - rows[bad_i], 0] - candidates = row[block_indices[bad_i]:last_prompt_block + 1] - valid_candidates = ((candidates >= 0) - & (candidates < max_blocks)).nonzero( - as_tuple=False) - if valid_candidates.numel() > 0: - values[bad_i] = candidates[valid_candidates[0, 0]] - - invalid_mask = (values < 0) | (values >= max_blocks) - if invalid_mask.any(): - bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) - req = requests[bad_i] - value = int(values[bad_i]) - raise RuntimeError( - f"Invalid recurrent state block index {value} " - f"(expected 0 <= index < {max_blocks}) for request {bad_i}, " - f"prompt_len={req.prompt_len}, " - f"is_context_finished={req.is_context_finished}, " - f"context_current_position={req.context_current_position}, " - f"prepopulated_token_num={req.prepopulated_prompt_len}, " - f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " - f"block_index for next step is {block_indices[bad_i]}, " - "\nblock_ids=" - f"{self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" - ) - self._host_state_indices[:n] = values - - self.cuda_state_indices.copy_(self._host_state_indices, - non_blocking=True) - is_dummy = [req.is_dummy for req in requests] - self._refresh_dummy_request_mask( - is_dummy if requests is - self.requests else [req.is_dummy for req in self.requests]) - - # Build request_id → pool block offset mapping so that - # get_state_indices can return indices in arbitrary request order. - # Bulk tolist avoids a per-request tensor-index + .item() round-trip. - state_values = self._host_state_indices[:n].tolist() - for req, value, dummy in zip(requests, state_values, is_dummy): - self._request_id_to_state_index[req.py_request_id] = value - self._request_id_to_is_dummy[req.py_request_id] = dummy - - def get_state_indices(self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None) -> list: - if self.local_num_mamba_layers == 0: - # Mamba metadata is prepared on every PP rank even when this rank - # owns only attention layers. No local kernel consumes these - # indices, so avoid consulting state that is intentionally absent. - return [0] * len(request_ids) if request_ids is not None else [] - if request_ids is not None: - # Return indices in the order of the caller's request_ids, - # not the internal self.requests order. This is critical when - # the batch is reordered after prepare_resources (e.g. disagg - # serving sorts generation_requests by py_batch_idx). - indices = [ - self._request_id_to_state_index[rid] for rid in request_ids - ] - if is_padding is None: - is_padding = [False] * len(request_ids) - assert len(request_ids) == len(is_padding) - is_dummy = [ - self._request_id_to_is_dummy.get(rid, False) or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - return self.cuda_state_indices - - def _setup_states(self) -> None: - # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) - pool: torch.Tensor = self.impl.get_recurrent_states_pool().view( - torch.uint8).reshape(self.local_num_mamba_layers, -1, - self.ssm_bytes + self.conv_bytes) - num_blocks_in_pool = pool.shape[1] - self.all_ssm_states = pool[:, :, :self.ssm_bytes].view( - self.ssm_state_dtype).view( - [self.local_num_mamba_layers, num_blocks_in_pool] + - self.ssm_state_shape) - self.all_conv_states = pool[:, :, self.ssm_bytes:self.ssm_bytes + - self.conv_bytes].view( - self.conv_state_dtype).view([ - self.local_num_mamba_layers, - num_blocks_in_pool - ] + self.conv_state_shape) - self.all_ssm_states.zero_() - self.all_conv_states.zero_() - - def _setup_replay_buffers(self, spec_config) -> None: - cache_size = self.all_ssm_states.shape[1] - device = self.all_ssm_states.device - if not self._allocate_pool_replay_buffers(spec_config, cache_size, - device): - return - - self._dummy_request_mask = torch.zeros(self.max_batch_size, - dtype=torch.bool, - device=device) - self._dummy_request_mask_host = torch.zeros( - self.max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned(), - ) - - -class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): - """Hybrid Mamba cache manager backed by KVCacheManagerV2. - - Attention KV pages and Mamba recurrent-state pages are both owned by the - Python V2 cache manager. Mamba layers are represented as V2 SSM layers, - while this wrapper exposes the state tensors and slot indices expected by - the PyTorch Mamba kernels. - """ - - _supports_additional_snapshot_offsets = True - # Recurrent-state snapshots use a specialized commit/history protocol, so - # keep main-like reuse endpoints and the existing unpaired draft path. - _supports_reuse_match_backoff = False - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - layer_mask: Optional[List[bool]] = None, - is_estimating_kv_cache: bool = False, - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", - **kwargs, - ) -> None: - if conv_state_layout not in ("x_b_c", "q_k_v"): - raise ValueError( - f"Unsupported convolution state layout: {conv_state_layout!r}") - if "model_type" in kwargs: - # The V1 managers select the conv-state layout by model_type; this - # class takes it explicitly. Silently absorbing model_type here - # means a caller's layout request would be dropped on the floor. - raise TypeError( - "MambaHybridCacheManagerV2 does not accept 'model_type' " - f"(got {kwargs['model_type']!r}); pass " - "conv_state_layout='x_b_c' or 'q_k_v' instead") - total_layers = len(mamba_layer_mask) - if layer_mask is None: - full_attention_layer_mask = [False] * total_layers - elif len(layer_mask) != total_layers: - raise ValueError( - f"layer_mask length ({len(layer_mask)}) must match " - f"mamba_layer_mask length ({total_layers})") - else: - full_attention_layer_mask = list(layer_mask) - - combined_layer_mask = [ - mamba_layer_mask[i] or full_attention_layer_mask[i] - for i in range(total_layers) - ] - - self._mamba_layer_mask = list(mamba_layer_mask) - self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = (use_replay_state_update - and conv_state_layout - == "q_k_v") - self._gdn_cached_replay_state_descriptors: Optional[torch.Tensor] = None - self._gdn_cached_replay_state_strides: Optional[Tuple[int, int, - int]] = None - self.replay_step_width: Optional[int] = ( - spec_config.tokens_per_gen_step - if spec_config is not None and use_replay_state_update else None) - self.replay_history_size: Optional[int] = ( - max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) - if self.replay_step_width is not None else None) - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - self._seed_request_counter = 0 - self._recurrent_evicted_blocks_total = 0 - self._recurrent_onboarded_blocks_total = 0 - self._recurrent_dropped_blocks_total = 0 - self._recurrent_status_logged = False - # Branch points keyed by request id. prepare_expect_snapshot_points() - # overwrites the per-request list on every scheduler pass, so the point - # has to be re-merged from here rather than stored only on the request. - self._branch_snapshot_points: Dict[int, int] = {} - self._snapshot_pruned_tokens_total = 0 - self._page_pruned_tokens_total = 0 - self._branch_snapshots_taken_total = 0 - self._branch_snapshots_skipped_total: Dict[str, int] = {} - num_cuda_graph_padding_dummy_slots = ( - _get_num_cuda_graph_padding_dummy_slots(spec_config, - max_batch_size)) - self._num_reserved_dummy_slots = (num_cuda_graph_padding_dummy_slots + - int(mapping.enable_attention_dp)) - self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype - is not None else mamba_cache_dtype) - self.conv_state_dtype = mamba_cache_dtype - - self.pp_layers, _ = get_pp_layers( - mamba_num_layers + num_layers, - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - self.mamba_pp_layers = [ - layer_idx for layer_idx in self.pp_layers - if mamba_layer_mask[layer_idx] - ] - self.local_num_mamba_layers = len(self.mamba_pp_layers) - self._use_gdn_cached_replay_all_layer_commit = ( - self._use_gdn_cached_replay_all_layer_commit - and self.local_num_mamba_layers > 0) - - if self.local_num_mamba_layers > 0: - tp_size = _mamba_effective_tp_size(mapping) - d_inner = mamba_head_dim * mamba_num_heads - grouped_state_dim = mamba_n_groups * mamba_d_state - conv_dim = d_inner + 2 * grouped_state_dim - nheads = mamba_num_heads - assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - if kwargs.get("is_disagg", - False) and grouped_state_dim % tp_size != 0: - raise ValueError( - "Disaggregated Mamba transfer requires each convolution " - "state section to be divisible by tp_size") - if use_replay_state_update: - assert mamba_n_groups % tp_size == 0, \ - "replay state update requires mamba_n_groups divisible by tp_size" - self._n_groups_per_rank = mamba_n_groups // tp_size - d_inner_local = d_inner // tp_size - grouped_state_dim_local = grouped_state_dim // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - self.conv_state_shape = [conv_dim, mamba_d_conv - 1] - self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] - # TP-mismatch disaggregated transfers must split the flat - # convolution state at its true semantic boundaries. Mamba2 stores - # [x | B | C], while GDN stores [Q | K | V]. The large section is - # therefore first for Mamba2 and last for GDN. - if conv_state_layout == "x_b_c": - self.conv_section_dims = [ - d_inner_local, - grouped_state_dim_local, - grouped_state_dim_local, - ] - else: - self.conv_section_dims = [ - grouped_state_dim_local, - grouped_state_dim_local, - d_inner_local, - ] - self.ssm_bytes = (math.prod(self.ssm_state_shape) * - self.ssm_state_dtype.itemsize) - self.conv_bytes = (math.prod(self.conv_state_shape) * - self.conv_state_dtype.itemsize) - else: - logger.info( - "No local mamba layers for this rank, skipping mamba state views" - ) - self._n_groups_per_rank = 0 - self.conv_state_shape = [] - self.ssm_state_shape = [] - self.conv_section_dims = [] - self.ssm_bytes = 0 - self.conv_bytes = 0 - - if isinstance(num_kv_heads, int): - per_layer_kv_heads = [num_kv_heads] * total_layers - else: - if len(num_kv_heads) != total_layers: - raise ValueError( - f"num_kv_heads list length ({len(num_kv_heads)}) does not " - f"match total layers ({total_layers})") - per_layer_kv_heads = list(num_kv_heads) - for i, is_mamba in enumerate(mamba_layer_mask): - if is_mamba: - per_layer_kv_heads[i] = 0 - - self._setup_mtp_intermediate_states(spec_config, max_batch_size) - - kv_cache_config = kv_cache_config.model_copy(deep=True) - if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: - block_reuse_config = kv_cache_config.block_reuse_config - block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) - if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: - # SSM reuse is valid only at explicit snapshot boundaries. - kv_cache_config.block_reuse_config = block_reuse_config.model_copy( - update={"policy": BlockReusePolicy.PER_REQUEST.value}) - self._attention_cache_layout_policy = None - pretrained_config = kwargs.get("pretrained_config") - if (dtype == DataType.NVFP4 and kv_cache_type == CacheTypeCpp.SELFKONLY - and getattr(pretrained_config, "kv_lora_rank", None) is not None - and getattr(pretrained_config, "qk_rope_head_dim", - None) is not None): - from tensorrt_llm._torch.attention_backend.fp4_mla.cache_manager import \ - Fp4MlaV2CacheLayoutPolicy - - kv_cache_config = Fp4MlaV2CacheLayoutPolicy.configure_manager( - self, - kv_cache_config, - kv_cache_type, - dtype=dtype, - tokens_per_block=tokens_per_block, - head_dim=head_dim, - pretrained_config=pretrained_config, - spec_config=spec_config, - is_disagg=kwargs.get("is_disagg", False), - ) - self.kv_cache_config = kv_cache_config - - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=mamba_num_layers + num_layers, - num_kv_heads=per_layer_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=combined_layer_mask, - is_draft=is_draft, - is_estimating_kv_cache=is_estimating_kv_cache, - num_reserved_index_slots=self._num_reserved_dummy_slots, - **kwargs, - ) - - if self._attention_cache_layout_policy is not None: - self._attention_cache_layout_policy.finalize_manager(self) - - self.mamba_layer_offsets = { - layer_id: offset - for offset, layer_id in enumerate(self.mamba_pp_layers) - } - self._request_id_to_state_index = {} - self._request_id_to_is_dummy = {} - - state_index_capacity = (self.max_batch_size + - self._num_reserved_dummy_slots) - self.cuda_state_indices = torch.zeros([state_index_capacity], - dtype=torch.int32, - device="cuda") - self._host_state_indices = torch.zeros([state_index_capacity], - dtype=torch.int32, - pin_memory=prefer_pinned()) - - if self.local_num_mamba_layers > 0: - first_mamba_local_layer = self.layer_offsets[ - self.mamba_pp_layers[0]] - self.ssm_layer_group_id = self.impl.get_layer_group_id( - LayerId(first_mamba_local_layer)) - self._ssm_page_index_scale = self.impl.get_page_index_scale( - LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) - num_ssm_pages = self.impl.get_page_index_upper_bound( - LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) - num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // - self._ssm_page_index_scale) - required_live_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - if num_ssm_slots < required_live_slots: - KVCacheManagerV2.shutdown(self) - raise ValueError( - "The V2 Mamba state pool has only " - f"{num_ssm_slots} slots but needs at least " - f"{required_live_slots} live/dummy slots. Increase the " - "KV cache budget or allocate a larger Mamba pool_ratio.") - self._setup_states() - self._setup_replay_buffers(spec_config) - if self._use_gdn_cached_replay_all_layer_commit: - state_layout = ("affine" - if self._gdn_cached_replay_state_strides - is not None else "indirect") - logger.info_once( - "Configured GDN cached replay commit mode for V2: " - "small-batch fused, large-batch all-layer; " - f"state layout: {state_layout}", - key="gdn_cached_replay_v2_commit_mode_fused", - ) - else: - self.ssm_layer_group_id = None - self._ssm_page_index_scale = 1 - self.all_ssm_states = [] - self.all_conv_states = [] - self._setup_replay_buffers(spec_config) - - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) - - def _commit_gdn_cached_replay_history_layers( - self, - attn_metadata: "AttentionMetadata", - num_decodes: int, - ) -> None: - """Advance all V2 GDN checkpoints through their state views.""" - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers) - - if (not getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) - or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): - return - if (not self.all_ssm_states - or (self._gdn_cached_replay_state_descriptors is None - and self._gdn_cached_replay_state_strides is None) - or self.old_x is None or self.old_B is None - or self.old_dt is None or self.replay_history_size is None): - raise RuntimeError( - "GDN cached replay all-layer commit requires V2 replay state buffers." - ) - - mamba_metadata = attn_metadata.mamba_metadata - if mamba_metadata.replay_num_decodes != num_decodes: - raise RuntimeError( - "GDN replay metadata contains " - f"{mamba_metadata.replay_num_decodes} decode requests, " - f"but state update received {num_decodes}.") - state_strides = self._gdn_cached_replay_state_strides - commit_gdn_cached_replay_history_layers( - ssm_states=self.all_ssm_states[0], - ssm_state_descriptors=(self._gdn_cached_replay_state_descriptors), - ssm_state_num_layers=(state_strides[0] - if state_strides is not None else None), - ssm_state_layer_stride=(state_strides[1] - if state_strides is not None else None), - ssm_state_slot_stride=(state_strides[2] - if state_strides is not None else None), - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=mamba_metadata.replay_work_items[:num_decodes], - n_writes=mamba_metadata.replay_n_writes, - history_size=self.replay_history_size, - ) - - def _create_kv_cache( - self, - request_id: int, - lora_task_id: Optional[int], - input_tokens: Optional[Sequence[TokenIdExt]], - *, - cache_salt: Optional[str] = None, - is_dummy: bool = False, - expected_prompt_length: Optional[int] = None, - ugpu_id: Optional[int] = None, - ) -> Optional[_KVCache]: - kv_cache = super()._create_kv_cache( - request_id, - lora_task_id, - input_tokens, - cache_salt=cache_salt, - is_dummy=is_dummy, - expected_prompt_length=expected_prompt_length, - ugpu_id=ugpu_id, - ) - if (self.mapping.rank == 0 and kv_cache is not None - and input_tokens is not None and not is_dummy - and self.local_num_mamba_layers > 0): - # Prefix lookup excludes the final prompt token because it must be - # recomputed by prefill. Restore it for the request-length metric. - request_total_tokens = len(input_tokens) + 1 - prefix_reuse_diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) - logger.debug( - f"[MambaHybridCacheManagerV2] prefix reuse rank={self.mapping.rank} " - f"request_id={request_id} " - f"request_total_tokens={request_total_tokens} " - f"longest_attention_match_tokens=" - f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_hybrid_pruning()} " - f"content_divergence_tokens=" - f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_pruning()} " - f"latest_recurrent_snapshot_tokens=" - f"{kv_cache.num_committed_tokens}") - return kv_cache - - def _skip_branch_snapshot(self, reason: str) -> None: - self._branch_snapshots_skipped_total[reason] = ( - self._branch_snapshots_skipped_total.get(reason, 0) + 1) - - def _format_branch_snapshot_counters(self) -> str: - """Attribute reuse loss between missing snapshots and evicted pages. - - A flat cache hit rate means something different depending on whether - branch points were never found or were found and not reused, so the - skip reasons are reported alongside the totals. - """ - skipped = self._branch_snapshots_skipped_total - return (f"snapshot_pruned_tokens={self._snapshot_pruned_tokens_total} " - f"page_pruned_tokens={self._page_pruned_tokens_total} " - f"branch_snapshots_taken={self._branch_snapshots_taken_total} " - f"branch_snapshots_skipped={sum(skipped.values())} " - f"branch_snapshots_skipped_by_reason=" - f"{dict(sorted(skipped.items()))}") - - def _record_branch_snapshot_point(self, req: LlmRequest, kv_cache: _KVCache, - num_lookup_tokens: Optional[int]) -> None: - """Snapshot where this request leaves the reuse tree, for its siblings.""" - if not self.kv_cache_config.mamba_state_config.enable_branch_snapshot: - return - if req.is_dummy_request: - return - if num_lookup_tokens is None: - self._skip_branch_snapshot("no_fresh_match") - # A fresh lookup can record a point before resume() fails. On the - # next attempt, reapply that point against the cache's authoritative - # reuse depth: the request cursor is not rewound until preparation - # succeeds and may still describe a deeper, abandoned attempt. - if req.py_request_id in self._branch_snapshot_points: - self._apply_branch_snapshot_point( - req, context_current_position=kv_cache.num_committed_tokens) - return - - diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) - divergence = diagnostics._get_num_reusable_tokens_before_pruning() - reused = kv_cache.num_committed_tokens - if self.local_num_mamba_layers > 0: - hybrid_depth = ( - diagnostics._get_num_reusable_tokens_before_hybrid_pruning()) - # Loss attributable to a missing recurrent snapshot, versus loss - # attributable to attention pages having been evicted. These call - # for different fixes, so they are counted apart. - self._snapshot_pruned_tokens_total += max(0, hybrid_depth - reused) - self._page_pruned_tokens_total += max(0, divergence - hybrid_depth) - - # The whole lookup range matched, so there is no fork here. - if divergence >= num_lookup_tokens: - self._skip_branch_snapshot("no_divergence") - return - # Align down: tokens past a block boundary come from a partial-match - # child rather than a confirmed shared prefix, and an unaligned commit - # takes the partial-block snapshot path. - point = (divergence // self.tokens_per_block) * self.tokens_per_block - if point == 0: - self._skip_branch_snapshot("aligned_to_zero") - return - # Reuse already reached the fork, so a snapshot there exists. - if point <= reused: - self._skip_branch_snapshot("already_reused") - return - # A branch point at or beyond the prompt end cannot improve sibling - # reuse. - if point >= req.prompt_len: - self._skip_branch_snapshot("at_prompt_end") - return - - self._branch_snapshot_points[req.py_request_id] = point - self._apply_branch_snapshot_point(req, context_current_position=reused) - - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - super().prepare_expect_snapshot_points(requests) - for request in requests: - self._align_fp4_mla_snapshot_points(request) - if (not self.enable_block_reuse or not self.kv_cache_config. - mamba_state_config.enable_branch_snapshot): - return - for request in requests: - self._apply_branch_snapshot_point(request) - - def _align_fp4_mla_snapshot_points(self, request: LlmRequest) -> None: - if self.dtype != DataType.NVFP4: - return - # NVFP4 KV is quantized in 16-token tiles, so FP4 context resumes only - # on a full cache block (partial-block reuse is disabled). Capture the - # recurrent state at that boundary; rounding only the reused length - # would restore a later token's state. - request.expect_snapshot_points = sorted({ - point // self.tokens_per_block * self.tokens_per_block - for point in request.expect_snapshot_points - if point >= self.tokens_per_block - }) - - def _apply_branch_snapshot_point( - self, - request: LlmRequest, - context_current_position: Optional[int] = None, - ) -> None: - # prompt_len is unioned in unconditionally: try_commit_blocks derives - # commit_limit from max(snapshot_points), so a lone branch point below - # the prompt end would silently stop commits there and lose the - # prompt-end snapshot. It also makes the flag usable on its own. - points = set(request.expect_snapshot_points) - # Keep the flag usable without any configured snapshot placement. If a - # placement is configured, preserve it: adding prompt_len can replace a - # nearby recurrent checkpoint in the same radix-tree block because a - # block stores one recurrent page per life cycle. - if not points: - points.add(request.prompt_len) - point = self._branch_snapshot_points.get(request.py_request_id) - current_position = (request.context_current_position - if context_current_position is None else - context_current_position) - if point is not None and point > current_position: - points.add(point) - request.expect_snapshot_points = sorted(points) - self._align_fp4_mla_snapshot_points(request) - - def get_iteration_stats(self) -> Optional[KVCacheV2IterationStatsReport]: - """Log recurrent-cache movement; this is KDA state for Kimi K3.""" - report = super().get_iteration_stats() - if report is None: - return None - - pool_group_ids = sorted({ - pool_group_id - for pool_group_id, _, kind in - self._stats_life_cycle_metadata().values() if kind == "ssm" - }) - if not pool_group_ids: - return report - pool_group_reports = [ - report.by_pool_group[pool_group_id] - for pool_group_id in pool_group_ids - ] - - stats = [ - pool_group_report.stats for pool_group_report in pool_group_reports - ] - evicted_blocks = sum(stat.iter_offload_blocks for stat in stats) - evicted_bytes = sum(stat.iter_offload_bytes for stat in stats) - onboarded_blocks = sum(stat.iter_onboard_blocks for stat in stats) - onboarded_bytes = sum(stat.iter_onboard_bytes for stat in stats) - dropped_blocks = sum(stat.iter_host_dropped_blocks for stat in stats) - dropped_bytes = sum(stat.iter_host_dropped_bytes for stat in stats) - has_movement = bool(evicted_blocks or onboarded_blocks - or dropped_blocks) - if has_movement: - self._recurrent_evicted_blocks_total += evicted_blocks - self._recurrent_onboarded_blocks_total += onboarded_blocks - self._recurrent_dropped_blocks_total += dropped_blocks - if (self.mapping.rank == 0 - and (has_movement or not self._recurrent_status_logged)): - logger.debug( - f"[MambaHybridCacheManagerV2] recurrent cache status " - f"rank={self.mapping.rank} pool_group_ids={pool_group_ids} " - f"evicted_recurrent_blocks={evicted_blocks} " - f"evicted_recurrent_bytes={evicted_bytes} " - f"onboarded_recurrent_blocks={onboarded_blocks} " - f"onboarded_recurrent_bytes={onboarded_bytes} " - f"dropped_recurrent_blocks={dropped_blocks} " - f"dropped_recurrent_bytes={dropped_bytes} " - f"total_evicted_recurrent_blocks=" - f"{self._recurrent_evicted_blocks_total} " - f"total_onboarded_recurrent_blocks=" - f"{self._recurrent_onboarded_blocks_total} " - f"total_dropped_recurrent_blocks=" - f"{self._recurrent_dropped_blocks_total} " - f"gpu_used_recurrent_blocks=" - f"{sum(stat.primary_used_num_blocks for stat in stats)} " - f"gpu_free_recurrent_blocks=" - f"{sum(stat.primary_free_num_blocks for stat in stats)} " - f"gpu_evictable_recurrent_blocks=" - f"{sum(stat.primary_evictable_num_blocks for stat in stats)} " - f"host_used_recurrent_blocks=" - f"{sum(stat.secondary_used_num_blocks for stat in stats)} " - f"host_free_recurrent_blocks=" - f"{sum(stat.secondary_free_num_blocks for stat in stats)} " - f"{self._format_branch_snapshot_counters()}") - self._recurrent_status_logged = True - return report - - @staticmethod - def get_cache_size_per_token(model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - num_layers: Optional[int] = None, - tokens_per_block: int = 32, - max_seq_len: Optional[int] = None, - **kwargs): - spec_config = kwargs.get("spec_config") - num_reserved_dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( - spec_config, max_batch_size) + int(mapping.enable_attention_dp)) - attention_slope_override = None - attention_intercept = 0 - attention_page_bytes_override = 0 - config = model_config.pretrained_config - quant_config = model_config.quant_config - if (getattr(config, "kv_lora_rank", None) is not None - and quant_config is not None - and quant_config.quant_mode.has_fp4_kv_cache()): - from tensorrt_llm._torch.attention_backend.fp4_mla.cache_manager import \ - Fp4MlaKVCacheManagerV2 - - _, _, local_attention_layers = _get_local_mamba_cache_layout( - model_config, - mapping, - spec_config=spec_config, - is_draft=kwargs.get("is_draft", False), - use_separate_draft_kv_cache=kwargs.get( - "use_separate_draft_kv_cache", False), - ) - attention_slope_override, attention_intercept = ( - Fp4MlaKVCacheManagerV2.get_cache_size_per_token( - model_config, - mapping, - num_layers=local_attention_layers, - tokens_per_block=tokens_per_block, - max_batch_size=max_batch_size, - spec_config=spec_config, - is_disagg=kwargs.get("is_disagg", False), - )) - max_resident_sequences = max_batch_size * mapping.pp_size - if max_resident_sequences > 0: - attention_page_bytes_override = (attention_intercept // - max_resident_sequences) - - slope, intercept = _estimate_mamba_hybrid_cache_cost( - model_config, - mapping, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - num_reserved_dummy_slots=num_reserved_dummy_slots, - include_explicit_snapshots=True, - cap_partial_attention_snapshots=True, - attention_slope_override=attention_slope_override, - attention_page_bytes_override=attention_page_bytes_override, - **kwargs, - ) - return slope, intercept + attention_intercept - - def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: - if self._is_standalone_draft_layer(local_layer_idx): - return False - layer_id = self.pp_layers[local_layer_idx] - return layer_id < len( - self._mamba_layer_mask) and self._mamba_layer_mask[layer_id] - - def _get_pool_roles(self, - pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: - layer_id = int(self.impl.layer_grouping[pool_id][0]) - if self._is_standalone_draft_layer(layer_id): - return KVCacheManagerV2._get_pool_roles(self, pool_id) - if layer_id >= self.num_local_layers: - if (getattr(self, "_attention_cache_layout_policy", - None) is not None - and LayerId(layer_id) in self._hp_manager_layer_ids): - return Role.MLA_HP_TAIL, None - raise ValueError(f"Unknown hybrid V2 virtual layer {layer_id}.") - if self._is_local_mamba_layer(layer_id): - return MambaRole.SSM_STATE, None - return super()._get_pool_roles(pool_id) - - def _max_resident_sequences(self) -> int: - return self.max_batch_size * self.mapping.pp_size - - def _mamba_state_bytes_per_slot(self) -> int: - return self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) - - def _num_ssm_snapshots_for_capacity( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> int: - if capacity <= 0 or not kv_cache_config.enable_block_reuse: - return 0 - - fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, - self.max_seq_len, - self.tokens_per_block) - interval = _mamba_regular_snapshot_interval(kv_cache_config, - self.max_seq_len) - regular_snapshots = capacity // interval if interval is not None else 0 - return (self._max_resident_sequences() * fixed_rules + - regular_snapshots) - - def _num_ssm_states_per_typical_request( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> int: - fixed_rules, _ = _mamba_snapshot_rule_counts( - kv_cache_config, - capacity, - self.tokens_per_block, - ) - # Additional snapshots are stable boundaries that must remain alive. - # Periodic snapshots are evictable cache entries and therefore do not - # increase the guaranteed state count represented by BatchDesc. - return 1 + fixed_rules - - def _typical_request_descs( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> List[KVCacheDesc]: - """Model one request with one descriptor per live SSM state.""" - num_states = self._num_ssm_states_per_typical_request( - capacity, kv_cache_config) - capacity_per_state, capacity_remainder = divmod(capacity, num_states) - capacities = [ - capacity_per_state + int(i < capacity_remainder) - for i in range(num_states) - ] - return [ - KVCacheDesc( - capacity=state_capacity, - history_length=max(0, state_capacity - 1), - ) for state_capacity in capacities - ] - - def _get_typical_request_capacity( - self, - kv_cache_config: KvCacheConfig, - ) -> int: - if kv_cache_config.avg_seq_len is not None: - return kv_cache_config.avg_seq_len - - fallback_capacity = max(1, self.max_seq_len // 2) - logger.warning( - "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " - "model using KV cache manager V2. Falling back to " - f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " - "'kv_cache_config.avg_seq_len' in the YAML configuration to the " - "workload's average total sequence length for an accurate KV/SSM " - "pool ratio.") - return fallback_capacity - - def _get_quota_from_max_tokens(self, max_tokens: int) -> int: - attention_quota = super()._get_quota_from_max_tokens(max_tokens) - num_request_lineages = self._max_resident_sequences() - snapshot_slots = self._num_ssm_snapshots_for_capacity( - max_tokens, self.kv_cache_config) - state_slots = (num_request_lineages + self._num_reserved_dummy_slots + - snapshot_slots) - state_quota = state_slots * self._mamba_state_bytes_per_slot() - # Once the plan contains any non-live SSM capacity, reserve one partial - # attention page per request lineage. This remains conservative when - # the plan contains fewer than one non-live slot per lineage. - extra_attention_quota = (num_request_lineages * - self._attention_cache_bytes_per_token() * - self.tokens_per_block - if snapshot_slots > 0 else 0) - return attention_quota + state_quota + extra_attention_quota - - def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: - draft_quota = super()._get_draft_quota_from_max_tokens(max_tokens) - snapshot_slots = self._num_ssm_snapshots_for_capacity( - max_tokens, self.kv_cache_config) - if snapshot_slots > 0: - # The snapshot plan's partial-page reserve also covers draft KV. - draft_quota += (self._max_resident_sequences() * - sum(self._draft_layer_sizes) * - self.tokens_per_block) - return draft_quota - - def _get_max_tokens_from_quota(self, quota: int) -> float: - if self._get_quota_from_max_tokens(0) > quota: - return 0 - - low = 0 - high = 1 - while self._get_quota_from_max_tokens(high) <= quota: - low = high - high *= 2 - if high >= 1 << 62: - return float("inf") - - while low + 1 < high: - mid = (low + high) // 2 - if self._get_quota_from_max_tokens(mid) <= quota: - low = mid - else: - high = mid - return low - - def _minimum_live_gpu_quota(self) -> int: - """Return the minimum quota for live states and one attention page.""" - attention_block_quota = (self._attention_cache_bytes_per_token() * - self.tokens_per_block) - num_state_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - state_quota = num_state_slots * self._mamba_state_bytes_per_slot() - return max( - self._get_quota_from_max_tokens(0), - state_quota + attention_block_quota, - ) - - def _build_cache_config( - self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: - kv_cache_config = self.kv_cache_config - cache_tiers = config.cache_tiers - gpu_quota = cache_tiers[0].quota - minimum_live_quota = self._minimum_live_gpu_quota() - if minimum_live_quota > gpu_quota: - raise ValueError( - "The V2 Mamba GPU cache quota is too small for live recurrent " - f"states and attention pages: got {gpu_quota} bytes, need at " - f"least {minimum_live_quota} bytes.") - # _build_base_config already constructed every attention layer, - # including dtype-specific scale and subclass-provided side buffers. - # Preserve those configs and replace only the local Mamba layers. - layers = list(config.layers) - for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): - if self._mamba_layer_mask[global_layer_idx]: - layer_id = LayerId(local_layer_idx) - layers[local_layer_idx] = SsmLayerConfig( - layer_id=layer_id, - buffers=[ - BufferConfig(role=MambaRole.SSM_STATE, - size=self.ssm_bytes), - BufferConfig(role=MambaRole.CONV_STATE, - size=self.conv_bytes), - ], - ) - - dummy_requests = [ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(self._num_reserved_dummy_slots) - ] - constraints = [ - replace( - batch, - kv_caches=[*batch.kv_caches, *dummy_requests], - ) for batch in config.constraints - ] - - typical_step = config.typical_step - if config.initial_pool_ratio is None: - typical_capacity = self._get_typical_request_capacity( - kv_cache_config) - request_descs = self._typical_request_descs(typical_capacity, - kv_cache_config) - typical_step = BatchDesc(request_descs * - self._max_resident_sequences() + - dummy_requests) - # The recurrent (SSM) state pool must hold one slot per resident - # sequence plus every reserved dummy slot. Unlike attention pages, a - # Mamba state is fixed-size per sequence, so this floor is independent - # of sequence length. The base config only emits constraints when - # ``avg_seq_len`` is set, and speculative decoding inflates the reserved - # dummy slots (CUDA-graph padding), so without an explicit floor the SSM - # pool can be undersized (see the live/dummy-slot check in _setup_states - # / __init__). Add a min-slots constraint of zero-capacity requests: - # these cost no attention pages but reserve one SSM slot each. - if any(isinstance(layer, SsmLayerConfig) for layer in layers): - ssm_floor_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - constraints = [ - *constraints, - BatchDesc([ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(ssm_floor_slots) - ]), - ] - result = replace( - config, - layers=layers, - typical_step=typical_step, - constraints=constraints, - # SSM lifecycles require minimum-snapshot commit semantics. The - # flag is harmless when reuse is disabled because no commits are - # attempted, while the runtime config still needs the invariant. - commit_min_snapshot=True, - ) - layout_policy = getattr(self, "_attention_cache_layout_policy", None) - if layout_policy is not None: - result = layout_policy._build_cache_config(self, result) - return result - - def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, - state_shape: List[int]) -> torch.Tensor: - addr = self.impl.get_mem_pool_base_address(LayerId(local_layer_idx), - role, PageIndexMode.SHARED) - num_pages = self.impl.get_page_index_upper_bound( - LayerId(local_layer_idx), role) - raw = convert_to_torch_tensor( - TensorWrapper(addr, dtype, [num_pages] + state_shape)) - page_index_scale = self.impl.get_page_index_scale( - LayerId(local_layer_idx), role) - num_slots = (num_pages + page_index_scale - 1) // page_index_scale - # V2 coalesces same-size per-layer buffers inside each slot. Kernels - # index Mamba states by logical slot id, so expose only this layer's - # sub-page from each coalesced slot instead of the raw page-index view. - return raw.as_strided( - [num_slots] + state_shape, - [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), - ) - - def _setup_states(self) -> None: - local_layer_ids = [ - self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers - ] - self.all_ssm_states = [ - self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, - self.ssm_state_dtype, self.ssm_state_shape) - for local_layer_idx in local_layer_ids - ] - self.all_conv_states = [ - self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, - self.conv_state_dtype, self.conv_state_shape) - for local_layer_idx in local_layer_ids - ] - if self._use_gdn_cached_replay_all_layer_commit: - expected_inner_strides = ( - self.ssm_state_shape[1] * self.ssm_state_shape[2], - self.ssm_state_shape[2], - 1, - ) - reference = self.all_ssm_states[0] - for state in self.all_ssm_states: - if (state.dtype != reference.dtype - or state.device != reference.device - or list(state.shape[1:]) != self.ssm_state_shape - or state.stride()[1:] != expected_inner_strides): - raise RuntimeError( - "GDN cached replay requires V2 SSM layers with matching " - "dtype, device, shape, and dense inner dimensions.") - pointers = [state.data_ptr() for state in self.all_ssm_states] - slot_stride = reference.stride(0) - layer_stride_bytes = (pointers[1] - - pointers[0] if len(pointers) > 1 else 0) - has_affine_layout = ( - layer_stride_bytes % reference.element_size() == 0 and all( - state.stride(0) == slot_stride - for state in self.all_ssm_states) - and all(pointer == pointers[0] + layer * layer_stride_bytes - for layer, pointer in enumerate(pointers))) - if has_affine_layout: - self._gdn_cached_replay_state_strides = ( - len(pointers), - layer_stride_bytes // reference.element_size(), - slot_stride, - ) - else: - logger.warning_once( - "V2 GDN state views are not affine; using indirect replay " - "checkpoint addressing", - key="gdn_cached_replay_v2_indirect_state_layout", - ) - self._gdn_cached_replay_state_descriptors = torch.tensor( - [(state.data_ptr(), state.stride(0)) - for state in self.all_ssm_states], - dtype=torch.int64, - device=reference.device, - ) - - def _setup_replay_buffers(self, spec_config) -> None: - cache_size = 0 - device = None - if self.local_num_mamba_layers > 0: - cache_size = self.all_ssm_states[0].shape[0] - assert all(t.shape[0] == cache_size for t in self.all_ssm_states) - device = self.all_ssm_states[0].device - if not self._allocate_pool_replay_buffers(spec_config, cache_size, - device): - return - - mask_capacity = self._host_state_indices.shape[0] - self._dummy_request_mask = torch.zeros(mask_capacity, - dtype=torch.bool, - device=device) - self._dummy_request_mask_host = torch.zeros( - mask_capacity, - dtype=torch.bool, - pin_memory=prefer_pinned(), - ) - - def _attention_cache_bytes_per_token(self) -> int: - layout_policy = getattr(self, "_attention_cache_layout_policy", None) - if layout_policy is not None: - return layout_policy.get_cache_bytes_per_token(self) - # Mamba layers have zero KV heads, so the generic calculation naturally - # returns only bytes owned by local attention layers. - return super().get_cache_bytes_per_token() - - def get_cache_bytes_per_token(self) -> int: - cache_bytes = self._attention_cache_bytes_per_token() - - interval = ( - self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) - if (self.kv_cache_config.enable_block_reuse and interval is not None - and interval > 0): - cache_bytes += self._mamba_state_bytes_per_slot() // interval - if cache_bytes == 0 and self.local_num_mamba_layers > 0: - cache_bytes = self._mamba_state_bytes_per_slot() - return max(1, cache_bytes) - - def get_num_free_blocks(self) -> int: - assert len(self.kv_cache_map) == 0, ( - "get_num_free_blocks is only used when the kv cache manager is empty" - ) - if getattr(self, "_attention_cache_layout_policy", None) is not None: - return self.blocks_in_primary_pool - attention_pages = [] - ssm_pages = [] - for local_layer_idx in range(self.num_local_layers): - layer_id = LayerId(local_layer_idx) - if self._is_local_mamba_layer(local_layer_idx): - ssm_pages.append( - self.impl.get_page_index_upper_bound( - layer_id, MambaRole.SSM_STATE) // - self._ssm_page_index_scale) - else: - attention_pages.append( - self._get_attention_pool_num_blocks(local_layer_idx)) - if attention_pages: - return max(attention_pages) - return max(ssm_pages) if ssm_pages else 0 - - @property - def blocks_in_primary_pool(self) -> int: - layout_policy = getattr(self, "_attention_cache_layout_policy", None) - if layout_policy is not None: - return layout_policy.blocks_in_primary_pool.fget(self) - for local_layer_idx in range(self.num_local_layers): - if self._is_local_mamba_layer(local_layer_idx): - continue - return self.impl.get_page_index_upper_bound( - LayerId(local_layer_idx), Role.KEY) - return 0 - - def get_buffers(self, - layer_idx: int, - kv_layout: str = "NHD") -> Optional[torch.Tensor]: - local_layer_idx = self.layer_offsets[layer_idx] - if self._is_local_mamba_layer(local_layer_idx): - return None - return super().get_buffers(layer_idx, kv_layout) - - def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: - layout_policy = getattr(self, "_attention_cache_layout_policy", None) - if layout_policy is not None: - yield from self._iter_fp4_mla_physical_pool_views() - else: - for global_layer_id, local_layer_id in self.layer_offsets.items(): - if self._is_local_mamba_layer(local_layer_id): - continue - # A layer group is a lifecycle, not a physical memory pool. - # Differently sized attention buffers can share one lifecycle, - # so scan every attention layer in this diagnostic path. - yield KVCacheManagerV2.get_buffers(self, global_layer_id) - - yield from self.all_ssm_states - yield from self.all_conv_states - - def add_dummy_requests( - self, - request_ids: List[int], - token_nums: Optional[List[int]] = None, - is_gen: bool = False, - prepare_resource: bool = True, - max_num_draft_tokens: int = 0, - kv_reserve_draft_tokens: Optional[int] = None, - use_mrope: bool = False, - max_beam_width: int = 1, - encoder_output_lens: Optional[List[int]] = None, - num_extra_decoding_steps: int = 0, - draft_kv_cache_manager: Optional[BaseResourceManager] = None, - ) -> List[LlmRequest]: - requests = super().add_dummy_requests( - request_ids=request_ids, - token_nums=token_nums, - is_gen=is_gen, - prepare_resource=prepare_resource, - max_num_draft_tokens=max_num_draft_tokens, - kv_reserve_draft_tokens=kv_reserve_draft_tokens, - use_mrope=use_mrope, - max_beam_width=max_beam_width, - encoder_output_lens=encoder_output_lens, - num_extra_decoding_steps=num_extra_decoding_steps, - draft_kv_cache_manager=draft_kv_cache_manager, - ) - if requests and prepare_resource: - self._setup_state_indices(requests) - return requests - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is not None and kv_cache.is_active: - self.try_commit_blocks(request, kv_cache) - self._request_id_to_state_index.pop(request.py_request_id, None) - self._request_id_to_is_dummy.pop(request.py_request_id, None) - self._branch_snapshot_points.pop(request.py_request_id, None) - super().free_resources(request, pin_on_release) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - super().prepare_resources(scheduled_batch) - if self.local_num_mamba_layers == 0: - return - requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - self._setup_state_indices(requests) - num_contexts = len(scheduled_batch.context_requests) - self._reset_context_mamba_slots(num_contexts) - - def _setup_state_indices(self, requests: List[LlmRequest]) -> None: - if self.local_num_mamba_layers == 0: - return - n = len(requests) - assert n <= self._host_state_indices.shape[0], ( - f"State-index batch size {n} exceeds max_batch_size " - f"{self._host_state_indices.shape[0]}") - self._host_state_indices.zero_() - if n > 0: - for i, req in enumerate(requests): - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - raise RuntimeError( - f"Missing V2 KV cache for request {req.py_request_id}") - base_index = kv_cache.get_ssm_block_base_index( - self.ssm_layer_group_id) - if base_index < 0: - raise RuntimeError( - f"Invalid SSM state block index {base_index} for " - f"request {req.py_request_id}") - self._host_state_indices[i] = base_index - - _copy_from_host_snapshot(self.cuda_state_indices, - self._host_state_indices) - is_dummy = [req.is_dummy for req in requests] - self._refresh_dummy_request_mask(is_dummy) - state_values = self._host_state_indices[:n].tolist() - for req, value, dummy in zip(requests, state_values, is_dummy): - self._request_id_to_state_index[req.py_request_id] = value - self._request_id_to_is_dummy[req.py_request_id] = dummy - - def get_state_indices(self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None): - if self.local_num_mamba_layers == 0: - # Mamba metadata is still prepared on attention-only PP ranks, - # but no local kernel consumes these indices. Return harmless - # placeholders instead of consulting an intentionally empty map. - if request_ids is not None: - return [0] * len(request_ids) - return self.cuda_state_indices - if request_ids is not None: - indices = [ - self._request_id_to_state_index[rid] for rid in request_ids - ] - if is_padding is None: - is_padding = [False] * len(request_ids) - assert len(request_ids) == len(is_padding) - is_dummy = [ - self._request_id_to_is_dummy.get(rid, False) or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - return self.cuda_state_indices - - def get_max_resource_count(self) -> int: - return self.max_batch_size - - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None, - accepted_leaf_positions: Optional[torch.Tensor] = None): - if self.local_num_mamba_layers == 0: - return - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - num_accepted_draft_tokens = ( - num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( - torch.int32) - # Dynamic tree selects a tree node rather than a linear draft depth. - accepted_positions = (accepted_leaf_positions.to(torch.int32) - if accepted_leaf_positions is not None else - num_accepted_draft_tokens) - if state_indices is None: - state_indices = self.get_state_indices() - state_indices_d = state_indices[num_contexts:num_contexts + - num_gens].to(torch.int32) - src_state_indices = self.intermediate_state_indices[:num_gens] - - if self._use_replay_state_update: - self._commit_gdn_cached_replay_history_layers( - attn_metadata, num_gens) - assert self._dummy_request_mask is not None - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - replay_metadata = self.get_replay_state_update_metadata() - assert replay_metadata is not None - _advance_replay_state( - replay_metadata, - state_indices_d, - num_accepted_tokens[num_contexts:num_contexts + num_gens], - is_dummy_request, - ) - else: - for layer_offset, dst in enumerate(self.all_ssm_states): - _promote_mamba_state_triton( - dst.unsqueeze(0), - self.intermediate_ssm_states[layer_offset:layer_offset + 1], - src_state_indices, - accepted_positions, - state_indices_d, - ) - - for layer_offset, dst in enumerate(self.all_conv_states): - _promote_mamba_state_triton( - dst.unsqueeze(0), - self.intermediate_conv_states[layer_offset:layer_offset + 1], - src_state_indices, - accepted_positions, - state_indices_d, - ) - - def _mark_context_position_as_history(self, request: LlmRequest, - kv_cache) -> None: - """Advance history without making later recurrent state reusable.""" - history_length = request.context_current_position - if history_length <= kv_cache.history_length: - return - capacity = max(kv_cache.capacity, history_length) - if not kv_cache.resize(capacity, history_length=history_length): - raise ValueError( - "Failed to resize history length of V2 Mamba cache for " - f"request {request.py_request_id} to {history_length} tokens") - - def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: - should_block_reuse = (self.enable_block_reuse and not self.is_draft - and not request.is_dummy_request) - if not should_block_reuse: - return - - if kv_cache is None: - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is None: - return - - snapshot_points = request.expect_snapshot_points - commit_limit = (min(max(snapshot_points), request.prompt_len) - if snapshot_points else request.prompt_len) - commit_end = min(request.context_current_position, commit_limit) - if (request.context_current_position in request.expect_snapshot_points - and commit_end > kv_cache.num_committed_tokens): - tokens = self._augment_tokens_for_block_reuse( - request.get_tokens(DEFAULT_BEAM_INDEX), - request, - start=kv_cache.num_committed_tokens, - end=commit_end, - ) - kv_cache.commit(tokens) - if (self.local_num_mamba_layers > 0 - and self._branch_snapshot_points.get( - request.py_request_id) == commit_end): - self._branch_snapshots_taken_total += 1 - if request.context_current_position >= commit_limit: - self._mark_context_position_as_history(request, kv_cache) - if request.context_remaining_length == 0: - kv_cache.stop_committing() - - def update_context_resources(self, - scheduled_batch: ScheduledRequests) -> None: - for request in scheduled_batch.context_requests: - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is None or not kv_cache.is_active: - continue - - should_block_reuse = (self.enable_block_reuse and not self.is_draft - and not request.is_dummy_request) - is_all_reusable = ( - self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE) - is_snapshot_boundary = (request.context_current_position - in request.expect_snapshot_points) - has_pending_snapshot = any( - point > request.context_current_position - for point in request.expect_snapshot_points) - # NVFP4: the BF16 HP-tail ring is positioned by the history - # length, so it needs the completed chunk's watermark even while a - # later recurrent-state snapshot is still pending. - should_resize = (self.dtype == DataType.NVFP4 - or not should_block_reuse or - (not is_all_reusable and not has_pending_snapshot)) - should_commit = (is_all_reusable or is_snapshot_boundary - or request.context_remaining_length == 0) - - if should_resize and not kv_cache.resize( - None, request.context_current_position): - raise ValueError( - "Failed to resize history length of V2 Mamba cache for " - f"request {request.py_request_id} to " - f"{request.context_current_position} tokens at context " - "update") - if should_commit: - self.try_commit_blocks(request, kv_cache) - if request.context_remaining_length == 0: - if self.conversation_manager is not None: - self.conversation_manager.save_drop_plan(request, kv_cache) - kv_cache.enable_swa_scratch_reuse = False - - def shutdown(self): - self._gdn_cached_replay_state_descriptors = None - self._gdn_cached_replay_state_strides = None - self.all_ssm_states = [] - self.all_conv_states = [] - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - self._branch_snapshot_points.clear() - super().shutdown() - - -def _is_disagg_context_server(is_disagg: bool) -> bool: - return is_disagg and os.getenv("TRTLLM_DISAGG_ROLE") == "context" - - -def _kda_replay_layout( - num_spec: int, num_heads: int, head_dim: int, - conv_window: int) -> dict[str, tuple[torch.dtype, list[int]]]: - """Physical per-slot shapes; beta rows retain the CuTe alignment.""" - dim = num_heads * head_dim - heads_per_alignment = _KDA_BETA_CACHE_ALIGNMENT_BYTES // torch.bfloat16.itemsize - aligned_heads = ((num_heads + heads_per_alignment - 1) // - heads_per_alignment) * heads_per_alignment - return { - "kda_conv_q": (torch.bfloat16, [conv_window + num_spec, dim]), - "kda_conv_k": (torch.bfloat16, [conv_window + num_spec, dim]), - "kda_conv_v": (torch.bfloat16, [conv_window + num_spec, dim]), - "kda_k_cache": (torch.bfloat16, [num_spec, dim]), - "kda_g_cache": (torch.float32, [num_spec, dim]), - "kda_v_cache": (torch.bfloat16, [num_spec, dim]), - "kda_beta_cache": (torch.bfloat16, [num_spec, aligned_heads]), - } - - -class KDAHybridCacheManagerV2(MambaHybridCacheManagerV2): - """V2 hybrid cache manager with Kimi KDA replay state. - - The live convolution pool stores the committed ``W - 1`` input window. - Fused verification extends that window with one column per draft token and - carries the accepted draft prefix between verification rounds. - """ - - def __init__( - self, - *args, - kda_replay_num_spec: Optional[int] = None, - **kwargs, - ) -> None: - spec_config = kwargs.get("spec_config") - use_replay_state_update = kwargs.get("use_replay_state_update", False) - conv_state_layout = kwargs.get("conv_state_layout", "x_b_c") - - self.kda_context_only = _is_disagg_context_server( - kwargs.get("is_disagg", False)) - if self.kda_context_only: - kda_replay_num_spec = None - self._kda_replay_layers: list[dict[str, torch.Tensor]] = [] - self._kda_replay_num_spec = kda_replay_num_spec - self._use_kda_replay_update = kda_replay_num_spec is not None - if self._use_kda_replay_update: - if use_replay_state_update: - raise ValueError( - "kda_replay_num_spec and use_replay_state_update are " - "mutually exclusive") - if spec_config is None: - raise ValueError( - "KDA replay caches require speculative decoding") - expected_num_spec = spec_config.tokens_per_gen_step - 1 - if kda_replay_num_spec != expected_num_spec: - raise ValueError( - "KDA replay cache width " - f"({kda_replay_num_spec}) must match the draft length " - f"({expected_num_spec})") - if conv_state_layout != "q_k_v": - raise ValueError( - "KDA replay requires conv_state_layout='q_k_v'") - self._record_kda_replay_in_update_resources = ( - self._use_kda_replay_update - and getattr(spec_config, "decoding_type", None) == "NGram") - - super().__init__(*args, **kwargs) - - @staticmethod - def get_cache_size_per_token(model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - tokens_per_block: int = 32, - max_seq_len: Optional[int] = None, - **kwargs): - from ..modules.kimi_kda._kda_kernels import is_kda_mtp_verify_available - - slope, intercept = MambaHybridCacheManagerV2.get_cache_size_per_token( - model_config, - mapping, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - **kwargs) - spec_config = kwargs.get("spec_config") - if (_is_disagg_context_server(kwargs.get("is_disagg", False)) - or spec_config is None or not is_kda_mtp_verify_available()): - return slope, intercept - - params, local_layers, _ = _get_local_mamba_cache_layout( - model_config, - mapping, - spec_config=spec_config, - is_draft=kwargs.get("is_draft", False), - use_separate_draft_kv_cache=kwargs.get( - "use_separate_draft_kv_cache", False)) - if not local_layers: - return slope, intercept - - tp_size = _mamba_effective_tp_size(mapping) - layout = _kda_replay_layout(spec_config.tokens_per_gen_step - 1, - params.num_heads // tp_size, - params.head_dim, params.conv_kernel - 1) - replay_bytes_per_slot = ( - local_layers * sum(dtype.itemsize * math.prod(shape) - for dtype, shape in layout.values()) + - _KDA_REPLAY_COUNT_BYTES) - resident_sequences = max_batch_size * mapping.pp_size - dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( - spec_config, max_batch_size) + int(mapping.enable_attention_dp)) - fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, - max_seq_len, - tokens_per_block) - intercept += (resident_sequences * - (1 + fixed_rules) + dummy_slots) * replay_bytes_per_slot - - interval = _mamba_regular_snapshot_interval(kv_cache_config, - max_seq_len) - if interval is not None: - # Round the combined snapshot cost once, preserving the parent's - # attention cost and avoiding a second rounding of replay bytes. - state_bytes_per_slot = (local_layers * - params.get_states_bytes_per_layer(mapping)) - slope += (math.ceil( - (state_bytes_per_slot + replay_bytes_per_slot) / interval) - - math.ceil(state_bytes_per_slot / interval)) - return slope, intercept - - def _replay_layout(self) -> dict[str, tuple[torch.dtype, list[int]]]: - if not self._use_kda_replay_update or not self.local_num_mamba_layers: - return {} - if (len(self.conv_section_dims) != 3 - or len(set(self.conv_section_dims)) != 1): - raise ValueError("KDA replay caches require equal [Q | K | V] " - "convolution-state sections") - return _kda_replay_layout(self._kda_replay_num_spec, - self.ssm_state_shape[0], - self.ssm_state_shape[1], - self.conv_state_shape[1]) - - def _mamba_state_bytes_per_slot(self) -> int: - replay_bytes = sum(dtype.itemsize * math.prod(shape) - for dtype, shape in self._replay_layout().values()) - count_bytes = _KDA_REPLAY_COUNT_BYTES if replay_bytes else 0 - return (super()._mamba_state_bytes_per_slot() + - self.local_num_mamba_layers * replay_bytes + count_bytes) - - def _build_cache_config( - self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: - config = super()._build_cache_config(config) - replay_buffers = [ - BufferConfig(role=DataRole(name), - size=dtype.itemsize * math.prod(shape)) - for name, (dtype, shape) in self._replay_layout().items() - ] - if not replay_buffers: - return config - first_layer_id = LayerId(self.layer_offsets[self.mamba_pp_layers[0]]) - layers = [] - for layer in config.layers: - if isinstance(layer, SsmLayerConfig): - buffers = [*layer.buffers, *replay_buffers] - if layer.layer_id == first_layer_id: - # All SSM layers share a slot lifecycle. Store the request's - # count once so KVCM moves it with the replay histories. - buffers.append( - BufferConfig(role=DataRole("prev_num_accepted_tokens"), - size=_KDA_REPLAY_COUNT_BYTES)) - layer = SsmLayerConfig(layer_id=layer.layer_id, buffers=buffers) - layers.append(layer) - return replace(config, layers=layers) - - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size: int) -> None: - if self.kda_context_only: - super()._setup_mtp_intermediate_states(None, max_batch_size) - return - if not self._use_kda_replay_update: - super()._setup_mtp_intermediate_states(spec_config, max_batch_size) - return - - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if spec_config is None or self.local_num_mamba_layers == 0: - return - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - - def _validate_replay_buffer_colocation(self, - replay_roles: Iterable[str]) -> None: - """Verify replay and recurrent buffers share a logical slot lifecycle.""" - # All layers reuse the same logical slot index. Check the core's - # allocated layout, not just equal capacities: buffers in one pool - # group and lifecycle are allocated and moved together. - required_buffers = {(self.layer_offsets[layer_id], str(role)) - for layer_id in self.mamba_pp_layers - for role in (MambaRole.SSM_STATE, - MambaRole.CONV_STATE, *replay_roles)} - required_buffers.add((self.layer_offsets[self.mamba_pp_layers[0]], - "prev_num_accepted_tokens")) - if not any( - required_buffers.issubset( - {(int(buffer_id.layer_id), str(buffer_id.role)) - for coalesced in variant.coalesced_buffers - for buffer_id in coalesced.buffer_ids}) - for pool_group in self.impl.pool_group_descs - for variant in pool_group.slot_desc.variants): - raise RuntimeError( - "KDA replay and recurrent-state buffers must share one KVCM " - "pool group and lifecycle to reuse the same logical slot index") - - def _allocate_pool_replay_buffers( - self, - spec_config, - cache_size: int, - device: Optional[torch.device], - ) -> bool: - allocated = super()._allocate_pool_replay_buffers( - spec_config, cache_size, device) - self._kda_replay_layers = [] - layout = self._replay_layout() - if not layout: - return allocated - self._validate_replay_buffer_colocation(layout) - for layer_id in self.mamba_pp_layers: - layer = {} - for name, (dtype, shape) in layout.items(): - view = self._get_state_buffer(self.layer_offsets[layer_id], - DataRole(name), dtype, shape) - if name in ("kda_conv_q", "kda_conv_k", "kda_conv_v"): - view = view.transpose(-1, -2) - elif name == "kda_beta_cache": - view = view[..., :self.ssm_state_shape[0]] - # Pool capacity can exceed physically resident memory. Only - # initialize rows after the request locks its pages. - layer[name] = view - self._kda_replay_layers.append(layer) - self.prev_num_accepted_tokens = self._get_state_buffer( - self.layer_offsets[self.mamba_pp_layers[0]], - DataRole("prev_num_accepted_tokens"), torch.int32, - [_KDA_REPLAY_COUNT_BYTES // torch.int32.itemsize])[..., 0] - return True - - def _reset_context_mamba_slots(self, num_contexts: int) -> None: - super()._reset_context_mamba_slots(num_contexts) - if num_contexts and self._use_kda_replay_update: - self._reset_kda_replay_slots( - self._host_state_indices[:num_contexts].tolist()) - - def _reset_kda_replay_slots(self, slots: Union[List[int], - torch.Tensor]) -> None: - if not self._kda_replay_layers: - return - device = self.prev_num_accepted_tokens.device - if isinstance(slots, torch.Tensor) and slots.device.type != "cpu": - slots = slots.to(device=device, dtype=torch.long, non_blocking=True) - else: - if isinstance(slots, torch.Tensor): - slots = slots.tolist() - if not slots: - return - slots = _host_slots_to_device(slots, device) - if slots.numel(): - self.prev_num_accepted_tokens.index_fill_(0, slots, 0) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - KVCacheManagerV2.prepare_resources(self, scheduled_batch) - if self.local_num_mamba_layers == 0: - return - requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - num_contexts = len(scheduled_batch.context_requests) - self._setup_state_indices(requests, num_contexts=num_contexts) - self._reset_context_mamba_slots(num_contexts) - - def _setup_state_indices(self, - requests: List[LlmRequest], - num_contexts: Optional[int] = None) -> None: - if not self._use_kda_replay_update: - super()._setup_state_indices(requests) - return - - replay_requests = (requests - if num_contexts is None else requests[num_contexts:]) - old_state_values = [ - self._request_id_to_state_index.get(request.py_request_id, -1) - for request in replay_requests - ] - super()._setup_state_indices(requests) - if self.local_num_mamba_layers == 0: - return - - new_state_values = [ - self._request_id_to_state_index[request.py_request_id] - for request in replay_requests - ] - self._reset_kda_replay_slots([ - new for old, new in zip(old_state_values, new_state_values) - if old < 0 - ]) - - def _record_kda_replay_acceptance( - self, - state_indices: torch.Tensor, - num_accepted_draft_tokens: torch.Tensor, - is_dummy_request: Optional[torch.Tensor] = None, - ) -> None: - """Record the draft prefix that the next KDA verify must replay.""" - slots = state_indices.to(torch.long) - accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) - counts = self.prev_num_accepted_tokens - values = (torch.where(is_dummy_request, counts[slots], accepted) - if is_dummy_request is not None else accepted) - counts[slots] = values - - @torch.inference_mode() - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - """Seed replay caches from transferred ``W - 1`` conv states.""" - if (not self._use_kda_replay_update - or self.local_num_mamba_layers == 0): - return - slots = [ - self._request_id_to_state_index[request_id] - for request_id in request_ids - if request_id in self._request_id_to_state_index - ] - if not slots: - return - - device = self.prev_num_accepted_tokens.device - state_indices = _host_slots_to_device(sorted(set(slots)), device) - self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0) - committed_window = self.conv_state_shape[1] - section_offsets = [0] - for section_dim in self.conv_section_dims: - section_offsets.append(section_offsets[-1] + section_dim) - - for layer, conv_state in zip(self._kda_replay_layers, - self.all_conv_states): - for replay_buffer in layer.values(): - replay_buffer.index_fill_(0, state_indices, 0) - selected_conv = conv_state.index_select(0, state_indices) - for section_idx, name in enumerate( - ("kda_conv_q", "kda_conv_k", "kda_conv_v")): - replay_layer = layer[name] - section = selected_conv[:, section_offsets[section_idx]: - section_offsets[section_idx + 1], :] - replay_layer[:, :, :committed_window].index_copy_( - 0, state_indices, section.to(replay_layer.dtype)) - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - if self.kda_context_only: - return PythonMambaCacheManager.State( - conv=self.get_conv_states(layer_idx), - temporal=self.get_ssm_states(layer_idx)) - if not self._use_kda_replay_update: - return super().mamba_layer_cache(layer_idx) - - layer_offset = self.mamba_layer_offsets[layer_idx] - spec_kwargs = dict(self._kda_replay_layers[layer_offset]) - spec_kwargs["prev_num_accepted_tokens"] = self.prev_num_accepted_tokens - if self.mamba_ssm_rand_seed is not None: - spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed - return PythonMambaCacheManager.SpeculativeState( - conv=self.get_conv_states(layer_idx), - temporal=self.get_ssm_states(layer_idx), - intermediate_conv_window=None, - **spec_kwargs, - ) - - def update_resources(self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None) -> None: - KVCacheManagerV2.update_resources(self, scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - if (not self._record_kda_replay_in_update_resources - or self.local_num_mamba_layers == 0): - return - - generation_requests = scheduled_batch.generation_requests - drafted_requests = [ - request for request in generation_requests - if request.py_draft_tokens is not None - and len(request.py_draft_tokens) > 0 - ] - if not drafted_requests: - return - if len(drafted_requests) != len(generation_requests): - raise RuntimeError( - "Mixed drafted/undrafted generation batch is not supported " - "for KDA replay bookkeeping") - - active_requests = [ - request for request in generation_requests - if request.py_request_id in self._request_id_to_state_index - ] - if not active_requests: - return - device = self.prev_num_accepted_tokens.device - state_indices = torch.tensor([ - self._request_id_to_state_index[request.py_request_id] - for request in active_requests - ], - dtype=torch.int32, - device=device) - accepted_drafts = torch.tensor([ - request.py_num_accepted_draft_tokens for request in active_requests - ], - dtype=torch.int32, - device=device) - is_dummy_request = torch.tensor([ - self._request_id_to_is_dummy.get(request.py_request_id, False) - for request in active_requests - ], - dtype=torch.bool, - device=device) - self._record_kda_replay_acceptance(state_indices, accepted_drafts, - is_dummy_request) - - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None, - accepted_leaf_positions: Optional[torch.Tensor] = None): - if self.kda_context_only: - return - if not self._use_kda_replay_update: - return super().update_mamba_states(attn_metadata, - num_accepted_tokens, - state_indices, - accepted_leaf_positions) - if self.local_num_mamba_layers == 0: - return - - num_contexts = attn_metadata.num_contexts - num_gens = attn_metadata.num_seqs - num_contexts - num_accepted_draft_tokens = ( - num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( - torch.int32) - if state_indices is None: - state_indices = self.get_state_indices() - state_indices_d = state_indices[num_contexts:num_contexts + - num_gens].to(torch.int32) - assert self._dummy_request_mask is not None - is_dummy_request = self._dummy_request_mask[num_contexts:num_contexts + - num_gens] - self._record_kda_replay_acceptance( - state_indices_d, - num_accepted_draft_tokens, - is_dummy_request, - ) - - def shutdown(self): - self._kda_replay_layers.clear() - self.prev_num_accepted_tokens = None - super().shutdown() diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b54cb6f5122c..40789830cb7d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -37,25 +37,14 @@ from ..distributed.communicator import init_pp_comm from ..distributed.ops import MNNVLAllReduce from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader -from ..models.modeling_multimodal_encoder import MultimodalEncoderMixin -from ..models.modeling_multimodal_mixin import (MultimodalEncoderContractError, - MultimodalModelMixin, - _build_request_multimodal_input) -from ..models.modeling_multimodal_utils import filter_mm_token_from_input_ids -from ..models.modeling_speculative import SpecDecOneEngineForCausalLM -from ..models.modeling_utils import DecoderModelForCausalLM -from ..modules.fused_moe.moe_load_balancer import (MoeLoadBalancer, - MoeLoadBalancerIterContext) -from ..modules.mamba.mamba2_metadata import Mamba2Metadata -from ..peft.lora.cuda_graph_lora_manager import CudaGraphLoraManager -from ..speculative import (SpecMetadata, get_draft_kv_cache_manager, - get_num_extra_kv_tokens, get_spec_metadata, - prepare_attn_metadata_for_draft_replay, - restore_attn_metadata_after_draft_replay, - update_spec_config_from_loaded_model) -from ..speculative.drafting_loops import BaseDraftingLoopWrapper -from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata -from ..speculative.spec_sampler_base import SampleStateTensorsSpec +from ..models.modeling_multimodal_mixin import MultimodalModelMixin +from ..models.modeling_utils import DecoderModelForCausalLM, timing_metric +from ..moe.expert_statistic import ExpertStatistic +from ..moe.fused_moe.moe_comm_timeout_guard import set_moe_comm_warmup +from ..moe.fused_moe.moe_load_balancer import MoeLoadBalancer +from ..nccl_window_tensor_scope import install_eager_nccl_window_tensor_scopes +from ..route_capture import ROUTE_CAPTURE_ATTR, RouteCapture +from ..speculative import SpecMetadata, update_spec_config_from_loaded_model from ..tensor_lru_cache import TensorLRUCache from ..utils import set_torch_compiling from .config_utils import is_hybrid_linear @@ -574,100 +563,6 @@ def _init_execution_services(self) -> None: self._runner = self._initialize_runner(runner_cls) - # Once encoder CUDA graphs are usable, enable mixed decoder graphs by - # default unless the user explicitly opts out. - encoder_decoder_cuda_graph_enabled = ( - self.encoder_cuda_graph_runner.enabled - and self.encoder_cuda_graph_runner.is_encoder_decoder - and bool(self.encoder_cuda_graph_runner.capture_keys)) - enable_encoder_decoder_mixed_cuda_graph = ( - encoder_decoder_cuda_graph_enabled - and self.cuda_graph_config is not None - and self.llm_args.enable_encoder_decoder_mixed_cuda_graph) - - # Create decoder CUDA graph config and runner. - cuda_graph_runner_config = CUDAGraphRunnerConfig( - use_cuda_graph=(not self._is_encode_only - and self.cuda_graph_config is not None), - cuda_graph_padding_enabled=self._cuda_graph_padding_enabled, - cuda_graph_batch_sizes=self._cuda_graph_batch_sizes, - max_cuda_graph_batch_size=self._max_cuda_graph_batch_size, - max_beam_width=self.max_beam_width, - spec_config=self.spec_config, - cuda_graph_mem_pool=self._cuda_graph_mem_pool, - dynamic_draft_len_mapping=self._dynamic_draft_len_mapping, - max_num_tokens=self.max_num_tokens, - use_mrope=self.use_mrope, - original_max_draft_len=self.original_max_draft_len, - original_max_total_draft_tokens=self. - original_max_total_draft_tokens, - is_draft_model=self.is_draft_model, - enable_attention_dp=self.enable_attention_dp, - is_encoder_decoder=self._is_encoder_decoder_model(), - batch_size=self.batch_size, - mapping=self.mapping, - dist=self.dist, - kv_cache_manager_key=self.kv_cache_manager_key, - sparse_attention_config=self.sparse_attention_config, - enable_encoder_decoder_mixed_cuda_graph=( - enable_encoder_decoder_mixed_cuda_graph), - ) - self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) - self.breakable_cuda_graph_runner = None - if self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.BREAKABLE: - decoder_model = (self.model if isinstance( - self.model, DecoderModelForCausalLM) else getattr( - self.model, "llm", None)) - if not isinstance(decoder_model, DecoderModelForCausalLM): - raise ValueError( - "breakable prefill CUDA graph requires a decoder model body" - ) - self.breakable_cuda_graph_runner = BreakableCUDAGraphRunner( - decoder_model.model) - - # Initialize CUDA Graph LoRA manager if LoRA is enabled - self.cuda_graph_lora_manager: Optional[CudaGraphLoraManager] = None - - # Setup the local cache indirection buffer only once and reuse it. - # This way it can also be used for CUDA graphs. - if self.use_beam_search: - self.cache_indirection_attention = torch.zeros( - (self.batch_size, self.max_beam_width, self.max_seq_len), - device="cuda", - dtype=torch.int32) - else: - self.cache_indirection_attention = None - - self.kv_cache_dtype_byte_size = self.get_kv_cache_dtype_byte_size() - - self._prepare_inputs_event: Optional[torch.cuda.Event] = None - - # Cache for enc-dec cross-attention stable generation steps. - # Populated on the first CUDA-graph generation step; cleared whenever - # the batch composition changes (new encoder request arrives). - self._cross_attn_stable_cached_tokens: Optional[List[int]] = None - self._cross_attn_stable_request_ids: Optional[List[int]] = None - - def register_forward_pass_callable(self, callable: Callable): - self.forward_pass_callable = callable - - def get_speculative_model(self) -> Optional[SpecDecOneEngineForCausalLM]: - """Return the one-engine speculative model, including inside a VLM.""" - if not self.is_spec_decode: - return None - model = self.model - if isinstance(model, MultimodalModelMixin): - model = model.language_model - if not isinstance(model, SpecDecOneEngineForCausalLM): - # Kimi's legacy multimodal wrapper exposes its decoder as llm - # without inheriting MultimodalModelMixin. - model = getattr(model, "llm", None) - return model if isinstance(model, SpecDecOneEngineForCausalLM) else None - - def get_kv_cache_dtype_byte_size(self) -> float: - """ - Returns the size (in bytes) occupied by kv cache type. - """ layer_quant_mode = self.model.model_config.quant_config.layer_quant_mode if layer_quant_mode.has_fp4_kv_cache(): self.kv_cache_dtype_byte_size = 1 / 2 @@ -1117,107 +1012,11 @@ def get_mm_encoder_item_cache_keys( return None return self._mm_item_scheduler.item_cache_keys(request) - def _resolve_bytes_per_mm_encoder_embedding(self) -> int: - """Bytes occupied by one multimodal encoder output embedding. - - Prefers the mixin's explicit `embedding_dim`/`embedding_dtype` - contract, then the text embedding layer's weight, then the - pretrained config's hidden size with the loaded weights' dtype — - both mixin properties are optional and most VLMs implement - neither. - """ - model = self.model - try: - return (model.embedding_dim * torch.empty( - (), dtype=model.embedding_dtype).element_size()) - except NotImplementedError: - pass - try: - weight = model.text_embedding_layer.weight - return weight.shape[-1] * weight.element_size() - except NotImplementedError: - pass - pretrained = model.model_config.pretrained_config - hidden_size = getattr(pretrained, "hidden_size", None) - if hidden_size is None: - hidden_size = getattr(getattr(pretrained, "text_config", None), - "hidden_size", None) - if hidden_size is None: - raise ValueError( - "Cannot derive the MM embedding row size: the model " - "implements neither embedding_dim/embedding_dtype nor " - "text_embedding_layer, and its pretrained config exposes " - "no (text_config.)hidden_size") - element_size = next(model.parameters()).dtype.itemsize - return hidden_size * element_size - - def _set_up_multimodal_encoder_attn_metadata(self) -> None: - """Construct AttentionMetadata for any multimodal encoders inside the - loaded model, using the engine's encoder token budget - (`encoder_max_num_tokens`, falling back to the LLM-side - `max_num_tokens`). - - Mirrors `_set_up_attn_metadata` for the LLM backbone: encoders opt in - by inheriting `MultimodalEncoderMixin`, and the engine drives the construction - so the sizes match the resolved encoder token budget rather than being - hardcoded inside each encoder's `__init__`. The optional - per-segment capacity combines the encoder token budget with the input - processor's largest supported item. - """ - max_seq_len = self.encoder_max_num_tokens - if isinstance(self.input_processor, BaseMultimodalDummyInputsBuilder): - max_tokens_per_item = ( - self.input_processor.get_mm_max_tokens_per_item()) - max_seq_len = max(max_seq_len, - max(max_tokens_per_item.values(), default=0)) - - for module in self.model.modules(): - if isinstance(module, MultimodalEncoderMixin): - setup_kwargs: Dict[str, Any] = dict( - max_num_tokens=self.encoder_max_num_tokens) - if self.mm_encoder_attention_metadata_capacity is not None: - setup_kwargs["attention_metadata_capacity"] = ( - self.mm_encoder_attention_metadata_capacity) - module.setup_attn_metadata(**setup_kwargs) - module.set_attn_max_seq_len(max_seq_len) - - def _set_up_spec_metadata( - self, - spec_resource_manager: Optional[BaseResourceManager], - no_cache=False): - spec_config = self.spec_config if self.enable_spec_decode else None - # The disaggregated attention-DP overlap path opts into larger metadata - # buffers. Passing None preserves the established max_num_requests - # fallback for other configurations, including PP. - # Size slot-indexed spec buffers (e.g. rejection-sampling draft_probs) - # by the same pool that sizes the SeqSlotManager - # (py_executor_creator passes model_engine.max_num_seq_slots), so every - # py_seq_slot has a row. With max_num_requests only, DSpark + rejection - # sampling on disaggregated attention-DP gen indexed past the buffer. - num_seq_slots = self.max_num_seq_slots - if no_cache: - return get_spec_metadata( - spec_config, - self.model.config, - self.batch_size, - max_num_tokens=self.max_num_tokens, - spec_resource_manager=spec_resource_manager, - is_draft_model=self.is_draft_model, - max_seq_len=self.max_seq_len, - num_seq_slots=num_seq_slots) - - if self.spec_metadata is not None: - return self.spec_metadata - self.spec_metadata = get_spec_metadata( - spec_config, - self.model.config, - self.batch_size, - max_num_tokens=self.max_num_tokens, - spec_resource_manager=spec_resource_manager, - is_draft_model=self.is_draft_model, - max_seq_len=self.max_seq_len, - num_seq_slots=num_seq_slots) - return self.spec_metadata + def invalidate_multimodal_encoder_cache(self) -> None: + """Clear cached MM encoder outputs when no request is using them.""" + encoder_cache = self.mm_encoder_cache + if encoder_cache is not None: + encoder_cache.clear() def cleanup(self) -> None: """Release resources owned by this model engine. @@ -1565,660 +1364,24 @@ def forward(self, resource_manager: Optional[ResourceManager] = None, new_tensors_device: Optional[SampleStateTensors] = None, cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None): - if not self._disable_overlap_scheduler: - # Do not refill reusable host staging while the previous - # iteration's asynchronous H2D copies still consume it. This - # event precedes the model forward, so its synchronization retains - # GPU forward overlap while protecting every forward entry, - # including speculative/draft paths. - self.wait_for_input_copy() - - kv_cache_manager = resource_manager.get_resource_manager( - self.kv_cache_manager_key) - draft_kv_cache_manager = self._get_draft_kv_cache_manager( - resource_manager) - - attn_metadata = self._set_up_attn_metadata(kv_cache_manager, - draft_kv_cache_manager) - if isinstance(attn_metadata, TrtllmAttentionMetadata): - attn_metadata.trtllm_gen_jit_warmup = self._trtllm_gen_jit_warmup - if self.enable_spec_decode: - spec_resource_manager = resource_manager.get_resource_manager( - ResourceManagerType.SPEC_RESOURCE_MANAGER) - spec_tree_manager = None - if spec_resource_manager is not None and hasattr( - spec_resource_manager, 'spec_tree_manager'): - spec_tree_manager = spec_resource_manager.spec_tree_manager - spec_metadata = self._set_up_spec_metadata(spec_resource_manager, - no_cache=kv_cache_manager - is None) - # attn_metadata now depends on spec_metadata since it determines the shape/content of spec_dec parameter Tensors - is_spec_dec_mode = spec_metadata.spec_dec_mode.attention_need_spec_dec_mode( - spec_resource_manager, self.is_draft_model, self.attn_backend, - self.model_is_wrapped) - # Propagate runtime_draft_len (already set on self by py_executor) - # to spec_metadata so downstream code (eagle3, interface, trtllm) can read it. - spec_metadata.runtime_draft_len = self.runtime_draft_len - spec_metadata.runtime_tokens_per_gen_step = ( - self.get_runtime_tokens_per_gen_step(self.runtime_draft_len)) - - # Parallel-draft modes advertise a per-gen-step width via - # tokens_per_gen_step (PARD: 2K, DFlash: K+1). Pass - # (tokens_per_gen_step - 1) so generation_lengths = tokens_per_gen_step - # and the XQA kernel computes the correct past_kv_len. - if spec_metadata.spec_dec_mode.is_parallel_draft(): - sd_max_draft_len = self.original_max_total_draft_tokens - sd_max_total = self.original_max_total_draft_tokens - else: - sd_max_draft_len = self.original_max_draft_len - sd_max_total = self._spec_dec_max_total_draft_tokens - - # Fill slot-ID buffer for update_spec_dec_param - if (spec_tree_manager is not None - and spec_tree_manager.use_dynamic_tree - and not self.is_draft_model): - spec_tree_manager.slot_storage.fill_all_slot_ids( - scheduled_requests.context_requests, - scheduled_requests.generation_requests, - ) - - attn_metadata.update_spec_dec_param( - batch_size=scheduled_requests.batch_size, - is_spec_decoding_enabled=is_spec_dec_mode, - is_spec_dec_tree=spec_metadata.is_spec_dec_tree, - is_spec_dec_dynamic_tree=spec_metadata.is_spec_dec_dynamic_tree, - max_draft_len=sd_max_draft_len, - max_total_draft_tokens=sd_max_total, - model_is_wrapped=self.model_is_wrapped, - spec_metadata=spec_metadata, - spec_tree_manager=spec_tree_manager, - num_contexts=scheduled_requests.num_context_requests) - else: - spec_resource_manager = None - spec_metadata = None - - moe_load_balancer: MoeLoadBalancer = getattr(self, 'moe_load_balancer', - None) - if kv_cache_manager is None: - inputs, gather_ids = self._prepare_tp_inputs_no_cache( - scheduled_requests, attn_metadata, spec_metadata, - resource_manager) - - with MoeLoadBalancerIterContext(moe_load_balancer): - # Special handling for multimodal encoder only mode - if self.llm_args.mm_encoder_only: - return self._forward_step_mm_encoder_only( - inputs, scheduled_requests) - else: - return self._forward_step( - inputs, - gather_ids=gather_ids, - gather_context_logits=gather_context_logits) - - graph_requests = scheduled_requests - promoted_context_request_ids: frozenset[int] = frozenset() - # Non-linear tree input preparation expands runtime_draft_len to the - # total tree width after graph selection. Only linear-tree zero-draft - # iterations can therefore safely reuse a zero-draft graph. - can_promote_spec_decode = (not self.enable_spec_decode - or (not self.is_draft_model - and self.runtime_draft_len == 0 - and self.spec_config is not None - and self.spec_config.is_linear_tree)) - # TODO: Generalize these conservative gates as actual-draft, beam, and - # context-parallel providers for decoder-only LLMs gain support for - # promoted final-context rows. Each relaxation must preserve whole-batch - # fallback on graph miss and prove parity with the provider's native - # q_len=1 path. Encoder-decoder and non-LLM engines remain out of scope. - if (scheduled_requests.num_context_requests > 0 - and self.cuda_graph_runner.enabled and can_promote_spec_decode - and not self.use_beam_search - and not self._is_encoder_decoder_model() - and not self._is_encode_only - and not self.llm_args.mm_encoder_only - and self.mapping.cp_size == 1): - graph_requests, promoted_context_request_ids = \ - _make_single_token_context_graph_batch( - scheduled_requests, - self._is_final_multimodal_context_decode_compatible) - - with self.cuda_graph_runner.pad_batch( - graph_requests, resource_manager, - self.runtime_draft_len) as padded_graph_requests: - # Callee already no-ops when use_mrope=False, but the Python call / - # frame setup itself is non-trivial under high concurrency. Gating - # at the caller avoids that overhead for non-mrope models. - if self.use_mrope: - self._pad_batch_seed_mrope_delta_cache(padded_graph_requests) - - # Refresh is_all_greedy_sample for the *current* batch BEFORE the - # CUDA graph key is built below. The key includes this flag to pick - # the argmax vs advanced-sampling graph variant; populate (inside - # _prepare_inputs) runs later and fills the matching GPU buffers. - # Without this pre-scan the key would use the previous iteration's - # stale value and could replay the advanced graph against - # unpopulated (greedy) buffers, hanging the run (e.g. MTP nextn>=2). - if spec_metadata is not None: - spec_metadata.update_is_all_greedy_sample( - padded_graph_requests.all_requests()) - self._sync_group_all_greedy_sample(spec_metadata) - - peft_cache_data_type = None - if getattr(self, "cuda_graph_lora_manager", None) is not None: - peft_cache_manager = resource_manager.get_resource_manager( - ResourceManagerType.PEFT_CACHE_MANAGER) - peft_cache_data_type = peft_cache_manager.data_type - - maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( - padded_graph_requests, - enable_spec_decode=self.enable_spec_decode, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - draft_tokens_cuda=self.draft_tokens_cuda - if self.is_spec_decode else None, - new_tensors_device=new_tensors_device, - spec_resource_manager=spec_resource_manager, - promoted_context_request_ids=promoted_context_request_ids, - peft_cache_data_type=peft_cache_data_type, - ) - - can_run_graph = key is not None - if can_run_graph: - attn_metadata = maybe_attn_metadata - spec_metadata = maybe_spec_metadata - execution_requests = padded_graph_requests - execution_promoted_context_ids = promoted_context_request_ids - else: - attn_metadata = self.attn_metadata - if self.enable_spec_decode: - spec_metadata = self.spec_metadata - else: - spec_metadata = None - execution_requests = scheduled_requests - execution_promoted_context_ids = frozenset() - - # Fill slot-ID buffer for scatter inside draft loop - if (self.enable_spec_decode and spec_tree_manager is not None - and spec_tree_manager.use_dynamic_tree - and not self.is_draft_model): - spec_tree_manager.slot_storage.fill_all_slot_ids( - execution_requests.context_requests, - execution_requests.generation_requests, - ) - - inputs, gather_ids = self._prepare_inputs( - execution_requests, kv_cache_manager, attn_metadata, - spec_metadata, new_tensors_device, cache_indirection_buffer, - num_accepted_tokens_device, req_id_to_old_request, - resource_manager, can_run_graph, execution_promoted_context_ids) - speculative_model = self.get_speculative_model() - spec_worker = (speculative_model.spec_worker - if speculative_model is not None else None) - if spec_worker is not None: - spec_worker.prepare_managed_draft_cache( - speculative_model.draft_model, spec_metadata, attn_metadata, - resource_manager) - if execution_promoted_context_ids: - self.iter_states[ - 'num_ctx_requests'] = scheduled_requests.num_context_requests - self.iter_states['num_ctx_tokens'] = sum( - request.context_chunk_size - for request in scheduled_requests.context_requests) - self.iter_states[ - 'num_generation_tokens'] = scheduled_requests.num_generation_requests - self._prepare_inputs_event = torch.cuda.Event() - self._prepare_inputs_event.record() - - breakable_runner = self.breakable_cuda_graph_runner - - with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): - - def forward_step(): - with MoeLoadBalancerIterContext(moe_load_balancer): - return self._forward_step( - inputs, - gather_ids=gather_ids, - gather_context_logits=gather_context_logits) - - if not can_run_graph: - if (breakable_runner is not None - and breakable_runner.is_capturing): - return breakable_runner.capture_model_body(forward_step) - - num_tokens = inputs['input_ids'].shape[0] - can_run_breakable_graph = ( - breakable_runner is not None - and get_per_request_prefill_cuda_graph_flag() - and not gather_context_logits - and breakable_runner.has_graph(num_tokens)) - if can_run_breakable_graph and not breakable_runner.is_warming_up: - outputs = breakable_runner.execute( - num_tokens, forward_step) - else: - # real eager or BCG warmup or PCG - outputs = forward_step() - else: - needs_capture = self.cuda_graph_runner.needs_capture(key) - if needs_capture: - - def capture_forward_fn(inputs: Dict[str, Any]): - with MoeLoadBalancerIterContext(moe_load_balancer): - return self._forward_step( - inputs, - gather_ids=gather_ids, - gather_context_logits=gather_context_logits) - - def capture_postprocess_fn(inputs: Dict[str, Any]): - self._postprocess_inputs(inputs) - - capture_outputs = self.cuda_graph_runner.capture( - key, - capture_forward_fn, - inputs, - enable_spec_decode=self.enable_spec_decode, - postprocess_fn=capture_postprocess_fn) - - if self.cuda_graph_runner.is_warmup_only: - outputs = capture_outputs - elif needs_capture: - # Refresh attention metadata for the current batch's - # draft cache before replaying the captured graph. - saved_draft = prepare_attn_metadata_for_draft_replay( - attn_metadata, draft_kv_cache_manager) - try: - outputs = self.cuda_graph_runner.replay(key, inputs) - finally: - restore_attn_metadata_after_draft_replay( - attn_metadata, saved_draft) - else: - saved_draft = prepare_attn_metadata_for_draft_replay( - attn_metadata, draft_kv_cache_manager) - try: - with MoeLoadBalancerIterContext(moe_load_balancer): - outputs = self.cuda_graph_runner.replay( - key, inputs) - finally: - restore_attn_metadata_after_draft_replay( - attn_metadata, saved_draft) - - if (spec_worker is not None and not self.is_warmup - and not self.cuda_graph_runner.is_warmup_only): - draft_history_update = spec_worker.snapshot_managed_draft_history( - ) - if draft_history_update is not None: - # Graph output dictionaries persist across replays. Each - # overlapped iteration must retain its own host readback. - outputs = dict(outputs) - outputs['draft_history_update'] = draft_history_update - - if self.forward_pass_callable is not None: - self.forward_pass_callable() - - self._execute_logit_post_processors(scheduled_requests, outputs) - - return outputs - - def _get_spec_worker(self): - """Access the spec_worker from DecoderModelForCausalLM (one-model spec dec).""" - return getattr(self.model, 'spec_worker', None) - - def model_forward(self, **kwargs): - attrs = get_model_extra_attrs() - assert attrs is not None, "Model extra attrs is not set" - attrs["attention_metadata"] = weakref.ref(kwargs['attn_metadata']) - attrs.update(self.model.model_config.extra_attrs) - attrs["spec_metadata"] = kwargs.get('spec_metadata', None) - - if self._torch_compile_backend is not None: - # Register aux streams and events to model extra attrs. - # The streams and events are list which could be updated during compilation. - attrs["aux_streams"] = weakref.ref(self.backend_num_streams) - attrs["events"] = weakref.ref(self._torch_compile_backend.events) - attrs["global_stream"] = torch.cuda.current_stream() - - if is_trace_enabled("TLLM_TRACE_MODEL_FORWARD"): - return trace_func(self.model.forward)(**kwargs) - else: - return self.model.forward(**kwargs) - - @nvtx_range("_forward_step") - def _forward_step(self, - inputs: Dict[str, Any], - *, - gather_ids: Optional[torch.Tensor] = None, - gather_context_logits: bool = False) -> Dict[str, Any]: - inputs = self._preprocess_inputs(inputs) - if inputs.get('spec_metadata', None): - gather_ids = inputs['spec_metadata'].gather_ids - - # For simplicity, just return all the the logits if we have special gather_ids - # from speculative decoding. - outputs = self.model_forward( - **inputs, - return_context_logits=gather_ids is not None - or gather_context_logits, - ) - - if self.without_logits: - return outputs - - if isinstance(outputs, dict): - # If the model returns a dict, get the logits from it. All other keys are kept. - logits = outputs.get('logits', None) - # If the logits are not found, no further processing is needed. - if logits is None: - return outputs - else: - # If the model returns a single tensor, assume it is the logits and wrap it in a dict. - logits = outputs - outputs = {'logits': logits} - - # If we have special gather_ids, gather the logits - if gather_ids is not None: - outputs['logits'] = logits[gather_ids] - - return outputs - - @nvtx_range("_forward_step_mm_encoder_only") - def _forward_step_mm_encoder_only( - self, inputs: Dict[str, Any], - scheduled_requests: ScheduledRequests) -> Dict[str, Any]: - """Forward step for multimodal encoder only mode - returns mm_embeddings instead of logits.""" - # Get multimodal parameters from inputs - multimodal_params = inputs.get("multimodal_params", []) - if not multimodal_params or len(multimodal_params) == 0: - # Return empty embeddings if no multimodal data - return { - 'mm_embeddings': [], - 'mm_embedding_request_indices': [], - 'mm_embedding_lengths': [], - } - # Some ctx requests carry only mrope metadata (no actual vision - # content). Skip them so the encoder only runs on real image payloads. - mm_context_requests = [(request_idx, request) for request_idx, request - in enumerate(scheduled_requests.context_requests) - if request.py_multimodal_data is not None] - if len(mm_context_requests) != len(multimodal_params): - raise ValueError( - "mm_encoder_only expects one multimodal payload per context " - "request carrying py_multimodal_data") - mm_request_indices_with_payload = [] - mm_params_with_payload = [] - mm_embedding_lengths = [] - for (request_idx, - request), multimodal_param in zip(mm_context_requests, - multimodal_params): - if not _has_mm_payload_keys(request.py_multimodal_data): - # mrope-only warmup request (no actual vision content) -> skip. - continue - multimodal_embedding_lengths = get_multimodal_embedding_lengths( - request) - if multimodal_embedding_lengths is None: - # Vision payload keys present but no pre-computed embedding - # lengths — skip to avoid a downstream sum(None) TypeError. - continue - mm_request_indices_with_payload.append(request_idx) - mm_params_with_payload.append(multimodal_param) - mm_embedding_lengths.append(multimodal_embedding_lengths) - if not mm_params_with_payload: - return { - 'mm_embeddings': [], - 'mm_embedding_request_indices': [], - 'mm_embedding_lengths': [], - } - # For mm_encoder_only mode, we only run the vision encoder part - # The model should be a vision encoder (e.g., Qwen2VisionModelBase) - mm_embeddings = self.model.forward(mm_params_with_payload) - assert len( - mm_embeddings - ) == 1, "mm_embeddings should be a 1-element list, mix modality (video+image) is not supported" - - split_lengths = [sum(lengths) for lengths in mm_embedding_lengths] - mm_embeddings = list(torch.split(mm_embeddings[0], split_lengths, - dim=0)) - if len(mm_embeddings) != len(mm_embedding_lengths): - raise ValueError( - "mm_encoder_only produced an embedding batch that does not " - "match mm_embedding_lengths") - - # Extract mrope position data from multimodal_params if available - mrope_position_ids_list = [] - mrope_position_deltas_list = [] - for multimodal_param in mm_params_with_payload: - mrope_config = multimodal_param.multimodal_data.get( - 'mrope_config', {}) - mrope_position_ids = mrope_config.get('mrope_position_ids') - mrope_position_deltas = mrope_config.get('mrope_position_deltas') - if mrope_position_ids is not None: - mrope_position_ids_list.append(mrope_position_ids) - if mrope_position_deltas is not None: - mrope_position_deltas_list.append(mrope_position_deltas) - - # mrope lists must align 1:1 with multimodal_params (or be empty); - # the sampler indexes them by per-MM-result position into mm_embeddings. - assert (len(mrope_position_ids_list) == len(mrope_position_deltas_list) - and len(mrope_position_ids_list) - in (0, len(mm_params_with_payload))), ( - f"mrope alignment: got {len(mrope_position_ids_list)} ids, " - f"{len(mrope_position_deltas_list)} deltas, " - f"{len(mm_params_with_payload)} mm params") - - result = { - 'mm_embeddings': mm_embeddings, - 'logits': None, - 'mm_embedding_request_indices': mm_request_indices_with_payload, - 'mm_embedding_lengths': mm_embedding_lengths, - } - if mrope_position_ids_list: - result['mrope_position_ids'] = mrope_position_ids_list - if mrope_position_deltas_list: - result['mrope_position_deltas'] = mrope_position_deltas_list - - return result - - def _make_encoder_attn_metadata( - self, - sequence_lengths: List[int], - request_ids: List[int], - ): - """Build fresh, no-cache attention metadata for one packed encoder - batch. ``self.attn_metadata`` is not reused because that object is - bound to the decoder's KV-cache manager.""" - if len(sequence_lengths) != len(request_ids): - raise ValueError("Encoder sequence lengths and request IDs must " - "have the same length.") - sparse_metadata_params = ( - self.sparse_attention_config.to_sparse_metadata_params( - pretrained_config=self.model.model_config.pretrained_config) - if self.sparse_attention_config is not None else None) - encoder_attn_metadata = self.attn_backend.Metadata( - max_num_requests=self.encoder_batch_size, - max_num_tokens=self.encoder_max_num_tokens, - max_num_sequences=self.encoder_batch_size * self.max_beam_width, - kv_cache_manager=None, - mapping=self.mapping, - runtime_features=self.attn_runtime_features, - enable_flash_mla=self.model.model_config.enable_flash_mla, - enable_context_mla_with_cached_kv=False, - cache_indirection=None, - sparse_metadata_params=sparse_metadata_params, - num_heads_per_kv=1, - ) - assert isinstance( - encoder_attn_metadata, - (VanillaAttentionMetadata, TrtllmAttentionMetadata) - ), "Only vanilla and trtllm attention metadata are supported for the encoder pass" - - encoder_attn_metadata.seq_lens = torch.tensor( - sequence_lengths, - dtype=torch.int, - pin_memory=prefer_pinned(), - ) - encoder_attn_metadata.num_contexts = len(sequence_lengths) - encoder_attn_metadata.max_seq_len = self.max_seq_len - encoder_attn_metadata.request_ids = request_ids - encoder_attn_metadata.prepare_encoder_only() - return encoder_attn_metadata - - def _prepare_encoder_decoder_encoder_inputs( - self, - encoder_input_ids: List[int], - encoder_position_ids: List[int], - sequence_lengths: List[int], - request_ids: List[int], - resource_manager: Optional[ResourceManager] = None, - ) -> Dict[str, Any]: - num_tokens = len(encoder_input_ids) - if num_tokens != len(encoder_position_ids): - raise ValueError("Encoder input IDs and position IDs must have " - "the same length.") - assert num_tokens <= self.encoder_max_num_tokens, ( - f"encoder packed length ({num_tokens}) exceeds " - f"encoder_max_num_tokens ({self.encoder_max_num_tokens})") - - encoder_attn_metadata = self._make_encoder_attn_metadata( - sequence_lengths, request_ids) - encoder_input_ids_t = torch.tensor(encoder_input_ids, - dtype=torch.int, - pin_memory=prefer_pinned()) - encoder_position_ids_t = torch.tensor(encoder_position_ids, - dtype=torch.int, - pin_memory=prefer_pinned()) - encoder_graph_runner = self.encoder_cuda_graph_runner - encoder_batch_size = len(sequence_lengths) - use_graph_staging = ( - encoder_graph_runner.enabled and - (encoder_batch_size in encoder_graph_runner.supported_batch_sizes or - (encoder_graph_runner.padding_enabled and encoder_batch_size - <= encoder_graph_runner.max_supported_batch_size))) - - return { - 'encoder_input_ids': - (encoder_input_ids_t if use_graph_staging else - encoder_input_ids_t.to('cuda', non_blocking=True)), - 'encoder_position_ids': - ((encoder_position_ids_t - if use_graph_staging else encoder_position_ids_t.to( - 'cuda', non_blocking=True)).unsqueeze(0)), - 'encoder_attn_metadata': - encoder_attn_metadata, - 'encoder_seq_lens': - sequence_lengths, - 'encoder_input_ids_host': - encoder_input_ids_t, - 'encoder_position_ids_host': - encoder_position_ids_t, - 'resource_manager': - resource_manager, - } - - @nvtx_range("_prepare_tp_inputs_encoder_features") - def _prepare_tp_inputs_encoder_features( - self, - encoder_requests: List[LlmRequest], - resource_manager: Optional[ResourceManager] = None, - ): - """Pack encoder inputs for feature-driven audio encoders (Whisper). - - The encoder input is a per-request feature tensor (an opaque audio - tensor, e.g. Whisper's 30 s-padded waveform) rather than token ids, and - the packed sequence lengths are the post-encoder position counts - (``encoder_output_len``), not the raw feature length. - """ - features: List[torch.Tensor] = [] - sequence_lengths: List[int] = [] - request_ids: List[int] = [] - - for request in encoder_requests: - request_features = request.py_encoder_input_features - if request_features is None: - raise ValueError( - f"Encoder request {request.py_request_id} has no " - "encoder_input_features; feature- and token-driven " - "encoder requests cannot share one batch.") - features.append(request_features) - sequence_lengths.append(int(request.encoder_output_len)) - request_ids.append(request.py_request_id) - - num_tokens = sum(sequence_lengths) - assert num_tokens <= self.encoder_max_num_tokens, ( - f"encoder packed length ({num_tokens}) exceeds " - f"encoder_max_num_tokens ({self.encoder_max_num_tokens})") - - encoder_attn_metadata = self._make_encoder_attn_metadata( - sequence_lengths, request_ids) - - inputs = { - 'input_features': - torch.cat(features, dim=0).to('cuda', non_blocking=True), - 'encoder_attn_metadata': - encoder_attn_metadata, - 'encoder_seq_lens': - sequence_lengths, - 'resource_manager': - resource_manager, - } - return inputs - - @nvtx_range("_prepare_tp_inputs_encoder") - def _prepare_tp_inputs_encoder( - self, - encoder_requests: List[LlmRequest], - resource_manager: Optional[ResourceManager] = None, - ): - """Pack encoder-side inputs for an encoder-decoder forward pass. - - Mirrors the no-cache path used by ``mm_encoder_only`` and the - legacy ``EncoderBuffers`` shape contract: ``encoder_input_ids`` - and ``encoder_position_ids`` are concatenated across requests - into a single ``[sum(encoder_output_len)]`` tensor, with one - non-causal :class:`AttentionMetadata` describing the packed - encoder batch. - - The encoder pass does not touch any KV-cache pool. The cross pool is - only written by the decoder's cross-attention on the first context - step. Self-pool blocks for the decoder are reserved on the next - scheduler iteration when the request transitions to ``CONTEXT_INIT``. - """ - if not encoder_requests: - raise ValueError( - "_prepare_tp_inputs_encoder called with no encoder requests") - - # Feature-driven audio encoders (Whisper) carry a tensor instead of - # encoder token ids; they take a dedicated prep path (which rejects - # mixed feature/token batches). - if any( - getattr(request, "py_encoder_input_features", None) is not None - for request in encoder_requests): - return self._prepare_tp_inputs_encoder_features( - encoder_requests, resource_manager=resource_manager) - - encoder_input_ids: List[int] = [] - encoder_position_ids: List[int] = [] - sequence_lengths: List[int] = [] - request_ids: List[int] = [] - - for request in encoder_requests: - tokens = request.encoder_tokens - if tokens is None: - raise ValueError( - f"Encoder request {request.py_request_id} has no " - "encoder_tokens; encoder_input_token_ids must be wired " - "through executor_request_to_llm_request.") - seq_len = len(tokens) - encoder_input_ids.extend(tokens) - encoder_position_ids.extend( - self._apply_position_id_offset(list(range(seq_len)))) - sequence_lengths.append(seq_len) - request_ids.append(request.py_request_id) - - return self._prepare_encoder_decoder_encoder_inputs( - encoder_input_ids=encoder_input_ids, - encoder_position_ids=encoder_position_ids, - sequence_lengths=sequence_lengths, - request_ids=request_ids, + *, + enable_spec_decode: bool = False, + runtime_draft_len: int = 0): + if isinstance(batch, PackedInputs): + assert isinstance(self._runner, PackedModelRunner), ( + "a packed batch requires a packed-batch runner") + return self._runner.forward(batch) + assert resource_manager is not None, ( + "scheduled execution requires a resource manager") + inputs = make_scheduled_inputs(batch, + new_tensors_device, + cache_indirection_buffer, + enable_spec_decode=enable_spec_decode, + runtime_draft_len=runtime_draft_len) + # Only a warmup scope makes the pass dummy; padding requests in a + # serving batch do not. + return self._forward_scheduled( + inputs, resource_manager=resource_manager, is_dummy=self._warmup_depth > 0, ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c52321250b31..873a469bd652 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -9084,7 +9084,11 @@ def _sample_async(self, scheduled_batch, if sample_state is not None: sample_state.draft_history_update = batch_outputs.get( 'draft_history_update') + self._maybe_record_hang_diagnostic_phase( + "sampling_returned", scheduled_batch) return sample_state + self._maybe_record_hang_diagnostic_phase("sampling_skipped", + scheduled_batch) except Exception as e: self._maybe_record_hang_diagnostic_phase("sampling_error", scheduled_batch) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 6fcedbea0505..397d67a33c41 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -43,13 +43,16 @@ from ..resource_manager import ResourceManager from ..scheduler import ScheduledRequests from ..standalone_draft_cache import DraftHistoryUpdate -from .beam_search import BeamHistoryBuilder, BeamSearchHandler, finalize_beam, prepare_beam_search -from .finish_reasons import FinishReasonsHandler -from .logprobs import ( - LogProbsState, - LogProbsStateList, - LogProbsStore, - store_logprobs_list_to_request, +from .beam_search import ( + BeamHistory, + BeamHistoryBuilder, + BeamSearchHandler, + BeamSearchParams, + BeamSearchStep, + BeamSearchStore, + PendingBeamHistories, + finalize_beam, + prepare_beam_search, ) from .finish_reasons import FinishReasonsHandler from .logprobs import LogProbsHandler, LogProbsState, LogProbsStateList, LogProbsStore diff --git a/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py b/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py index ce224f3a4649..3e049ef8eb23 100644 --- a/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py +++ b/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py @@ -9,7 +9,7 @@ import torch if TYPE_CHECKING: - from .kv_cache_manager_v2 import KVCacheManagerV2 + from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @dataclass(frozen=True) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 71b2cad0af3e..8a8ce92796b8 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -31,10 +31,10 @@ from ..attention.backends import AttentionMetadata from ..modules.logits_processor import LogitsProcessor +from ..pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID -from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager -from ..pyexecutor.resource_manager import BaseResourceManager, ResourceManagerType +from ..pyexecutor.resource_manager import BaseResourceManager, ResourceManager, ResourceManagerType from ..pyexecutor.standalone_draft_cache import DraftHistoryUpdate from .accept_stats import maybe_create_recorder from .dflash_attention import ( @@ -57,7 +57,6 @@ if TYPE_CHECKING: from ...llmapi.llm_args import DFlashDecodingConfig - from ..pyexecutor.resource_manager import ResourceManager def dflash_context_dtype(draft_model: nn.Module) -> torch.dtype: @@ -617,24 +616,41 @@ def prepare(self): # prefilling, so _store_prefill_context -- the only other assigner -- # never runs and every request would land on the dummy slot. num_contexts = max(0, len(self.request_ids) - self.num_generations) - for rid in self.request_ids[num_contexts:]: + assigned = {} + for request_idx, rid in enumerate(self.request_ids[num_contexts:], num_contexts): if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - if worker._assign_slot(rid) is None: + slot = worker._assign_slot(rid, updates=assigned) + if slot is None: # Not fatal -- spec dec is lossless, so the dummy-slot # fallback only costs acceptance. Warn so it is not silent. logger.warning( f"DFlash: no free drafter slot for generation request {rid}; " "it shares the dummy slot and will draft poorly." ) + elif worker._ctx_unified_cache: + # New requests have no in-flight predecessor. Subsequent + # accepted tokens advance the device length during replay. + token_offset = sum(self.seq_lens[:request_idx]) + assigned[slot] = min( + int(self.host_position_ids[token_offset]), worker._max_ctx + ) + worker._write_ctx_len(assigned) # Route unknown request IDs (cuda graph padding or warmup dummies) - # to dummy slot to avoid corrupting real request's context + # to dummy slot to avoid corrupting real request's context. A + # generation row with the idle id is a dummy even when a context + # warmup stored a prompt under that id: its length is not this row's. mapping = torch.tensor( - [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], + [ + worker._dummy_slot + if i >= num_contexts and rid == ATTENTION_DP_DUMMY_REQUEST_ID + else worker._req_to_slot.get(rid, worker._dummy_slot) + for i, rid in enumerate(self.request_ids) + ], dtype=torch.long, ) worker._batch_to_slot_buffer.update(mapping) @@ -805,6 +821,9 @@ def __init__( # the table width on the default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None self._ctx_pool_idx = 0 + self._ctx_layer_ids = None + self._ctx_unified_cache = False + self._ctx_block_offset = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index self._ctx_kv_indptr = None self._ctx_kv_last_page_len = None @@ -868,7 +887,7 @@ def _validate_draft_attention_backend(self, draft_model) -> None: draft_model, "dflash_attention_backend", None ) if draft_model_dflash_attention_backend is None: - raise ValueError("DFlash draft model is missing dflash_attention_backend.") + raise ValueError("DFlash draft model is missing dflash_attention.backends.") if self._dflash_attention_backend == "AUTO": # The draft model resolved AUTO against its own family and this # build; adopt that, because _lazy_init_ctx_buffers keys paging and @@ -1016,142 +1035,6 @@ def _argmax_block_tokens( local_block.reshape(-1, shard), out_dtype=torch.int32 ).view(num_gens, block_len) - def get_draft_kv_cache_manager( - self, resource_manager: Optional["ResourceManager"] - ) -> Optional[BaseResourceManager]: - if resource_manager is not None: - manager = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) - if bool(getattr(manager, "draft_layer_ids", ())): - return manager - return super().get_draft_kv_cache_manager(resource_manager) - - def _has_unified_draft_cache(self) -> bool: - return bool(getattr(self._ctx_kv_manager, "draft_layer_ids", ())) - - def _prepare_managed_history(self, request_ids: list[int]) -> list[int]: - """Restore newly resident histories without rewinding overlapping iterations.""" - updates = {self._dummy_slot: 0} - offsets = {self._dummy_slot: 0} - restore_rows = [] - real_request_ids = { - request_id - for request_id in request_ids - if request_id != ATTENTION_DP_DUMMY_REQUEST_ID - and request_id < self._graph_dummy_id_floor - and request_id not in self._ctx_kv_manager._draft_dummy_request_ids - } - for request_id in list(self._req_to_slot): - if request_id not in real_request_ids: - slot = self._req_to_slot.pop(request_id) - updates[slot] = 0 - offsets[slot] = 0 - self._req_ctx_pos.pop(request_id, None) - self._managed_cache_bindings.pop(request_id, None) - self._free_slots.append(slot) - for row, request_id in enumerate(request_ids): - if request_id in real_request_ids: - self._assign_slot(request_id) - slot = self._req_to_slot.get(request_id) - if slot is not None: - cache = self._ctx_kv_manager.kv_cache_map[request_id] - if self._managed_cache_bindings.get(request_id) is cache: - continue - history = self._ctx_kv_manager.get_draft_history(request_id) - updates[slot] = history.valid_length if history is not None else 0 - self._req_ctx_pos[request_id] = history.position if history is not None else 0 - offsets[slot] = self._req_ctx_pos[request_id] - updates[slot] - self._managed_cache_bindings[request_id] = cache - restore_rows.append(row) - self._write_ctx_len(updates) - self._write_slot_values(self._ctx_position_offset, offsets) - self._batch_to_slot[: len(request_ids)].copy_( - torch.tensor( - [self._req_to_slot.get(request_id, self._dummy_slot) for request_id in request_ids], - dtype=torch.long, - device="cpu", - pin_memory=prefer_pinned(), - ), - non_blocking=True, - ) - self._managed_request_ids = tuple( - request_id for request_id in request_ids if request_id in real_request_ids - ) - # Each upload owns fresh host storage. PyTorch's pinned allocator defers - # recycling it until the asynchronous copy completes. - self._managed_snapshot_slots[: len(self._managed_request_ids)].copy_( - torch.tensor( - [self._req_to_slot[request_id] for request_id in self._managed_request_ids], - dtype=torch.long, - device="cpu", - pin_memory=prefer_pinned(), - ), - non_blocking=True, - ) - return restore_rows - - def _gather_managed_context(self, request_ids: list[int], restore_rows: list[int]) -> None: - """Refresh VANILLA's dense staging; manager pages remain authoritative.""" - if self._dflash_attention_backend != "VANILLA": - return - for row in restore_rows: - request_id = request_ids[row] - slot = self._req_to_slot.get(request_id) - if slot is None: - continue - length = self._ctx_len_host[slot] - positions = torch.arange(length, device=self._ctx_k_buf.device) - pages = self._ctx_block_tables[row, positions // self._ctx_page_size].long() - offsets = positions % self._ctx_page_size - for layer_idx, pool in enumerate(self._ctx_kv_buf): - self._ctx_k_buf[slot, layer_idx, :length] = pool[pages, 0, :, offsets, :] - self._ctx_v_buf[slot, layer_idx, :length] = pool[pages, 1, :, offsets, :] - - def prepare_managed_draft_cache( - self, draft_model, spec_metadata, attn_metadata, resource_manager - ) -> None: - """Refresh managed inputs before eager execution, capture, or graph replay.""" - manager = self.get_draft_kv_cache_manager(resource_manager) - if not getattr(manager, "draft_layer_ids", ()): - return - self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata, manager) - spec_metadata._dflash_worker = self - restore_rows = self._prepare_managed_history(spec_metadata.request_ids) - self._refresh_ctx_block_tables( - attn_metadata, attn_metadata.num_seqs, spec_metadata.request_ids - ) - self._gather_managed_context(spec_metadata.request_ids, restore_rows) - - def _write_slot_values(self, target: torch.Tensor, values: dict[int, int]) -> None: - """Scatter a slot-to-value mapping into a device tensor with one upload.""" - if not values: - return - slots, vals = zip(*values.items()) - pinned = prefer_pinned() - target.index_copy_( - 0, - torch.tensor(slots, dtype=torch.long, pin_memory=pinned).to( - target.device, non_blocking=True - ), - torch.tensor(vals, dtype=target.dtype, pin_memory=pinned).to( - target.device, non_blocking=True - ), - ) - - def snapshot_managed_draft_history(self) -> DraftHistoryUpdate | None: - """Queue iteration-owned readback; the executor publishes after completion.""" - if not self._has_unified_draft_cache() or not self._managed_request_ids: - return None - # Preparation, forward, snapshot gathers and the next preparation run - # on the execution stream, so its next upload cannot overtake this read. - slots = self._managed_snapshot_slots[: len(self._managed_request_ids)] - lengths = self._ctx_len[slots] - positions = lengths + self._ctx_position_offset[slots] - return DraftHistoryUpdate.capture( - self._ctx_kv_manager, - self._managed_request_ids, - torch.stack((lengths, positions), dim=1), - ) - def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype, kv_factor=2): """Fail with the arithmetic before allocating the drafter context arena. @@ -1239,13 +1122,25 @@ def _managed_ctx_pool( # it: _util.py skips that bail for an external drafter, which is # what DSpark is, so this drafter does get a manager under DP. return None - unified = bool(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) - get_buffers = ( - draft_kv_cache_manager.get_draft_buffers - if unified - else draft_kv_cache_manager.get_buffers - ) - layers = [get_buffers(i, kv_layout="HND") for i in range(L)] + draft_layer_ids = tuple(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) + layer_ids = draft_layer_ids if draft_layer_ids else tuple(range(L)) + if len(layer_ids) != L: + logger.warning( + f"DFlash: KV cache manager exposes {len(layer_ids)} draft layers " + f"but the drafter has {L}; using the private ctx arena." + ) + return None + self._ctx_layer_ids = layer_ids + unified = bool(getattr(draft_kv_cache_manager, "_draft_layer_sizes", ())) + if unified: + layers = [ + draft_kv_cache_manager.get_draft_buffers(i, kv_layout="HND") for i in range(L) + ] + else: + layers = [ + draft_kv_cache_manager.get_buffers(layer_id, kv_layout="HND") + for layer_id in layer_ids + ] base = layers[0] page_size = draft_kv_cache_manager.tokens_per_block if unified else base.size(-2) # kv_factor 1 is the MLA drafter's SELFKONLY pool: one latent, no V. @@ -1272,8 +1167,10 @@ def _managed_ctx_pool( ) return layers - def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page_size) -> bool: - """Validate the legacy pool layout and resolve its offset decode constant. + def _init_ctx_block_tables( + self, draft_kv_cache_manager, pool, num_slots, L, nkv, hd, page_size + ) -> bool: + """Size the per-iteration block table and the offset decode constant. ``kv_cache_block_offsets`` entries are ``pool_block_index * num_pool_layers * kv_factor`` plus the K/V field index. What that has to @@ -1283,10 +1180,42 @@ def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page """ mgr = draft_kv_cache_manager pool_mapping = mgr.kv_cache_pool_mapping - layer_offset = mgr.layer_offsets[0] + first_layer_id = self._ctx_layer_ids[0] if self._ctx_layer_ids is not None else 0 + layer_offset = mgr.layer_offsets[first_layer_id] self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) + per_layer_tables = bool(getattr(mgr, "_use_per_layer_page_tables", False)) num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) - if num_pool_layers != L: + self._ctx_block_offset = 0 + if isinstance(mgr, KVCacheManagerV2): + converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) + layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] + unsupported = None + if mgr.enable_swa_scratch_reuse: + unsupported = "DFlash managed context does not support SWA scratch pages" + for layer_id in self._ctx_layer_ids: + local_id = mgr.layer_offsets[layer_id] + if ( + mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id + or mgr.get_layer_page_index_scale(layer_id) != converter.scale + ): + unsupported = ( + "DFlash's shared context block table requires draft layers " + "with the same lifecycle and page-index scale" + ) + break + if unsupported is not None: + # A unified manager was selected for this layout, so the drafter + # cannot move to private storage without breaking allocation, + # reuse and transfer. A separate manager keeps the private arena. + if self._ctx_unified_cache: + raise NotImplementedError(unsupported) + logger.warning(f"{unsupported}; using the private ctx arena.") + return False + if per_layer_tables: + # Attention tables address the physical pool origin, whereas + # get_buffers() starts at this layer's own K buffer. + self._ctx_block_offset = int(converter.layer_offset) + elif num_pool_layers != L: logger.warning( f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " f"the drafter has {L}; using the private ctx arena." @@ -1295,12 +1224,23 @@ def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page # V1 slices a page-major pool, so stride(0) spans every layer; V2 wraps # the layer base in a dense view over interleaved pages, so stride(0) is # one layer slot. Guessing wrong is silent: acceptance quietly drops. - layer_block_elems = mgr.kv_factor * nkv * page_size * hd + manager_kv_factor = ( + mgr.get_layer_kv_factor(first_layer_id) + if hasattr(mgr, "get_layer_kv_factor") + else mgr.kv_factor + ) + if manager_kv_factor != getattr(self, "_ctx_kv_factor", manager_kv_factor): + logger.warning( + "DFlash: draft KV layer factor does not match the drafter; " + "using the private ctx arena." + ) + return False + layer_block_elems = manager_kv_factor * nkv * page_size * hd stride0 = pool[0].stride(0) if stride0 == layer_block_elems: - self._ctx_block_divisor = mgr.kv_factor + self._ctx_block_divisor = manager_kv_factor elif stride0 == num_pool_layers * layer_block_elems: - self._ctx_block_divisor = num_pool_layers * mgr.kv_factor + self._ctx_block_divisor = num_pool_layers * manager_kv_factor else: logger.warning( f"DFlash: draft KV pool layer stride {stride0} matches neither V1's " @@ -1309,13 +1249,22 @@ def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page f"rather than guessing block indices." ) return False + max_blocks = mgr.max_blocks_per_seq + self._ctx_block_tables = torch.zeros( + (num_slots, max_blocks), dtype=torch.int32, device="cuda" + ) + self._ctx_block_indptr = torch.arange( + 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" + ) + self._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") logger.info( - f"DFlash: ctx block offsets pool_idx={self._ctx_pool_idx}, " - f"divisor={self._ctx_block_divisor}, view_pages={pool[0].size(0)}" + f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}, " + f"pool_idx={self._ctx_pool_idx}, divisor={self._ctx_block_divisor}, " + f"view_pages={pool[0].size(0)}" ) return True - def _init_ctx_block_tables(self, num_slots: int, max_blocks: int) -> None: + def _init_managed_ctx_block_tables(self, num_slots: int, max_blocks: int) -> None: self._ctx_block_tables = torch.zeros( (num_slots, max_blocks), dtype=torch.int32, device="cuda" ) @@ -1325,6 +1274,23 @@ def _init_ctx_block_tables(self, num_slots: int, max_blocks: int) -> None: self._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") logger.info(f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}") + @staticmethod + def _check_block_within_allocation( + num_ctx: torch.Tensor, allocated_ctx_limit: torch.Tensor, block_size: int + ) -> None: + """Raise when a request's drafter block would leave its allocated pages. + + The shared manager sizes the request, so a shortfall there would have the + block's scratch write land on the last context position, or on another + request's page. Asynchronous, so it adds no host sync and is safe to + capture in a CUDA graph; on the device it traps instead of raising. + """ + torch._assert_async( + (num_ctx <= allocated_ctx_limit).all(), + f"DFlash: the drafter's block of {block_size} positions falls outside the " + "pages the shared KV cache manager allocated for the request", + ) + def _refresh_ctx_block_tables( self, attn_metadata, num_seqs: int, request_ids: list[int] | None = None ) -> bool: @@ -1356,132 +1322,72 @@ def _refresh_ctx_block_tables( non_blocking=True, ) return True - src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None) + offsets_name = ( + "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + src = getattr(attn_metadata, offsets_name, None) if src is None: # The table is bound but unfillable. Leaving it at zeros would send # every request's every layer to pool block 0 -- silently. raise RuntimeError( "DFlash: bound to the draft KV pool but " f"{type(attn_metadata).__name__} carries no " - "draft_kv_cache_block_offsets this iteration." + f"{offsets_name} this iteration." ) # Field 0 of the K/V pair: both name the same block for a kv_factor=2 # pool, and the drafter addresses blocks, not planes. The block count # is taken before the clamp erases the placeholders; it is informative # only on the python V2 backend, the cpp one already mapped # BAD_PAGE_INDEX to 0. + encoded = src[self._ctx_pool_idx, :num_seqs, 0] + if self._ctx_block_offset: + encoded = encoded.to(torch.int64) - self._ctx_block_offset decode_draft_block_table( - src[self._ctx_pool_idx, :num_seqs, 0], + encoded, self._ctx_block_divisor, self._ctx_block_tables, self._ctx_block_counts, ) return True - def _init_unified_draft_tables(self, manager) -> None: - """Choose how the unified pool's block table is built each iteration.""" - self._device_draft_tables = False - layer = manager.draft_layer_ids[0] - pool = manager.layer_to_pool_mapping_dict[manager.layer_offsets[layer]] - if ( - getattr(manager, "_helix_cp_size", 1) != 1 - or getattr(manager, "enable_swa_scratch_reuse", False) - or manager._index_scale_ints[pool] != manager.get_layer_page_index_scale(layer) - ): + def _lazy_init_ctx_buffers( + self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None + ): + self._validate_draft_attention_backend(draft_model) + self._ctx_unified_cache = ( + draft_kv_cache_manager is not None + and draft_kv_cache_manager is getattr(attn_metadata, "kv_cache_manager", None) + ) + offsets_name = ( + "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" + ) + if draft_kv_cache_manager is not None and not hasattr(attn_metadata, offsets_name): + if self._ctx_unified_cache: + raise NotImplementedError( + f"DFlash unified KV requires {offsets_name} in attention metadata" + ) logger.warning( - "DFlash: the draft pool's block offsets cannot be decoded on the device; " - "building draft block tables on the host." + f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; " + "using the private ctx arena." ) - return - self._ctx_pool_idx = pool - self._ctx_block_divisor = manager.get_layer_kv_factor(layer) - self._device_draft_tables = True - self._draft_table_checks_left = 8 - logger.info( - f"DFlash: draft block tables decoded on the device from pool {pool} " - f"(divisor {self._ctx_block_divisor})" - ) + draft_kv_cache_manager = None - def _refresh_unified_block_tables( - self, attn_metadata, num_seqs: int, request_ids: list[int] - ) -> bool: - """Decode the draft pool's rows of the target's device block offsets. + unified = bool(getattr(draft_kv_cache_manager, "_draft_layer_sizes", ())) + if self._ctx_buf_inited: + if self._ctx_kv_manager is draft_kv_cache_manager: + return + # KV cache estimation swaps a probe manager for the real one. + # Holding the old pool's tensors across that writes freed memory. + logger.info("DFlash: draft KV cache manager replaced, rebinding the ctx cache.") + self._ctx_buf_inited = False - The target metadata already uploads every pool's offsets for this - batch, so the draft rows need one device divide plus the per-request - block counts, rather than a host table rebuilt from page lists. - Returns False when the host path must build this iteration's table. - """ - src = getattr(attn_metadata, "kv_cache_block_offsets", None) - if src is None or src.size(0) <= self._ctx_pool_idx: - return False - width = min(src.size(-1), self._ctx_block_tables.size(1)) - encoded = src[self._ctx_pool_idx, :num_seqs, 0, :width] - decoded = torch.div(encoded.clamp(min=0), self._ctx_block_divisor, rounding_mode="floor") - self._ctx_block_tables[:num_seqs, :width].copy_(decoded) - cache_map = self._ctx_kv_manager.kv_cache_map - counts = [] - for request_id in request_ids[:num_seqs]: - cache = cache_map.get(request_id) - counts.append(cache.num_blocks if cache is not None else 0) - self._ctx_block_counts[:num_seqs].copy_( - torch.tensor(counts, dtype=torch.long, pin_memory=prefer_pinned()), - non_blocking=True, - ) - if self._draft_table_checks_left > 0 and not torch.cuda.is_current_stream_capturing(): - self._draft_table_checks_left -= 1 - if not self._check_device_block_tables(num_seqs, request_ids, counts): - return False - return True - - def _check_device_block_tables( - self, num_seqs: int, request_ids: list[int], counts: list[int] - ) -> bool: - """Compare device-decoded rows with the manager's page lists (startup only).""" - manager = self._ctx_kv_manager - rows = [] - for row, request_id in enumerate(request_ids[:num_seqs]): - cache = manager.kv_cache_map.get(request_id) - if counts[row] > 0 and cache is not None and cache.is_active: - rows.append(row) - if not rows: - return True - expected = manager.get_batch_cache_indices( - [request_ids[row] for row in rows], manager.draft_layer_ids[0] - ) - device_rows = self._ctx_block_tables[rows].cpu() - for k, row in enumerate(rows): - got = device_rows[k, : counts[row]].tolist() - want = expected[k][: counts[row]] - if any(w >= 0 and g != w for g, w in zip(got, want)): - logger.error( - "DFlash: device-decoded draft block table disagrees with the manager " - f"for request {request_ids[row]}; building draft tables on the host." - ) - self._device_draft_tables = False - return False - return True - - def _lazy_init_ctx_buffers( - self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None - ): - self._validate_draft_attention_backend(draft_model) - unified = bool(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) - if self._ctx_buf_inited: - if self._ctx_kv_manager is draft_kv_cache_manager: - return - # KV cache estimation swaps a probe manager for the real one. - # Holding the old pool's tensors across that writes freed memory. - logger.info("DFlash: draft KV cache manager replaced, rebinding the ctx cache.") - self._ctx_buf_inited = False - - # Recorded even when the buffers stay private, so that the estimation - # pass falling back to the arena still rebinds once the real pool lands. - self._ctx_kv_manager = draft_kv_cache_manager - self._ctx_block_tables = None - self._ctx_block_counts = None - self._paged_latent_prepare = False - max_batch = spec_metadata.max_num_requests + # Recorded even when the buffers stay private, so that the estimation + # pass falling back to the arena still rebinds once the real pool lands. + self._ctx_kv_manager = draft_kv_cache_manager + self._ctx_block_tables = None + self._ctx_block_counts = None + self._paged_latent_prepare = False + max_batch = spec_metadata.max_num_requests # ctx_len is 1:1 with the target's positions, so max_seq_len bounds it. # max_position_embeddings does not: K3's drafter advertises 1048576, @@ -1536,8 +1442,13 @@ def _lazy_init_ctx_buffers( self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") self._ctx_position_offset = torch.zeros_like(self._ctx_len) self._ctx_len_host = [0] * num_slots - self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") - self._managed_snapshot_slots = torch.empty_like(self._batch_to_slot) if unified else None + self._batch_to_slot_buffer = DeviceCopyBuffer( + torch.zeros(max_batch, dtype=torch.long, device="cuda"), + pin_memory=prefer_pinned(), + ) + self._managed_snapshot_slots = ( + torch.empty_like(self._batch_to_slot_buffer.tensor) if unified else None + ) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} @@ -1570,6 +1481,7 @@ def _lazy_init_ctx_buffers( nkv = draft_model._num_kv_heads hd = draft_model._head_dim kv_factor = getattr(draft_model, "_kv_factor", 2) + self._ctx_kv_factor = kv_factor capacity = self._max_ctx + self._compute_block_size # Published before any drafter table is built, so a drafter may read it # instead of its config cap. _max_ctx is attn_metadata.max_seq_len, which @@ -1624,6 +1536,8 @@ def _lazy_init_ctx_buffers( ): pool = None if pool is None: + if self._ctx_unified_cache: + raise ValueError("DFlash could not bind the unified manager's draft KV layers") # Per-slot reservation is an arena-only notion; the pool path # takes its pages from the manager one iteration at a time. self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size @@ -1651,21 +1565,32 @@ def _lazy_init_ctx_buffers( # block table one iteration at a time, so the footprint follows # the sequences served rather than max_batch x max_seq_len. self._ctx_kv_buf = pool + self._ctx_kv_ptrs = torch.tensor( + [layer.data_ptr() for layer in pool], dtype=torch.int64, device="cuda" + ) + # The fused preparation stores single-latent context entries + # projected from the captured target features through fc and + # hidden_norm (which every loaded DFlash drafter owns), as the + # eager path does; LiLiCorr drafters keep the eager path. + # Unified caches also track retained-history positions and + # validate per-request allocation in the eager preparation path. + self._paged_latent_prepare = ( + kv_factor == 1 and not draft_model._uses_lilicorr and not unified + ) max_blocks = ( draft_kv_cache_manager.draft_max_blocks_per_seq if unified else draft_kv_cache_manager.max_blocks_per_seq ) - self._init_ctx_block_tables(num_slots, max_blocks) + self._init_managed_ctx_block_tables(num_slots, max_blocks) if unified: self._init_unified_draft_tables(draft_kv_cache_manager) # Only a manager that publishes and matches its own blocks can # hand back a reused prefix's drafter K/V; the private arena and # an unpaired draft pool both start every request at 0. - if not unified: - self._ctx_reuse_addressable = bool( - getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) - ) + self._ctx_reuse_addressable = self._ctx_unified_cache or bool( + getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) + ) self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) @@ -1736,8 +1661,7 @@ def clear(slot: int) -> None: if reset and req_id in self._req_to_slot: old_slot = self._req_to_slot.pop(req_id) - self._ctx_len[old_slot] = 0 - self._ctx_len_host[old_slot] = 0 + clear(old_slot) self._free_slots.append(old_slot) if req_id not in self._req_to_slot: if not self._free_slots: @@ -1749,8 +1673,7 @@ def clear(slot: int) -> None: if self._has_unified_draft_cache(): self._managed_cache_bindings.pop(req_id, None) self._ctx_position_offset[slot] = 0 - self._ctx_len[slot] = 0 - self._ctx_len_host[slot] = 0 + clear(slot) return self._req_to_slot[req_id] def _get_ctx_paged_append(self) -> Callable[..., None]: @@ -1777,6 +1700,7 @@ def _store_context_kv_paged( page_size = self._ctx_page_size pages = table[rows.long(), positions.long() // page_size].long() offsets = positions.long() % page_size + for layer_idx in range(k.size(1)): pool = self._ctx_kv_buf[layer_idx] if fp8_cache: @@ -1804,7 +1728,7 @@ def _store_context_kv_paged( kv_indices, kv_indptr = self._ctx_paged_index_args() for layer_idx in range(k.size(1)): pool = self._ctx_kv_buf[layer_idx] - trtllm_gen_ops.append_paged_kv_cache( + append_paged_kv_cache( append_key=k[:, layer_idx].to(pool.dtype).contiguous(), append_value=v[:, layer_idx].to(pool.dtype).contiguous(), batch_indices=rows_i32, @@ -1859,6 +1783,23 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount attn_metadata.kv_lens_cuda[nc:bs].clamp_(min=0) + def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> None: + """Add this step's accepted tokens to the context lengths of ``slots``, on the device. + + Padding, warmup and attention-DP dummy rows share the dummy slot, which no request owns + or frees. Emptying it right after the add keeps it at 0 between steps: those rows attend + no context, and their next context K / V starts at column 0, inside the padding + request's own pages (past them, its page-table row maps page 0, another request's). + The store is to a fixed slot, so a captured step replays it. + """ + self._ctx_len[slots] += num_accepted + self._ctx_len.clamp_(max=self._max_ctx) + self._ctx_len[self._dummy_slot].zero_() + + def _rebuild_batch_to_slot(self, request_ids: List[int]) -> None: + mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] + self._batch_to_slot_buffer.update(torch.tensor(mapping, dtype=torch.long)) + def _write_ctx_len(self, updates: dict[int, int]) -> None: """Apply a slot-to-length mapping to _ctx_len in one async scatter.""" if not updates: @@ -1926,6 +1867,15 @@ def _store_prefill_context( ctx_len_updates = {} offset = 0 num_contexts = attn_metadata.num_contexts + # A context write addresses pages through row i of the block table, so + # it must fit that row's allocation as well as the arena's length. + # One tolist() for the whole loop: the lengths themselves come from + # _ctx_len_host, so the loop body below stays sync-free. + ctx_alloc = ( + (self._ctx_block_counts[:num_contexts] * self._ctx_page_size).tolist() + if self._ctx_block_tables is not None + else None + ) # Host copies of each chunk's start avoid a device sync per request. cached_tokens = None kv_cache_params = getattr(attn_metadata, "kv_cache_params", None) @@ -1939,31 +1889,35 @@ def _store_prefill_context( if cached_tokens is not None and i < len(cached_tokens): first_pos = int(cached_tokens[i]) + elif slen == 0: + first_pos = 0 + elif host_pos is not None: + first_pos = int(host_pos[offset]) else: - first_pos = chunk_pos[0].item() if slen > 0 else 0 - - # Assign slots for new requests or reset reused IDs in private caches. - # Managed preparation already binds fresh histories and records their - # snapshot slots; keep those slots stable through the forward pass. - reset_slot = first_pos == 0 and not self._has_unified_draft_cache() - if self._assign_slot(req_id, reset=reset_slot) is None: + first_pos = chunk_pos[0].item() + + # A chunk continues what we already stored only if it starts exactly + # where the previous one ended. Everything else -- a fresh request, a + # request id reused after completion, a prefill restarted after + # preemption -- has to start from a clean slot. + reset = ( + not self._has_unified_draft_cache() and self._req_ctx_pos.get(req_id) != first_pos + ) + if self._assign_slot(req_id, reset=reset, updates=ctx_len_updates) is None: logger.warning("DFlash: no free slots, skipping context store") self._req_ctx_pos.pop(req_id, None) offset += slen continue slot = self._req_to_slot[req_id] - if self._has_unified_draft_cache(): - # Restore and prefill keep the host shadow current here. - cur = self._ctx_len_host[slot] - else: - cur = int(self._ctx_len[slot].item()) + cur = ctx_len_updates.get(slot, self._ctx_len_host[slot]) + cap = self._max_ctx if ctx_alloc is None else min(self._max_ctx, ctx_alloc[i]) if cur == 0 and first_pos > 0 and self._ctx_reuse_addressable: # Reuse let the target skip [0, first_pos), so only the tail # reaches the capture buffer. The prefix K/V are still in the # paired pool: precompute_context_kv is per-token in (hidden, p). cur = first_pos - if cur + slen > self._max_ctx: + if cur + slen > cap: if self._has_unified_draft_cache(): self._ctx_len[slot] = 0 self._ctx_len_host[slot] = 0 @@ -1972,17 +1926,13 @@ def _store_prefill_context( self._managed_cache_bindings[req_id] = self._ctx_kv_manager.kv_cache_map[req_id] offset += slen continue - # Truncating here keeps the OLDEST tokens and drops the newest — - # the arena is append-only, not a sliding window — so the drafter - # would go on drafting from a stale prefix with no error and only - # a quiet acceptance-rate loss to show for it. ctx_len tracks the - # target sequence 1:1, so overflow means max_seq_len disagrees - # with the workload, which is a configuration bug. - raise RuntimeError( - f"DFlash drafter context overflow on slot {slot}: " - f"{cur} + {slen} > max_ctx {self._max_ctx}. The drafter " - f"context tracks the target sequence 1:1, so raise " - f"max_seq_len to cover the sequences being served." + # Request-level, like the no-free-slots path above: truncating + # would silently draft from a stale prefix, but killing the + # forward would take every other in-flight request with it. + logger.warning( + f"DFlash: ctx overflow on slot {slot} " + f"({cur} + {slen} > {cap}); skipping its context " + "store, so this request drafts nothing." ) self._req_to_slot.pop(req_id, None) self._req_ctx_pos.pop(req_id, None) @@ -2005,8 +1955,7 @@ def _store_prefill_context( # the store converts its outputs to the cache dtype. cache_dtype = chunk_proj.dtype chunk_proj_cast = chunk_proj[:actual].to(cache_dtype) - self._ctx_len[slot] = end - self._ctx_len_host[slot] = end + ctx_len_updates[slot] = end # Precompute post-norm/post-RoPE K,V for this prefill chunk # so decode iters can read without re-projecting. # @@ -2030,7 +1979,7 @@ def _store_prefill_context( chunk_k, chunk_v, torch.full((actual,), row, dtype=torch.long, device="cuda"), - local_pos, + torch.arange(cur, end, dtype=torch.long, device="cuda"), ) if self._ctx_k_buf is not None: self._ctx_k_buf[slot, :, cur:end] = chunk_k.permute(1, 0, 2, 3) @@ -2106,28 +2055,26 @@ def _forward_impl( ) draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - if bool(getattr(draft_kv_cache_manager, "draft_layer_ids", ())): + if bool(getattr(draft_kv_cache_manager, "_draft_layer_sizes", ())): if not self._ctx_buf_inited or self._ctx_kv_manager is not draft_kv_cache_manager: raise RuntimeError("Unified DSpark draft inputs must be prepared before forward") else: + previous_ctx_len = self._ctx_len self._lazy_init_ctx_buffers( draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager ) spec_metadata._dflash_worker = self + if ( + self._ctx_unified_cache + and self._ctx_len is not previous_ctx_len + and spec_metadata.request_ids + ): + # New buffers clear request slots; seed them before generation. + spec_metadata.prepare() # Before any store: prefill and decode both address pages through it. - refreshed = self._refresh_ctx_block_tables( - attn_metadata, batch_size, spec_metadata.request_ids - ) - if self._ctx_block_tables is not None and not refreshed: - # False is normal only while no paged pool is bound (the private - # arena addresses slots directly). Once one IS bound, the counts - # drive the advertised context length, so a missed refresh leaves - # them at zero and every drafter attends an empty context -- no - # error, just acceptance quietly collapsing. - raise RuntimeError( - "DFlash: draft block tables are bound to a paged pool but this " - "iteration's attn_metadata carried no draft_kv_cache_block_offsets." - ) + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + self._refresh_ctx_block_tables(attn_metadata, batch_size, spec_metadata.request_ids) # Save context lengths so both warmup and a failed forward can roll # back the in-place _ctx_len updates made during drafting. @@ -2214,7 +2161,8 @@ def _forward_impl( if num_gens > 0: with self.draft_kv_cache_context( - attn_metadata, None if self._has_unified_draft_cache() else draft_kv_cache_manager + attn_metadata, + None if self._has_unified_draft_cache() else draft_kv_cache_manager, ): hidden_states_out = draft_model.dflash_forward( noise_embedding=inputs["noise_embedding"], @@ -2372,7 +2320,12 @@ def _draft_block_width(self, draft_model) -> int: this alongside :meth:`_draft_slot_ids` — the two must agree, since the width bounds the slots the gather is allowed to name. """ - return self.max_draft_len + 1 + return self.block_width_for(self.spec_config) + + @staticmethod + def block_width_for(spec_config) -> int: + """Block slots a worker of this family computes for ``spec_config``.""" + return spec_config.max_draft_len + 1 def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int @@ -2625,7 +2578,11 @@ def prepare_1st_drafter_inputs( ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) noise_embed_2d = dflash_noise_block_embedding( - embed_tokens, bonus, mask_token_id, query_tokens_per_req, trained_mask_embedding + embed_tokens, + bonus, + mask_token_id, + query_tokens_per_req, + trained_mask_embedding, ) # Accumulate new accepted features into context buffers @@ -2642,8 +2599,9 @@ def prepare_1st_drafter_inputs( projected_rows = projected_to_store.reshape(num_gens, K_plus_1, -1) lilicorr_anchor = last_accepted_hidden(projected_rows, gen_num_accepted) gen_num_accepted_long = gen_num_accepted.long() - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) + col_idx = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) + if self._ctx_block_tables is not None: # Bound by what the manager handed THIS request: the masked # entries below are still written (fixed-size, for graph @@ -2693,15 +2651,11 @@ def prepare_1st_drafter_inputs( if v_new is not None: self._ctx_v_buf[slot_long, :, col_long] = v_new - self._ctx_len[slots] += gen_num_accepted_long if self._has_unified_draft_cache(): self._ctx_position_offset[slots] += ( - self._ctx_len[slots] - self._max_ctx + self._ctx_len[slots] + gen_num_accepted_long - self._max_ctx ).clamp_min(0) - self._ctx_len.clamp_(max=self._max_ctx) - # Padding and ADP dummy rows share a slot that is never freed; keep its - # context empty. zero_() is graph-capturable, a scalar store is not. - self._ctx_len[self._dummy_slot].zero_() + self._advance_ctx_len(slots, gen_num_accepted_long) num_ctx_per_req_t = self._ctx_len[slots] if self._ctx_block_tables is not None: @@ -2712,6 +2666,10 @@ def prepare_1st_drafter_inputs( allocated = dflash_allocated_ctx_limit( self._ctx_block_counts, self._ctx_page_size, block_size ) + if self._ctx_unified_cache: + self._check_block_within_allocation( + num_ctx_per_req_t, allocated[gen_rows_out], block_size + ) num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) noise_embedding = noise_embed_2d query_positions = query_position_ids.long() @@ -2846,3 +2804,269 @@ def _prepare_paged_latent_gen_inputs( prep.blk_pages, prep.blk_offsets, prep.blk_seq_lens, prep.blk_kv_bounds ), } + + def get_draft_kv_cache_manager( + self, resource_manager: Optional["ResourceManager"] + ) -> Optional[BaseResourceManager]: + if resource_manager is not None: + manager = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) + if bool(getattr(manager, "draft_layer_ids", ())): + return manager + return super().get_draft_kv_cache_manager(resource_manager) + + def _has_unified_draft_cache(self) -> bool: + return bool(getattr(self._ctx_kv_manager, "_draft_layer_sizes", ())) + + def _prepare_managed_history(self, request_ids: list[int]) -> list[int]: + """Restore newly resident histories without rewinding overlapping iterations.""" + updates = {self._dummy_slot: 0} + offsets = {self._dummy_slot: 0} + restore_rows = [] + real_request_ids = { + request_id + for request_id in request_ids + if request_id != ATTENTION_DP_DUMMY_REQUEST_ID + and request_id < self._graph_dummy_id_floor + and request_id not in self._ctx_kv_manager._draft_dummy_request_ids + } + for request_id in list(self._req_to_slot): + if request_id not in real_request_ids: + slot = self._req_to_slot.pop(request_id) + updates[slot] = 0 + offsets[slot] = 0 + self._req_ctx_pos.pop(request_id, None) + self._managed_cache_bindings.pop(request_id, None) + self._free_slots.append(slot) + for row, request_id in enumerate(request_ids): + if request_id in real_request_ids: + self._assign_slot(request_id) + slot = self._req_to_slot.get(request_id) + if slot is not None: + cache = self._ctx_kv_manager.kv_cache_map[request_id] + if self._managed_cache_bindings.get(request_id) is cache: + continue + history = self._ctx_kv_manager.get_draft_history(request_id) + updates[slot] = history.valid_length if history is not None else 0 + self._req_ctx_pos[request_id] = history.position if history is not None else 0 + offsets[slot] = self._req_ctx_pos[request_id] - updates[slot] + self._managed_cache_bindings[request_id] = cache + restore_rows.append(row) + self._write_ctx_len(updates) + self._write_slot_values(self._ctx_position_offset, offsets) + self._batch_to_slot_buffer.update( + torch.tensor( + [self._req_to_slot.get(request_id, self._dummy_slot) for request_id in request_ids], + dtype=torch.long, + ) + ) + self._managed_request_ids = tuple( + request_id for request_id in request_ids if request_id in real_request_ids + ) + # Each upload owns fresh host storage. PyTorch's pinned allocator defers + # recycling it until the asynchronous copy completes. + self._managed_snapshot_slots[: len(self._managed_request_ids)].copy_( + torch.tensor( + [self._req_to_slot[request_id] for request_id in self._managed_request_ids], + dtype=torch.long, + device="cpu", + pin_memory=prefer_pinned(), + ), + non_blocking=True, + ) + return restore_rows + + def _gather_managed_context(self, request_ids: list[int], restore_rows: list[int]) -> None: + """Refresh VANILLA's dense staging; manager pages remain authoritative.""" + if self._dflash_attention_backend != "VANILLA": + return + for row in restore_rows: + request_id = request_ids[row] + slot = self._req_to_slot.get(request_id) + if slot is None: + continue + length = self._ctx_len_host[slot] + positions = torch.arange(length, device=self._ctx_k_buf.device) + pages = self._ctx_block_tables[row, positions // self._ctx_page_size].long() + offsets = positions % self._ctx_page_size + for layer_idx, pool in enumerate(self._ctx_kv_buf): + self._ctx_k_buf[slot, layer_idx, :length] = pool[pages, 0, :, offsets, :] + self._ctx_v_buf[slot, layer_idx, :length] = pool[pages, 1, :, offsets, :] + + def prepare_managed_draft_cache( + self, draft_model, spec_metadata, attn_metadata, resource_manager + ) -> None: + """Refresh managed inputs before eager execution, capture, or graph replay.""" + manager = self.get_draft_kv_cache_manager(resource_manager) + if not getattr(manager, "draft_layer_ids", ()): + return + self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata, manager) + spec_metadata._dflash_worker = self + restore_rows = self._prepare_managed_history(spec_metadata.request_ids) + self._refresh_ctx_block_tables( + attn_metadata, attn_metadata.num_seqs, spec_metadata.request_ids + ) + self._gather_managed_context(spec_metadata.request_ids, restore_rows) + + def _write_slot_values(self, target: torch.Tensor, values: dict[int, int]) -> None: + """Scatter a slot-to-value mapping into a device tensor with one upload.""" + if not values: + return + slots, vals = zip(*values.items()) + pinned = prefer_pinned() + target.index_copy_( + 0, + torch.tensor(slots, dtype=torch.long, pin_memory=pinned).to( + target.device, non_blocking=True + ), + torch.tensor(vals, dtype=target.dtype, pin_memory=pinned).to( + target.device, non_blocking=True + ), + ) + + def snapshot_managed_draft_history(self) -> DraftHistoryUpdate | None: + """Queue iteration-owned readback; the executor publishes after completion.""" + if not self._has_unified_draft_cache() or not self._managed_request_ids: + return None + # Preparation, forward, snapshot gathers and the next preparation run + # on the execution stream, so its next upload cannot overtake this read. + slots = self._managed_snapshot_slots[: len(self._managed_request_ids)] + lengths = self._ctx_len[slots] + positions = lengths + self._ctx_position_offset[slots] + return DraftHistoryUpdate.capture( + self._ctx_kv_manager, + self._managed_request_ids, + torch.stack((lengths, positions), dim=1), + ) + + def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page_size) -> bool: + """Validate the legacy pool layout and resolve its offset decode constant. + + ``kv_cache_block_offsets`` entries are ``pool_block_index * + num_pool_layers * kv_factor`` plus the K/V field index. What that has to + be divided by to index a ``get_buffers`` view depends on the manager -- + see the stride check below; attention_backend/utils.py hardcodes the V1 + divisor and is not a general contract. + """ + mgr = draft_kv_cache_manager + pool_mapping = mgr.kv_cache_pool_mapping + layer_offset = mgr.layer_offsets[0] + self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) + num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) + if num_pool_layers != L: + logger.warning( + f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " + f"the drafter has {L}; using the private ctx arena." + ) + return False + # The two managers hand back different index spaces and only the + # tensor's own stride distinguishes them. V1 slices one page-major + # allocation, so index i steps a whole block (stride = num_pool_layers x + # a layer's block) and the raw block index is what's wanted. V2 wraps + # the layer's base address in a dense view over the *interleaved* page + # space -- layer 0's view spans the entire pool -- so index i steps a + # single layer slot and the same block is num_pool_layers apart. + # Getting this backwards does not fail: the drafter's layers simply + # overwrite each other and acceptance quietly drops. + layer_block_elems = mgr.kv_factor * nkv * page_size * hd + stride0 = pool[0].stride(0) + if stride0 == layer_block_elems: + self._ctx_block_divisor = mgr.kv_factor + elif stride0 == num_pool_layers * layer_block_elems: + self._ctx_block_divisor = num_pool_layers * mgr.kv_factor + else: + logger.warning( + f"DFlash: draft KV pool layer stride {stride0} matches neither the per-layer " + f"({num_pool_layers * layer_block_elems}) nor the interleaved " + f"({layer_block_elems}) layout; using the private ctx arena rather than " + f"guessing block indices." + ) + return False + logger.info( + f"DFlash: ctx block offsets pool_idx={self._ctx_pool_idx}, " + f"divisor={self._ctx_block_divisor}, view_pages={pool[0].size(0)}" + ) + return True + + def _init_unified_draft_tables(self, manager) -> None: + """Choose how the unified pool's block table is built each iteration.""" + self._device_draft_tables = False + layer = manager.draft_layer_ids[0] + pool = manager.layer_to_pool_mapping_dict[manager.layer_offsets[layer]] + if ( + getattr(manager, "_helix_cp_size", 1) != 1 + or getattr(manager, "enable_swa_scratch_reuse", False) + or manager._index_scale_ints[pool] != manager.get_layer_page_index_scale(layer) + ): + logger.warning( + "DFlash: the draft pool's block offsets cannot be decoded on the device; " + "building draft block tables on the host." + ) + return + self._ctx_pool_idx = pool + self._ctx_block_divisor = manager.get_layer_kv_factor(layer) + self._device_draft_tables = True + self._draft_table_checks_left = 8 + logger.info( + f"DFlash: draft block tables decoded on the device from pool {pool} " + f"(divisor {self._ctx_block_divisor})" + ) + + def _refresh_unified_block_tables( + self, attn_metadata, num_seqs: int, request_ids: list[int] + ) -> bool: + """Decode the draft pool's rows of the target's device block offsets. + + The target metadata already uploads every pool's offsets for this + batch, so the draft rows need one device divide plus the per-request + block counts, rather than a host table rebuilt from page lists. + Returns False when the host path must build this iteration's table. + """ + src = getattr(attn_metadata, "kv_cache_block_offsets", None) + if src is None or src.size(0) <= self._ctx_pool_idx: + return False + width = min(src.size(-1), self._ctx_block_tables.size(1)) + encoded = src[self._ctx_pool_idx, :num_seqs, 0, :width] + decoded = torch.div(encoded.clamp(min=0), self._ctx_block_divisor, rounding_mode="floor") + self._ctx_block_tables[:num_seqs, :width].copy_(decoded) + cache_map = self._ctx_kv_manager.kv_cache_map + counts = [] + for request_id in request_ids[:num_seqs]: + cache = cache_map.get(request_id) + counts.append(cache.num_blocks if cache is not None else 0) + self._ctx_block_counts[:num_seqs].copy_( + torch.tensor(counts, dtype=torch.long, pin_memory=prefer_pinned()), + non_blocking=True, + ) + if self._draft_table_checks_left > 0 and not torch.cuda.is_current_stream_capturing(): + self._draft_table_checks_left -= 1 + if not self._check_device_block_tables(num_seqs, request_ids, counts): + return False + return True + + def _check_device_block_tables( + self, num_seqs: int, request_ids: list[int], counts: list[int] + ) -> bool: + """Compare device-decoded rows with the manager's page lists (startup only).""" + manager = self._ctx_kv_manager + rows = [] + for row, request_id in enumerate(request_ids[:num_seqs]): + cache = manager.kv_cache_map.get(request_id) + if counts[row] > 0 and cache is not None and cache.is_active: + rows.append(row) + if not rows: + return True + expected = manager.get_batch_cache_indices( + [request_ids[row] for row in rows], manager.draft_layer_ids[0] + ) + device_rows = self._ctx_block_tables[rows].cpu() + for k, row in enumerate(rows): + got = device_rows[k, : counts[row]].tolist() + want = expected[k][: counts[row]] + if any(w >= 0 and g != w for g, w in zip(got, want)): + logger.error( + "DFlash: device-decoded draft block table disagrees with the manager " + f"for request {request_ids[row]}; building draft tables on the host." + ) + self._device_draft_tables = False + return False + return True diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index fe0377e51dea..b75d97ae8655 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -34,7 +34,8 @@ from ..modules.linear import Linear, TensorParallelMode from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID from ..pyexecutor.resource_manager import ResourceManagerType -from .dflash import DFlashWorker, dflash_draft_slot_ids +from .dflash import DFlashWorker, VocabParallelDraftTokens, dflash_draft_slot_ids +from .dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT from .interface import SpecMetadata, SpecWorkerBase from .vocab_parallel import tp_pair_exchange @@ -284,30 +285,31 @@ class DSv4DSparkWorker(SpecWorkerBase): DSpark drafts a whole block of ``block_size`` tokens in one backbone forward (``DSv4DSparkDraftModel.forward``): it projects the captured target-layer hidden states (``main_proj`` + ``main_norm``) into the draft's captured-context - attention, runs the ``num_stages`` DSpark blocks over a captured + attention, runs the ``num_stages`` DSpark blocks over a rolling captured window, refines the per-position logits with the Markov head, and predicts a per-position acceptance confidence used by the verification scheduler. - Attention reads a sliding window of projected captured context directly - from manager-owned KV pages, alongside the current draft block. + Unlike DFlash, the draft does NOT use the paged KV cache or mask-token + cross-attention: its attention K/V come from the worker-owned rolling window + of projected captured context (one ``main_kv`` per decode step, per stage). Acceptance of the previous block goes through the unified :meth:`SpecWorkerBase.sample_and_accept_draft_tokens` (strict target-verify, or rejection sampling for a non-greedy batch), so greedy parity with no-spec is preserved regardless of draft quality. - The context window is kept consistent across the whole decode: it is seeded + The rolling window is kept consistent across the whole decode: it is seeded from the prompt's captured context at prefill and back-filled with the intermediate accepted tokens of a multi-accept step (both via - ``DSv4DSparkDraftModel.write_context_pages``), in addition to the per-step bonus + ``DSv4DSparkDraftModel.write_context_windows``), in addition to the per-step bonus write done by the generation path. These affect draft acceptance rate only, not correctness, which the standard target verify guarantees. Naming: workers are classified by *deployment form*, not by draft backbone (see :class:`DSparkWorker`). This one is form-specific - because it tracks captured-context history and drives the draft + because it owns a rolling captured-context window and drives the draft through attributes only an embedded DeepSeek-V4-Pro draft has -- - ``num_stages``, ``_attn_params``, ``write_context_pages`` and - ``forward``. A standalone + ``num_stages``, ``_attn_params``, ``write_context_windows``, + ``write_context_windows_batched`` and ``forward_batched``. A standalone drafter has none of them and is served by :class:`DSparkWorker`. Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). @@ -323,7 +325,11 @@ def __init__( self.spec_config = spec_config self.mapping = mapping - # Request progress is slot-indexed; managed KV lives only in its pages. + # Per-slot rolling captured-context KV windows, built lazily on the + # first forward (fixed-size for slot-indexed reads/writes). + self._win_inited = False + self._output_verify_lens: Optional[torch.Tensor] = None + self._attention_warmup_attempted = False self._ctx_len: Optional[torch.Tensor] = None # [max_batch] abs decode position self._valid_len: Optional[torch.Tensor] = None # [max_batch] written window entries self._win = 0 @@ -445,7 +451,7 @@ def max_draft_len(self) -> int: def _publish_position_ceiling(self, draft_model, attn_metadata, block_size: int) -> None: """Size the drafter's RoPE table from the runtime bound, and bound positions by it. - Refreshes the bound for every batch. The first forward can + Runs on EVERY forward, not once behind ``_win_inited``. The first forward can land on a KV-estimation probe manager whose ``max_seq_len`` is below the real one, and a table pinned to that probe would index out of range for the rest of the process; ``DFlashWorker._lazy_init_ctx_buffers`` re-publishes for that @@ -453,8 +459,9 @@ def _publish_position_ceiling(self, draft_model, attn_metadata, block_size: int) Two bounds exist and neither dominates: ``attn_metadata`` carries the KV manager's ``max_seq_len`` while ``_freqs_cap`` carries the user's, and - ``_create_cuda_graph_warmup_request`` sizes its dummy request at whichever is - larger. Covering both is what keeps warmup off the end of the table. + ``DummyRequestFactory.create_cuda_graph_warmup_request`` sizes its dummy + request at whichever is larger. Covering both is what keeps warmup off + the end of the table. """ inner_model = getattr(draft_model, "dspark_model", None) or draft_model config_cap = int(getattr(inner_model, "_freqs_cap", 0) or 0) @@ -490,6 +497,7 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: if self._ctx_len is None: max_batch = spec_metadata.max_num_requests + self._output_verify_lens = torch.empty(max_batch, dtype=torch.int32, device="cuda") draft_model.configure_generation_autotuner(max_batch) self._num_stages = draft_model.num_stages self._win = int(draft_model._attn_params["window_size"]) @@ -517,6 +525,82 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: ) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} + self._position_initialized = torch.zeros(num_rows, dtype=torch.bool, device="cuda") + if self.return_confidence: + # Neutral fill = large positive logit (sigmoid ~ 1.0): + # unscored slots fail safe by verifying the full block. The + # neutral row sits one past the scratch row and is never used + # as a draft-output destination. + self._neutral_conf_row = num_rows + self._confidence_logits = torch.full( + (num_rows + 1, block_size), + _NEUTRAL_CONFIDENCE_LOGIT, + dtype=torch.float32, + device="cuda", + ) + self._confidence_stamp = torch.zeros( + (num_rows + 1,), dtype=torch.int32, device="cuda" + ) + self._draft_seq_host = 0 + # The executor updates this scalar outside inference mode; + # inference tensors reject that in-place host-side update. + with torch.inference_mode(False): + self._draft_seq_cuda = torch.zeros((), dtype=torch.int32, device="cuda") + self._build_verify_planner(draft_model, block_size) + self._win_inited = True + + def _build_verify_planner(self, draft_model, block_size: int) -> None: + """Construct the host-side verify planner (once, at lazy init). + + ModelEngine owns and authenticates the deployment cost table. The + worker creates only the policy state and calibration, then receives + that exact runtime object before its first decision. + """ + from .dspark_schedule import DSparkScheduleConfig + from .dspark_verify import DSparkVerifyPlanner + + cfg = DSparkScheduleConfig(block_size=block_size, min_verify_len=1) + + table_path = self.spec_config.confidence_sps_table_path + # ModelEngine owns the actual CUDA graph ladder, so it is the only + # component that may read, authenticate, and construct the deployment + # cost object. Reopening this path here creates a TOCTOU window in + # which graph capture and scheduling can observe different files. + logger.info( + f"DSpark: deferring exact SPS table {table_path} to the validated " + "ModelEngine runtime object" + ) + + # Calibration lives on the draft's confidence head; fall back to a plain + # sigmoid when the head is absent (no confidence to calibrate anyway). + head = resolve_confidence_head(draft_model) + apply_calibration = head.apply_sts if head is not None else None + sts_path = self.spec_config.confidence_sts_path + if sts_path and head is None: + # Refuse rather than degrade: a configured calibration that cannot + # be attached is a wiring bug, never a preference. + raise ValueError( + f"confidence_sts_path is set but no confidence head was found " + f"on {type(draft_model).__name__}; calibration cannot be " + f"applied, so the planner would schedule on raw sigmoid " + f"survivals. This is a model-wiring bug, not a config choice." + ) + if sts_path and head is not None: + # Accepts either key spelling ("sts_temperatures" here, + # "temperatures" in SGLang); the vectors are interchangeable. + temps = load_sts_temperatures_from_path(sts_path) + head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) + logger.info(f"DSpark: loaded STS calibration from {sts_path}") + + self.verify_planner = DSparkVerifyPlanner( + cfg=cfg, + apply_calibration=apply_calibration, + device_windows=bool(self.spec_config.enable_fused_confidence_scheduler), + ) + logger.info( + f"DSpark verify planner: max_verify_len={self.verify_planner.max_verify_len}, " + f"device_windows={self.verify_planner.device_windows}" + ) def _assign_slot(self, req_id: int) -> int: """Get the persistent progress slot for a request.""" @@ -530,208 +614,14 @@ def _assign_slot(self, req_id: int) -> int: self._req_to_slot[req_id] = slot self._ctx_len[slot] = 0 self._valid_len[slot] = 0 + self._position_initialized[slot] = False + if self._confidence_logits is not None: + # A recycled slot still holds the previous occupant's scores; + # reset to neutral so the first scheduling decision verifies + # the full block. + self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] - def _release_inactive_slots(self, request_ids: list[int]) -> None: - current = set(request_ids) - for request_id in list(self._req_to_slot): - if request_id not in current: - slot = self._req_to_slot.pop(request_id) - self._managed_residency.pop(request_id, None) - self._ctx_len[slot] = 0 - self._valid_len[slot] = 0 - self._free_slots.append(slot) - - def _is_managed_request(self, request_id: int) -> bool: - return ( - request_id != ATTENTION_DP_DUMMY_REQUEST_ID - and request_id < self._graph_dummy_id_floor - and request_id not in self._draft_kv_manager._draft_dummy_request_ids - ) - - def _bind_managed_history(self, resource_manager) -> None: - manager = ( - resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) - if resource_manager is not None - else None - ) - if not getattr(manager, "draft_layer_ids", ()): - raise ValueError("Embedded DSpark requires manager-owned draft KV pages") - if manager is self._draft_kv_manager: - return - if ( - manager.draft_window_size != self._win - or len(manager.draft_layer_ids) != self._num_stages - or any( - manager.kv_factor_per_layer[manager.layer_offsets[layer_id]] != 1 - or manager.head_dim_per_layer[manager.layer_offsets[layer_id]] != self._head_dim - for layer_id in manager.draft_layer_ids - ) - ): - raise ValueError("Embedded DSpark draft cache does not match its attention dimensions") - buffers = tuple( - manager.get_draft_buffers(stage) for stage in range(len(manager.draft_layer_ids)) - ) - self._draft_kv_manager = manager - self._draft_kv_buffers = buffers - self._managed_residency.clear() - self._draft_page_indices.clear() - self._draft_block_tables = torch.zeros( - (self._batch_to_slot.shape[0], manager.draft_max_blocks_per_seq), - dtype=torch.int32, - device=self._ctx_len.device, - ) - self._draft_capacities = torch.zeros( - self._batch_to_slot.shape[0], dtype=torch.long, device=self._ctx_len.device - ) - - def prepare_managed_draft_cache( - self, draft_model, spec_metadata, attn_metadata, resource_manager - ) -> None: - """Refresh persistent draft inputs before eager execution or graph replay.""" - self._lazy_init(draft_model, spec_metadata, attn_metadata) - self._bind_managed_history(resource_manager) - self._prepare_managed_history(spec_metadata.request_ids, attn_metadata.num_contexts) - - def _prepare_managed_history(self, request_ids: list[int], num_contexts: int) -> None: - """Bind live page tables and initialize progress for newly resident histories.""" - # Startup probes and graph/ADP padding never publish synthetic history. - real_rows = [ - row - for row, request_id in enumerate(request_ids) - if self._is_managed_request(request_id) - ] - real_ids = [request_ids[row] for row in real_rows] - self._release_inactive_slots(real_ids) - batch_indices = self._draft_kv_manager.get_draft_block_indices(real_ids) - mappings = {row: tuple(indices) for row, indices in zip(real_rows, batch_indices)} - updates = [] - for row in sorted(self._draft_page_indices.keys() | mappings.keys()): - previous = self._draft_page_indices.get(row, ()) - current = mappings.get(row, ()) - if current == previous: - continue - # Clear removed mappings as well as uploading new/reassigned pages. - width = max(len(previous), len(current)) - old = previous + (0,) * (width - len(previous)) - new = current + (0,) * (width - len(current)) - first = next((i for i in range(width) if old[i] != new[i]), width) - if first == width: - continue - last = width - while old[last - 1] == new[last - 1]: - last -= 1 - updates.append((row, first, new[first:last])) - self._draft_page_indices = mappings - self._ctx_len[self._scratch_slot] = 0 - self._valid_len[self._scratch_slot] = 0 - batch_slots = [self._scratch_slot] * len(request_ids) - for row in real_rows: - request_id = request_ids[row] - cache = self._draft_kv_manager.kv_cache_map[request_id] - slot = self._assign_slot(request_id) - batch_slots[row] = slot - if self._managed_residency.get(request_id) is cache: - continue - position = cache.history_length - if row >= num_contexts and position == 0: - raise ValueError( - f"Embedded DSpark generation request {request_id} has no committed KV history" - ) - self._ctx_len[slot] = position - self._valid_len[slot] = min(self._win, position) - self._managed_residency[request_id] = cache - # Staging may still be read by an earlier asynchronous H2D. Reuse it - # only after its copy event, independently of overlapped model execution. - count = sum(len(indices) for _, _, indices in updates) - host = next( - (buffers for buffers in self._draft_metadata_host_buffers if buffers["event"].query()), - None, - ) - if host is None: - host = { - "tables": torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()), - "capacities": torch.empty( - self._draft_capacities.shape, dtype=torch.long, pin_memory=prefer_pinned() - ), - "slots": torch.empty( - self._batch_to_slot.shape, dtype=torch.long, pin_memory=prefer_pinned() - ), - "event": torch.cuda.Event(), - } - self._draft_metadata_host_buffers.append(host) - elif host["tables"].numel() < count: - host["tables"] = torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()) - offset = 0 - for row, first, indices in updates: - size = len(indices) - source = host["tables"][offset : offset + size] - source.copy_(torch.tensor(indices, dtype=torch.int32)) - self._draft_block_tables[row, first : first + size].copy_(source, non_blocking=True) - offset += size - host["capacities"].zero_() - for row, request_id in zip(real_rows, real_ids): - host["capacities"][row] = self._draft_kv_manager.kv_cache_map[request_id].capacity - self._draft_capacities.copy_(host["capacities"], non_blocking=True) - host["slots"].fill_(self._scratch_slot) - host["slots"][: len(batch_slots)] = torch.tensor(batch_slots, dtype=torch.long) - self._batch_to_slot.copy_(host["slots"], non_blocking=True) - host["event"].record() - - def _seed_context_pages( - self, - draft_model, - spec_metadata: "DSparkSpecMetadata", - attn_metadata, - position_ids: torch.Tensor, - total_target_tokens: int, - ) -> None: - """Seed context chunks using their absolute positions. - - A request can arrive in multiple prefill chunks; continuation chunks - append to the same request's pages. - """ - captured = spec_metadata.get_hidden_states(total_target_tokens) - capture_lens = getattr(spec_metadata, "context_capture_lens", None) - if capture_lens is not None and len(capture_lens) != attn_metadata.num_contexts: - raise ValueError("DSpark capture lengths must match the context requests") - flat_position_ids = position_ids.reshape(-1) - context_offset = 0 - for i in range(attn_metadata.num_contexts): - chunk_len = int(attn_metadata._seq_lens[i]) - capture_len = chunk_len if capture_lens is None else capture_lens[i] - if not 0 <= capture_len <= chunk_len: - raise ValueError("DSpark capture length must be a suffix of its context chunk") - chunk_positions = flat_position_ids[context_offset : context_offset + chunk_len].long() - if chunk_len == 0: - context_offset += chunk_len - continue - - req_id = spec_metadata.request_ids[i] - slot = self._req_to_slot.get(req_id, self._scratch_slot) - if captured is None and slot != self._scratch_slot: - raise RuntimeError( - "Embedded DSpark requires captured context to initialize KV history" - ) - if captured is not None: - keep = min(self._win, chunk_len) - # Prefix reuse can publish any complete page in this chunk. - count = chunk_len if self._draft_kv_manager.enable_block_reuse else keep - positions = (chunk_positions[-count:] + 1).unsqueeze(0) - draft_model.write_context_pages( - captured[ - context_offset + chunk_len - count : context_offset + chunk_len - ].unsqueeze(0), - positions, - torch.ones_like(positions, dtype=torch.bool), - self._draft_kv_buffers, - self._draft_block_tables[i : i + 1], - self._draft_capacities[i : i + 1], - ) - self._valid_len[slot] = torch.clamp(self._valid_len[slot] + keep, max=self._win) - self._ctx_len[slot] = chunk_positions[-1] + 1 - context_offset += chunk_len - def _advance_generation_state( self, slots: torch.Tensor, @@ -745,9 +635,15 @@ def _advance_generation_state( scratch = self._scratch_slot self._ctx_len[scratch].zero_() self._valid_len[scratch].zero_() - old = torch.where(slots == scratch, input_positions, self._ctx_len[slots]) + self._position_initialized[scratch].zero_() + old = torch.where(self._position_initialized[slots], self._ctx_len[slots], input_positions) start_pos = old + num_accepted_tokens - # Keep dummy graph positions within the RoPE table. + # Bound both positions by the RoPE table they index. Warmup reaches this + # through a captured graph replay, and prepare() cannot tell a synthetic + # warmup request from a real one, so the bound has to live in the graph. + # Only a warmup row climbs -- no completion frees its slot, and it ran 53 + # entries past the table [measured job 3097207] -- while _position_cap sits + # exactly at what a full-length real request reaches. if self._position_cap is not None: old = torch.clamp(old, max=self._position_cap) start_pos = torch.clamp(start_pos, max=self._position_cap) @@ -796,7 +692,11 @@ def _draft_gen_block_batched( return None captured = spec_metadata.get_hidden_states(total_target_tokens) if captured is None: - raise RuntimeError("Embedded DSpark requires captured context to advance KV history") + raise RuntimeError( + "DSpark generation requires captured target hidden states; " + "check target_layer_ids and the target model's capture hooks. " + "Refusing to substitute zero draft tokens for a missing capture buffer." + ) # gen-only graph batches have num_ctx_tokens == 0; mixed eager batches put # the gen tokens after the context tokens. @@ -831,10 +731,14 @@ def _draft_gen_block_batched( bonus_row = torch.where(is_scratch, torch.zeros_like(base), base + gidx) main_hidden = captured.index_select(0, bonus_row) - # Back-fill intermediate accepted tokens (everything but the bonus) at - # frames old+1 .. old+nacc-1. Fixed [G, K] tensors keep this capture-safe; - # j >= nacc-1 is masked before any page store. - input_positions = position_ids.reshape(-1)[base].long() + # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens + # (everything but the bonus) into the rolling window — same frames as the + # eager path (old+1 .. old+nacc-1), with j >= nacc-1 masked out. + # A disaggregated generation worker never sees prompt prefill, so a new + # slot has no absolute decode position. Bootstrap it once from the first + # target input position; locally-prefilled and existing slots keep their + # monotonically advanced position. + input_positions = position_ids.reshape(-1).index_select(0, base).long() old, start_pos = self._advance_generation_state(slots, nacc, input_positions) j = torch.arange(K, device=device) # [K] interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1) # [G, K] @@ -846,13 +750,18 @@ def _draft_gen_block_batched( tables = self._draft_block_tables[num_contexts:batch_size] capacities = self._draft_capacities[num_contexts:batch_size] draft_model.write_context_pages( - interim_hidden, interim_pos, interim_valid, self._draft_kv_buffers, tables, capacities + interim_hidden, + interim_pos, + interim_valid, + self._draft_kv_buffers, + tables, + capacities, ) # Surface the per-position corrected block logits ([num_gens, K, vocab]) # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) # sampling + TP gather + draft_probs scatter, rather than argmaxing here. - _toks, _num_proposed, block_logits = draft_model.forward( + _toks, confidence, block_logits = draft_model.forward( main_hidden, bonus, start_pos, @@ -948,6 +857,33 @@ def _forward_impl( total_target_tokens = input_ids.shape[0] + # CUDA-graph warmup guard: the warmup forwards (is_cuda_graph set, stream + # NOT yet capturing) run synthetic gen batches that would otherwise advance + # the persistent rolling-window state. Snapshot and restore it so warmup is + # side-effect-free. (During the capture pass itself the stream IS capturing, + # so we skip the save/restore and let the ops be recorded. Capture-time + # mutation stays on dummy/scratch state: a locally-prefilled request + # resets its slot at position zero, while a disaggregated generation-only + # request receives a freshly zeroed slot.) + is_warmup = ( + getattr(spec_metadata, "is_cuda_graph", False) + and not torch.cuda.is_current_stream_capturing() + ) + if is_warmup: + saved_ctx_len = self._ctx_len.clone() + saved_valid_len = self._valid_len.clone() + saved_position_initialized = self._position_initialized.clone() + # The confidence buffer is persistent state too; warmup scores + # must not reach the verification planner. + saved_confidence = ( + None if self._confidence_logits is None else self._confidence_logits.clone() + ) + saved_stamp = None if self._confidence_stamp is None else self._confidence_stamp.clone() + + # Assign / reset window slots for context (prefill) requests and seed each + # request's rolling KV window from its prompt's captured context, so the + # first generation step drafts against real context instead of an all-zero + # window (acceptance-rate only; verified decoding keeps output correct). if num_contexts > 0: self._seed_context_pages( draft_model, @@ -1052,7 +988,16 @@ def _forward_impl( num_accepted_tokens, ) - return { + if is_warmup: + self._ctx_len.copy_(saved_ctx_len) + self._valid_len.copy_(saved_valid_len) + self._position_initialized.copy_(saved_position_initialized) + if saved_confidence is not None: + self._confidence_logits.copy_(saved_confidence) + if saved_stamp is not None: + self._confidence_stamp.copy_(saved_stamp) + + outputs = { "logits": raw_logits, "new_tokens": accepted_tokens, "new_tokens_lens": num_accepted_tokens, @@ -1068,6 +1013,207 @@ def _forward_impl( ) return outputs + def _release_inactive_slots(self, request_ids: list[int]) -> None: + current = set(request_ids) + for request_id in list(self._req_to_slot): + if request_id not in current: + slot = self._req_to_slot.pop(request_id) + self._managed_residency.pop(request_id, None) + self._ctx_len[slot] = 0 + self._valid_len[slot] = 0 + self._free_slots.append(slot) + + def _is_managed_request(self, request_id: int) -> bool: + return ( + request_id != ATTENTION_DP_DUMMY_REQUEST_ID + and request_id < self._graph_dummy_id_floor + and request_id not in self._draft_kv_manager._draft_dummy_request_ids + ) + + def _bind_managed_history(self, resource_manager) -> None: + manager = ( + resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) + if resource_manager is not None + else None + ) + if not getattr(manager, "draft_layer_ids", ()): + raise ValueError("Embedded DSpark requires manager-owned draft KV pages") + if manager is self._draft_kv_manager: + return + if ( + manager.draft_window_size != self._win + or len(manager.draft_layer_ids) != self._num_stages + or any( + manager.kv_factor_per_layer[manager.layer_offsets[layer_id]] != 1 + or manager.head_dim_per_layer[manager.layer_offsets[layer_id]] != self._head_dim + for layer_id in manager.draft_layer_ids + ) + ): + raise ValueError("Embedded DSpark draft cache does not match its attention dimensions") + buffers = tuple( + manager.get_draft_buffers(stage) for stage in range(len(manager.draft_layer_ids)) + ) + self._draft_kv_manager = manager + self._draft_kv_buffers = buffers + self._managed_residency.clear() + self._draft_page_indices.clear() + self._draft_block_tables = torch.zeros( + (self._batch_to_slot.shape[0], manager.draft_max_blocks_per_seq), + dtype=torch.int32, + device=self._ctx_len.device, + ) + self._draft_capacities = torch.zeros( + self._batch_to_slot.shape[0], dtype=torch.long, device=self._ctx_len.device + ) + + def prepare_managed_draft_cache( + self, draft_model, spec_metadata, attn_metadata, resource_manager + ) -> None: + """Refresh persistent draft inputs before eager execution or graph replay.""" + self._lazy_init(draft_model, spec_metadata, attn_metadata) + self._bind_managed_history(resource_manager) + self._prepare_managed_history(spec_metadata.request_ids, attn_metadata.num_contexts) + + def _prepare_managed_history(self, request_ids: list[int], num_contexts: int) -> None: + """Bind live page tables and initialize progress for newly resident histories.""" + # Startup probes and graph/ADP padding never publish synthetic history. + real_rows = [ + row + for row, request_id in enumerate(request_ids) + if self._is_managed_request(request_id) + ] + real_ids = [request_ids[row] for row in real_rows] + self._release_inactive_slots(real_ids) + batch_indices = self._draft_kv_manager.get_draft_block_indices(real_ids) + mappings = {row: tuple(indices) for row, indices in zip(real_rows, batch_indices)} + updates = [] + for row in sorted(self._draft_page_indices.keys() | mappings.keys()): + previous = self._draft_page_indices.get(row, ()) + current = mappings.get(row, ()) + if current == previous: + continue + # Clear removed mappings as well as uploading new/reassigned pages. + width = max(len(previous), len(current)) + old = previous + (0,) * (width - len(previous)) + new = current + (0,) * (width - len(current)) + first = next((i for i in range(width) if old[i] != new[i]), width) + if first == width: + continue + last = width + while old[last - 1] == new[last - 1]: + last -= 1 + updates.append((row, first, new[first:last])) + self._draft_page_indices = mappings + self._ctx_len[self._scratch_slot] = 0 + self._valid_len[self._scratch_slot] = 0 + batch_slots = [self._scratch_slot] * len(request_ids) + for row in real_rows: + request_id = request_ids[row] + cache = self._draft_kv_manager.kv_cache_map[request_id] + slot = self._assign_slot(request_id) + batch_slots[row] = slot + if self._managed_residency.get(request_id) is cache: + continue + position = cache.history_length + if row >= num_contexts and position == 0: + raise ValueError( + f"Embedded DSpark generation request {request_id} has no committed KV history" + ) + self._ctx_len[slot] = position + self._valid_len[slot] = min(self._win, position) + self._managed_residency[request_id] = cache + # Staging may still be read by an earlier asynchronous H2D. Reuse it + # only after its copy event, independently of overlapped model execution. + count = sum(len(indices) for _, _, indices in updates) + host = next( + (buffers for buffers in self._draft_metadata_host_buffers if buffers["event"].query()), + None, + ) + if host is None: + host = { + "tables": torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()), + "capacities": torch.empty( + self._draft_capacities.shape, dtype=torch.long, pin_memory=prefer_pinned() + ), + "slots": torch.empty( + self._batch_to_slot.shape, dtype=torch.long, pin_memory=prefer_pinned() + ), + "event": torch.cuda.Event(), + } + self._draft_metadata_host_buffers.append(host) + elif host["tables"].numel() < count: + host["tables"] = torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()) + offset = 0 + for row, first, indices in updates: + size = len(indices) + source = host["tables"][offset : offset + size] + source.copy_(torch.tensor(indices, dtype=torch.int32)) + self._draft_block_tables[row, first : first + size].copy_(source, non_blocking=True) + offset += size + host["capacities"].zero_() + for row, request_id in zip(real_rows, real_ids): + host["capacities"][row] = self._draft_kv_manager.kv_cache_map[request_id].capacity + self._draft_capacities.copy_(host["capacities"], non_blocking=True) + host["slots"].fill_(self._scratch_slot) + host["slots"][: len(batch_slots)] = torch.tensor(batch_slots, dtype=torch.long) + self._batch_to_slot.copy_(host["slots"], non_blocking=True) + host["event"].record() + + def _seed_context_pages( + self, + draft_model, + spec_metadata: "DSparkSpecMetadata", + attn_metadata, + position_ids: torch.Tensor, + total_target_tokens: int, + ) -> None: + """Seed context chunks using their absolute positions. + + A request can arrive in multiple prefill chunks; continuation chunks + append to the same request's pages. + """ + captured = spec_metadata.get_hidden_states(total_target_tokens) + capture_lens = getattr(spec_metadata, "context_capture_lens", None) + if capture_lens is not None and len(capture_lens) != attn_metadata.num_contexts: + raise ValueError("DSpark capture lengths must match the context requests") + flat_position_ids = position_ids.reshape(-1) + context_offset = 0 + for i in range(attn_metadata.num_contexts): + chunk_len = int(attn_metadata._seq_lens[i]) + capture_len = chunk_len if capture_lens is None else capture_lens[i] + if not 0 <= capture_len <= chunk_len: + raise ValueError("DSpark capture length must be a suffix of its context chunk") + chunk_positions = flat_position_ids[context_offset : context_offset + chunk_len].long() + if chunk_len == 0: + context_offset += chunk_len + continue + + req_id = spec_metadata.request_ids[i] + slot = self._req_to_slot.get(req_id, self._scratch_slot) + if captured is None and slot != self._scratch_slot: + raise RuntimeError( + "Embedded DSpark requires captured context to initialize KV history" + ) + if captured is not None and capture_len: + keep = min(self._win, capture_len) + # Prefix reuse can publish any complete page in this chunk. + count = capture_len if self._draft_kv_manager.enable_block_reuse else keep + positions = (chunk_positions[-count:] + 1).unsqueeze(0) + draft_model.write_context_pages( + captured[ + context_offset + chunk_len - count : context_offset + chunk_len + ].unsqueeze(0), + positions, + torch.ones_like(positions, dtype=torch.bool), + self._draft_kv_buffers, + self._draft_block_tables[i : i + 1], + self._draft_capacities[i : i + 1], + ) + self._valid_len[slot] = torch.clamp(self._valid_len[slot] + keep, max=self._win) + self._ctx_len[slot] = chunk_positions[-1] + 1 + self._position_initialized[slot] = True + context_offset += chunk_len + class DSparkWorker(DFlashWorker): """Worker for a *standalone* DSpark drafter (DFlash lineage). @@ -1101,8 +1247,9 @@ class DSparkWorker(DFlashWorker): subclass. Naming workers by backbone would produce N classes with identical bodies. - Embedded DSpark uses :class:`DSv4DSparkWorker` to manage captured-context - progress and the draft block forward. + Deployment form is the axis the runtime state actually splits on: paged + draft K/V here, a worker-owned rolling window in + :class:`DSv4DSparkWorker`. """ def set_draft_model(self, draft_model) -> None: @@ -1145,6 +1292,16 @@ def bind_draft_lm_head(self, draft_model) -> None: ): self._markov_tp_exchange = tp_pair_exchange(mapping, AllReduceStrategy.NCCL) + @staticmethod + def block_width_for(spec_config) -> int: + """Block slots a DSpark worker computes for ``spec_config``. + + The ``shift_label`` convention, which the DSpark drafter defaults to, + reads slots 0..K-1; a checkpoint that turns it off needs K+1 and is + caught by the block-within-allocation check. + """ + return spec_config.max_draft_len + def _draft_block_width(self, draft_model) -> int: """Block width under the dspark ``shift_label`` convention. @@ -1154,8 +1311,8 @@ def _draft_block_width(self, draft_model) -> int: drafters are meant to run. """ if getattr(draft_model, "_dspark_shift_label", False): - return self.max_draft_len - return super()._draft_block_width(draft_model) + return self.block_width_for(self.spec_config) + return DFlashWorker.block_width_for(self.spec_config) def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 1f9f5747ebb0..0f012b000373 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1128,8 +1128,10 @@ def write_padding_onehot_draft_probs(self, padding_slot_ids, draft_len): return onehot_vocab = (self.draft_probs_last_dim if self.draft_probs_last_dim > 0 else self.draft_probs_vocab_size) - bad = [s for s in padding_slot_ids - if s is None or s < 0 or s >= self.draft_probs.shape[0]] + bad = [ + s for s in padding_slot_ids + if s is None or s < 0 or s >= self.draft_probs.shape[0] + ] if bad: # Writing a different row would leave the real slot's distribution # stale and break lossless acceptance, so fail loudly instead. @@ -1177,14 +1179,11 @@ def create_cuda_graph_metadata(self, max_batch_size: int): if self.is_cuda_graph: return self - # Allocate the slot-indexed rejection-sampling buffers on the base - # metadata BEFORE the shallow copy so every CUDA-graph view shares - # them. A copy taken while draft_probs is still None allocates its own - # buffer sized by the graph batch size: the drafter then scatters into - # one buffer while the captured rejection kernel reads another, and - # slot ids >= the graph batch size index out of bounds (seen with - # DSpark + attention DP, write_padding_onehot_draft_probs). + # Slot-indexed state persists across steps that may run eagerly or under + # different graphs, so every view must share one allocation: allocate it + # here before the shallow copy rather than lazily in each view. self.prepare_rejection_sampling_buffers() + self.prepare_penalty_buffers() cuda_graph_metadata = copy.copy(self) cuda_graph_metadata.is_cuda_graph = True cuda_graph_metadata.max_num_requests = max_batch_size diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index a4d1064aa04d..feb34389991b 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -456,7 +456,6 @@ def _build_spec_metadata(spec_config, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, - num_seq_slots=num_seq_slots, ) if spec_config.spec_dec_mode.is_dspark(): target_layer_ids = getattr(spec_config, 'target_layer_ids', None) @@ -477,7 +476,6 @@ def _build_spec_metadata(spec_config, use_rejection_sampling=use_rejection_sampling, vocab_size=vocab_size, draft_vocab_size=draft_vocab_size, - num_seq_slots=num_seq_slots, ) if spec_config.spec_dec_mode.is_draft_target_one_model(): return DraftTargetOneModelSpecMetadata( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index bbab9e7b77ec..97f1b146017a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -117,6 +117,8 @@ class _BatchDescFieldSpec: # field), matching the binding's constructor. import dataclasses as _dataclasses # noqa: E402 +CacheDomain = _cpp.CacheDomain + @_dataclasses.dataclass class _KVCacheManagerConfigFieldSpec: @@ -137,17 +139,6 @@ class _KVCacheManagerConfigFieldSpec: pool_rebalance: object = None cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY - # Symbols added on main that are not yet ported to the C++ backend. - # TODO(kvCacheManagerV2-cpp): port these and replace the fallbacks. - AttnLifeCycle = getattr(_cpp, "AttnLifeCycle", None) - OutOfMemoryError = getattr(_cpp, "OutOfMemoryError", MemoryError) - PageIndexConverter = getattr(_cpp, "PageIndexConverter", None) - ReuseScope = getattr(_cpp, "ReuseScope", ReuseScope) - ScratchDesc = getattr(_cpp, "ScratchDesc", None) - SsmSnapshotIterationStatsDelta = _cpp.SsmSnapshotIterationStatsDelta - SwaScratchReuseConfig = getattr(_cpp, "SwaScratchReuseConfig", None) - UniqueToken = _cpp.UniqueToken - CacheDomain = _cpp.CacheDomain KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ del _KVCacheManagerConfigFieldSpec, _dataclasses @@ -319,5 +310,4 @@ def typed_range(*args: int) -> range: "typed_range", ] -if _BACKEND == "cpp": - __all__.append("CacheDomain") +__all__.append("CacheDomain") diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 9458848c6340..ed7828864a56 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -32,7 +32,7 @@ from typing import ( Union, ) -from ._utils import HalfOpenRange as HalfOpenRange +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig # From _common.py NDEBUG: Final[int] @@ -97,12 +97,14 @@ LifeCycleId = NewType("LifeCycleId", int) LayerGroupId: TypeAlias = LifeCycleId class AttnLifeCycle: - """The attention life cycle, keyed by window, sink-token shape, and cache domain.""" + """The attention life cycle, keyed by window, sink tokens, sparsity, and reuse policy.""" def __init__( self, window_size: int | None, num_sink_blocks: int, + is_sparse: bool = False, + reuse_policy: AttentionReusePolicy = ..., cache_domain: CacheDomain = CacheDomain.TARGET, ) -> None: ... @staticmethod @@ -110,6 +112,8 @@ class AttnLifeCycle: window_size: int | None, num_sink_tokens: int | None, tokens_per_block: int, + is_sparse: bool = False, + reuse_policy: AttentionReusePolicy = ..., cache_domain: CacheDomain = CacheDomain.TARGET, ) -> "AttnLifeCycle": ... @property @@ -117,6 +121,10 @@ class AttnLifeCycle: @property def num_sink_blocks(self) -> int: ... @property + def is_sparse(self) -> bool: ... + @property + def reuse_policy(self) -> AttentionReusePolicy: ... + @property def cache_domain(self) -> CacheDomain: ... def get_stale_range(self, history_length: int, tokens_per_block: int) -> HalfOpenRange: ... @@ -276,6 +284,7 @@ class AttentionLayerConfig: buffers: list[BufferConfig] sliding_window_size: int | None = None num_sink_tokens: int | None = None + reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED cache_domain: CacheDomain = CacheDomain.TARGET @property def window_size(self) -> int | None: ... diff --git a/tests/unittest/_torch/attention/test_fp4_mla_packed.py b/tests/unittest/_torch/attention/test_fp4_mla_packed.py index e03e3cc1db91..50b4a50b40f3 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla_packed.py +++ b/tests/unittest/_torch/attention/test_fp4_mla_packed.py @@ -387,13 +387,15 @@ def test_packed_plan_wave_plan_env_rolls_back(monkeypatch) -> None: monkeypatch.setattr( _plan_mod(), "launch_packed_plan", lambda *args, **kwargs: calls.append(kwargs["wave_plan"]) ) - geometry = fp4_mla_backend._PackedPlanGeometry(2, 4, 3, 17, 64, 6, 1024, None, None, ("key",)) + geometry = fp4_mla_packed_plan._PackedPlanGeometry( + 2, 4, 3, 17, 64, 6, 1024, None, None, ("key",) + ) kv_lens = torch.zeros(2, dtype=torch.int32, device="cuda") - metadata = SimpleNamespace(_fp4_mla_packed_plan_key=None) - monkeypatch.delenv(fp4_mla_backend.FP4_MLA_PACKED_WAVE_PLAN_ENV, raising=False) - fp4_mla_backend._launch_fp4_mla_packed_plan(geometry, kv_lens, metadata) - monkeypatch.setenv(fp4_mla_backend.FP4_MLA_PACKED_WAVE_PLAN_ENV, "0") - fp4_mla_backend._launch_fp4_mla_packed_plan(geometry, kv_lens, metadata) + metadata = SimpleNamespace(fp4_mla_state=SimpleNamespace(packed_plan_key=None)) + monkeypatch.delenv(fp4_mla_config.FP4_MLA_PACKED_WAVE_PLAN_ENV, raising=False) + fp4_mla_packed_plan._launch_fp4_mla_packed_plan(geometry, kv_lens, metadata) + monkeypatch.setenv(fp4_mla_config.FP4_MLA_PACKED_WAVE_PLAN_ENV, "0") + fp4_mla_packed_plan._launch_fp4_mla_packed_plan(geometry, kv_lens, metadata) assert calls == [True, False] diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index bf18029588ac..cc7cb074e74b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -37,11 +37,13 @@ KVCacheLayerSpec, KVCacheManagerV2, Role, + _get_draft_layer_sizes, _get_static_cache_size_layer_components, ) from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings import DataType from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( @@ -2059,3 +2061,45 @@ def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) assert manager_config is kv_cache_config assert manager_config.max_attention_window == ([1024] if guard == "explicit_window" else None) + + +@pytest.mark.parametrize( + "attention_backend,dtype,bytes_per_layer", + [("TRTLLM", torch.float8_e4m3fn, 2048), ("VANILLA", torch.bfloat16, 4096)], +) +def test_standalone_draft_kv_dtype_follows_the_drafter_attention( + attention_backend, dtype, bytes_per_layer +) -> None: + """The TRTLLM drafter attention stores the drafter's history in FP8; VANILLA keeps BF16.""" + creator = object.__new__(KvCacheCreator) + draft = SimpleNamespace( + fc=SimpleNamespace(weight=torch.empty(1, dtype=torch.bfloat16)), + dflash_attention_backend=attention_backend, + config=SimpleNamespace(max_position_embeddings=4096), + ) + creator._model_engine = SimpleNamespace(model=SimpleNamespace(draft_model=draft)) + creator._speculative_config = SimpleNamespace( + spec_dec_mode=SimpleNamespace(is_dspark=lambda: True), + draft_is_embedded_in_target=False, + _use_shared_kv_cache=False, + attention_backend="AUTO", + max_draft_len=3, + ) + # Drafter shapes remain unsharded under attention DP. + creator._draft_config = SimpleNamespace( + pretrained_config=SimpleNamespace( + num_attention_heads=64, + num_key_value_heads=16, + head_dim=64, + hidden_size=4096, + num_hidden_layers=5, + ) + ) + creator._mapping = SimpleNamespace( + has_cp_helix=lambda: False, enable_attention_dp=True, tp_size=16 + ) + + kwargs = creator._get_unified_draft_cache_kwargs() + + assert kwargs["draft_dtype"] == torch_dtype_to_binding(dtype) + assert _get_draft_layer_sizes(**kwargs) == (bytes_per_layer,) * 5 diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py deleted file mode 100644 index cf0caff78d65..000000000000 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ /dev/null @@ -1,1010 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for KV cache token estimation in KvCacheCreator._get_token_num_for_estimation. - -Guards the ADP (Attention Data Parallelism) cache-block reduction: when -enable_attention_dp is True and tp_size > 1, _create_dummy_context_requests -produces tp_size duplicate requests, but the scheduler distributes them -1-per-rank. Each rank's KV cache therefore only needs capacity for its own -share, not all copies. -""" - -from types import SimpleNamespace -from unittest.mock import Mock, patch - -import pytest -import torch - -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin -from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator -from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( - KVCacheManagerV2, - _get_draft_layer_sizes, -) -from tensorrt_llm._utils import torch_dtype_to_binding -from tensorrt_llm.inputs.multimodal import MultimodalParams -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MultimodalConfig - -pytestmark = pytest.mark.cpu_only - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_mock_request(num_input_tokens, beam_width=1): - """Create a mock request with the fields _get_token_num_for_estimation reads.""" - req = Mock() - req.input_token_ids = list(range(num_input_tokens)) - req.sampling_config.beam_width = beam_width - return req - - -class _TextModel: - def __init__(self, encoder_cache_max_bytes): - self.model_config = ModelConfig( - multimodal_config=MultimodalConfig(encoder_cache_max_bytes=encoder_cache_max_bytes) - ) - - -class _MultimodalModel(MultimodalModelMixin): - def __init__(self, encoder_cache_max_bytes): - self.model_config = ModelConfig( - multimodal_config=MultimodalConfig(encoder_cache_max_bytes=encoder_cache_max_bytes) - ) - - -class _EncoderCacheMultimodalModel(_MultimodalModel): - supports_encoder_cache = True - - -def test_encoder_profiling_uses_full_budget_independent_of_llm_limit( - monkeypatch, -): - class _InputProcessor: - def __init__(self): - self.calls = [] - - def get_mm_max_tokens_per_item(self, max_num_encoder_tokens=None): - del max_num_encoder_tokens - return {"image": 4096} - - def get_dummy_mm_data( - self, - *, - max_num_encoder_tokens, - mm_counts, - dtype, - ): - self.calls.append((max_num_encoder_tokens, mm_counts, dtype)) - item_count = mm_counts["image"] - return {"image": {"item_count": item_count}} - - class _Model(MultimodalModelMixin): - dtype = torch.float16 - mm_encoder = object() - - def __init__(self): - self.forwarded_item_counts = [] - self.last_output = None - - def encode_multimodal_inputs(self, multimodal_params): - count = multimodal_params[0].multimodal_data["image"]["item_count"] - self.forwarded_item_counts.append(count) - self.last_output = torch.arange(count * 4).reshape(count, 4) - return self.last_output - - input_processor = _InputProcessor() - model = _Model() - model.model_config = SimpleNamespace(pretrained_config=SimpleNamespace(vocab_size=128)) - creator = object.__new__(KvCacheCreator) - creator._model_engine = SimpleNamespace( - model=model, - input_processor=input_processor, - encoder_batch_size=4, - encoder_max_num_tokens=8192, - use_mrope=False, - ) - creator._profiling_stage_data = {"enable_mm_reqs": True} - creator._mapping = SimpleNamespace(enable_attention_dp=False) - creator._max_num_tokens = 18 - creator._max_beam_width = 1 - - # The LLM dummy remains text-only and is not allowed to shrink the - # independent encoder profiling budget. - requests = creator._create_dummy_context_requests(input_seq_len=18) - assert sum(len(request.input_token_ids) for request in requests) == 18 - assert input_processor.calls == [] - assert all(getattr(request, "py_multimodal_data", None) is None for request in requests) - - creator._dummy_encoder_inputs = creator._create_dummy_encoder_inputs() - assert input_processor.calls == [(8192, {"image": 2}, torch.float16)] - - monkeypatch.setattr(MultimodalParams, "to_device", lambda self, *args, **kwargs: self) - retained_output = creator._encode_dummy_inputs() - assert model.forwarded_item_counts == [2] - assert retained_output.data_ptr() != model.last_output.data_ptr() - assert creator._dummy_encoder_inputs == [] - - -def _make_creator( - tokens_per_block, - dummy_reqs, - enable_attention_dp, - tp_size, - batch_size=1, - model_max_seq_len=1, - max_cuda_graph_batch_size=1, - layer_types=None, - sliding_window=None, - use_sliding_window=None, - max_attention_window=None, -): - """Build a minimal KvCacheCreator (bypasses __init__) wired up for - _get_token_num_for_estimation only.""" - c = object.__new__(KvCacheCreator) - - c._tokens_per_block = tokens_per_block - c._net_max_seq_len = 2048 - c._speculative_config = None - c._dummy_reqs = dummy_reqs - - c._mapping = Mock(enable_attention_dp=enable_attention_dp, tp_size=tp_size, cp_config={}) - - c._llm_args = Mock(disable_overlap_scheduler=True) - - pretrained = SimpleNamespace( - layer_types=layer_types, - num_hidden_layers=(len(layer_types) if isinstance(layer_types, (list, tuple)) else None), - sliding_window=sliding_window, - use_sliding_window=use_sliding_window, - ) - - model_config = Mock() - model_config.pretrained_config = pretrained - - c._model_engine = Mock( - batch_size=batch_size, - max_seq_len=model_max_seq_len, - _max_cuda_graph_batch_size=max_cuda_graph_batch_size, - ) - c._model_engine.model = Mock(model_config=model_config) - - c._kv_cache_config = Mock( - free_gpu_memory_fraction=0.9, max_attention_window=max_attention_window - ) - - # _get_token_num_for_estimation gates pool-group scaling on V2 (the - # split-pool layout that motivates the scaling). Mamba hybrid uses - # MambaHybridCacheManager and must NOT scale. These tests target V2 - # behavior; production sets this in __init__ via - # _get_model_kv_cache_manager_cls(), which we bypass here. - c._kv_cache_manager_cls = KVCacheManagerV2 - c._is_kv_cache_manager_v2 = True - - return c - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - -@pytest.fixture(autouse=True) -def _no_gpu(): - """Stub out CUDA memory queries and per-token KV size so the test runs on - any machine and the memory cap never constrains the result. - - _get_kv_size_per_token now returns a CacheCost; using slope=1 + zero - intercept keeps the legacy ``budget // bytes_per_token`` arithmetic - untouched downstream.""" - huge = 100 * (1 << 30) - with ( - patch("torch.cuda.mem_get_info", return_value=(huge, huge)), - patch.object(KvCacheCreator, "_get_kv_size_per_token", return_value=CacheCost(slope=1)), - ): - yield - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -def test_cache_cost_tokens_for_budget_not_int_division(): - """Regression: V2 KV cache capacity must derive max_tokens via - ``CacheCost.tokens_for_budget``, not ``budget // CacheCost``.""" - budget = 10_000 - cost = CacheCost(slope=512, intercept=128) - assert cost.tokens_for_budget(budget) == (budget - 128) // 512 - with pytest.raises(TypeError): - _ = budget // cost - - -def test_adp_reduces_blocks_to_per_rank_share(): - """With ADP + tp_size duplicated requests the result must equal a single - rank's share, not the sum across all duplicates.""" - tpb = 64 - tp = 4 - n_in = 128 # ceil((128+1)/64) = 3 blocks per request - - baseline = _make_creator(tpb, [_make_mock_request(n_in)], enable_attention_dp=False, tp_size=1) - adp = _make_creator( - tpb, [_make_mock_request(n_in) for _ in range(tp)], enable_attention_dp=True, tp_size=tp - ) - - assert adp._get_token_num_for_estimation() == baseline._get_token_num_for_estimation() - - -def test_without_adp_all_blocks_counted(): - """Without ADP every request's blocks contribute to the total.""" - tpb = 64 - n_in = 128 # 3 blocks each - n_reqs = 4 - - c = _make_creator( - tpb, [_make_mock_request(n_in) for _ in range(n_reqs)], enable_attention_dp=False, tp_size=1 - ) - - # 4 reqs * 3 blocks * 64 tokens/block = 768 - assert c._get_token_num_for_estimation() == n_reqs * 3 * tpb - - -@pytest.mark.parametrize("tp_size", [2, 4, 8]) -def test_adp_various_tp_sizes(tp_size): - """ADP division must hold for several representative tp_size values.""" - tpb = 64 - n_in = 128 # 3 blocks per request - - c = _make_creator( - tpb, - [_make_mock_request(n_in) for _ in range(tp_size)], - enable_attention_dp=True, - tp_size=tp_size, - ) - - total = tp_size * 3 - expected_blocks = (total + tp_size - 1) // tp_size - assert c._get_token_num_for_estimation() == expected_blocks * tpb - - -def test_regression_without_fix_would_overcount(): - """If the ADP ceil-division fix were removed, the returned - value would be tp_size times too large. This test guards that fix.""" - tpb = 64 - tp = 4 - n_in = 128 - - c = _make_creator( - tpb, [_make_mock_request(n_in) for _ in range(tp)], enable_attention_dp=True, tp_size=tp - ) - - result = c._get_token_num_for_estimation() - - correct = 3 * tpb # 192 (per-rank share) - wrong = tp * 3 * tpb # 768 (all duplicates summed) - assert result == correct - assert result != wrong - - -@pytest.mark.parametrize( - ("model_cls", "encoder_cache_max_bytes", "expected_reserve"), - [ - (_TextModel, 64, 0), - (_MultimodalModel, 0, 0), - (_MultimodalModel, 64, 0), - (_EncoderCacheMultimodalModel, 0, 0), - (_EncoderCacheMultimodalModel, 64, 64), - ], -) -def test_kv_cache_estimation_reserves_multimodal_encoder_cache( - model_cls, - encoder_cache_max_bytes, - expected_reserve, -): - creator = object.__new__(KvCacheCreator) - creator._model_engine = SimpleNamespace(model=model_cls(encoder_cache_max_bytes)) - - assert creator._get_multimodal_encoder_memory_reserve() == expected_reserve - - -def test_reserve_adds_only_unprofiled_output_capacity(): - creator = object.__new__(KvCacheCreator) - creator._model_engine = SimpleNamespace( - mm_encoder_output_budget_bytes=512, - model=_MultimodalModel(0), - ) - assert creator._get_multimodal_encoder_memory_reserve(profiled_output_bytes=400) == 112 - - -# --------------------------------------------------------------------------- -# VSWA hybrid attention pool-group scaling (Gemma4 hybrid MMMU Pro hang fix) -# --------------------------------------------------------------------------- -# -# KVCacheManagerV2 creates one pool group per distinct attention-window size. -# The quota derived from max_tokens is split proportionally across pool -# groups, so each pool ends up with roughly num_cache_blocks / num_pool_groups -# blocks. A single long-context request then overflows the full-attention -# pool and the scheduler livelocks on suspend/retry. The fix scales -# num_cache_blocks by the number of distinct attention-window sizes inferred -# effective per-layer sliding windows on the pretrained config (preferred) or -# an explicit ``max_attention_window`` list on kv_cache_config (fallback). - - -def test_uniform_layer_types_no_scaling(): - """All-sliding or all-full layers stay a single pool group.""" - tpb = 32 - max_seq_len = 4096 - uniform = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=["sliding_attention"] * 26, - ) - # num_pool_groups = 1 -> behaviour unchanged from legacy ADP-only case. - baseline = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - ) - assert uniform._get_token_num_for_estimation() == baseline._get_token_num_for_estimation() - - -def test_get_layer_attention_window_honors_max_window_layers(): - config = SimpleNamespace( - use_sliding_window=True, - sliding_window=4096, - max_window_layers=2, - ) - - assert [get_layer_attention_window(config, layer_idx) for layer_idx in range(4)] == [ - None, - None, - 4096, - 4096, - ] - - -def test_gemma4_hybrid_scales_by_num_pool_groups(): - """Gemma4 hybrid attention (mixed sliding/full layers) must scale the - estimated block count by the number of distinct layer types. Otherwise - the per-pool quota is too small to hold a single max_seq_len request, - which is the MMMU Pro livelock reproducer.""" - tpb = 32 - max_seq_len = 12288 - layer_types = ["sliding_attention"] * 28 + ["full_attention"] * 7 - # Mixed sliding/full attention must include the model's actual window size; - # Gemma4-E2B uses a 512-token sliding window. - sliding_window = 512 - assert len(set(layer_types)) == 2 - - hybrid = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1), _make_mock_request(1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=layer_types, - sliding_window=sliding_window, - ) - uniform = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1), _make_mock_request(1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=["full_attention"] * 35, - ) - - hybrid_tokens = hybrid._get_token_num_for_estimation() - uniform_tokens = uniform._get_token_num_for_estimation() - assert hybrid_tokens == 2 * uniform_tokens, ( - f"Expected 2x scaling for 2 pool groups, got " - f"hybrid={hybrid_tokens}, uniform={uniform_tokens}" - ) - - -def test_hybrid_linear_attention_scales_by_num_pool_groups(): - """Hybrid linear/attention V2 managers retain both estimation pools.""" - tpb = 32 - max_seq_len = 4096 - layer_types = ["linear_attention"] * 30 + ["full_attention"] * 10 - - hybrid = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=layer_types, - ) - uniform = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=["full_attention"] * 40, - ) - - assert hybrid._get_token_num_for_estimation() == 2 * uniform._get_token_num_for_estimation() - - -@pytest.mark.parametrize( - ("sliding_window", "use_sliding_window"), - [ - ([512, 1024], None), - (None, True), - ], - ids=["multiple_window_sizes", "missing_window"], -) -def test_v2_pool_estimation_falls_back_for_unsupported_window_metadata( - sliding_window, - use_sliding_window, -): - tpb = 32 - max_seq_len = 4096 - hybrid = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=["sliding_attention", "full_attention"], - sliding_window=sliding_window, - use_sliding_window=use_sliding_window, - ) - uniform = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=["full_attention", "full_attention"], - ) - - assert hybrid._get_token_num_for_estimation() == 2 * uniform._get_token_num_for_estimation() - - -def test_vswa_max_attention_window_fallback_scales(): - """When layer_types is absent but kv_cache_config.max_attention_window is - a heterogeneous list (VSWA), we still scale by the number of distinct - windows.""" - tpb = 32 - max_seq_len = 12288 - max_attention_window = [1024] * 28 + [max_seq_len] * 7 - assert len(set(max_attention_window)) == 2 - - c = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=None, - max_attention_window=max_attention_window, - ) - uniform = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - ) - assert c._get_token_num_for_estimation() == 2 * uniform._get_token_num_for_estimation() - - -def test_pool_scaling_prevents_mmmu_pro_underestimation(): - """Regression: with max_seq_len=12288 and max_num_tokens=12288 (MMMU Pro - config), hybrid estimation must produce enough capacity to hold one full - max_seq_len request per pool, with max_util_for_resume headroom.""" - tpb = 32 - max_seq_len = 12288 - layer_types = ["sliding_attention"] * 28 + ["full_attention"] * 7 - # Mixed sliding/full attention must include the model's actual window size; - # Gemma4-E2B uses a 512-token sliding window. - sliding_window = 512 - - c = _make_creator( - tpb, - [_make_mock_request(max_seq_len - 1), _make_mock_request(1)], - enable_attention_dp=False, - tp_size=1, - model_max_seq_len=max_seq_len, - max_cuda_graph_batch_size=4, - layer_types=layer_types, - sliding_window=sliding_window, - ) - - total_tokens = c._get_token_num_for_estimation() - per_pool_tokens = total_tokens // 2 # 2 pool groups - # Must be enough to hold a max_seq_len request per pool. - assert per_pool_tokens >= max_seq_len - - -def test_v2_cache_size_per_token_models_generation_swa_cost(): - class FakeModelConfig: - quant_config = None - pretrained_config = SimpleNamespace( - hidden_size=32, - num_attention_heads=4, - num_key_value_heads=2, - ) - - def get_num_attention_layers(self): - return 3 - - mapping = Mock(enable_attention_dp=False, tp_size=1) - mapping.pp_layers.return_value = [0, 1, 2] - - no_scratch_size_per_token = CacheCost.from_raw( - KVCacheManagerV2.get_cache_size_per_token( - FakeModelConfig(), - mapping, - tokens_per_block=64, - max_seq_len=4096, - max_batch_size=3, - kv_cache_config=KvCacheConfig(max_attention_window=[2048, 2048, 4096]), - ) - ) - scratch_size_per_token = CacheCost.from_raw( - KVCacheManagerV2.get_cache_size_per_token( - FakeModelConfig(), - mapping, - tokens_per_block=64, - max_seq_len=4096, - max_batch_size=3, - kv_cache_config=KvCacheConfig(max_attention_window=[2048, 2048, 4096]), - enable_swa_scratch_reuse=True, - ) - ) - - # Per layer: K+V * kv_heads * head_dim * bf16 bytes = 2 * 2 * 8 * 2. - expected = CacheCost(slope=64, intercept=3 * 2 * 2048 * 64) - assert no_scratch_size_per_token == expected - assert scratch_size_per_token == expected - - -def test_creator_uses_v2_affine_cache_cost(): - class FakeV2Manager(KVCacheManagerV2): - @staticmethod - def get_cache_size_per_token(model_config, mapping, **kwargs): - return 20, 21 - - creator = object.__new__(KvCacheCreator) - creator._mapping = Mock() - creator._tokens_per_block = 64 - creator._max_seq_len = 1024 - creator._max_batch_size = 3 - creator._kv_cache_config = KvCacheConfig() - creator._speculative_config = None - - cost = creator._per_manager_cache_cost(FakeV2Manager, Mock()) - - assert cost == CacheCost(slope=20, intercept=21) - - -def test_v2_quota_from_max_tokens_models_context_swa_scratch(): - manager = object.__new__(KVCacheManagerV2) - manager._has_cp_helix = False - manager.num_local_layers = 3 - manager.pp_layers = [0, 1, 2] - manager.max_attention_window_vec = [128, 128, None] - manager.tokens_per_block = 64 - manager.max_batch_size = 4 - manager.max_num_tokens = 1000 - manager.get_layer_bytes_per_token = lambda local_layer_idx, data_role: [10, 10, 20][ - local_layer_idx - ] - - max_tokens = 1200 - - manager.enable_swa_scratch_reuse = False - no_scratch_quota = manager._get_quota_from_max_tokens(max_tokens) - assert no_scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 20 + 4 * 2 * 128 * 10) - assert manager._get_max_tokens_from_quota(no_scratch_quota) == max_tokens - - manager.enable_swa_scratch_reuse = True - scratch_quota = manager._get_quota_from_max_tokens(max_tokens) - assert scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 10 + 4 * 2 * 128 * 10) - assert manager._get_max_tokens_from_quota(scratch_quota) == max_tokens - - -# --------------------------------------------------------------------------- -# KVCacheManagerV2 clamp_max_seq_len_for_mem float-to-int cast regression -# --------------------------------------------------------------------------- -# -# ``KVCacheManagerV2.get_num_available_tokens`` returns -# ``clamp_max_seq_len_for_mem(...) - extra_tokens``, where the clamp helper -# does a floating-point memory-budget division and returns a float. When -# ``max_seq_len > max_num_tokens`` the ``__init__`` body assigned the float -# directly to ``self.max_seq_len``, which then propagated into -# ``_util.py::_create_dummy_context_requests``'s ``torch.randint(size=(...))`` -# call and crashed on a float size tuple. -# -# resource_manager.py:1820 now casts ``int(max_num_tokens)`` at the assign -# site. The tests below lock in that invariant: (a) the replay test -# demonstrates the fix applied to the scenario that previously crashed, -# and (b) the propagation test shows that downstream arithmetic in -# ``_util.py:610-620`` stays int-safe. - - -def test_kv_cache_manager_v2_clamp_casts_float_to_int(): - """Replay of the V2 clamp block with a float clamp result. - - Reproduces the exact arithmetic from ``resource_manager.py`` - ``KVCacheManagerV2.__init__`` lines 1813-1825. Pre-fix, the ``if`` - branch stored a float on ``self.max_seq_len``. Post-fix, it casts - to int. A broken fix (e.g. dropping the cast) will make this test - fail on the ``isinstance(..., int)`` assertion.""" - # Scenario pulled from the MMMU Pro 26B + s=131K crash log: - # clamp_max_seq_len_for_mem returned 60160.0 when requested - # 131072 tokens of max_seq_len. - initial_max_seq_len = 131072 - # Simulate the float return from clamp_max_seq_len_for_mem. - clamp_float_result = 60160.0 - assert isinstance(clamp_float_result, float) - # Match the production code block (resource_manager.py:1816-1825). - self_max_seq_len = initial_max_seq_len - if self_max_seq_len > clamp_float_result: - self_max_seq_len = int(clamp_float_result) - - assert isinstance(self_max_seq_len, int), ( - "KVCacheManagerV2 must cast the clamp result to int — float " - "propagates into torch.randint(size=...) and crashes." - ) - assert self_max_seq_len == 60160 - - -def test_kv_cache_manager_v2_float_max_seq_len_would_crash_torch_randint(): - """Pre-fix behaviour: a float max_seq_len propagating into - torch.randint(size=(...,)) raises. This test documents WHY the cast - is necessary — if the cast is dropped, the following code crashes.""" - - float_seq_len = 60160.0 - with pytest.raises((TypeError, RuntimeError)): - # torch.randint only accepts int-valued size tuples. - torch.randint(low=0, high=32000, size=(float_seq_len,)) - - -def test_kv_manager_int_max_seq_len_stays_int_through_util_expression(): - """The downstream expression in ``_util.py:615-620`` (which builds the - dummy context request input_seq_len) must stay int when the V2 - manager's max_seq_len is int. Covers the full propagation chain - from the root-cause fix to the downstream consumer.""" - net_max_seq_len = 131072 # from ModelEngine - creator_max_seq_len = 131072 # KvCacheCreator.self._max_seq_len - kv_manager_max_seq_len = 60160 # post-fix int (was 60160.0 pre-fix) - - # Replay _util.py:616-619 arithmetic. - input_seq_len = max( - 1, - net_max_seq_len - 1 - (creator_max_seq_len - kv_manager_max_seq_len), - ) - assert isinstance(input_seq_len, int) - assert input_seq_len >= 1 - - # Sanity: the same expression with a float would yield float (the - # original bug), ensuring this test would notice regression. - buggy_input = max( - 1, - net_max_seq_len - 1 - (creator_max_seq_len - float(kv_manager_max_seq_len)), - ) - assert isinstance(buggy_input, float), ( - "Sanity check: the pre-fix float path must still be demonstrable " - "in the test, otherwise this regression guard is hollow." - ) - - -# --------------------------------------------------------------------------- -# _create_kv_cache_manager: MLA branch must forward max_num_tokens -# --------------------------------------------------------------------------- - - -def test_mla_branch_forwards_max_num_tokens_to_manager() -> None: - """The is_mla branch of ``_create_kv_cache_manager`` must pass - ``max_num_tokens`` to the manager, like the generic branch does. - - Regression guard: when the argument is dropped, - ``DeepseekV4CacheManager._max_num_tokens`` stays ``None`` and the - profiling-phase context extra quota is sized by the full ``max_tokens`` - estimate instead of the runtime chunk size (observed as a 27.59 GiB vs - 11.92 GiB temp-quota inflation on DeepSeek-V4-Pro, raising peak memory - and OOM risk during KV cache estimation). - """ - - from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager - - captured_kwargs = {} - - class _RecordingManager: - def __init__(self, *args, **kwargs) -> None: - captured_kwargs.update(kwargs) - - # Minimal MLA pretrained config: is_mla() keys on kv_lora_rank and - # qk_rope_head_dim being set. - pretrained = SimpleNamespace( - kv_lora_rank=512, - qk_rope_head_dim=64, - hidden_size=1024, - num_attention_heads=8, - num_key_value_heads=8, - num_hidden_layers=2, - vocab_size=32000, - ) - model_config = Mock() - model_config.pretrained_config = pretrained - model_config.quant_config = None - - _create_kv_cache_manager( - model_engine=None, - kv_cache_manager_cls=_RecordingManager, - mapping=Mock(), - kv_cache_config=KvCacheConfig(), - tokens_per_block=32, - max_seq_len=2048, - max_batch_size=8, - spec_config=None, - sparse_attention_config=None, - max_num_tokens=333, - max_beam_width=1, - kv_connector_manager=None, - model_config=model_config, - dtype=torch.bfloat16, - is_draft=False, - ) - - assert captured_kwargs.get("max_num_tokens") == 333, ( - "is_mla branch dropped max_num_tokens: DeepseekV4CacheManager then " - "falls back to _max_num_tokens=None and over-sizes the estimation " - "temp quota." - ) - - -def test_estimation_temporarily_uses_inferred_pool_sizing() -> None: - pool_ratio = [0.2, 0.3, 0.5] - avg_seq_len = 128 - max_seq_len = 4096 - user_max_tokens = 1024 - estimation_max_tokens = 256 - kv_cache_config = KvCacheConfig( - max_tokens=user_max_tokens, - pool_ratio=pool_ratio, - avg_seq_len=avg_seq_len, - ) - model_engine = Mock() - model_engine.model.model_config.attn_backend = "TRTLLM" - # Explicit False: try_prepare_estimation skips estimation for - # encoder-decoder models, and a bare Mock attribute is truthy. - model_engine.model.model_config.is_encoder_decoder = False - # A bare Mock would auto-create the attribute; real engines set it to - # None unless the model opted into MM item scheduling. - model_engine.mm_encoder_output_budget_bytes = None - llm_args = Mock(cache_transceiver_config=None) - - with patch.object( - KvCacheCreator, - "_get_model_kv_cache_manager_cls", - return_value=KVCacheManagerV2, - ): - creator = KvCacheCreator( - model_engine=model_engine, - draft_model_engine=None, - mapping=Mock(cp_config={}), - net_max_seq_len=max_seq_len, - kv_connector_manager=None, - max_num_tokens=256, - max_beam_width=1, - tokens_per_block=128, - max_seq_len=max_seq_len, - max_batch_size=8, - kv_cache_config=kv_cache_config, - llm_args=llm_args, - speculative_config=None, - sparse_attention_config=None, - profiling_stage_data=None, - is_disagg=False, - ) - - with ( - patch.object( - creator, - "_get_token_num_for_estimation", - return_value=estimation_max_tokens, - ), - patch.object(creator, "_cal_max_memory", return_value=512), - patch.object(torch.cuda, "mem_get_info", return_value=(768, 1024)), - patch.object(torch.cuda, "memory_stats", return_value={"allocated_bytes.all.current": 128}), - patch.object(torch.cuda, "empty_cache"), - patch.object(torch.cuda, "reset_peak_memory_stats"), - # This test exercises inferred pool sizing, not the backend workspace reserve; the mock - # model_config would otherwise walk into backend resolution and the MLA byte-cost path. - # Neutralize it so no reserve applies. - patch( - "tensorrt_llm._torch.pyexecutor._util.get_attention_workspace_bytes_per_token", - return_value=0, - ), - ): - assert creator.try_prepare_estimation() - assert kv_cache_config.max_tokens == estimation_max_tokens - assert kv_cache_config.pool_ratio is None - assert kv_cache_config.avg_seq_len == max_seq_len - - creator.configure_kv_cache_capacity() - - assert kv_cache_config.max_tokens == user_max_tokens - assert kv_cache_config.pool_ratio == pool_ratio - assert kv_cache_config.avg_seq_len == avg_seq_len - - -@pytest.mark.parametrize( - ("estimating_kv_cache", "expected_avg_seq_len"), - [(True, 2045), (False, 2055)], -) -def test_manager_estimation_clamps_only_temporary_avg_seq_len( - estimating_kv_cache, - expected_avg_seq_len, -) -> None: - import torch - - from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager - - captured_configs = [] - - class _RecordingKVCacheManagerV2(KVCacheManagerV2): - def __init__(self, kv_cache_config, _kv_cache_type, **kwargs) -> None: - captured_configs.append(kv_cache_config) - self.max_seq_len = kwargs["max_seq_len"] - - pretrained = SimpleNamespace( - hidden_size=1024, - num_attention_heads=8, - num_key_value_heads=8, - num_hidden_layers=2, - vocab_size=32000, - ) - model_config = Mock() - model_config.pretrained_config = pretrained - model_config.quant_config = None - kv_cache_config = KvCacheConfig( - max_tokens=2048, - avg_seq_len=2055, - ) - - _create_kv_cache_manager( - model_engine=None, - kv_cache_manager_cls=_RecordingKVCacheManagerV2, - mapping=Mock(), - kv_cache_config=kv_cache_config, - tokens_per_block=32, - max_seq_len=2045, - max_batch_size=4, - spec_config=None, - sparse_attention_config=None, - max_num_tokens=2048, - max_beam_width=1, - kv_connector_manager=None, - estimating_kv_cache=estimating_kv_cache, - model_config=model_config, - dtype=torch.bfloat16, - is_draft=False, - ) - - assert captured_configs[0].avg_seq_len == expected_avg_seq_len - assert kv_cache_config.avg_seq_len == 2055 - - -def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: - creator = object.__new__(KvCacheCreator) - target_pool_ratio = [0.32, 0.68] - creator._kv_cache_config = KvCacheConfig( - pool_ratio=target_pool_ratio, - max_attention_window=None, - ) - creator._max_seq_len = 9472 - creator._tokens_per_block = 32 - creator._max_batch_size = 1024 - creator._max_num_tokens = 9472 - creator._max_beam_width = 1 - creator._kv_connector_manager = None - creator._skip_est = False - creator._execution_stream = None - creator._is_disagg = False - creator._mapping = Mock() - creator._speculative_config = Mock() - - effective_draft_config = Mock() - effective_draft_config.pretrained_config.torch_dtype = "bfloat16" - effective_draft_config.sparse_attention_config = None - - with ( - patch.object(creator, "_get_num_draft_layers", return_value=1), - patch.object(creator, "_get_one_model_draft_layer_mask", return_value=[True]), - patch.object( - creator, - "_get_effective_draft_config", - return_value=effective_draft_config, - ), - patch.object(creator, "_enable_kv_cache_stats", return_value=False), - patch.object( - creator, - "_validate_or_fallback_kv_cache_manager_v2", - return_value=KVCacheManagerV2, - ), - patch( - "tensorrt_llm._torch.pyexecutor._util._derive_draft_max_attention_window", - return_value=None, - ), - patch( - "tensorrt_llm._torch.pyexecutor._util.get_kv_cache_manager_cls", - return_value=KVCacheManagerV2, - ), - patch( - "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", - return_value=Mock(), - ) as create_manager, - ): - creator._create_one_model_draft_kv_cache_manager(creator._max_seq_len) - - draft_config = create_manager.call_args.kwargs["kv_cache_config"] - assert draft_config.pool_ratio == [1.0] - assert creator._kv_cache_config.pool_ratio == target_pool_ratio - - -@pytest.mark.parametrize( - "attention_backend,dtype,bytes_per_layer", - [("TRTLLM", torch.float8_e4m3fn, 2048), ("VANILLA", torch.bfloat16, 4096)], -) -def test_standalone_draft_kv_dtype_follows_the_drafter_attention( - attention_backend, dtype, bytes_per_layer -) -> None: - """The TRTLLM drafter attention stores the drafter's history in FP8; VANILLA keeps BF16.""" - creator = object.__new__(KvCacheCreator) - draft = SimpleNamespace( - fc=SimpleNamespace(weight=torch.empty(1, dtype=torch.bfloat16)), - dflash_attention_backend=attention_backend, - config=SimpleNamespace(max_position_embeddings=4096), - ) - creator._model_engine = SimpleNamespace(model=SimpleNamespace(draft_model=draft)) - creator._speculative_config = SimpleNamespace( - spec_dec_mode=SimpleNamespace(is_dspark=lambda: True), - draft_is_embedded_in_target=False, - attention_backend="AUTO", - max_draft_len=3, - ) - # Drafter shapes remain unsharded under attention DP. - creator._draft_config = SimpleNamespace( - pretrained_config=SimpleNamespace( - num_attention_heads=64, - num_key_value_heads=16, - head_dim=64, - hidden_size=4096, - num_hidden_layers=5, - ) - ) - creator._mapping = SimpleNamespace( - has_cp_helix=lambda: False, enable_attention_dp=True, tp_size=16 - ) - - kwargs = creator._get_unified_draft_cache_kwargs() - - assert kwargs["draft_dtype"] == torch_dtype_to_binding(dtype) - assert _get_draft_layer_sizes(**kwargs) == (bytes_per_layer,) * 5 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py index 31355f231135..6c0da73da58e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py @@ -87,6 +87,85 @@ def test_trtllm_backend_builds_private_paged_context_cache(monkeypatch): ] +def test_generation_step_keeps_the_dummy_slot_context_empty(): + """Padding and ADP dummy rows must not accumulate drafter context. + + They all map to the shared dummy slot, which no request ever frees. If the + generation step advanced its length like a real slot's, the dummies' drafter + context would grow every step until it reached ``_max_ctx``. + """ + max_draft_len = 3 + num_slots = 3 # two request slots plus the dummy slot + hidden, num_layers, num_kv_heads, head_dim = 8, 1, 1, 4 + max_ctx = 64 + config = DFlashDecodingConfig( + max_draft_len=max_draft_len, attention_backend="TRTLLM", mask_token_id=0 + ) + worker = DFlashWorker(config, Mapping()) + worker._dummy_slot = num_slots - 1 + worker._max_ctx = max_ctx + worker._compute_block_size = max_draft_len + 1 + worker._ctx_len = torch.tensor([5, 0, 0], dtype=torch.long, device="cuda") + # Gen request 0 owns slot 0; gen request 1 is padding on the dummy slot. + worker._batch_to_slot_buffer = DeviceCopyBuffer( + torch.tensor([0, worker._dummy_slot], dtype=torch.long, device="cuda") + ) + buf_shape = (num_slots, num_layers, max_ctx + max_draft_len + 1, num_kv_heads, head_dim) + worker._ctx_k_buf = torch.zeros(buf_shape, device="cuda") + worker._ctx_v_buf = torch.zeros(buf_shape, device="cuda") + + num_gens = 2 + tokens_per_req = max_draft_len + 1 + num_target_tokens = num_gens * tokens_per_req + + def precompute_context_kv(projected, positions): + rows = projected.shape[0] + k = torch.ones(rows, num_layers, num_kv_heads, head_dim, device="cuda") + return k, k.clone() + + embed_tokens = torch.nn.Embedding(16, hidden).cuda() + draft_model = SimpleNamespace( + fc=object(), + hidden_norm=object(), + project_target_hidden=lambda hs: hs, + precompute_context_kv=precompute_context_kv, + draft_model_full=SimpleNamespace(model=SimpleNamespace(embed_tokens=embed_tokens)), + ) + captured = torch.randn(num_target_tokens, hidden, device="cuda") + spec_metadata = SimpleNamespace( + hidden_size=hidden, + runtime_draft_len=max_draft_len, + get_hidden_states=lambda num_tokens: captured[:num_tokens], + ) + attn_metadata = SimpleNamespace( + num_contexts=0, + num_seqs=num_gens, + num_ctx_tokens=0, + _seq_lens_cuda=torch.zeros(num_gens, dtype=torch.int32, device="cuda"), + _seq_lens=torch.zeros(num_gens, dtype=torch.int32), + ) + accepted_tokens = torch.arange(num_gens * tokens_per_req, dtype=torch.long, device="cuda").view( + num_gens, tokens_per_req + ) + num_accepted_tokens = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + + for step in range(2): + inputs = worker.prepare_1st_drafter_inputs( + input_ids=torch.zeros(num_target_tokens, dtype=torch.long, device="cuda"), + position_ids=torch.zeros(num_target_tokens, dtype=torch.long, device="cuda"), + hidden_states=torch.zeros(num_target_tokens, hidden, device="cuda"), + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + total_target_tokens=num_target_tokens, + ) + real_len = 5 + 2 * (step + 1) + assert worker._ctx_len.tolist() == [real_len, 0, 0] + assert inputs["num_ctx_per_req"].tolist() == [real_len, 0] + + def _make_worker(monkeypatch): monkeypatch.setattr( "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index 098e42770acf..b620f71e19d8 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -12,7 +12,13 @@ # 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. -"""DSpark attention math and managed-page addressing tests.""" +"""Unit tests for the DSpark captured-context attention primitives (CPU). + +The two primitives are validated against fully independent computations: +``dspark_sparse_attn`` vs (a) a per-element loop reference of the kernel formula +and (b) ``torch.nn.functional.scaled_dot_product_attention`` for the no-sink, +all-valid case; ``get_dspark_topk_idxs`` vs the reference index formula. +""" import types from unittest.mock import Mock @@ -24,12 +30,15 @@ import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark from tensorrt_llm._torch.models.modeling_dspark import ( DSv4DSparkDraftModel, + _output_projection, + _qkv_projections, apply_dspark_rotary, + dspark_attention_forward, + dspark_sparse_attn, + get_dspark_topk_idxs, precompute_dspark_freqs_cis, ) - -__extra_import_path__ = [".."] - +from tensorrt_llm._torch.modules.linear import Linear # The captured-context attention primitives were folded into # modeling_dspark; keep the historical alias so monkeypatch targets @@ -37,6 +46,9 @@ dspark_attention = modeling_dspark +__extra_import_path__ = [".."] + + def test_rmsnorm_rope_fallback_applies_weight_without_rmsnorm(monkeypatch): monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_AVAILABLE", False) torch.manual_seed(11) @@ -93,6 +105,426 @@ def test_batched_attention_rejects_mismatched_window_size(): forward(**args) +def test_rope_table_is_cached_once_per_device(): + model = types.SimpleNamespace( + _attn_params={"rope_head_dim": 16}, + _freqs_cap=64, + _rope_theta=10000.0, + _freqs_table_cache={}, + ) + + first = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + second = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + + assert first.data_ptr() == second.data_ptr() + assert len(model._freqs_table_cache) == 1 + positions = torch.tensor([1, 17, 63]) + expected = precompute_dspark_freqs_cis(16, 64, rope_theta=10000.0) + torch.testing.assert_close(first[positions], expected[positions]) + + +def test_dspark_block_uses_stage_id_as_attention_layer_idx(monkeypatch): + captured = {} + + def fake_decoder_layer_init( + self, + model_config, + layer_idx, + aux_stream_dict, + attention_layer_idx=None, + mapping_with_cp=None, + disable_post_moe_fusion=False, + ): + torch.nn.Module.__init__(self) + self.model_config = model_config + self.config = model_config.pretrained_config + self.layer_idx = layer_idx + captured.update( + layer_idx=layer_idx, + attention_layer_idx=attention_layer_idx, + aux_stream_dict=aux_stream_dict, + mapping_with_cp=mapping_with_cp, + disable_post_moe_fusion=disable_post_moe_fusion, + ) + + monkeypatch.setattr( + modeling_dspark.DeepseekV4DecoderLayer, + "__init__", + fake_decoder_layer_init, + ) + model_config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2), + spec_config=None, + ) + + block = modeling_dspark.DSv4DSparkBlock( + model_config, + layer_idx=10, + aux_stream_dict={}, + stage_id=1, + num_stages=3, + num_capture_layers=0, + ) + + assert block.layer_idx == captured["layer_idx"] == 10 + assert captured["attention_layer_idx"] == block.stage_id == 1 + assert captured["disable_post_moe_fusion"] is True + + +@pytest.mark.parametrize("enable_fused_hc", [True, False]) +def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): + """The draft stage must use the inherited fused-HC rollback setting.""" + torch.manual_seed(71) + num_requests, block_size, hc_mult, hidden_size = 1, 2, 2, 3 + h = torch.randn(num_requests, block_size, hc_mult, hidden_size) + attention_input = torch.randn(num_requests, block_size, hidden_size) + attention_output = torch.randn_like(attention_input) + mid_residual = torch.randn_like(h) + attention_post_mix = torch.randn(num_requests, block_size, hc_mult, 1) + attention_comb_mix = torch.randn(num_requests, block_size, hc_mult, hc_mult) + ffn_post_mix = torch.randn_like(attention_post_mix) + ffn_comb_mix = torch.randn_like(attention_comb_mix) + raw_ffn_input = torch.randn_like(attention_input) + normed_ffn_input = torch.randn_like(attention_input) + moe_output = torch.randn(num_requests * block_size, hidden_size) + final_h = torch.randn_like(h) + events = [] + + def record(name, result): + def call(*args, **kwargs): + events.append(name) + return result + + return call + + monkeypatch.setattr( + modeling_dspark, + "dspark_attention_forward", + Mock(return_value=attention_output), + ) + + hc_attn = types.SimpleNamespace( + pre_mapping=Mock(return_value=(attention_post_mix, attention_comb_mix, attention_input)), + post_mapping=Mock(side_effect=record("attention_post", mid_residual)), + ) + hc_ffn = types.SimpleNamespace( + fused_hc=Mock( + side_effect=record( + "fused", + (mid_residual, ffn_post_mix, ffn_comb_mix, normed_ffn_input), + ) + ), + pre_mapping=Mock( + side_effect=record("ffn_pre", (ffn_post_mix, ffn_comb_mix, raw_ffn_input)) + ), + post_mapping=Mock(side_effect=record("ffn_post", final_h)), + ) + post_attention_layernorm = Mock(side_effect=record("ffn_norm", normed_ffn_input)) + post_attention_layernorm.weight = torch.ones(hidden_size) + post_attention_layernorm.variance_epsilon = 1e-6 + stage = types.SimpleNamespace( + enable_fused_hc=enable_fused_hc, + hc_attn=hc_attn, + hc_ffn=hc_ffn, + input_layernorm=Mock(side_effect=lambda tensor: tensor), + post_attention_layernorm=post_attention_layernorm, + mlp=Mock(return_value=moe_output), + _dspark_attn={}, + ) + model = types.SimpleNamespace( + use_real_mla=False, + _attn_params={"window_size": 2, "head_dim": 1}, + model_config=types.SimpleNamespace( + mapping=types.SimpleNamespace(enable_attention_dp=False, tp_size=8) + ), + ) + + actual = DSv4DSparkDraftModel._forward_stage( + model, + stage, + h, + torch.randn(num_requests, hidden_size), + torch.ones(num_requests, dtype=torch.long), + torch.empty(0), + torch.zeros(num_requests, block_size, dtype=torch.long), + torch.empty(1, 128, 1), + torch.zeros(num_requests, 1, dtype=torch.int32), + torch.full((num_requests,), 128, dtype=torch.long), + torch.ones(num_requests, dtype=torch.long), + ) + + assert actual is final_h + torch.testing.assert_close( + stage.mlp.call_args.args[0], + normed_ffn_input.reshape(num_requests * block_size, hidden_size), + ) + # Non-attention-DP multi-GPU (enable_attention_dp=False, tp_size>1): the draft + # MoE must all-reduce its TP-sharded output, mirroring the target MoE. The + # attention-DP and single-GPU paths keep it disabled. + assert stage.mlp.call_args.kwargs["final_all_reduce_params"].enable_allreduce is True + if enable_fused_hc: + assert events == ["fused", "ffn_post"] + hc_ffn.fused_hc.assert_called_once() + hc_attn.post_mapping.assert_not_called() + hc_ffn.pre_mapping.assert_not_called() + post_attention_layernorm.assert_not_called() + fused_kwargs = hc_ffn.fused_hc.call_args.kwargs + assert fused_kwargs["norm_weight"] is post_attention_layernorm.weight + assert fused_kwargs["norm_eps"] == post_attention_layernorm.variance_epsilon + else: + assert events == ["attention_post", "ffn_pre", "ffn_norm", "ffn_post"] + hc_ffn.fused_hc.assert_not_called() + hc_attn.post_mapping.assert_called_once() + hc_ffn.pre_mapping.assert_called_once_with(mid_residual) + post_attention_layernorm.assert_called_once_with(raw_ffn_input) + + +def _ref_precompute_freqs_cis(dim, seqlen, base): + """DeepSpec precompute_freqs_cis with original_seq_len == 0 (no YaRN).""" + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): + """DeepSpec apply_rotary_emb (returns a fresh tensor instead of in-place).""" + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if xc.ndim == 3: + fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) + else: + fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) + return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) + + +def _loop_reference(q, kv, attn_sink, topk_idxs, scale): + """Obvious, slow per-(b,m,h) reference of the exact kernel math.""" + b, m, h, d = q.shape + out = torch.zeros(b, m, h, d, dtype=torch.float32) + qf, kvf, sink = q.float(), kv.float(), attn_sink.float() + for bi in range(b): + for mi in range(m): + idxs = topk_idxs[bi, mi].tolist() + for hi in range(h): + scores, vecs = [], [] + for j in idxs: + if j < 0: + continue + k = kvf[bi, j] + scores.append(torch.dot(qf[bi, mi, hi], k) * scale) + vecs.append(k) + if not scores: + continue + s = torch.stack(scores) + smax = s.max() + p = torch.exp(s - smax) + denom = p.sum() + torch.exp(sink[hi] - smax) + num = (p.unsqueeze(-1) * torch.stack(vecs)).sum(0) + out[bi, mi, hi] = num / denom + return out + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_sparse_attn_matches_loop_reference(seed): + torch.manual_seed(seed) + b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, n, d) + attn_sink = torch.randn(h) + idx = torch.stack( + [torch.stack([torch.randperm(n)[:topk] for _ in range(m)]) for _ in range(b)] + ).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + ref = _loop_reference(q, kv, attn_sink, idx, scale) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_no_sink_matches_sdpa(): + """With a -inf sink (no sink mass) and all-valid contiguous indices, the + primitive must equal standard scaled-dot-product attention over the gathered + KV — an independent implementation.""" + torch.manual_seed(0) + b, m, h, d, topk = 2, 4, 5, 16, 9 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, topk, d) # n == topk, attend to all + attn_sink = torch.full((h,), float("-inf")) + idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + + # SDPA: q [b,h,m,d], k/v [b,h,topk,d] (broadcast the shared KV over heads). + qh = q.permute(0, 2, 1, 3) + kvh = kv.unsqueeze(1).expand(b, h, topk, d) + ref = F.scaled_dot_product_attention(qh, kvh, kvh, scale=scale).permute(0, 2, 1, 3) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_sink_reduces_mass(): + """A finite sink must strictly shrink the attention output magnitude vs an + infinitely-negative (disabled) sink, because it adds denominator mass only.""" + torch.manual_seed(0) + b, m, h, d, topk = 1, 2, 2, 16, 6 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, topk, d) + idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() + scale = d**-0.5 + no_sink = dspark_sparse_attn(q, kv, torch.full((h,), float("-inf")), idx, scale) + with_sink = dspark_sparse_attn(q, kv, torch.zeros(h), idx, scale) + assert with_sink.abs().sum() < no_sink.abs().sum() + + +def test_sparse_attn_masked_indices_excluded(): + """An index of -1 must be excluded exactly (equiv. to dropping that column).""" + torch.manual_seed(0) + b, m, h, d = 1, 1, 2, 16 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, 5, d) + sink = torch.full((h,), float("-inf")) + scale = d**-0.5 + full = torch.tensor([[[0, 1, 2, 3]]]).int() + masked = torch.tensor([[[0, 1, 2, -1]]]).int() + drop3 = torch.tensor([[[0, 1, 2]]]).int() + got_masked = dspark_sparse_attn(q, kv, sink, masked, scale) + got_drop = dspark_sparse_attn(q, kv, sink, drop3, scale) + torch.testing.assert_close(got_masked, got_drop, rtol=1e-5, atol=1e-5) + # And masking genuinely changes the result vs attending to position 3. + got_full = dspark_sparse_attn(q, kv, sink, full, scale) + assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize( + "start_pos,window,block", [(1, 128, 5), (3, 128, 5), (10, 4, 5), (200, 128, 6)] +) +def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): + bsz = 3 + got = get_dspark_topk_idxs(window, bsz, block, start_pos) + # Reference formula (DeepSpec get_dspark_topk_idxs). + ctx = torch.arange(min(window, start_pos + 1)) + blk = window + torch.arange(block) + ref_row = torch.cat([ctx, blk]).int() + assert got.shape == (bsz, block, ref_row.numel()) + for bi in range(bsz): + for mi in range(block): + torch.testing.assert_close(got[bi, mi], ref_row) + + +def test_get_dspark_topk_idxs_requires_generation(): + with pytest.raises(AssertionError): + get_dspark_topk_idxs(128, 1, 5, 0) + + +@pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) +def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): + got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) + ref = _ref_precompute_freqs_cis(rope_head_dim, seqlen, 10000.0) + torch.testing.assert_close(got, ref) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_apply_rotary_matches_reference(ndim): + torch.manual_seed(0) + b, s, h, rd = 2, 5, 4, 64 + x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) + fc = precompute_dspark_freqs_cis(rd, s) + got = apply_dspark_rotary(x, fc) + ref = _ref_apply_rotary_emb(x, fc) + torch.testing.assert_close(got, ref) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_apply_rotary_inverse_roundtrip(ndim): + """De-rotation (inverse) must undo the forward rotation (property test).""" + torch.manual_seed(1) + b, s, h, rd = 2, 6, 3, 64 + x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) + fc = precompute_dspark_freqs_cis(rd, s) + roundtrip = apply_dspark_rotary(apply_dspark_rotary(x, fc), fc, inverse=True) + torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) + + +def _make_attn_inputs(seed=0): + """Small synthetic DSpark attention inputs/weights (CPU bf16).""" + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block, start_pos = 8, 3, 5 + b = 2 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + start_pos=start_pos, + b=b, + eps=1e-6, + softmax_scale=head_dim**-0.5, + ) + bf = torch.bfloat16 + g["x"] = torch.randn(b, block, dim, dtype=bf) + g["main_x"] = torch.randn(b, 1, dim, dtype=bf) + g["kv_cache0"] = torch.randn(b, window, head_dim, dtype=bf) + g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 + g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 + g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 + g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 + g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 + g["q_norm"] = torch.ones(q_lora) + g["kv_norm"] = torch.ones(head_dim) + g["attn_sink"] = torch.randn(n_heads) + g["freqs"] = precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) + return g + + +def _run(g): + return dspark_attention_forward( + g["x"], + g["main_x"], + g["start_pos"], + g["kv_cache0"], + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def test_attention_forward_shape_and_determinism(): + g = _make_attn_inputs() + o = _run(g) + assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) + assert torch.isfinite(o.float()).all() + torch.testing.assert_close(o, _run(g)) # deterministic + + +def test_attention_forward_does_not_mutate_kv_cache(): + """The rolling window write must be functional (cache cloned, not mutated).""" + g = _make_attn_inputs() + before = g["kv_cache0"].clone() + _run(g) + torch.testing.assert_close(g["kv_cache0"], before) + + @pytest.mark.parametrize( "value,tp_size,enable_attention_dp,expected", [ @@ -276,7 +708,7 @@ def fake_bmm_quantize( packed_scale.zero_() monkeypatch.setattr(dspark_attention, "get_sm_version", lambda: 107) - monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_INTERNAL_AVAILABLE", True) + monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", True) monkeypatch.setattr( torch.ops.trtllm, "fp8_batched_quantize_1x128_permute102", @@ -429,72 +861,6 @@ def output_projection(grouped_o_fp8, grouped_o_scale, **kwargs): assert calls == ["prepared_fp8_attention", "fp8_output_consumer"] -def test_rope_table_is_cached_once_per_device(): - model = types.SimpleNamespace( - _attn_params={"rope_head_dim": 16}, - _freqs_cap=64, - _rope_theta=10000.0, - _freqs_table_cache={}, - ) - - first = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - second = DSv4DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - - assert first.data_ptr() == second.data_ptr() - assert len(model._freqs_table_cache) == 1 - positions = torch.tensor([1, 17, 63]) - expected = precompute_dspark_freqs_cis(16, 64, rope_theta=10000.0) - torch.testing.assert_close(first[positions], expected[positions]) - - -def test_dspark_block_uses_stage_id_as_attention_layer_idx(monkeypatch): - captured = {} - - def fake_decoder_layer_init( - self, - model_config, - layer_idx, - aux_stream_dict, - attention_layer_idx=None, - mapping_with_cp=None, - disable_post_moe_fusion=False, - ): - torch.nn.Module.__init__(self) - self.model_config = model_config - self.config = model_config.pretrained_config - self.layer_idx = layer_idx - captured.update( - layer_idx=layer_idx, - attention_layer_idx=attention_layer_idx, - aux_stream_dict=aux_stream_dict, - mapping_with_cp=mapping_with_cp, - disable_post_moe_fusion=disable_post_moe_fusion, - ) - - monkeypatch.setattr( - modeling_dspark.DeepseekV4DecoderLayer, - "__init__", - fake_decoder_layer_init, - ) - model_config = types.SimpleNamespace( - pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2), - spec_config=None, - ) - - block = modeling_dspark.DSv4DSparkBlock( - model_config, - layer_idx=10, - aux_stream_dict={}, - stage_id=1, - num_stages=3, - num_capture_layers=0, - ) - - assert block.layer_idx == captured["layer_idx"] == 10 - assert captured["attention_layer_idx"] == block.stage_id == 1 - assert captured["disable_post_moe_fusion"] is True - - def test_context_wkv_autotune_bucket_contract(): from tensorrt_llm._torch.custom_ops.dspark_attention_custom_op import ( _gen_dspark_context_wkv_m_profiles, @@ -713,7 +1079,7 @@ def test_bind_attn_projections_reuses_loaded_mla_modules(monkeypatch): mtp_layers=[stage], ) monkeypatch.setattr(dspark_attention, "get_sm_version", lambda: 107) - monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_INTERNAL_AVAILABLE", True) + monkeypatch.setattr(dspark_attention, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", True) DSv4DSparkDraftModel._bind_attn_projections(model) @@ -967,158 +1333,28 @@ def fake_gemm(rows, weight, scale, weight_scale): assert len(quant_calls) == len(gemm_calls) == 1 -@pytest.mark.parametrize("enable_fused_hc", [True, False]) -def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): - """The draft stage must use the inherited fused-HC rollback setting.""" - torch.manual_seed(71) - num_requests, block_size, hc_mult, hidden_size = 1, 2, 2, 3 - h = torch.randn(num_requests, block_size, hc_mult, hidden_size) - attention_input = torch.randn(num_requests, block_size, hidden_size) - attention_output = torch.randn_like(attention_input) - mid_residual = torch.randn_like(h) - attention_post_mix = torch.randn(num_requests, block_size, hc_mult, 1) - attention_comb_mix = torch.randn(num_requests, block_size, hc_mult, hc_mult) - ffn_post_mix = torch.randn_like(attention_post_mix) - ffn_comb_mix = torch.randn_like(attention_comb_mix) - raw_ffn_input = torch.randn_like(attention_input) - normed_ffn_input = torch.randn_like(attention_input) - moe_output = torch.randn(num_requests * block_size, hidden_size) - final_h = torch.randn_like(h) - events = [] - - def record(name, result): - def call(*args, **kwargs): - events.append(name) - return result - - return call - - monkeypatch.setattr( - modeling_dspark, - "dspark_attention_forward", - Mock(return_value=attention_output), - ) - - hc_attn = types.SimpleNamespace( - pre_mapping=Mock(return_value=(attention_post_mix, attention_comb_mix, attention_input)), - post_mapping=Mock(side_effect=record("attention_post", mid_residual)), - ) - hc_ffn = types.SimpleNamespace( - fused_hc=Mock( - side_effect=record( - "fused", - (mid_residual, ffn_post_mix, ffn_comb_mix, normed_ffn_input), - ) - ), - pre_mapping=Mock( - side_effect=record("ffn_pre", (ffn_post_mix, ffn_comb_mix, raw_ffn_input)) - ), - post_mapping=Mock(side_effect=record("ffn_post", final_h)), - ) - post_attention_layernorm = Mock(side_effect=record("ffn_norm", normed_ffn_input)) - post_attention_layernorm.weight = torch.ones(hidden_size) - post_attention_layernorm.variance_epsilon = 1e-6 - stage = types.SimpleNamespace( - enable_fused_hc=enable_fused_hc, - hc_attn=hc_attn, - hc_ffn=hc_ffn, - input_layernorm=Mock(side_effect=lambda tensor: tensor), - post_attention_layernorm=post_attention_layernorm, - mlp=Mock(return_value=moe_output), - _dspark_attn={}, - ) - model = types.SimpleNamespace( - use_real_mla=False, - _attn_params={"window_size": 2, "head_dim": 1}, - model_config=types.SimpleNamespace( - mapping=types.SimpleNamespace(enable_attention_dp=False, tp_size=8) - ), - ) - - actual = DSv4DSparkDraftModel._forward_stage( - model, - stage, - h, - torch.randn(num_requests, hidden_size), - torch.ones(num_requests, dtype=torch.long), - torch.empty(0), - torch.zeros(num_requests, block_size, dtype=torch.long), - torch.empty(1, 128, 1), - torch.zeros(num_requests, 1, dtype=torch.int32), - torch.full((num_requests,), 128, dtype=torch.long), - torch.ones(num_requests, dtype=torch.long), - ) +@pytest.mark.cpu_only +def test_zero_argument_post_load_transforms_non_attention_linears(monkeypatch): + from tensorrt_llm._torch.models import modeling_dspark - assert actual is final_h - torch.testing.assert_close( - stage.mlp.call_args.args[0], - normed_ffn_input.reshape(num_requests * block_size, hidden_size), - ) - # Non-attention-DP multi-GPU (enable_attention_dp=False, tp_size>1): the draft - # MoE must all-reduce its TP-sharded output, mirroring the target MoE. The - # attention-DP and single-GPU paths keep it disabled. - assert stage.mlp.call_args.kwargs["final_all_reduce_params"].enable_allreduce is True - if enable_fused_hc: - assert events == ["fused", "ffn_post"] - hc_ffn.fused_hc.assert_called_once() - hc_attn.post_mapping.assert_not_called() - hc_ffn.pre_mapping.assert_not_called() - post_attention_layernorm.assert_not_called() - fused_kwargs = hc_ffn.fused_hc.call_args.kwargs - assert fused_kwargs["norm_weight"] is post_attention_layernorm.weight - assert fused_kwargs["norm_eps"] == post_attention_layernorm.variance_epsilon - else: - assert events == ["attention_post", "ffn_pre", "ffn_norm", "ffn_post"] - hc_ffn.fused_hc.assert_not_called() - hc_attn.post_mapping.assert_called_once() - hc_ffn.pre_mapping.assert_called_once_with(mid_residual) - post_attention_layernorm.assert_called_once_with(raw_ffn_input) - - -def _ref_precompute_freqs_cis(dim, seqlen, base): - """DeepSpec precompute_freqs_cis with original_seq_len == 0 (no YaRN).""" - freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - t = torch.arange(seqlen) - freqs = torch.outer(t, freqs) - return torch.polar(torch.ones_like(freqs), freqs) - - -def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): - """DeepSpec apply_rotary_emb (returns a fresh tensor instead of in-place).""" - xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) - if inverse: - freqs_cis = freqs_cis.conj() - if xc.ndim == 3: - fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) - else: - fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) - return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) + class TestLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.transforms = 0 + def transform_weights(self): + self.transforms += 1 -@pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) -def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): - got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) - ref = _ref_precompute_freqs_cis(rope_head_dim, seqlen, 10000.0) - torch.testing.assert_close(got, ref) + monkeypatch.setattr(modeling_dspark, "Linear", TestLinear) + draft = object.__new__(modeling_dspark.DSv4DSparkDraftModel) + torch.nn.Module.__init__(draft) + draft._use_bf16_linear = False + stage = torch.nn.Module() + stage.self_attn = torch.nn.Sequential(TestLinear()) + draft.mtp_layers = torch.nn.ModuleList([stage]) + draft.main_proj = TestLinear() + draft.post_load_weights() -@pytest.mark.parametrize("ndim", [3, 4]) -def test_apply_rotary_matches_reference(ndim): - torch.manual_seed(0) - b, s, h, rd = 2, 5, 4, 64 - x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) - fc = precompute_dspark_freqs_cis(rd, s) - got = apply_dspark_rotary(x, fc) - ref = _ref_apply_rotary_emb(x, fc) - torch.testing.assert_close(got, ref) - - -@pytest.mark.parametrize("ndim", [3, 4]) -def test_apply_rotary_inverse_roundtrip(ndim): - """De-rotation (inverse) must undo the forward rotation (property test).""" - torch.manual_seed(1) - b, s, h, rd = 2, 6, 3, 64 - x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) - fc = precompute_dspark_freqs_cis(rd, s) - roundtrip = apply_dspark_rotary(apply_dspark_rotary(x, fc), fc, inverse=True) - torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) + assert draft.main_proj.transforms == 1 + assert stage.self_attn[0].transforms == 0 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index 2e47dc49014f..5d77d8cff968 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -12,7 +12,17 @@ # 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. -"""Per-request rotary equivalence for batched DSpark graph inputs.""" +"""CUDA-graph-safety unit tests for the DSpark batched draft-attention path. + +The load-bearing invariant for the DSpark batched draft path (the default, used +whenever ``cuda_graph_config`` is set, since the one-engine drafter is captured in +the target's graph) is that the batched, sync-free primitives are **numerically +identical, per request**, to the validated scalar path — only the host-int +``start_pos`` and the per-request window indexing are tensorized. These tests assert +that equivalence on CPU (so they run in pre-merge CI without a GPU), plus a +GPU-gated capture+replay smoke test that proves the batched attention is actually +graph-capturable. +""" import pytest import torch @@ -20,10 +30,188 @@ from tensorrt_llm._torch.models.modeling_dspark import ( apply_dspark_rotary, apply_dspark_rotary_batched, + dspark_attention_forward, + dspark_attention_forward_batched, + get_dspark_topk_idxs, + get_dspark_topk_idxs_batched, precompute_dspark_freqs_cis, ) +def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): + """Per-request DSpark attention inputs/weights (shared weights, distinct pos). + + Mirrors ``test_dspark_attention._make_attn_inputs`` but builds ``G`` requests + each with its own ``start_pos`` (small => partial context, large => full + rolling window) and its own pre-seeded window, so the batched-vs-scalar + comparison exercises the per-request RoPE gather + windowed context read. + """ + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block = 8, 3 + G = len(start_positions) + bf = torch.bfloat16 + # A single fixed RoPE table covering every request's positions (both paths + # index/gather the same values, so freqs are identical across paths). + maxlen = max(start_positions) + 1 + block + 4 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + G=G, + start_positions=list(start_positions), + eps=1e-6, + softmax_scale=head_dim**-0.5, + x=torch.randn(G, block, dim, dtype=bf), + main_x=torch.randn(G, 1, dim, dtype=bf), + # Distinct, non-zero seeded window per request to exercise context reads. + kv_cache=torch.randn(G, window, head_dim, dtype=bf), + wq_a=torch.randn(q_lora, dim, dtype=bf) * 0.1, + wq_b=torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1, + wkv=torch.randn(head_dim, dim, dtype=bf) * 0.1, + wo_a=torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1, + wo_b=torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1, + q_norm=torch.ones(q_lora), + kv_norm=torch.ones(head_dim), + attn_sink=torch.randn(n_heads), + freqs=precompute_dspark_freqs_cis(rd, maxlen), + ) + return g + + +def _attn_kwargs(g): + return dict( + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def _scalar_reference(g): + """Run the validated scalar attention once per request and stack the outputs.""" + outs = [] + for i, sp in enumerate(g["start_positions"]): + out_i = dspark_attention_forward( + g["x"][i : i + 1], + g["main_x"][i : i + 1], + int(sp), + g["kv_cache"][i : i + 1].clone(), + **_attn_kwargs(g), + ) + outs.append(out_i) + return torch.cat(outs, dim=0) + + +def _batched(g, persist=False): + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + return dspark_attention_forward_batched( + g["x"], + g["main_x"], + start_pos, + g["kv_cache"].clone(), + slots, + persist=persist, + **_attn_kwargs(g), + ) + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_batched_attention_matches_scalar_per_request(seed): + """Batched attention == per-request scalar attention at distinct start_pos. + + This is the invariant that lets the batched path replace the per-request loop + under CUDA graphs without changing draft quality / greedy parity. + """ + g = _make_batched_inputs(seed=seed) + ref = _scalar_reference(g) + got = _batched(g) + assert tuple(got.shape) == (g["G"], g["block"], g["dim"]) + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) + + +def test_batched_attention_persist_writes_through_window(): + """persist=True writes main_kv into the shared window at start_pos%window.""" + g = _make_batched_inputs(seed=3) + G, win = g["G"], g["window"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + cache = g["kv_cache"].clone() + before = cache.clone() + dspark_attention_forward_batched( + g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) + ) + # Exactly the start_pos%win row of each request changed. + for i, sp in enumerate(g["start_positions"]): + changed = (cache[i] != before[i]).any(dim=-1) + expected = torch.zeros(win, dtype=torch.bool) + expected[sp % win] = True + assert torch.equal(changed, expected), f"req {i}: wrong window row written" + + +def test_batched_attention_no_persist_keeps_window(): + """persist=False must not mutate the caller's window (functional).""" + g = _make_batched_inputs(seed=4) + before = g["kv_cache"].clone() + _batched(g, persist=False) + torch.testing.assert_close(g["kv_cache"], before) + + +@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) +def test_batched_topk_matches_scalar(start_positions): + """Fixed-size masked batched topk == scalar topk per request (valid slots).""" + window, block = 8, 3 + start_pos = torch.tensor(start_positions, dtype=torch.long) + batched = get_dspark_topk_idxs_batched(window, block, start_pos) + assert tuple(batched.shape) == (len(start_positions), block, window + block) + for i, sp in enumerate(start_positions): + scalar = get_dspark_topk_idxs(window, 1, block, int(sp))[0] # [block, topk_i] + # The batched row drops the -1-masked context slots to recover the scalar + # (variable-width) index list; both must then be identical. + for m in range(block): + valid = batched[i, m][batched[i, m] >= 0] + torch.testing.assert_close(valid, scalar[m].to(valid.dtype)) + + +def test_batched_topk_respects_partial_window_valid_len(): + """Only actually written circular slots are visible for partial seeds.""" + window, block = 8, 3 + start_pos = torch.tensor([3, 10, 20], dtype=torch.long) + valid_len = torch.tensor([0, 3, 8], dtype=torch.long) + batched = get_dspark_topk_idxs_batched(window, block, start_pos, valid_len) + assert tuple(batched.shape) == (3, block, window + block) + + for i, (pos, count) in enumerate(zip(start_pos.tolist(), valid_len.tolist())): + expected_context = sorted((pos - age) % window for age in range(count)) + row = batched[i, 0] + context = row[(row >= 0) & (row < window)].tolist() + assert context == expected_context + block_indices = row[row >= window].tolist() + assert block_indices == list(range(window, window + block)) + + @pytest.mark.parametrize("ndim", [3, 4]) def test_batched_rotary_matches_scalar_per_row(ndim): torch.manual_seed(0) @@ -36,3 +224,44 @@ def test_batched_rotary_matches_scalar_per_row(ndim): for i, sp in enumerate(starts): ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) torch.testing.assert_close(got[i : i + 1], ref_i) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_batched_attention_cuda_graph_capture_replay(): + """The batched attention captures + replays and matches eager output. + + Proves the path is free of capture-illegal ops (host syncs, dynamic shapes). + """ + g = _make_batched_inputs(seed=0) + dev = "cuda" + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) + slots = torch.arange(G, dtype=torch.long, device=dev) + # Static input tensors the graph reads/writes. + x = g["x"].to(dev) + main_x = g["main_x"].to(dev) + cache = g["kv_cache"].to(dev) + kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} + + def run(persist): + return dspark_attention_forward_batched( + x, main_x, start_pos, cache, slots, persist=persist, **kw + ) + + eager = run(persist=False) + + # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so + # the comparison isn't perturbed by the window write-through. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run(persist=False) + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = run(persist=False) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 1d846c73ed6a..4e5d2a647f88 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -16,12 +16,13 @@ Covers the framework-side logic that does NOT need the full draft model: ``DSparkSpecMetadata`` hidden-state capture (incl. the mHC hc-mean reduction) -and ``DSv4DSparkWorker`` progress-slot / managed-page management. The end-to-end block +and ``DSv4DSparkWorker`` slot / rolling-KV-window management. The end-to-end block draft and acceptance path is covered by the DSpark test in ``integration/defs/accuracy/test_llm_api_pytorch.py``. """ import types +from unittest.mock import Mock import pytest import torch @@ -122,60 +123,10 @@ def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): "softmax_scale": 1.0, "eps": 1e-6, }, + configure_generation_autotuner=Mock(), ) -def _bind_manager(worker, request_ids=()): - from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType - - manager = types.SimpleNamespace( - draft_layer_ids=tuple(range(worker._num_stages)), - draft_window_size=worker._win, - kv_factor_per_layer=[1] * worker._num_stages, - head_dim_per_layer=[worker._head_dim] * worker._num_stages, - layer_offsets=list(range(worker._num_stages)), - draft_max_blocks_per_seq=8192, - tokens_per_block=128, - _draft_dummy_request_ids=set(), - enable_block_reuse=False, - kv_cache_map={}, - ) - buffers = tuple( - torch.zeros(65, 1, 1, 128, worker._head_dim, device="cuda", dtype=torch.bfloat16) - for _ in range(worker._num_stages) - ) - manager.get_draft_buffers = lambda stage: buffers[stage] - manager.get_draft_block_indices = lambda ids: [manager.kv_cache_map[rid].indices for rid in ids] - resource_manager = types.SimpleNamespace( - get_resource_manager=lambda kind: manager - if kind == ResourceManagerType.KV_CACHE_MANAGER - else None - ) - worker._bind_managed_history(resource_manager) - for i, rid in enumerate(request_ids): - manager.kv_cache_map[rid] = types.SimpleNamespace( - history_length=1, capacity=4096, indices=[i + 1] * 32 - ) - return manager - - -def _prepare(worker, meta, num_contexts=0): - if worker._draft_kv_manager is None: - _bind_manager(worker) - manager = worker._draft_kv_manager - for rid in meta.request_ids: - if worker._is_managed_request(rid) and rid not in manager.kv_cache_map: - manager.kv_cache_map[rid] = types.SimpleNamespace( - history_length=1, capacity=4096, indices=[len(manager.kv_cache_map) + 1] * 32 - ) - # A previously assigned slot represents already initialized resident history. - for rid in worker._req_to_slot: - if rid in manager.kv_cache_map: - worker._managed_residency.setdefault(rid, manager.kv_cache_map[rid]) - meta.prepare() - worker._prepare_managed_history(meta.request_ids, num_contexts) - - def test_worker_lazy_init_window_buffers(): worker = _make_worker() dm = _fake_draft_model(num_stages=3, window_size=128, head_dim=64) @@ -248,78 +199,6 @@ def test_worker_slot_exhaustion_preserves_live_request(): assert int(worker._valid_len[slot]) == 1 -def test_seed_context_pages_preserves_state_across_prefill_chunks(): - class DraftModel: - num_stages = 1 - block_size = 5 - _attn_params = { - "window_size": 8, - "head_dim": 4, - "n_heads": 128, - "softmax_scale": 1.0, - "eps": 1e-6, - } - - def __init__(self): - self.written_positions = [] - - def configure_generation_autotuner(self, _max_batch): - pass - - def write_context_pages(self, hidden, positions, mask, pools, tables, capacities): - from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context - - self.written_positions.append(positions.clone()) - write_dspark_context( - torch.ones(*positions.shape, 4, device="cuda", dtype=torch.bfloat16), - positions, - mask, - pools[0][:, 0, 0], - tables, - capacities, - ) - - worker = _make_worker() - draft_model = DraftModel() - metadata = types.SimpleNamespace( - max_num_requests=1, - request_ids=[100], - get_hidden_states=lambda _num_tokens: torch.zeros( - 3, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 - ), - ) - worker._lazy_init(draft_model, metadata) - manager = _bind_manager(worker, [100]) - manager.kv_cache_map[100].history_length = 0 - worker._prepare_managed_history([100], 1) - - first_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[3]) - worker._seed_context_pages( - draft_model, metadata, first_chunk, torch.tensor([[0, 1, 2]], device="cuda"), 3 - ) - slot = worker._req_to_slot[100] - assert int(worker._ctx_len[slot]) == 3 - assert int(worker._valid_len[slot]) == 3 - assert worker._managed_residency[100] is manager.kv_cache_map[100] - - metadata.get_hidden_states = lambda _num_tokens: torch.zeros( - 2, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 - ) - second_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[2]) - worker._seed_context_pages( - draft_model, metadata, second_chunk, torch.tensor([[3, 4]], device="cuda"), 2 - ) - - assert int(worker._ctx_len[slot]) == 5 - assert int(worker._valid_len[slot]) == 5 - assert [positions.tolist() for positions in draft_model.written_positions] == [ - [[1, 2, 3]], - [[4, 5]], - ] - assert torch.all(worker._draft_kv_buffers[0][1, 0, 0, :5] == 1.0) - assert torch.count_nonzero(worker._draft_kv_buffers[0][1, 0, 0, 5:]) == 0 - - @pytest.mark.parametrize( "chunks", [ @@ -332,21 +211,25 @@ def write_context_pages(self, hidden, positions, mask, pools, tables, capacities ) @pytest.mark.parametrize("window", [64, 128, 256]) @pytest.mark.parametrize("final_window", [False, True]) -def test_seed_context_windows_reads_only_each_chunk_tail(chunks, window, final_window): +def test_seed_context_pages_reads_only_each_chunk_tail(chunks, window, final_window): """Never consume dropped captures, including an entirely skipped chunk.""" worker = _make_worker() draft = _fake_draft_model() draft._attn_params = dict(draft._attn_params, window_size=window) written = [] - def write(hidden, positions, windows): + def write(hidden, positions, valid, pages, tables, capacities): assert torch.isfinite(hidden).all() - torch.testing.assert_close(hidden[:, 0].float(), positions.float() - 1) - written.append(positions.tolist()) + torch.testing.assert_close(hidden[..., 0].float(), positions.float() - 1) + written.append(positions.squeeze(0).tolist()) - draft.write_context_windows = write + draft.write_context_pages = write meta = types.SimpleNamespace(max_num_requests=2, request_ids=[100, 101]) worker._lazy_init(draft, meta) + manager = _bind_manager(worker, meta.request_ids) + for request in manager.kv_cache_map.values(): + request.history_length = 0 + worker._prepare_managed_history(meta.request_ids, 2) expected = [] total_kept = 0 prompt_end = chunks[-1][0] + chunks[-1][1] @@ -372,7 +255,7 @@ def write(hidden, positions, windows): packed_positions = torch.cat( (positions, positions.new_tensor([0, 1, 2, 50, 51, 52, 53, 54, 55])) ) - worker._seed_context_windows(draft, meta, attn, packed_positions, len(captured)) + worker._seed_context_pages(draft, meta, attn, packed_positions, len(captured)) end = start + length if kept: expected.append(list(range(end - min(window, kept) + 1, end + 1))) @@ -544,9 +427,7 @@ def test_prepare_keeps_cuda_graph_slots_across_replays(): meta = _make_metadata(max_num_requests=4) worker._lazy_init(_fake_draft_model(), meta) meta._dspark_worker = worker - # Not 0: that is ATTENTION_DP_DUMMY_REQUEST_ID, which prepare() routes to - # the scratch row by design, so it would never take a persistent slot. - meta.request_ids = [11, 12, 13] + meta.request_ids = [0, 1, 2] meta.num_generations = 3 meta.is_cuda_graph = True _prepare(worker, meta) @@ -1016,7 +897,8 @@ def test_worker_and_draft_kv_follow_the_draft_form(embedded, worker_cls, uses_se worker = get_spec_worker(spec_config, None, Mapping()) assert type(worker) is worker_cls - # Embedded drafts share the target manager; standalone drafts can use a separate manager. + # The embedded draft owns a rolling window and never reads the paged draft + # KV cache; the standalone one is DFlash lineage and does. assert should_use_separate_draft_kv_cache(spec_config) is uses_separate_draft_kv @@ -1055,10 +937,9 @@ def test_lazy_init_publishes_the_runtime_position_ceiling(): worker._lazy_init(dm, _make_metadata(max_num_requests=2), attn_metadata) - # start_pos is a frame index (position + 1); a full verification then accepts - # K+1 tokens and DSpark indexes K block positions from that start. Length is - # max index + 1. - expected = 4096 + 1 + 5 + 1 + 5 + 1 + # A full target verification accepts K+1 tokens, then DSpark indexes K + # block positions from that new start. The table length is max index + 1. + expected = 4096 + 5 + 1 + 5 + 1 assert dm._runtime_position_ceiling == expected assert dm.dspark_model._runtime_position_ceiling == expected @@ -1182,3 +1063,208 @@ def test_prepare_resets_the_scratch_row_every_iteration(): assert int(worker._ctx_len[live]) == 17 assert int(worker._valid_len[live]) == 5 assert worker._managed_residency[100] is worker._draft_kv_manager.kv_cache_map[100] + + +def test_ragged_draft_gathers_each_requests_hidden_state(): + worker = _make_worker() + draft_model = _RecordingDraftModel() + metadata = types.SimpleNamespace( + max_num_requests=2, + qo_indptr=torch.tensor([0, 4, 6], dtype=torch.int32, device="cuda"), + verify_lens=torch.tensor([4, 2], dtype=torch.int32, device="cuda"), + ) + worker._lazy_init(draft_model, metadata) + worker._batch_to_slot[:2] = torch.tensor([0, 1], device="cuda") + captured = torch.stack( + [torch.full((HIDDEN * NCAP,), float(index), device="cuda") for index in range(6)] + ) + metadata.get_hidden_states = lambda _num_tokens: captured + accepted = torch.zeros((2, 6), dtype=torch.int32, device="cuda") + + worker._draft_gen_block_batched( + draft_model, + metadata, + types.SimpleNamespace(num_ctx_tokens=0), + accepted, + torch.ones(2, dtype=torch.int32, device="cuda"), + num_contexts=0, + batch_size=2, + total_target_tokens=6, + position_ids=torch.arange(6, device="cuda").unsqueeze(0), + ) + + assert draft_model.forward_calls[-1]["main_hidden"][:, 0].tolist() == [0.0, 4.0] + + +def test_confidence_rows_are_slot_indexed_and_fail_open(): + worker = _make_worker(enable_confidence_scheduling=True) + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) + + buffer = worker.staged_confidence_buffer() + assert buffer is not None + neutral = worker.confidence_row_for(12345) + assert neutral == worker._neutral_conf_row + assert neutral > worker._scratch_slot + assert torch.sigmoid(buffer[neutral]).min().item() > 0.999 + + worker._req_to_slot = {100: 5, 101: 2, 102: 7} + for slot in range(worker._scratch_slot): + buffer[slot] = float(slot) + rows = [worker.confidence_row_for(request_id) for request_id in (100, 101, 102)] + assert rows == [5, 2, 7] + assert buffer[rows, 0].tolist() == [5.0, 2.0, 7.0] + + +def test_recycled_confidence_slot_is_reset_to_neutral(): + worker = _make_worker(enable_confidence_scheduling=True) + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) + worker._confidence_logits[: worker._scratch_slot] = -20.0 + + slot = worker._assign_slot(100, reset=False) + + assert torch.sigmoid(worker._confidence_logits[slot]).min().item() > 0.999 + + +def test_confidence_state_is_absent_when_scheduling_is_disabled(): + worker = _make_worker(enable_confidence_scheduling=False) + worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=8)) + + assert worker.staged_confidence_buffer() is None + assert worker.confidence_stamp_buffer() is None + assert worker.verify_planner is None + + +def _bind_manager(worker, request_ids=()): + from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType + + manager = types.SimpleNamespace( + draft_layer_ids=tuple(range(worker._num_stages)), + draft_window_size=worker._win, + kv_factor_per_layer=[1] * worker._num_stages, + head_dim_per_layer=[worker._head_dim] * worker._num_stages, + layer_offsets=list(range(worker._num_stages)), + draft_max_blocks_per_seq=8192, + tokens_per_block=128, + _draft_dummy_request_ids=set(), + enable_block_reuse=False, + kv_cache_map={}, + ) + buffers = tuple( + torch.zeros(65, 1, 1, 128, worker._head_dim, device="cuda", dtype=torch.bfloat16) + for _ in range(worker._num_stages) + ) + manager.get_draft_buffers = lambda stage: buffers[stage] + manager.get_draft_block_indices = lambda ids: [manager.kv_cache_map[rid].indices for rid in ids] + resource_manager = types.SimpleNamespace( + get_resource_manager=lambda kind: manager + if kind == ResourceManagerType.KV_CACHE_MANAGER + else None + ) + worker._bind_managed_history(resource_manager) + for i, rid in enumerate(request_ids): + manager.kv_cache_map[rid] = types.SimpleNamespace( + history_length=1, capacity=4096, indices=[i + 1] * 32 + ) + return manager + + +def _prepare(worker, meta, num_contexts=0): + if worker._draft_kv_manager is None: + _bind_manager(worker) + manager = worker._draft_kv_manager + for rid in meta.request_ids: + if worker._is_managed_request(rid) and rid not in manager.kv_cache_map: + manager.kv_cache_map[rid] = types.SimpleNamespace( + history_length=1, capacity=4096, indices=[len(manager.kv_cache_map) + 1] * 32 + ) + # A previously assigned slot represents already initialized resident history. + for rid in worker._req_to_slot: + if rid in manager.kv_cache_map: + worker._managed_residency.setdefault(rid, manager.kv_cache_map[rid]) + meta.prepare() + worker._prepare_managed_history(meta.request_ids, num_contexts) + + +def test_worker_lazy_init_configures_generation_autotuner_once(): + worker = _make_worker() + draft_model = _fake_draft_model() + draft_model.configure_generation_autotuner = Mock() + metadata = _make_metadata(max_num_requests=8) + + worker._lazy_init(draft_model, metadata) + worker._lazy_init(draft_model, metadata) + + draft_model.configure_generation_autotuner.assert_called_once_with(8) + + +def test_seed_context_pages_preserves_state_across_prefill_chunks(): + class DraftModel: + num_stages = 1 + block_size = 5 + _attn_params = { + "window_size": 8, + "head_dim": 4, + "n_heads": 128, + "softmax_scale": 1.0, + "eps": 1e-6, + } + + def __init__(self): + self.written_positions = [] + + def configure_generation_autotuner(self, _max_batch): + pass + + def write_context_pages(self, hidden, positions, mask, pools, tables, capacities): + from tensorrt_llm._torch.attention.backends.dspark import write_dspark_context + + self.written_positions.append(positions.clone()) + write_dspark_context( + torch.ones(*positions.shape, 4, device="cuda", dtype=torch.bfloat16), + positions, + mask, + pools[0][:, 0, 0], + tables, + capacities, + ) + + worker = _make_worker() + draft_model = DraftModel() + metadata = types.SimpleNamespace( + max_num_requests=1, + request_ids=[100], + get_hidden_states=lambda _num_tokens: torch.zeros( + 3, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 + ), + ) + worker._lazy_init(draft_model, metadata) + manager = _bind_manager(worker, [100]) + manager.kv_cache_map[100].history_length = 0 + worker._prepare_managed_history([100], 1) + + first_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[3]) + worker._seed_context_pages( + draft_model, metadata, first_chunk, torch.tensor([[0, 1, 2]], device="cuda"), 3 + ) + slot = worker._req_to_slot[100] + assert int(worker._ctx_len[slot]) == 3 + assert int(worker._valid_len[slot]) == 3 + assert worker._managed_residency[100] is manager.kv_cache_map[100] + assert bool(worker._position_initialized[slot]) + + metadata.get_hidden_states = lambda _num_tokens: torch.zeros( + 2, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 + ) + second_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[2]) + worker._seed_context_pages( + draft_model, metadata, second_chunk, torch.tensor([[3, 4]], device="cuda"), 2 + ) + + assert int(worker._ctx_len[slot]) == 5 + assert int(worker._valid_len[slot]) == 5 + assert [positions.tolist() for positions in draft_model.written_positions] == [ + [[1, 2, 3]], + [[4, 5]], + ] + assert torch.all(worker._draft_kv_buffers[0][1, 0, 0, :5] == 1.0) + assert torch.count_nonzero(worker._draft_kv_buffers[0][1, 0, 0, 5:]) == 0 diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 878a70531840..7fa009b9cc80 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU correctness tests for paged DSpark CuteDSL attention and preparation.""" +"""GPU correctness tests for the DSpark CuteDSL attention op.""" from unittest.mock import Mock @@ -13,7 +13,7 @@ from tensorrt_llm._torch.custom_ops import dspark_rmsnorm_rope_custom_op as preparation_op from tensorrt_llm._torch.cute_dsl_utils import ( IS_CUTLASS_DSL_AVAILABLE, - IS_CUTLASS_DSL_INTERNAL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, ) from tensorrt_llm._torch.models import modeling_dspark as dspark from tensorrt_llm._utils import get_sm_version @@ -23,28 +23,11 @@ pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or not IS_CUTLASS_DSL_AVAILABLE - or get_sm_version() not in (100, 103), - reason="DSpark CuteDSL attention requires an SM100 or SM103 CUDA GPU", + or get_sm_version() not in (100, 103, 107), + reason="DSpark CuteDSL attention requires an SM100, SM103 or SM107 CUDA GPU", ) -@pytest.fixture(params=[False, True], ids=["native", "unsupported_arch_fallback"]) -def force_fallback(request, monkeypatch): - if request.param: - monkeypatch.setattr(preparation_op, "_get_dspark_arch_str", lambda: None) - monkeypatch.setattr(dspark, "_get_dspark_arch_str", lambda: None) - - def unsupported_kernel(*args, **kwargs): - pytest.fail("Unsupported architecture entered a fused DSpark kernel") - - monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope", unsupported_kernel) - monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope_page_write", unsupported_kernel) - monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope_draft_block", unsupported_kernel) - monkeypatch.setattr(torch.ops.trtllm, "fused_dsv4_dspark_attention", unsupported_kernel) - assert preparation_op._get_dspark_arch_str() is None - return request.param - - def _make_inputs( seed=0, batch=2, @@ -98,6 +81,24 @@ def _make_inputs( ) +def _legacy_valid_len(start_pos): + return (start_pos.long() + 1).clamp(max=128) + + +def _prepare_attention_inputs(main_kv, block_kv, kv_cache, slots, start_pos): + kv_cache[slots.long(), (start_pos % kv_cache.shape[1]).long()] = main_kv + draft_block = block_kv.new_zeros((block_kv.shape[0], 8, block_kv.shape[2])) + draft_block[:, : block_kv.shape[1]].copy_(block_kv) + return draft_block, slots.to(torch.int32), start_pos.to(torch.int32) + + +def _with_row_padding(x: torch.Tensor, padding: int = 128) -> torch.Tensor: + storage = torch.empty((*x.shape[:-1], x.shape[-1] + padding), device=x.device, dtype=x.dtype) + padded = storage[..., : x.shape[-1]] + padded.copy_(x) + return padded + + def _reference( q, draft_block, @@ -129,26 +130,48 @@ def _reference( probs = torch.cat((scores, sink), dim=-1).softmax(-1)[..., :-1] outputs.append(torch.einsum("sht,td->shd", probs, kv.float()).to(q.dtype)) return dspark._rope_last_dims_batched( - torch.stack(outputs), _ROPE_DIM, torch.view_as_complex(inverse_rope_freqs), inverse=True + torch.stack(outputs), + _ROPE_DIM, + torch.view_as_complex(inverse_rope_freqs), + inverse=True, ) -def _norm_rope(x, weight, freqs): - normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6) - normalized = (normalized * weight.float()).to(x.dtype) - return dspark._rope_last_dims_batched(normalized, _ROPE_DIM, freqs) +@pytest.mark.parametrize( + ("invalid_case", "reason"), + [ + ("q_dtype", "q dtype must be BF16"), + ("block_size", "draft block size must be 5 or 6"), + ("valid_len_dtype", "valid_len must use INT64"), + ("freqs_shape", "inverse_rope_freqs shape must be"), + ], +) +def test_fused_dsv4_dspark_attention_support_gate_logs_rejection(monkeypatch, invalid_case, reason): + import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op + q, _, _, kv_cache, _, start_pos, sink, _, freqs = _make_inputs() + inputs = [q, kv_cache, _legacy_valid_len(start_pos), sink, freqs] + if invalid_case == "q_dtype": + inputs[0] = q.float() + elif invalid_case == "block_size": + inputs[0] = q[:, :4].contiguous() + inputs[4] = freqs[:, :4].contiguous() + elif invalid_case == "valid_len_dtype": + inputs[2] = inputs[2].to(torch.int32) + else: + inputs[4] = freqs[:, :, :-1].contiguous() -def _write_reference(cache, main_kv, inputs): - cache = cache.clone() - for row, (end, capacity) in enumerate( - zip(inputs["positions"].tolist(), inputs["capacities"].tolist()) - ): - token = end - 1 - if 0 <= token < capacity: - page = int(inputs["block_tables"][row, token // cache.shape[1]]) - cache[page, token % cache.shape[1]] = main_kv[row, 0] - return cache + messages = [] + monkeypatch.setattr( + dspark_attention_op.logger, + "debug_once", + lambda *message, key: messages.append((" ".join(map(str, message)), key)), + ) + + assert not dspark_attention_op.is_fused_dsv4_dspark_attention_supported(*inputs) + assert len(messages) == 1 + assert reason in messages[0][0] + assert messages[0][1][0] == "fused_dsv4_dspark_attention_unsupported" @pytest.mark.parametrize( @@ -229,7 +252,11 @@ def test_fused_dsv4_dspark_attention_matches_reference(heads, block, page_size): @pytest.mark.parametrize("heads,page_size", [(64, 32), (128, 64), (128, 256)]) def test_fused_dsv4_dspark_attention_cuda_graph_replay(heads, page_size): inputs = _make_inputs( - 3, heads=heads, page_size=page_size, positions=[257, 5, 390], capacities=[512, 512, 0] + 3, + heads=heads, + page_size=page_size, + positions=[257, 5, 390], + capacities=[512, 512, 0], ) main_x = torch.randn(3, 1, 512, device="cuda", dtype=torch.bfloat16) block_x = torch.randn(3, 6, 512, device="cuda", dtype=torch.bfloat16) @@ -282,92 +309,42 @@ def forward(): assert torch.count_nonzero(draft[:, 6:]) == 0 -@pytest.mark.parametrize("heads", (64, 128)) -@pytest.mark.parametrize("page_size", (32, 256)) -@pytest.mark.parametrize("block", (3, 5, 6)) -def test_dspark_attention_forward_matches_reference( - heads, page_size, block, force_fallback, monkeypatch -): - inputs = _make_inputs( - 17, - block=block, - heads=heads, - page_size=page_size, - positions=[1, 33, 129, 390, 0], - valid_lengths=[1, 3, 128, 3, 0], - capacities=[512, 512, 512, 512, 0], - ) - batch, hidden, rank, groups, o_rank = 5, 64, 64, 8, 32 +@pytest.mark.parametrize("page_size", (128, 256)) +def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, page_size): + g = _make_attn_inputs(seed=17, device="cuda", page_size=page_size) + for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): + g[name].mul_(0.2) + for name in ("x", "main_x", "kv_cache0"): + g[name].mul_(0.1) + before = g["kv_cache0"].clone() + actual = _run(g) + actual_cache = g["kv_cache0"].clone() + g["kv_cache0"].copy_(before) + with monkeypatch.context() as patch: + patch.setattr(torch.ops.trtllm, "fused_dsv4_dspark_attention", _reference) + patch.setattr(dspark, "is_fused_dspark_rmsnorm_rope_supported", lambda *a: False) + expected = _run(g) + torch.testing.assert_close(actual, expected, rtol=8e-2, atol=1e-2) + torch.testing.assert_close(actual_cache, g["kv_cache0"], rtol=2e-2, atol=2e-2) - def weight(*shape): - return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.02 - x, main_x = weight(batch, block, hidden), weight(batch, 1, hidden) - kwargs = dict( - wq_a=weight(rank, hidden), - q_norm_w=torch.ones(rank, device="cuda", dtype=x.dtype), - wq_b=weight(heads * 512, rank), - wkv=weight(512, hidden), - kv_norm_w=torch.ones(512, device="cuda", dtype=x.dtype), - wo_a=weight(groups * o_rank, heads * 512 // groups), - wo_b=weight(hidden, groups * o_rank), - attn_sink=inputs["attn_sink"], - n_heads=heads, - head_dim=512, - rope_head_dim=_ROPE_DIM, - n_groups=groups, - o_lora_rank=o_rank, - window_size=128, - eps=1e-6, - softmax_scale=inputs["softmax_scale"], - freqs_cis=dspark.precompute_dspark_freqs_cis(_ROPE_DIM, 1024, device="cuda"), - ) - positions = inputs["positions"] - block_freqs = torch.view_as_complex(inputs["inverse_rope_freqs"]) - q = F.linear( - dspark._rmsnorm(F.linear(x, kwargs["wq_a"]), kwargs["q_norm_w"], 1e-6), kwargs["wq_b"] - ).unflatten(-1, (heads, 512)) - q = _norm_rope(q, torch.ones(512, device="cuda", dtype=x.dtype), block_freqs) - main = _norm_rope( - F.linear(main_x, kwargs["wkv"]), - kwargs["kv_norm_w"], - kwargs["freqs_cis"][positions][:, None], - ) - draft = _norm_rope(F.linear(x, kwargs["wkv"]), kwargs["kv_norm_w"], block_freqs) - expected_cache = _write_reference(inputs["kv_pages"], main, inputs) - output = _reference(**(inputs | {"q": q, "draft_block": draft, "kv_pages": expected_cache})) - output = torch.einsum( - "bsgd,grd->bsgr", - output.reshape(batch, block, groups, -1), - kwargs["wo_a"].view(groups, o_rank, -1), - ) - expected = F.linear(output.flatten(2), kwargs["wo_b"]) - fused_calls = [] - for module, name in ( - (dspark, "cute_dsl_dspark_rmsnorm_rope_page_write"), - (dspark, "cute_dsl_dspark_rmsnorm_rope_draft_block"), - (torch.ops.trtllm, "fused_dsv4_dspark_attention"), - ): - call = Mock(wraps=getattr(module, name)) - monkeypatch.setattr(module, name, call) - fused_calls.append(call) - actual = dspark.dspark_attention_forward( - x, - main_x, - positions, - inputs["kv_pages"], - inputs["block_tables"], - inputs["capacities"], - inputs["valid_lengths"], - **kwargs, +def test_warmup_is_best_effort(monkeypatch): + import tensorrt_llm._torch.custom_ops.dspark_attention_custom_op as dspark_attention_op + + warnings = [] + monkeypatch.setattr(dspark_attention_op, "_get_dspark_arch_str", lambda: "sm_100") + monkeypatch.setattr( + dspark_attention_op, + "cute_dsl_dspark_rmsnorm_rope_cache_write", + lambda *args: (_ for _ in ()).throw(RuntimeError("synthetic compile failure")), ) - for call in fused_calls: - if not force_fallback and block in (5, 6): - call.assert_called_once() - else: - call.assert_not_called() - torch.testing.assert_close(actual, expected, rtol=8e-2, atol=1e-2) - torch.testing.assert_close(inputs["kv_pages"], expected_cache, rtol=2e-2, atol=2e-2) + monkeypatch.setattr(dspark_attention_op.logger, "warning", warnings.append) + + dspark_attention_op.warmup_fused_dsv4_dspark_attention(5, 1e-6) + + assert len(warnings) == 1 + assert "self-JIT on first use" in warnings[0] + assert "synthetic compile failure" in warnings[0] def test_self_jit_reuses_one_kernel_across_runtime_shapes_and_scales(monkeypatch): @@ -409,7 +386,9 @@ def unexpected_compile_op(*args, **kwargs): patch.setattr(attention_op.cute.runtime, "from_dlpack", unexpected_compile_op) compiled = attention_op._compile_dspark_attention(5, arch_str, 32, None, 128) monkeypatch.setattr( - attention_op, "_dspark_attention_kernel_cache", {(5, 128, 32, arch_str, None): compiled} + attention_op, + "_dspark_attention_kernel_cache", + {(5, 128, 32, arch_str, None): compiled}, ) inputs = _make_inputs(41, block=5, positions=[257, 9]) @@ -490,23 +469,127 @@ def test_preparation_self_jit_covers_dynamic_batches(): assert preparation_op._compile_dspark_rmsnorm_rope_draft_block.cache_info().misses == 1 -@pytest.mark.parametrize("page_size", (128, 256)) -def test_dspark_attention_forward_batched_matches_fallback(monkeypatch, page_size): - g = _make_attn_inputs(seed=17, device="cuda", page_size=page_size) - for name in ("wq_a", "wq_b", "wkv", "wo_a", "wo_b"): - g[name].mul_(0.2) - for name in ("x", "main_x", "kv_cache0"): - g[name].mul_(0.1) - before = g["kv_cache0"].clone() - actual = _run(g) - actual_cache = g["kv_cache0"].clone() - g["kv_cache0"].copy_(before) - with monkeypatch.context() as patch: - patch.setattr(torch.ops.trtllm, "fused_dsv4_dspark_attention", _reference) - patch.setattr(dspark, "is_fused_dspark_rmsnorm_rope_supported", lambda *a: False) - expected = _run(g) +@pytest.fixture(params=[False, True], ids=["native", "unsupported_arch_fallback"]) +def force_fallback(request, monkeypatch): + if request.param: + monkeypatch.setattr(preparation_op, "_get_dspark_arch_str", lambda: None) + monkeypatch.setattr(dspark, "_get_dspark_arch_str", lambda: None) + + def unsupported_kernel(*args, **kwargs): + pytest.fail("Unsupported architecture entered a fused DSpark kernel") + + monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope", unsupported_kernel) + monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope_page_write", unsupported_kernel) + monkeypatch.setattr(dspark, "cute_dsl_dspark_rmsnorm_rope_draft_block", unsupported_kernel) + monkeypatch.setattr(torch.ops.trtllm, "fused_dsv4_dspark_attention", unsupported_kernel) + assert preparation_op._get_dspark_arch_str() is None + return request.param + + +def _norm_rope(x, weight, freqs): + normalized = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + 1e-6) + normalized = (normalized * weight.float()).to(x.dtype) + return dspark._rope_last_dims_batched(normalized, _ROPE_DIM, freqs) + + +def _write_reference(cache, main_kv, inputs): + cache = cache.clone() + for row, (end, capacity) in enumerate( + zip(inputs["positions"].tolist(), inputs["capacities"].tolist()) + ): + token = end - 1 + if 0 <= token < capacity: + page = int(inputs["block_tables"][row, token // cache.shape[1]]) + cache[page, token % cache.shape[1]] = main_kv[row, 0] + return cache + + +@pytest.mark.parametrize("heads", (64, 128)) +@pytest.mark.parametrize("page_size", (32, 256)) +@pytest.mark.parametrize("block", (3, 5, 6)) +def test_dspark_attention_forward_matches_reference( + heads, page_size, block, force_fallback, monkeypatch +): + inputs = _make_inputs( + 17, + block=block, + heads=heads, + page_size=page_size, + positions=[1, 33, 129, 390, 0], + valid_lengths=[1, 3, 128, 3, 0], + capacities=[512, 512, 512, 512, 0], + ) + batch, hidden, rank, groups, o_rank = 5, 64, 64, 8, 32 + + def weight(*shape): + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) * 0.02 + + x, main_x = weight(batch, block, hidden), weight(batch, 1, hidden) + kwargs = dict( + wq_a=weight(rank, hidden), + q_norm_w=torch.ones(rank, device="cuda", dtype=x.dtype), + wq_b=weight(heads * 512, rank), + wkv=weight(512, hidden), + kv_norm_w=torch.ones(512, device="cuda", dtype=x.dtype), + wo_a=weight(groups * o_rank, heads * 512 // groups), + wo_b=weight(hidden, groups * o_rank), + attn_sink=inputs["attn_sink"], + n_heads=heads, + head_dim=512, + rope_head_dim=_ROPE_DIM, + n_groups=groups, + o_lora_rank=o_rank, + window_size=128, + eps=1e-6, + softmax_scale=inputs["softmax_scale"], + freqs_cis=dspark.precompute_dspark_freqs_cis(_ROPE_DIM, 1024, device="cuda"), + ) + positions = inputs["positions"] + block_freqs = torch.view_as_complex(inputs["inverse_rope_freqs"]) + q = F.linear( + dspark._rmsnorm(F.linear(x, kwargs["wq_a"]), kwargs["q_norm_w"], 1e-6), kwargs["wq_b"] + ).unflatten(-1, (heads, 512)) + q = _norm_rope(q, torch.ones(512, device="cuda", dtype=x.dtype), block_freqs) + main = _norm_rope( + F.linear(main_x, kwargs["wkv"]), + kwargs["kv_norm_w"], + kwargs["freqs_cis"][positions][:, None], + ) + draft = _norm_rope(F.linear(x, kwargs["wkv"]), kwargs["kv_norm_w"], block_freqs) + expected_cache = _write_reference(inputs["kv_pages"], main, inputs) + output = _reference(**(inputs | {"q": q, "draft_block": draft, "kv_pages": expected_cache})) + output = torch.einsum( + "bsgd,grd->bsgr", + output.reshape(batch, block, groups, -1), + kwargs["wo_a"].view(groups, o_rank, -1), + ) + expected = F.linear(output.flatten(2), kwargs["wo_b"]) + fused_calls = [] + for module, name in ( + (dspark, "cute_dsl_dspark_rmsnorm_rope_page_write"), + (dspark, "cute_dsl_dspark_rmsnorm_rope_draft_block"), + (torch.ops.trtllm, "fused_dsv4_dspark_attention"), + ): + call = Mock(wraps=getattr(module, name)) + monkeypatch.setattr(module, name, call) + fused_calls.append(call) + actual = dspark.dspark_attention_forward( + x, + main_x, + positions, + inputs["kv_pages"], + inputs["block_tables"], + inputs["capacities"], + inputs["valid_lengths"], + **kwargs, + ) + for call in fused_calls: + if not force_fallback and block in (5, 6): + call.assert_called_once() + else: + call.assert_not_called() torch.testing.assert_close(actual, expected, rtol=8e-2, atol=1e-2) - torch.testing.assert_close(actual_cache, g["kv_cache0"], rtol=2e-2, atol=2e-2) + torch.testing.assert_close(inputs["kv_pages"], expected_cache, rtol=2e-2, atol=2e-2) def test_warmup_propagates_compile_failure(monkeypatch): @@ -577,7 +660,7 @@ def test_sparse_attn_masked_indices_excluded(): "start_pos,page_size,block", [(1, 128, 5), (3, 64, 5), (10, 16, 5), (200, 256, 6)] ) def test_page_mapping_matches_reference(start_pos, page_size, block): - from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context + from tensorrt_llm._torch.attention.backends.dspark import write_dspark_context inputs = _make_inputs(block=block, positions=[start_pos] * 3, page_size=page_size) values = inputs["draft_block"][:, :1].contiguous() @@ -594,7 +677,7 @@ def test_page_mapping_matches_reference(start_pos, page_size, block): def test_page_mapping_masks_non_generation_rows(): - from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context + from tensorrt_llm._torch.attention.backends.dspark import write_dspark_context values = torch.randn(1, 1, 16, device="cuda") pool = torch.randn(1, 128, 16, device="cuda") @@ -760,7 +843,7 @@ def test_batched_attention_dummy_rows_keep_pages(): @pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) def test_batched_page_writes_match_per_request(start_positions): - from tensorrt_llm._torch.attention_backend.dspark import write_dspark_context + from tensorrt_llm._torch.attention.backends.dspark import write_dspark_context g = _make_batched_inputs(start_positions=start_positions) values = torch.randn(g["b"], 1, g["head_dim"], device="cuda", dtype=torch.bfloat16) @@ -873,7 +956,7 @@ def assert_quantized_close(actual_fp8, actual_scale, expected_fp8, expected_scal @pytest.mark.skipif( - get_sm_version() != 107 or not IS_CUTLASS_DSL_INTERNAL_AVAILABLE, + get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The fused DSpark FP8 input producers require SM107 Rubin CuTe DSL", ) def test_dspark_fp8_input_producers_match_reference_and_reuse_dynamic_m_compile(): @@ -972,7 +1055,7 @@ def active_scales(scale: torch.Tensor, rows: int) -> torch.Tensor: @pytest.mark.skipif( - get_sm_version() != 107 or not IS_CUTLASS_DSL_INTERNAL_AVAILABLE, + get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The fused DSpark output projection requires SM107 Rubin CuTe DSL", ) def test_dspark_fused_fp8_output_projection_matches_separate_and_reuses_cache(): From 6f765ba33c2193c0cb0ba62129a8a3a84d61c3ca Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Wed, 28 Oct 2026 17:54:16 -0700 Subject: [PATCH 10/16] Fix DSpark unsupported-architecture fallback fixture Disable the independent RMSNorm/RoPE architecture gate when simulating an unsupported GPU. Preserve the reference comparisons and kernel-launch assertions. Signed-off-by: Mingyang Hao --- .../_torch/speculative/test_dspark_cute_dsl_attention.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py index 6dc0c3e986b2..bd36327f8f1e 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_attention.py @@ -479,6 +479,7 @@ def test_preparation_self_jit_covers_dynamic_batches(): def force_fallback(request, monkeypatch): if request.param: monkeypatch.setattr(preparation_op, "_get_dspark_arch_str", lambda: None) + monkeypatch.setattr(preparation_op, "_is_dspark_rmsnorm_rope_sm_supported", lambda: False) monkeypatch.setattr(dspark, "_get_dspark_arch_str", lambda: None) def unsupported_kernel(*args, **kwargs): From 646c479056063478c599f476d5460828e5c6f1bb Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Wed, 28 Oct 2026 17:57:10 -0700 Subject: [PATCH 11/16] Declare absent draft layers in disaggregation test manager The mocked target-only manager must expose an empty draft-layer list so layer-id mapping does not iterate an unconstrained Mock. Keep pool-layout and role assertions unchanged. Signed-off-by: Mingyang Hao --- tests/unittest/disaggregated/test_extractor.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index b8d1f656bc2d..a6db3d7a83cf 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -313,6 +313,7 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l get_disagg_ignored_roles=lambda: frozenset(), impl=impl, pp_layers=[0, 1], + draft_layer_ids=(), num_kv_heads_per_layer=[1, 1], cp_as_tp_layer_ids=frozenset(), get_disagg_role_mapper_kinds=lambda: role_mapper_kinds, From f4587daabcce4d6d06561a478d040a835028f7aa Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Wed, 28 Oct 2026 18:13:40 -0700 Subject: [PATCH 12/16] Fix unified-cache warmup and restore DSpark SM107 routing Avoid duplicate dummy-cache allocation when target and draft share a manager. Restore the SM107 allocator mapping present in the GitLab source and retain the independent RMSNorm/RoPE support gate. Mark the inherited Mamba quota override and complete partial Helix/Mamba manager fixtures for the unified draft cache interface. Signed-off-by: Mingyang Hao --- .../_torch/custom_ops/dspark_attention_custom_op.py | 4 +++- .../_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py | 5 ++--- .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 3 +++ .../kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py | 1 + .../kv_cache/test_kv_cache_manager_v2_helix_superblock.py | 2 ++ .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 7 ++++++- 6 files changed, 17 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py index a149d781204e..6f235ae5a8cf 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py @@ -1019,7 +1019,9 @@ def is_fused_dsv4_dspark_attention_supported( ) -> bool: """Return whether tensors can use the fused DSV4 DSpark attention kernel.""" if _get_dspark_arch_str() is None: - return _log_unsupported(f"SM {get_sm_version()} is not SM100 or SM103", "sm_version") + return _log_unsupported( + f"SM {get_sm_version()} is not SM100, SM103, or SM107", "sm_version" + ) if not all(t.is_cuda for t in (q, kv_cache, valid_len, attn_sink, inverse_rope_freqs)): return _log_unsupported("all inputs must be CUDA tensors", "device") if q.dtype != torch.bfloat16: diff --git a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py index 0fc68189f14c..c9ea0c5acd0f 100644 --- a/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py @@ -27,10 +27,9 @@ _DSV4_DSPARK_ARCH_BY_SM = { 100: "sm_100", 103: "sm_103", + 107: "sm_107", } -# The RMSNorm/RoPE kernels are plain SIMT CuTe kernels with no architecture-specific -# instructions, so they run on more GPUs than the DSV4 DSpark attention kernel gated -# by _DSV4_DSPARK_ARCH_BY_SM. +# RMSNorm/RoPE support is checked independently of attention allocator selection. _DSPARK_RMSNORM_ROPE_SM_VERSIONS = frozenset((100, 103, 107)) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 9dc02d51c2ff..cafc0ee08105 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6961,6 +6961,9 @@ def add_dummy_requests( on this manager and the supplied draft manager when it is no longer needed, including when subsequent warmup or graph capture fails. """ + # A unified manager creates target and draft pools in one cache. + if draft_kv_cache_manager is self: + draft_kv_cache_manager = None _kv_draft = ( kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 968af8a521d3..e978406cc15c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -579,6 +579,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: ) return attention_quota + state_quota + extra_attention_quota + @override def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: draft_quota = super()._get_draft_quota_from_max_tokens(max_tokens) snapshot_slots = self._num_ssm_snapshots_for_capacity(max_tokens, self.kv_cache_config) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index 7f9d7ef36e06..47c0df353fd9 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -229,6 +229,7 @@ def test_update_resources_leaves_history_untouched_under_helix() -> None: # rewind falls back to the reserve width. _allocated_draft_lens={}, _inflight_draft_lens={}, + _standalone_draft_layer_ids=(), kv_cache_map={7: kv}, kv_compression_manages_history=False, _has_cp_helix=True, @@ -387,6 +388,7 @@ def _context_resize_mgr(is_disagg: bool, resizes: list) -> SimpleNamespace: is_disagg=is_disagg, kv_cache_map={7: kv}, num_extra_kv_tokens=0, + _standalone_draft_reserve=0, _connector_reservations_enabled=lambda: False, _fill_fresh_kv_pages=lambda req_id: None, _log_window_crossing=lambda *args: None, diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 51321fbb476b..58259fbf57ca 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2745,10 +2745,15 @@ def allocated_memory(pool_ratio): _install_default_v2_state(mgr) mgr._num_reserved_dummy_slots = 1 mgr.dtype = DataType.HALF + mgr.dtype_per_layer = [mgr.dtype] * mgr.num_local_layers + mgr.cp_as_tp_layer_ids = set() + mgr._helix_cp_size = 1 + mgr._standalone_draft_layer_ids = () + mgr._draft_layer_sizes = [] mgr.enable_swa_scratch_reuse = False mgr.enable_stats = False mgr.num_extra_kv_tokens = 0 - mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + mgr.get_layer_bytes_per_token = lambda local_layer_idx, data_role: 8 kv_cache_config = KvCacheConfig( pool_ratio=pool_ratio, From 36ae25e34ea5a88210e2d4245a74bcc3e5ab8a6b Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Wed, 28 Oct 2026 21:19:36 -0700 Subject: [PATCH 13/16] Remove legacy unified DSpark draft-cache transfer Remove the superseded standalone draft KV ownership, lifecycle and CTX-to-GEN transfer path, including its later integration adaptations. Restore the main worker/cache-manager contracts in preparation for the layer_specs based design. Keep the independent FP8 projections, Rubin GEMMs, paged attention and page-write kernels. Let the restored rolling-window path reuse the loaded projection modules, with CPU regressions for attention and context-window writes. Preserve the original GitLab replay commits, including 272f7bf274, unchanged. Signed-off-by: Mingyang Hao --- .../kv_cache_manager_v2/common.h | 7 - .../kv_cache_manager_v2/config.h | 2 - .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 4 +- .../kv_cache_manager_v2/lifeCycleRegistry.h | 14 +- .../kv_cache_manager_v2/storage/config.cpp | 9 +- .../batch_manager/kvCacheManagerV2.cpp | 21 +- .../kvCacheManagerV2ConfigTest.cpp | 52 -- .../backends/fp4_mla/cache_manager.py | 43 +- .../sparse/deepseek_v4/cache_manager.py | 172 ++-- .../backends/sparse/dsa/cache_manager.py | 48 +- .../_torch/disaggregation/native/auxiliary.py | 137 +-- .../disaggregation/native/bounce/impl.py | 10 +- .../native/mixers/attention/peer.py | 26 +- .../_torch/disaggregation/native/peer.py | 72 +- .../_torch/disaggregation/native/rank_info.py | 6 - .../_torch/disaggregation/native/transfer.py | 78 +- .../disaggregation/resource/cache_reuse.py | 16 +- .../disaggregation/resource/kv_extractor.py | 76 +- .../_torch/disaggregation/resource/page.py | 7 - .../_torch/disaggregation/transceiver.py | 85 +- tensorrt_llm/_torch/models/modeling_dspark.py | 141 ++- tensorrt_llm/_torch/pyexecutor/_util.py | 155 +--- .../engine/runners/decoder/runner.py | 30 - .../kv_cache/kv_cache_manager_v2.py | 812 ++---------------- .../kv_cache/mamba_cache_manager/common.py | 22 - .../mamba_cache_manager_v2.py | 24 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 +- .../_torch/pyexecutor/sampler/sampler.py | 2 - .../pyexecutor/standalone_draft_cache.py | 66 -- tensorrt_llm/_torch/speculative/dflash.py | 700 ++------------- tensorrt_llm/_torch/speculative/dspark.py | 451 ++++------ tensorrt_llm/_torch/speculative/interface.py | 15 - tensorrt_llm/_torch/speculative/utils.py | 6 +- .../runtime/kv_cache_manager_v2/__init__.py | 4 - .../runtime/kv_cache_manager_v2/__init__.pyi | 9 - .../sparse/glm_kpool/test_glm_kpool.py | 4 - .../kv_cache/test_kv_cache_budget_split.py | 7 - .../kv_cache/test_kv_cache_estimation.py | 46 - ...st_kv_cache_manager_v2_helix_superblock.py | 2 - .../kv_cache/test_mamba_cache_manager.py | 7 +- .../executor/test_disagg_receive_ordering.py | 1 - .../hw_agnostic/test_dflash_copy_skip.py | 1 - .../hw_agnostic/test_dflash_dummy_slot.py | 1 - .../hw_agnostic/test_dflash_worker.py | 121 --- .../hw_agnostic/test_dspark_attention.py | 108 +++ .../hw_agnostic/test_dspark_cuda_graph.py | 4 +- .../hw_agnostic/test_dspark_v41.py | 4 - .../hw_agnostic/test_dspark_worker.py | 343 +++----- .../test_dspark_worker_confidence_state.py | 16 +- .../hw_agnostic/test_ragged_spec_sampler.py | 1 - .../unittest/disaggregated/region/test_aux.py | 22 - .../unittest/disaggregated/test_extractor.py | 2 - 52 files changed, 858 insertions(+), 3164 deletions(-) delete mode 100644 tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h index 79ae6c456226..4d867eb906a4 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h @@ -51,13 +51,6 @@ enum class PageStatus : int DROPPABLE = 2, // Allow eviction and dropping. }; -// Standalone draft pools precede target pools in storage layout. -enum class CacheDomain : int -{ - kStandaloneDraft = 0, - kTarget = 1, -}; - enum class CacheTier : int { GPU_MEM = 0, diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index 983073986b11..19600e0fd52a 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -207,8 +207,6 @@ struct AttentionLayerConfig // Private working pages never constrain prefix matching or enter the tree. AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; - CacheDomain cacheDomain = CacheDomain::kTarget; - [[nodiscard]] std::optional windowSize() const noexcept { return slidingWindowSize; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index 1d798cccd7a7..a4c57cacea07 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -41,8 +41,8 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) else { bool const isSparse = !cfg.buffers.empty() && cfg.buffers.front().isSparse; - return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, - cfg.reusePolicy, cfg.cacheDomain); + return AttnLifeCycle::make( + cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, cfg.reusePolicy); } }, layer); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index 6dbecf343e68..bdf550e6fcc1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -25,7 +25,6 @@ #include #include #include -#include #include #include #include @@ -43,7 +42,6 @@ struct AttnLifeCycle int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) bool isSparse = false; AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; - CacheDomain cacheDomain = CacheDomain::kTarget; bool isPrivate() const noexcept { @@ -77,7 +75,7 @@ struct AttnLifeCycle bool operator==(AttnLifeCycle const& o) const noexcept { return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse - && reusePolicy == o.reusePolicy && cacheDomain == o.cacheDomain; + && reusePolicy == o.reusePolicy; } bool operator<(AttnLifeCycle const& o) const noexcept @@ -94,19 +92,18 @@ struct AttnLifeCycle { return isSparse < o.isSparse; } - return std::tie(reusePolicy, cacheDomain) < std::tie(o.reusePolicy, o.cacheDomain); + return reusePolicy < o.reusePolicy; } static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock, - bool isSparse = false, AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED, - CacheDomain cacheDomain = CacheDomain::kTarget) + bool isSparse = false, AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED) { TLLM_CHECK_DEBUG(tokensPerBlock > 0); TLLM_CHECK_DEBUG(!ws.has_value() || *ws > 0); TLLM_CHECK_DEBUG(!numSinkTokens.has_value() || *numSinkTokens >= 0); TLLM_CHECK_DEBUG((!numSinkTokens.has_value() || *numSinkTokens == 0) || ws.has_value()); int sinkBlocks = divUp(numSinkTokens.value_or(0), tokensPerBlock); - return AttnLifeCycle{ws, sinkBlocks, isSparse, reusePolicy, cacheDomain}; + return AttnLifeCycle{ws, sinkBlocks, isSparse, reusePolicy}; } }; @@ -166,8 +163,7 @@ inline HalfOpenRange computeScratchRange( LifeCycle const& lc, int historyLength, int capacity, int tokensPerBlock, int maxRewindLen) { auto const* attn = std::get_if(&lc); - // Draft history must retain private pages for later drafting or transfer. - if (!attn || !attn->windowSize.has_value() || attn->cacheDomain == CacheDomain::kStandaloneDraft) + if (!attn || !attn->windowSize.has_value()) { return {BlockOrdinal{0}, BlockOrdinal{0}}; } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index d2bbad0920c9..a5d17a7817e1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -212,12 +212,11 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; std::vector slotSizes; bool isSparse = false; - CacheDomain cacheDomain = CacheDomain::kTarget; bool operator<(PoolGroupKey const& other) const { - return std::tie(reusePolicy, slotSizes, isSparse, cacheDomain) - < std::tie(other.reusePolicy, other.slotSizes, other.isSparse, other.cacheDomain); + return std::tie(reusePolicy, slotSizes, isSparse) + < std::tie(other.reusePolicy, other.slotSizes, other.isSparse); } }; @@ -228,9 +227,7 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) auto const* attn = std::get_if(®istry.getLifeCycle(sg.lifeCycleId)); bool const isSparse = attn != nullptr && attn->isSparse; auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; - CacheDomain const domain = attn ? attn->cacheDomain : CacheDomain::kTarget; - poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .isSparse = isSparse, .cacheDomain = domain}] - .push_back(std::move(sg)); + poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .isSparse = isSparse}].push_back(std::move(sg)); } StorageConfig out; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 3ac482b6f064..c503e0250648 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -901,10 +901,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("DROPPABLE", kv::PageStatus::DROPPABLE) .export_values(); - nb::enum_(m, "CacheDomain") - .value("STANDALONE_DRAFT", kv::CacheDomain::kStandaloneDraft) - .value("TARGET", kv::CacheDomain::kTarget); - nb::enum_(m, "CacheTier") .value("GPU_MEM", kv::CacheTier::GPU_MEM) .value("HOST_MEM", kv::CacheTier::HOST_MEM) @@ -1411,21 +1407,18 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .value("OPTIONAL", kv::AttentionReusePolicy::OPTIONAL); nb::class_(m, "AttnLifeCycle") - .def(nb::init, int, bool, kv::AttentionReusePolicy, kv::CacheDomain>(), - nb::arg("window_size").none(), nb::arg("num_sink_blocks"), nb::arg("is_sparse") = false, - nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, - nb::arg("cache_domain") = kv::CacheDomain::kTarget) + .def(nb::init, int, bool, kv::AttentionReusePolicy>(), nb::arg("window_size").none(), + nb::arg("num_sink_blocks"), nb::arg("is_sparse") = false, + nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) // Sink tokens round up to whole blocks. Bound rather than repeated in Python so the // connector's view of a life cycle is built by the same code as the allocator's. .def_static("make", &kv::AttnLifeCycle::make, nb::arg("window_size").none(), nb::arg("num_sink_tokens").none(), nb::arg("tokens_per_block"), nb::arg("is_sparse") = false, - nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, - nb::arg("cache_domain") = kv::CacheDomain::kTarget) + nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) .def_prop_ro("window_size", [](kv::AttnLifeCycle const& self) { return self.windowSize; }) .def_ro("num_sink_blocks", &kv::AttnLifeCycle::numSinkBlocks) .def_ro("is_sparse", &kv::AttnLifeCycle::isSparse) .def_ro("reuse_policy", &kv::AttnLifeCycle::reusePolicy) - .def_ro("cache_domain", &kv::AttnLifeCycle::cacheDomain) .def("get_stale_range", &kv::AttnLifeCycle::getStaleRange, nb::arg("history_length"), nb::arg("tokens_per_block")) .def("__eq__", &kv::AttnLifeCycle::operator==); @@ -1634,16 +1627,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, - kv::AttentionReusePolicy, kv::CacheDomain>(), + kv::AttentionReusePolicy>(), nb::arg("layer_id"), nb::arg("buffers"), nb::arg("sliding_window_size") = std::nullopt, - nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED, - nb::arg("cache_domain") = kv::CacheDomain::kTarget) + nb::arg("num_sink_tokens") = std::nullopt, nb::arg("reuse_policy") = kv::AttentionReusePolicy::REQUIRED) .def_rw("layer_id", &kv::AttentionLayerConfig::layerId) .def_rw("buffers", &kv::AttentionLayerConfig::buffers) .def_rw("sliding_window_size", &kv::AttentionLayerConfig::slidingWindowSize) .def_rw("num_sink_tokens", &kv::AttentionLayerConfig::numSinkTokens) .def_rw("reuse_policy", &kv::AttentionLayerConfig::reusePolicy) - .def_rw("cache_domain", &kv::AttentionLayerConfig::cacheDomain) .def_prop_ro("window_size", &kv::AttentionLayerConfig::windowSize) DEF_COPY(kv::AttentionLayerConfig); nb::enum_(m, "LayerType") diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index bb43f5c618e8..99756c07befc 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp @@ -162,56 +162,4 @@ TEST(KvCacheManagerV2ConfigTest, SparseBuffersCoalesceByExpandedSize) EXPECT_EQ(second.expansion, 1); } -TEST(KvCacheManagerV2ConfigTest, AttentionLifecycleFactoryPreservesCacheDomain) -{ - auto const target = AttnLifeCycle::make(128, 5, 4); - auto const draft - = AttnLifeCycle::make(128, 5, 4, false, AttentionReusePolicy::REQUIRED, CacheDomain::kStandaloneDraft); - EXPECT_EQ(target.numSinkBlocks, 2); - EXPECT_EQ(target.cacheDomain, CacheDomain::kTarget); - EXPECT_EQ(draft.cacheDomain, CacheDomain::kStandaloneDraft); - EXPECT_FALSE(target == draft); - EXPECT_TRUE(target < draft || draft < target); - EXPECT_TRUE( - draft == AttnLifeCycle::make(128, 8, 4, false, AttentionReusePolicy::REQUIRED, CacheDomain::kStandaloneDraft)); -} - -TEST(KvCacheManagerV2ConfigTest, MatchingSlotShapesSharePoolsOnlyWithinEachCacheDomain) -{ - KVCacheManagerConfig config; - config.tokensPerBlock = 32; - config.cacheTiers - = {GpuCacheTierConfig{4 << 20}, HostCacheTierConfig{4 << 20}, DiskCacheTierConfig{4 << 20, "/tmp"}}; - for (LayerId layerId = 0; layerId < 4; ++layerId) - { - AttentionLayerConfig layer; - layer.layerId = layerId; - layer.buffers = {BufferConfig{"key", 4096, std::nullopt}}; - layer.slidingWindowSize = layerId % 2 == 0 ? 128 : 256; - layer.cacheDomain = layerId < 2 ? CacheDomain::kTarget : CacheDomain::kStandaloneDraft; - config.layers.push_back(layer); - } - ASSERT_NO_THROW(config.validate()); - - LifeCycleRegistry const registry(config); - auto const storage = createStorageConfig(config); - auto const attributes = storage.bufferAttributes(); - auto const grouping = storage.lifeCycleGrouping(); - auto lifeCycle = [&](LayerId layerId) { return attributes.at(BufferId{layerId, "key"}).lifeCycleId; }; - auto poolGroup = [&](LayerId layerId) { return grouping[lifeCycle(layerId)]; }; - EXPECT_EQ(registry.size(), LifeCycleId{4}); - EXPECT_EQ(storage.slotDescList.size(), PoolGroupIndex{2}); - for (LayerId first = 0; first < 4; ++first) - { - auto const& attention = std::get(registry[lifeCycle(first)]); - EXPECT_EQ(attention.cacheDomain, first < 2 ? CacheDomain::kTarget : CacheDomain::kStandaloneDraft); - for (LayerId second = first + 1; second < 4; ++second) - { - EXPECT_NE(lifeCycle(first), lifeCycle(second)); - EXPECT_FALSE(registry[lifeCycle(first)] == registry[lifeCycle(second)]); - EXPECT_EQ(poolGroup(first) == poolGroup(second), first / 2 == second / 2); - } - } -} - } // namespace diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py index 76d9d5806be0..b6dcc2d963c4 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py @@ -288,8 +288,6 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. dtype and logical shape as an all-BF16 cache manager. """ local_layer = self.layer_offsets[layer_idx] - if self._is_standalone_draft_layer(local_layer): - return self.get_draft_buffers(self.draft_layer_ids.index(layer_idx), kv_layout) is_linear_layer = getattr(self, "_is_local_mamba_layer", None) if callable(is_linear_layer) and is_linear_layer(local_layer): return None @@ -324,8 +322,6 @@ def get_kv_cache_dtype(self, layer_idx: Optional[int] = None) -> DataType: if layer_idx is None: return self.dtype local_layer = self.layer_offsets[layer_idx] - if self._is_standalone_draft_layer(local_layer): - return self.dtype_per_layer[local_layer] if local_layer in self._bf16_mla_local_layer_indices(): return DataType.BF16 return self.dtype @@ -973,15 +969,10 @@ def get_disagg_global_layer_ids(self, layer_group_id: int) -> list[int]: int(layer_id): self._fp4_mla_compact_to_local[compact_layer] for compact_layer, layer_id in enumerate(self._hp_manager_layer_ids) } - draft_global = { - self.layer_offsets[layer_id]: 2 * layer_id for layer_id in self.draft_layer_ids - } result = [] for layer_id in local_layer_ids: internal_layer = int(layer_id) - if internal_layer in draft_global: - result.append(draft_global[internal_layer]) - elif internal_layer in cache_local: + if internal_layer in cache_local: result.append(2 * int(self.pp_layers[cache_local[internal_layer]])) elif internal_layer in hp_local: result.append(2 * int(self.pp_layers[hp_local[internal_layer]]) + 1) @@ -1007,8 +998,6 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[list[int], list[Opti for local_layer in fp4_local_layers ) windows.extend([self._fp4_mla_hp_pool_size] * len(fp4_local_layers)) - sizes.extend(self._draft_layer_sizes) - windows.extend([self.draft_retention_window_size] * len(self._draft_layer_sizes)) return sizes, windows def _get_generation_request_capacity(self) -> int: @@ -1018,7 +1007,7 @@ def _get_generation_request_capacity(self) -> int: def get_cache_bytes_per_token(self) -> int: sizes, _ = self._get_runtime_cache_size_layer_components() - return sum(sizes) + (self._helix_cp_size - 1) * sum(self._draft_layer_sizes) + return sum(sizes) @staticmethod def get_cache_size_per_token( @@ -1072,33 +1061,7 @@ def get_cache_size_per_token( * torch.empty((), dtype=torch.bfloat16).element_size() * num_bf16_layers ) - fixed_bytes = hp_bytes * max_batch_size * mapping.pp_size - from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - _estimate_draft_cache_size_components, - _get_draft_layer_sizes, - _get_generation_kv_capacity, - ) - - draft_sizes = _get_draft_layer_sizes(**kwargs) - if draft_sizes: - scratch = kwargs.get("draft_scratch_tokens", 0) - _, headroom = _get_generation_kv_capacity(spec_config, is_draft=False) - draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( - draft_sizes, - tokens_per_block, - scratch_tokens=scratch, - window_size=kwargs.get("draft_window_size"), - generation_capacity_headroom=headroom + scratch, - helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, - ) - # The shared allocation envelope reserves space in target pools too. - draft_fixed += scratch * (cache_bytes + hp_bytes / (HP_BLOCK_SIZE + rewind)) - fixed_bytes += int(draft_fixed * max_batch_size * mapping.pp_size) - fixed_bytes += (draft_context - draft_generation) * int( - kwargs.get("max_num_tokens") or 0 - ) - cache_bytes += draft_generation - return cache_bytes, fixed_bytes + return cache_bytes, hp_bytes * max_batch_size * mapping.pp_size class Fp4MlaKVCacheManagerV2(Fp4MlaV2CacheLayoutPolicy, KVCacheManagerV2): diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index 784b02992fcd..1d97a1f2f1cd 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -27,10 +27,7 @@ CUDA_GRAPH_DUMMY_REQUEST_ID, GPU_LEVEL, KVCacheManagerV2, - _estimate_draft_cache_size_components, _fill_kv_pages, - _get_draft_layer_sizes, - _get_generation_kv_capacity, ) from tensorrt_llm._torch.pyexecutor.resource_manager import scan_kv_buffer_for_invalid_values from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -952,8 +949,7 @@ def _get_extra_quota_padding(self) -> int: """Ensure each attention type has minimal space when max_tokens is small.""" return len(DeepseekV4AttentionType) * (2 << 20) - def _get_cache_cost_components(self) -> tuple[int, int, int, int]: - """Return context/generation bytes per token, then fixed bytes per request.""" + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] has_fp8_kv_cache = self.dtype == DataType.FP8 non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( @@ -968,7 +964,7 @@ def _get_cache_cost_components(self) -> tuple[int, int, int, int]: ) ( context_swa_size_per_token, - context_swa_size_per_request, + _, ) = _estimate_swa_cache_size( self.head_dim, self.index_head_dim, @@ -996,60 +992,65 @@ def _get_cache_cost_components(self) -> tuple[int, int, int, int]: indexer_k_dtype=self._indexer_k_dtype, use_fp8_ds_mla=self.use_fp8_ds_mla, ) - context_size_per_token = non_sliding_attn_size_per_token + context_swa_size_per_token - generation_size_per_token = non_sliding_attn_size_per_token + generation_swa_size_per_token - if self._draft_layer_sizes: - draft_context, draft_generation, draft_per_request = ( - _estimate_draft_cache_size_components( - self._draft_layer_sizes, - self.tokens_per_block, - scratch_tokens=self._standalone_draft_reserve, - window_size=self.draft_window_size, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - ) - ) - draft_per_request += self._standalone_draft_reserve * context_size_per_token - context_size_per_token += draft_context - generation_size_per_token += draft_generation - context_swa_size_per_request += draft_per_request - generation_swa_size_per_request += draft_per_request - return ( - context_size_per_token, - generation_size_per_token, - context_swa_size_per_request, - generation_swa_size_per_request, - ) - - def _get_quota_from_max_tokens(self, max_tokens: int) -> int: - ( - context_size_per_token, - generation_size_per_token, - _, - generation_size_per_request, - ) = self._get_cache_cost_components() max_context_tokens = ( self._max_num_tokens if self._max_num_tokens is not None else max_tokens ) context_tokens = min(max_tokens, max_context_tokens) generation_tokens = max_tokens - context_tokens - return int( - context_tokens * context_size_per_token - + generation_tokens * generation_size_per_token - + self.max_batch_size * generation_size_per_request - + self._get_extra_quota_padding() + generation_quota = ( + max_tokens * non_sliding_attn_size_per_token + + generation_tokens * generation_swa_size_per_token + + self.max_batch_size * generation_swa_size_per_request ) + context_extra_quota = context_tokens * context_swa_size_per_token + padding = self._get_extra_quota_padding() + return int(generation_quota + context_extra_quota + padding) def _get_max_tokens_from_quota(self, quota: int) -> float: + compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + has_fp8_kv_cache = self.dtype == DataType.FP8 + non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + has_nvfp4_compress=self._use_nvfp4_compress, + nvfp4_residual_dim=self.nvfp4_residual_dim, + ) + context_swa_size_per_token, _ = _estimate_swa_cache_size( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + self.tokens_per_block, + self._swa_window_size, + context=True, + scratch=self.enable_swa_scratch_reuse, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + ) ( - context_size_per_token, - generation_size_per_token, - _, - generation_size_per_request, - ) = self._get_cache_cost_components() + generation_swa_size_per_token, + generation_swa_size_per_request, + ) = _estimate_swa_cache_size( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + self.tokens_per_block, + self._swa_window_size, + context=False, + scratch=False, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + ) padding = self._get_extra_quota_padding() - size_per_batch = self.max_batch_size * generation_size_per_request + padding + size_per_batch = self.max_batch_size * generation_swa_size_per_request + padding if quota < size_per_batch: return 0 + context_size_per_token = non_sliding_attn_size_per_token + context_swa_size_per_token if self._max_num_tokens is None: return (quota - size_per_batch) / context_size_per_token @@ -1057,6 +1058,7 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: if quota <= context_limit_quota: return (quota - size_per_batch) / context_size_per_token + generation_size_per_token = non_sliding_attn_size_per_token + generation_swa_size_per_token if generation_size_per_token <= 0: return float("inf") return self._max_num_tokens + (quota - context_limit_quota) / generation_size_per_token @@ -1247,14 +1249,6 @@ def _add_layer( constraints=constraints, ) - def _append_standalone_draft_layers( - self, config: KVCacheManagerConfigPy, **kwargs - ) -> KVCacheManagerConfigPy: - # Preserve target virtual-layer indices when registering draft layers. - return super()._append_standalone_draft_layers( - config, register_model_layers=False, **kwargs - ) - def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None: # Indexer compressor cache layout. Two modes are supported: # - "fp8" (FP8 blockwise): 1 byte per value + 1 fp32 scale per 128 @@ -1439,15 +1433,34 @@ def _get_generation_bytes(self, request: llm_request.LlmRequest) -> int: return self._get_cache_bytes_for_tokens(total_tokens, context=False) def _get_cache_bytes_for_tokens(self, total_tokens: int, *, context: bool) -> int: - ( - context_size_per_token, - generation_size_per_token, - context_size_per_request, - generation_size_per_request, - ) = self._get_cache_cost_components() - if context: - return int(total_tokens * context_size_per_token + context_size_per_request) - return int(total_tokens * generation_size_per_token + generation_size_per_request) + has_fp8_kv_cache = self.dtype == DataType.FP8 + compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + non_sliding_attn_size_per_token = _estimate_non_sliding_attn_size_per_token( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + has_nvfp4_compress=self._use_nvfp4_compress, + nvfp4_residual_dim=self.nvfp4_residual_dim, + ) + swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + self.tokens_per_block, + self._swa_window_size, + context=context, + scratch=self.enable_swa_scratch_reuse, + indexer_k_dtype=self._indexer_k_dtype, + use_fp8_ds_mla=self.use_fp8_ds_mla, + ) + return int( + total_tokens * (non_sliding_attn_size_per_token + swa_size_per_token) + + swa_size_per_request + ) def get_needed_resource_to_completion(self, request: llm_request.LlmRequest) -> int: if self._is_generation_request(request): @@ -1461,8 +1474,6 @@ def get_layer_bytes_per_token( local_layer_idx: int, data_role: DataRole, ) -> int: - if local_layer_idx in self.cp_as_tp_layer_ids: - return super().get_layer_bytes_per_token(local_layer_idx, data_role) # The generic layers in the base config are replaced by # _build_cache_config, so their buffer sizes are only placeholders. return 1 @@ -2047,33 +2058,6 @@ def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwar use_fp8_ds_mla=use_fp8_ds_mla, ) max_batch_size = int(kwargs.get("max_batch_size") or 0) - draft_layer_sizes = _get_draft_layer_sizes(**kwargs) - draft_scratch_tokens = kwargs.get("draft_scratch_tokens", 0) - if draft_layer_sizes: - _, headroom = _get_generation_kv_capacity(kwargs.get("spec_config"), is_draft=False) - target_context_swa_bytes, _ = _estimate_swa_cache_size( - head_dim, - index_head_dim, - compress_ratios, - has_fp8_kv_cache, - kwargs["tokens_per_block"], - model_config.sparse_attention_config.window_size, - context=True, - scratch=False, - indexer_k_dtype=indexer_k_dtype, - ) - _, draft_generation, draft_per_request = _estimate_draft_cache_size_components( - draft_layer_sizes, - kwargs["tokens_per_block"], - scratch_tokens=draft_scratch_tokens, - window_size=kwargs.get("draft_window_size"), - generation_capacity_headroom=headroom + draft_scratch_tokens, - ) - draft_per_request += draft_scratch_tokens * ( - non_sliding_attn_size_per_token + target_context_swa_bytes - ) - swa_size_per_token += draft_generation - swa_size_per_request += draft_per_request return ( non_sliding_attn_size_per_token + swa_size_per_token, swa_size_per_request * max_batch_size, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py index 9a1c8820758a..dc653921e1eb 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py @@ -612,15 +612,9 @@ def get_unique_primary_pool(self) -> torch.Tensor: first_head_dim = self.head_dim_per_layer[0] first_num_heads = self.num_kv_heads_per_layer[0] - if any( - head_dim != first_head_dim - for head_dim in self.head_dim_per_layer[: self.num_local_layers] - ): + if any(head_dim != first_head_dim for head_dim in self.head_dim_per_layer): raise ValueError("DSA requires a uniform KV head dimension") - if any( - num_heads != first_num_heads - for num_heads in self.num_kv_heads_per_layer[: self.num_local_layers] - ): + if any(num_heads != first_num_heads for num_heads in self.num_kv_heads_per_layer): raise ValueError("DSA requires a uniform KV head count") first_converter = self.impl.get_page_index_converter(0, Role.KEY) @@ -663,22 +657,13 @@ def get_unique_primary_pool(self) -> torch.Tensor: ) return self._unique_primary_pool - def _append_standalone_draft_layers(self, config, **kwargs): - # DSA masks and fused primary views describe target layers only. - return super()._append_standalone_draft_layers( - config, register_model_layers=False, **kwargs - ) - def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) -> int: if data_role == Role.INDEX_KEY: return _get_indexer_k_cache_bytes_per_token( self.index_head_dim, self.quant_block_size, self.use_fp4 ) cache_bytes = super().get_layer_bytes_per_token(local_layer_idx, data_role) - if ( - self.dtype_per_layer[local_layer_idx] == DataType.NVFP4 - and self.mla_kv_cache_residual_dim > 0 - ): + if self.dtype == DataType.NVFP4 and self.mla_kv_cache_residual_dim > 0: if data_role == Role.KEY: cache_bytes += self.mla_kv_cache_residual_dim // 2 elif data_role == Role.KEY_BLOCK_SCALE: @@ -688,15 +673,12 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) - cache_bytes += self.mla_kv_cache_residual_dim // 16 if ( data_role == Role.ALL - and not self._is_standalone_draft_layer(local_layer_idx) and self._indexer_k_cache_layer_mask[self.pp_layers[local_layer_idx]] ): cache_bytes += self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY) return cache_bytes def get_cache_bytes_per_token(self) -> int: - if self._draft_layer_sizes: - return super().get_cache_bytes_per_token() return sum( self.get_layer_bytes_per_token(local_layer_idx, Role.ALL) for local_layer_idx in range(self.num_local_layers) @@ -727,30 +709,6 @@ def get_cache_size_per_token( model_config, mapping, num_layers ) cache_bytes += num_attention_layers * (residual_dim // 2 + math.ceil(residual_dim / 16)) - from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - _estimate_draft_cache_size_components, - _get_draft_layer_sizes, - _get_generation_kv_capacity, - ) - - draft_sizes = _get_draft_layer_sizes(**kwargs) - if draft_sizes: - scratch = kwargs.get("draft_scratch_tokens", 0) - _, headroom = _get_generation_kv_capacity(kwargs.get("spec_config"), is_draft=False) - context, generation, fixed = _estimate_draft_cache_size_components( - draft_sizes, - kwargs["tokens_per_block"], - scratch_tokens=scratch, - window_size=kwargs.get("draft_window_size"), - generation_capacity_headroom=headroom + scratch, - helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, - ) - fixed += scratch * cache_bytes - return ( - cache_bytes + generation, - fixed * int(kwargs.get("max_batch_size") or 0) - + (context - generation) * int(kwargs.get("max_num_tokens") or 0), - ) return cache_bytes def shutdown(self) -> None: diff --git a/tensorrt_llm/_torch/disaggregation/native/auxiliary.py b/tensorrt_llm/_torch/disaggregation/native/auxiliary.py index 04da6a4c56c2..e7a3c411a7a2 100644 --- a/tensorrt_llm/_torch/disaggregation/native/auxiliary.py +++ b/tensorrt_llm/_torch/disaggregation/native/auxiliary.py @@ -114,63 +114,6 @@ def build_aux_transfer_layout( AuxSlot = namedtuple("AuxSlot", ["id", "buffer"]) -_DRAFT_HISTORY_VERSION = 2 -_DRAFT_HISTORY_FIELDS = 10 -_DRAFT_DTYPE_CODES = {"torch.float16": 1, "torch.bfloat16": 2, "torch.float8_e4m3fn": 3} -_DRAFT_BACKEND_CODES = {"VANILLA": 1, "TRTLLM": 2, "DSv4": 3} - - -def _encode_draft_history(history: dict[str, Any]) -> list[int]: - """Serialize the manager's validated history and rank-local storage identity.""" - if history is None: - raise ValueError("Standalone draft transfer requires draft history metadata") - layout = history["layout"] - return [ - _DRAFT_HISTORY_VERSION, - history["valid_length"], - history["position"], - layout["num_layers"], - layout["num_kv_heads"], - layout["head_dim"], - _DRAFT_DTYPE_CODES[layout["dtype"]], - _DRAFT_BACKEND_CODES[layout["attention_backend"]], - layout["kv_factor"], - layout["window_size"] or 0, - ] - - -def _decode_draft_history(values: list[int]) -> dict[str, Any]: - if values[0] != _DRAFT_HISTORY_VERSION: - raise ValueError("Missing or unsupported standalone draft history metadata version") - ( - _, - valid_length, - position, - num_layers, - num_kv_heads, - head_dim, - dtype_code, - backend_code, - kv_factor, - window_size, - ) = values - dtypes = {code: name for name, code in _DRAFT_DTYPE_CODES.items()} - backends = {code: name for name, code in _DRAFT_BACKEND_CODES.items()} - layout = { - "num_layers": num_layers, - "num_kv_heads": num_kv_heads, - "head_dim": head_dim, - "dtype": dtypes.get(dtype_code), - "attention_backend": backends.get(backend_code), - "kv_factor": kv_factor, - "window_size": window_size or None, - } - return { - "valid_length": valid_length, - "position": position, - "layout": layout, - } - class AuxBufferBase(ABC): """ @@ -226,15 +169,8 @@ def get_slot_data(self, slot: int) -> tuple[list[int], list[int], tuple[int, int class AuxBuffer(AuxBufferBase): - def __init__( - self, - max_slot_num: int, - beam_width: int, - max_draft_len: int, - device: str = "cpu", - *, - draft_history: bool = False, - ) -> None: + def __init__(self, max_slot_num: int, beam_width: int, max_draft_len: int, device: str = "cpu"): + # public constructor args remain the same, internals are private self._max_slot_num = int(max_slot_num) self._beam_width = int(beam_width) self._max_draft_len = int(max_draft_len) @@ -242,6 +178,9 @@ def __init__( self._free_slots = deque(list(range(self._max_slot_num))) self._occupied_slots: set[int] = set() + self._slot_token_counts: dict[ + int, tuple[int, int] + ] = {} # slot -> (first_tokens_len, draft_tokens_len) data_type = torch.int32 self._first_tokens_buffer = torch.empty( @@ -260,32 +199,35 @@ def __init__( self._prompt_token_counts_buffer = torch.zeros( self._max_slot_num, 2, dtype=data_type, device=self._device ) - # Version zero denotes an unfilled or recycled slot. - self._draft_history_buffer = ( - torch.zeros( - self._max_slot_num, _DRAFT_HISTORY_FIELDS, dtype=torch.int64, device=self._device - ) - if draft_history - else None - ) - - buffers = [ - self._first_tokens_buffer, - self._draft_tokens_buffer, - self._token_counts_buffer, - self._prompt_token_counts_buffer, - ] - if self._draft_history_buffer is not None: - buffers.append(self._draft_history_buffer) self._meta = AuxBufferMeta( - ptrs=np.array([buffer.data_ptr() for buffer in buffers], dtype=np.int64), + ptrs=np.array( + [ + self._first_tokens_buffer.data_ptr(), + self._draft_tokens_buffer.data_ptr(), + self._token_counts_buffer.data_ptr(), + self._prompt_token_counts_buffer.data_ptr(), + ], + dtype=np.int64, + ), size=np.array( - [buffer.numel() * buffer.element_size() for buffer in buffers], + [ + self._first_tokens_buffer.numel() * self._first_tokens_buffer.element_size(), + self._draft_tokens_buffer.numel() * self._draft_tokens_buffer.element_size(), + self._token_counts_buffer.numel() * self._token_counts_buffer.element_size(), + self._prompt_token_counts_buffer.numel() + * self._prompt_token_counts_buffer.element_size(), + ], dtype=np.int64, ), item_sizes=np.array( - [buffer[0].numel() * buffer.element_size() for buffer in buffers], + [ + self._first_tokens_buffer[0].numel() * self._first_tokens_buffer.element_size(), + self._draft_tokens_buffer[0].numel() * self._draft_tokens_buffer.element_size(), + self._token_counts_buffer[0].numel() * self._token_counts_buffer.element_size(), + self._prompt_token_counts_buffer[0].numel() + * self._prompt_token_counts_buffer.element_size(), + ], dtype=np.int64, ), device=self._device, @@ -305,8 +247,7 @@ def alloc_slot(self) -> AuxSlot: "This indicates a bug in slot management." ) self._occupied_slots.add(slot_id) - if self._draft_history_buffer is not None: - self._draft_history_buffer[slot_id].zero_() + self._slot_token_counts[slot_id] = (0, 0) return AuxSlot(slot_id, self) def free_slot(self, slot: int) -> None: @@ -320,17 +261,13 @@ def free_slot(self, slot: int) -> None: f"Invalid slot id {slot}. Valid slot indices are in the range 0..{self._max_slot_num - 1}." ) self._occupied_slots.remove(slot) + self._slot_token_counts.pop(slot, None) self._free_slots.append(slot) @property def meta(self) -> AuxBufferMeta: return self._meta - @property - def has_draft_history(self) -> bool: - """Whether this buffer transfers standalone drafter history metadata.""" - return self._draft_history_buffer is not None - def fill_slot(self, slot: int, request: LlmRequest) -> None: if slot not in self._occupied_slots: raise ValueError( @@ -358,6 +295,7 @@ def fill_slot(self, slot: int, request: LlmRequest) -> None: self._draft_tokens_buffer[slot][: len(draft_tokens)].copy_( torch.tensor(draft_tokens, dtype=torch.int32, device=self._device) ) + self._slot_token_counts[slot] = (len(first_gen_tokens), len(draft_tokens)) self._token_counts_buffer[slot].copy_( torch.tensor( [len(first_gen_tokens), len(draft_tokens)], dtype=torch.int32, device=self._device @@ -367,11 +305,6 @@ def fill_slot(self, slot: int, request: LlmRequest) -> None: self._prompt_token_counts_buffer[slot].copy_( torch.tensor([prompt_tokens, cached_tokens], dtype=torch.int32, device=self._device) ) - if self._draft_history_buffer is not None: - values = _encode_draft_history(request.py_draft_transfer_history) - self._draft_history_buffer[slot].copy_( - torch.tensor(values, dtype=torch.int64, device=self._device) - ) @staticmethod def _resolve_prompt_token_counts(request: LlmRequest) -> tuple[int, int]: @@ -402,11 +335,3 @@ def get_slot_data(self, slot: int) -> tuple[list[int], list[int], tuple[int, int first_gen_tokens, draft_tokens = self.get_slot_tokens(slot) prompt_tokens, cached_tokens = self._prompt_token_counts_buffer[slot].tolist() return first_gen_tokens, draft_tokens, (int(prompt_tokens), int(cached_tokens)) - - def get_slot_draft_history(self, slot: int) -> dict[str, Any]: - """Decode transferred history for validation against the receiving request/cache.""" - if slot not in self._occupied_slots: - raise ValueError(f"Cannot read slot {slot}: slot is not currently allocated.") - if self._draft_history_buffer is None: - raise ValueError("Standalone draft history transfer is not enabled for this buffer") - return _decode_draft_history(self._draft_history_buffer[slot].tolist()) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py index 87b1daca257c..b6bf6e81addc 100644 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py @@ -631,10 +631,7 @@ def block_bytes_per_group(page_table: KVCachePageTable) -> list[int | None]: remains aligned with receive-request layer-group indices. """ from tensorrt_llm._torch.disaggregation.resource.page import CacheKind - from tensorrt_llm._torch.disaggregation.resource.utils import ( - get_physical_pool, - get_pool_view_num_layers, - ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool assert page_table is not None out: list[int | None] = [] @@ -642,11 +639,6 @@ def block_bytes_per_group(page_table: KVCachePageTable) -> list[int | None]: if lg.kind != CacheKind.PAGED: out.append(None) continue - if lg.cp_as_tp: - out.append( - sum(get_pool_view_num_layers(pv) * pv.bytes_per_layer for pv in lg.pool_views) - ) - continue pool_indices = {pool_view.pool_idx for pool_view in lg.pool_views} out.append( sum( diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index 8cace14f89d6..f141cce93ff4 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -26,16 +26,8 @@ SpecRegionPair, ) from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo -from tensorrt_llm._torch.disaggregation.resource.page import ( - AttentionLayerGroup, - CacheKind, - KVCachePageTable, - MapperKind, -) -from tensorrt_llm._torch.disaggregation.resource.utils import ( - find_replicated_role_mismatch, - get_layer_to_layer_group, -) +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind +from tensorrt_llm._torch.disaggregation.resource.utils import find_replicated_role_mismatch from tensorrt_llm._utils import nvtx_range @@ -581,20 +573,6 @@ def validate_peer_compatible( "AttentionPolicy.validate_peer_compatible: replicated roles differ on " f"overlapping layers: {differing}" ) - if self_page_table is not None and peer_page_table is not None: - local_layers = get_layer_to_layer_group(self_page_table) - peer_layers = get_layer_to_layer_group(peer_page_table) - for layer in local_layers.keys() & peer_layers.keys(): - local = self_page_table.layer_groups[local_layers[layer]] - peer = peer_page_table.layer_groups[peer_layers[layer]] - if not isinstance(local, AttentionLayerGroup) or not isinstance( - peer, AttentionLayerGroup - ): - continue - if local.cp_as_tp != peer.cp_as_tp: - raise ValueError("Attention peers disagree on full-sequence CP-to-TP storage") - if local.cp_as_tp and local.total_kv_head_num != peer.total_kv_head_num: - raise ValueError("Standalone draft peers have different global KV head counts") def check_peer_compatible(self, peer_ri: RankInfo) -> bool: a = self._ri.attention diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 035a5314ed8f..d8d37658195b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -14,7 +14,7 @@ # limitations under the License. from collections import Counter -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple, Union import numpy as np @@ -26,12 +26,7 @@ from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import ( - AttentionLayerGroup, - CacheKind, - MapperKind, - PoolView, -) +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, MapperKind, PoolView from tensorrt_llm._torch.disaggregation.resource.utils import ( get_layer_byte_ranges, get_layer_to_layer_group, @@ -172,10 +167,6 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: if not self._get_policy(CacheKind.PAGED).check_peer_compatible(peer_ri): return False - AttentionPolicy.validate_peer_compatible( - self._self_ext_cache.page_table, peer_ri.page_table - ) - # Recurrent-state (Mamba/KDA) layout gate. Raises ValueError with a # field-level diagnostic instead of returning False, so the precise # mismatch reaches the caller of register(). @@ -189,29 +180,6 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: peer_ri.page_table, ) AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table) - if self._ri.num_draft_layers != peer_ri.num_draft_layers: - # The layer-count tolerance below would transfer only the target - # layers and leave the drafter's prompt KV unset on the receiver. - raise ValueError( - "PeerRegistrar: unified draft KV layer count differs " - f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); " - "both peers must use the same speculative decoding KV cache layout; " - "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers." - ) - if self._ri.num_draft_layers: - if self._ri.tp_size != peer_ri.tp_size: - raise ValueError( - "PeerRegistrar: a shared draft KV cache needs equal tp_size " - f"(local={self._ri.tp_size}, peer={peer_ri.tp_size}); " - "unset TRTLLM_DISAGG_SHARED_DRAFT_KV or use the same TP on both workers." - ) - if self._ri.attention.enable_attention_dp != peer_ri.attention.enable_attention_dp: - raise ValueError( - "PeerRegistrar: a shared draft KV cache needs the same attention DP " - f"setting (local={self._ri.attention.enable_attention_dp}, " - f"peer={peer_ri.attention.enable_attention_dp}); unset " - "TRTLLM_DISAGG_SHARED_DRAFT_KV or match it on both workers." - ) self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) @@ -455,12 +423,8 @@ def get_kv_map( peer_overlap_positions[0] if peer_overlap_positions else 0 ) - mapper_peer_ri = peer_ri - if getattr(self_lg, "cp_as_tp", False): - policy = AttentionPolicy(self._group_rank_info(self._ri, self_lg)) - mapper_peer_ri = self._group_rank_info(peer_ri, peer_lg) mapper = policy.build_mapper( - peer_ri=mapper_peer_ri, + peer_ri=peer_ri, mapper_kind=self_pv.mapper_kind, self_layer_offsets=self_layer_offsets, peer_layer_offsets=peer_layer_offsets, @@ -623,41 +587,11 @@ def should_send_pool( """ layer_group = self._self_ext_cache.page_table.layer_groups[layer_group_id] pool_view = layer_group.pool_views[pool_idx] - if getattr(layer_group, "cp_as_tp", False): - peer_lg_id, _ = self.get_pool_mapping(peer_rank_info)[(layer_group_id, pool_idx)] - peer_lg = peer_rank_info.page_table.layer_groups[peer_lg_id] - local = self._group_rank_info(self._ri, layer_group) - peer = self._group_rank_info(peer_rank_info, peer_lg) - if local.tp_size <= peer.tp_size: - if peer.tp_rank // (peer.tp_size // local.tp_size) != local.tp_rank: - return False - elif local.tp_rank // (local.tp_size // peer.tp_size) != peer.tp_rank: - return False - duplicate, _ = AttentionPolicy(local).duplicate_head_factors(peer) - return duplicate <= 1 or peer.dp_rank % duplicate == local.tp_rank % duplicate # Delegate to the policy's should_send. None means fan-in election. policy = self._get_policy(layer_group.kind) result = policy.should_send(peer_overlap, peer_rank_info, mapper_kind=pool_view.mapper_kind) return self._owns_tp_fan_in(peer_rank_info) if result is None else result - @staticmethod - def _group_rank_info(rank_info: RankInfo, layer_group: AttentionLayerGroup) -> RankInfo: - tp_size, tp_rank = MambaPolicy._mamba_tp(rank_info) - return replace( - rank_info, - tp_size=tp_size, - tp_rank=tp_rank, - cp_size=1, - cp_rank=0, - dp_size=1, - attention=replace( - rank_info.attention, - kv_heads_per_rank=layer_group.kv_head_num_per_rank, - is_mla=False, - enable_attention_dp=False, - ), - ) - def should_send_aux(self, peer_rank_info: RankInfo) -> bool: # to ensure the transfer aux is not duplicated diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 358b1b576412..87731d303204 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -49,9 +49,6 @@ class RankInfo: aux_meta: Optional[AuxBufferMeta] = None page_table: Optional[KVCachePageTable] = None bounded_replay_on_generation: bool = False - # Layers a unified KV cache manager owns for the drafter (layer specs with - # a DRAFT owner). They ride the target's pages, so both peers must agree. - num_draft_layers: int = 0 @property def tp_size_per_dp_group(self) -> int: @@ -117,9 +114,6 @@ def from_kv_cache_manager( aux_meta=aux_buffer_meta, page_table=build_page_table_from_manager(kvm), bounded_replay_on_generation=getattr(kvm, "bounded_replay_on_generation", False), - num_draft_layers=( - len(kvm.draft_layer_ids) if getattr(kvm, "layer_specs", None) is not None else 0 - ), ) @classmethod diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 82278c14d278..e28ca9d91dd9 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -367,13 +367,7 @@ class AgentResult(Enum): def _make_kv_result_msg( - instance_rank, - unique_rid, - slice_id, - is_last_slice, - agent_result, - transfer_size=0, - tail=None, + instance_rank, unique_rid, slice_id, is_last_slice, agent_result, transfer_size=0, tail=None ): """Build a KV_AGENT_RESULT message. ALL result sends (success AND failed/cancelled) must go through this single binary frame so the receiver's _KV_RESULT_PREFIX.unpack never hits a stale @@ -969,9 +963,7 @@ def _poll_in_doubt_transfers(self, thread_idx: int) -> None: pending.send_slot_id = None if meta.meta_type == WriteMetaType.AUX: message = _make_aux_result_msg( - self._instance_rank, - meta.unique_rid, - AgentResult.FAILED_QUIESCED, + self._instance_rank, meta.unique_rid, AgentResult.FAILED_QUIESCED ) else: message = _make_kv_result_msg( @@ -1015,11 +1007,7 @@ def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferReque f"_make_agent_request: dst_device_id is None for KV transfer " f"unique_rid={write_meta.unique_rid}" ) - src_dev, dst_dev, mem_type = ( - device_id, - write_meta.dst_device_id, - MemoryType.VRAM, - ) + src_dev, dst_dev, mem_type = device_id, write_meta.dst_device_id, MemoryType.VRAM if n == 0: src_memory_descs = MemoryDescs(mem_type, []) @@ -1036,11 +1024,7 @@ def _make_agent_request(write_meta: WriteMeta, device_id: int) -> "TransferReque # After this call, src_memory_descs and dst_memory_descs are in a moved-from # state and must NOT be accessed again. return TransferRequest( - TransferOp.WRITE, - src_memory_descs, - dst_memory_descs, - write_meta.peer_name, - None, + TransferOp.WRITE, src_memory_descs, dst_memory_descs, write_meta.peer_name, None ) @nvtx_range("_deliver_kv_to_agent") @@ -1428,6 +1412,7 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write lg_info = extractor.page_table.layer_groups[self_lg] src_block_ids = src_block_ids_per_groups[self_lg] dst_block_ids = dst_block_ids_per_groups[peer_lg] + # Extract regions: STATE = per-layer ptrs, PAGED = per-block ptrs if lg_info.kind == CacheKind.STATE: if src_block_ids.size == 0 or dst_block_ids.size == 0: @@ -1439,11 +1424,7 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write else: src_block_ids = np.asarray(src_block_ids, dtype=np.int64) dst_block_ids = np.asarray(dst_block_ids, dtype=np.int64) - if ( - peer_ri.cp_size > 1 - and self_ri.cp_size == 1 - and not getattr(extractor.page_table.layer_groups[self_lg], "cp_as_tp", False) - ): + if peer_ri.cp_size > 1 and self_ri.cp_size == 1: # Helix: the receiver owns global blocks [cp_rank::cp_size] # (same protocol as partition_context_for_helix), so its # table is the strided subset of ours; block reuse is @@ -1542,9 +1523,7 @@ def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> Wri if timer: timer.record_prepare_args_end(peer_ri.instance_rank) timer.record_transfer_sizes( - peer_ri.instance_rank, - int(sizes.sum()) if sizes.size > 0 else 0, - src_ptrs.size, + peer_ri.instance_rank, int(sizes.sum()) if sizes.size > 0 else 0, src_ptrs.size ) return WriteMeta( @@ -2018,9 +1997,7 @@ def __init__( self._timeout_s = timeout_s self._overall_timeout_s = overall_timeout_s self._deadline_monotonic_s: Optional[float] = None - self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST or ( - aux_buffer is not None and aux_buffer.has_draft_history - ) + self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST self._enforce_physical_ownership = getattr(sender, "_enforce_physical_ownership", False) self._sender: Sender # narrow base class type for Pylance self.request_id = request_id @@ -3105,10 +3082,7 @@ def send_cancel_to_senders(self, unique_rid: int, sender_endpoints: set[str]) -> if self._enforce_physical_ownership: rank_info = self._registrar.self_rank_info message.extend( - [ - rank_info.instance_name.encode(), - str(rank_info.instance_rank).encode("ascii"), - ] + [rank_info.instance_name.encode(), str(rank_info.instance_rank).encode("ascii")] ) sent = set() for endpoint in sender_endpoints: @@ -3183,14 +3157,9 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): f"_process_kv_agent_result: unexpected msg_type={message[0]!r}, expected KV_AGENT_RESULT" ) return - ( - peer_rank, - unique_rid, - receiver_slice_id, - is_last_slice, - status_code, - transfer_size, - ) = _KV_RESULT_PREFIX.unpack(message[1]) + peer_rank, unique_rid, receiver_slice_id, is_last_slice, status_code, transfer_size = ( + _KV_RESULT_PREFIX.unpack(message[1]) + ) from .bounce import decode_result_tail dst_ptrs, sizes, src_base = decode_result_tail(message) @@ -3259,9 +3228,7 @@ def __init__( ): super().__init__(receiver, SessionArgsBase(params, prompt_len=prompt_len)) self._timeout_s = timeout_s - self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST or ( - aux_buffer is not None and aux_buffer.has_draft_history - ) + self._need_aux = params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST self._enforce_physical_ownership = getattr(receiver, "_enforce_physical_ownership", False) self._receiver: Receiver # narrow base class type for Pylance self.request_id = request_id @@ -3838,8 +3805,6 @@ def unpack_aux(self, request: LlmRequest) -> None: """Read token data from the aux buffer slot into the given request.""" assert self._aux_buffer is not None, "No aux_buffer set for this session" assert self.aux_slot is not None, "No aux_slot set for this session" - if self._aux_buffer.has_draft_history: - self.unpack_draft_history(request) first_gen_tokens, draft_tokens, (prompt_tokens, cached_tokens) = ( self._aux_buffer.get_slot_data(self.aux_slot) ) @@ -4084,20 +4049,9 @@ def __del__(self): except Exception as e: logger.warning(f"RxSession.__del__: exception during close: {e}") - def unpack_draft_history(self, request: LlmRequest) -> None: - """Read standalone history without changing context-first token fields.""" - if self._aux_buffer is None or self.aux_slot is None: - raise ValueError("Standalone draft transfer requires an auxiliary buffer slot") - request.py_draft_transfer_history = self._aux_buffer.get_slot_draft_history(self.aux_slot) - class RankInfoServer: - def __init__( - self, - rank_info: RankInfo, - addr: Optional[str] = None, - port: Optional[int] = None, - ): + def __init__(self, rank_info: RankInfo, addr: Optional[str] = None, port: Optional[int] = None): self._rank_info = rank_info self._shutdown = False # must be set before _start_listener() so __del__ is safe if addr is None and port is None: @@ -4178,10 +4132,7 @@ def _create_nixl_agent( def _make_aux_buffer( kvm: KVCacheManagerV2, max_slots: int, max_draft_len: Optional[int] = None ) -> Optional[AuxBuffer]: - draft_history = bool(getattr(kvm, "_draft_layer_sizes", ())) if max_slots <= 0: - if draft_history: - raise ValueError("Standalone draft transfer requires auxiliary buffer slots") return None if max_draft_len is None: max_draft_len = max(0, int(getattr(kvm, "max_draft_len", 0))) @@ -4190,7 +4141,6 @@ def _make_aux_buffer( beam_width=max(1, int(getattr(kvm, "max_beam_width", 1))), max_draft_len=max_draft_len, device="cpu", - draft_history=draft_history, ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py index 7945d3395435..b97f16d52a53 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py +++ b/tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py @@ -19,9 +19,8 @@ import numpy as np -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode from .page import AttentionLayerGroup @@ -145,9 +144,6 @@ def get_block_ids(self, req, group_idx, lg): # noqa: ARG002 # IDs), and active sequences GPU-lock their pages (_UniqPageLock enforces # cache_level==GPU), so the slot_ids yielded here are already the right # offsets for primary-pool pointer arithmetic. No translation is needed. - if getattr(lg, "cp_as_tp", False): - ordinals = self.get_block_ordinals(req, group_idx, lg) - return ordinals[ordinals >= 0] return np.fromiter( self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices( group_idx, valid_only=True @@ -159,20 +155,12 @@ def get_block_ordinals(self, req, group_idx, lg): # noqa: ARG002 # valid_only=False yields one entry per block ordinal, with -1 # (BAD_PAGE_INDEX) for out-of-window (SWA-evicted) and unbound blocks. # Position is the index; no length arithmetic needed. - ordinals = np.fromiter( + return np.fromiter( self._mgr.kv_cache_map[req.py_request_id].get_aggregated_page_indices( group_idx, valid_only=False ), dtype=np.int64, ) - if lg.cp_as_tp: - converter = self._mgr.impl.get_page_index_converter( - lg.local_layers[0].local_layer_id, Role.KEY - ) - return np.asarray( - converter(ordinals.tolist(), index_mode=PageIndexMode.SHARED), dtype=np.int64 - ) - return ordinals def commit_blocks_for_reuse(self, req: LlmRequest) -> None: self._mgr.try_commit_blocks(req) diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 629af5872c8c..b1132c10ddc4 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -46,7 +46,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole, PoolGroupDesc, SlotDescVariant +from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole # Mapper kinds a V2 manager may declare via get_disagg_role_mapper_kinds(). # A physical pool may mix kinds (V2 storage coalesces buffers purely by @@ -201,12 +201,7 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: if not hasattr(manager, "_layer_attn_to_layer_id"): # Standard: local_layer_id is index into pp_layers - local_to_global = dict(enumerate(manager.pp_layers)) - local_to_global.update( - (manager.layer_offsets[layer_id], layer_id) - for layer_id in getattr(manager, "draft_layer_ids", ()) - ) - return [local_to_global[lid] for lid in local_layer_ids] + return [manager.pp_layers[lid] for lid in local_layer_ids] # Virtual layers: build inverse mapping internal_layer_id -> (model_layer, attn_type) # and encode as model_layer * num_attn_types + attn_type_value @@ -220,8 +215,6 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: inverse = {} for (model_layer, attn_type), layer_id in manager._layer_attn_to_layer_id.items(): inverse.setdefault(layer_id, (model_layer, attn_type.value)) - for draft_layer in getattr(manager, "draft_layer_ids", ()): - inverse[manager.layer_offsets[draft_layer]] = (draft_layer, 0) # Use the full enum range for consistent encoding across all PP ranks. # Different PP ranks may have different subsets of attention types (e.g., @@ -334,56 +327,6 @@ def _build_pool_views_for_variant( return pool_views -def _build_draft_pool_views( - manager: KVCacheManagerV2, variant: SlotDescVariant, pg_desc: PoolGroupDesc -) -> tuple[PhysicalPoolGroup, list[PoolView]]: - """Expose expanded draft buffers as physical-token pages for transfer.""" - if len(variant.coalesced_buffers) != 1: - raise ValueError("Standalone draft layers must share one equal-sized K/V pool") - buffer = variant.coalesced_buffers[0] - first_layer = buffer.buffer_ids[0].layer_id - expansion = manager.impl.get_page_index_converter(first_layer, Role.KEY).expansion - buffer_size = int(buffer.single_buffer_size) - if buffer_size % expansion: - raise ValueError("Draft buffer size is not divisible by its page expansion") - page_bytes = buffer_size // expansion - physical_pool = pg_desc.pools[0] - slot_bytes = int(physical_pool.slot_bytes) - if slot_bytes % page_bytes: - raise ValueError("Draft pool slots must contain whole physical-token pages") - pool_group = PhysicalPoolGroup( - pools=[ - PhysicalPool( - base_address=int(physical_pool.base_address), - slot_bytes=page_bytes, - num_slots=int(pg_desc.num_slots) * (slot_bytes // page_bytes), - slot_stride_bytes=page_bytes, - ) - ] - ) - entries_by_role = defaultdict(list) - for index, buffer_id in enumerate(buffer.buffer_ids): - if buffer_id.role not in (Role.KEY, Role.VALUE): - raise ValueError("Standalone draft transfer expects only K/V buffers") - entries_by_role[buffer_id.role].append( - (int(buffer_id.layer_id), index * buffer_size, page_bytes) - ) - # K and V occupy separate expanded runs. Keeping one role per view lets - # the existing head mapper copy each physical-token page without treating - # the intervening subpages as part of that page's payload. - views = [ - PoolView( - pool_idx=0, - buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE), - pool_role=frozenset({str(role)}), - mapper_kind=MapperKind.INDEXED, - bytes_per_layer=page_bytes, - ) - for role, entries in entries_by_role.items() - ] - return pool_group, views - - def _build_page_table_v2(manager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManagerV2. @@ -485,21 +428,6 @@ def _window_size_for_layer(internal_layer_id: int): for iid, gid in zip(all_internal_layer_ids, all_global_layer_ids) ] - layer_id = int(all_internal_layer_ids[0]) - if layer_id in manager.cp_as_tp_layer_ids: - draft_pool_group, pool_views = _build_draft_pool_views(manager, variant, pg_desc) - draft_pool_group_idx = len(pool_groups) - pool_groups.append(draft_pool_group) - layer_groups_by_id[layer_group_id] = AttentionLayerGroup( - pool_group_idx=draft_pool_group_idx, - local_layers=local_layers, - pool_views=pool_views, - kv_head_num_per_rank=manager.num_kv_heads_per_layer[layer_id], - sliding_window_size=_window_size_for_layer(all_internal_layer_ids[0]), - cp_as_tp=True, - total_kv_head_num=manager._draft_total_num_kv_heads, - ) - continue # A life cycle is recurrent when its layers are the manager's # mamba layers. Recurrent and attention layers never share a life # cycle: their buffer configs differ, so V2 groups them apart. diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 837ce31ce20d..1254c06001f6 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -293,9 +293,6 @@ class AttentionLayerGroup(LayerGroup): kv_head_num_per_rank: int = 0 sliding_window_size: Optional[int] = None - # Full-sequence KV whose heads use the flattened TP x CP rank grid. - cp_as_tp: bool = False - total_kv_head_num: Optional[int] = None def to_dict(self) -> dict: return { @@ -303,8 +300,6 @@ def to_dict(self) -> dict: "pool_group_idx": int(self.pool_group_idx), "kv_head_num_per_rank": int(self.kv_head_num_per_rank), "sliding_window_size": self.sliding_window_size, - "cp_as_tp": self.cp_as_tp, - "total_kv_head_num": self.total_kv_head_num, "local_layers": [ll.to_dict() for ll in self.local_layers], "pool_views": [pv.to_dict() for pv in self.pool_views], } @@ -317,8 +312,6 @@ def from_dict(cls, data: dict) -> "AttentionLayerGroup": pool_views=[PoolView.from_dict(pv) for pv in data.get("pool_views", [])], kv_head_num_per_rank=int(data["kv_head_num_per_rank"]), sliding_window_size=data.get("sliding_window_size"), - cp_as_tp=bool(data.get("cp_as_tp", False)), - total_kv_head_num=data.get("total_kv_head_num"), ) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 51ae7033ae8c..759d408fa362 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -633,7 +633,6 @@ def _create_cache_extent(self, req: LlmRequest) -> CacheExtent: registered by.""" rid = get_unique_rid(req) assert rid is not None - self._pack_draft_history(req) return CacheExtent(name=rid, local=self._describe_local(req)) def _describe_local(self, req: LlmRequest) -> Chunk: @@ -654,9 +653,7 @@ def _describe_local(self, req: LlmRequest) -> Chunk: assert self._page_table is not None layer_groups = self._page_table.layer_groups remote_tail_start = getattr(req, "py_csa2_remote_tail_start", None) - transfer_len = ( - remote_tail_start if remote_tail_start is not None else self._global_prompt_len(req) - ) + transfer_len = remote_tail_start if remote_tail_start is not None else req.prompt_len prompt_blocks = (transfer_len + tpb - 1) // tpb remote_tail_split = getattr(req, "py_csa2_remote_tail_split", None) excluded_pool_views = set() @@ -685,14 +682,7 @@ def _describe_local(self, req: LlmRequest) -> Chunk: # prompt positions through beam 0's block table), so the # positional path applies to every beam width. ordinals = adapter.get_block_ordinals(req, idx, lg) - local_prompt_blocks = prompt_blocks - if not getattr(lg, "cp_as_tp", False): - cp_size = getattr(self._kv_cache_manager, "_helix_cp_size", 1) - cp_rank = getattr(self._kv_cache_manager, "_helix_cp_rank", 0) - local_prompt_blocks = len(range(cp_rank, prompt_blocks, cp_size)) - group = self._positional_window( - ordinals, local_prompt_blocks, cached_per_lg[idx] // tpb - ) + group = self._positional_window(ordinals, prompt_blocks, cached_per_lg[idx] // tpb) groups.append(group) return Chunk( @@ -719,69 +709,10 @@ def _positional_window( window[: min(cached_blocks, prompt_blocks)] = -1 return window - def _need_aux_transfer(self, req: LlmRequest) -> bool: + @staticmethod + def _need_aux_transfer(req: LlmRequest) -> bool: params = req.py_disaggregated_params - return bool(getattr(self._kv_cache_manager, "_draft_layer_sizes", ())) or ( - params is not None and params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST - ) - - def _global_prompt_len(self, req: LlmRequest) -> int: - if getattr(self._kv_cache_manager, "_has_cp_helix", False): - return req.total_input_len_cp - return req.prompt_len - - def _validate_draft_history_range(self, req: LlmRequest, history: dict) -> None: - # Full-attention drafters can skip an oversized context store; rolling - # drafters retain their complete live suffix. Neither includes scratch. - window_size = history["layout"]["window_size"] - max_positions = self._kv_cache_manager.draft_max_position_embeddings - prompt_len = self._global_prompt_len(req) - expected_length = prompt_len - if window_size is not None: - if type(window_size) is not int or window_size <= 0: - raise ValueError("Invalid draft history window size") - expected_length = min(expected_length, window_size) - elif max_positions is not None and prompt_len > max_positions: - expected_length = min(history["valid_length"], max_positions) - if history["valid_length"] != expected_length or history["position"] != prompt_len: - raise ValueError( - "DSpark transfer requires valid draft history and sequence position " - f"covering the complete prompt ({prompt_len} tokens, " - f"{expected_length} retained)." - ) - - def _pack_draft_history(self, req: LlmRequest) -> None: - manager = getattr(self, "_kv_cache_manager", None) - if not getattr(manager, "_draft_layer_sizes", ()): - return - history = self._kv_cache_manager.export_draft_history(req.py_request_id) - self._validate_draft_history_range(req, history) - req.py_draft_transfer_history = history - - def _restore_draft_history(self, req: LlmRequest) -> None: - history = getattr(req, "py_draft_transfer_history", None) - manager = getattr(self, "_kv_cache_manager", None) - has_draft = bool(getattr(manager, "_draft_layer_sizes", ())) - if history is None: - if has_draft: - raise ValueError( - "Standalone DSpark generation requires draft history from a prefill worker " - "with matching speculative configuration; draft history metadata is missing." - ) - return - if not has_draft: - raise ValueError( - "Received standalone DSpark draft history without a manager-owned draft cache." - ) - self._validate_draft_history_range(req, history) - # K/V already occupies receiver-local pages; restore only validity and position. - self._kv_cache_manager.restore_draft_history(req.py_request_id, history) - - def _prepare_received_history(self, session: RxSessionBase, req: LlmRequest) -> None: - if self._need_aux_transfer(req): - self._apply_aux(session, req) - self._assert_disagg_history_declared(req) - self._restore_draft_history(req) + return params is not None and params.schedule_style == DisaggScheduleStyle.GENERATION_FIRST def _validate_bridge_req(self, req: LlmRequest, synchronous: bool = False) -> bool: if not getattr(self, "_fp4_mla_bridge_enabled", False): @@ -1084,10 +1015,6 @@ def _close_session_or_raise(self, session: object, rid: int, outcome: str) -> No def _apply_aux(self, session, req: LlmRequest): """Unpack aux tokens from session into request's context_phase_params.""" - params = req.py_disaggregated_params - if params is not None and params.schedule_style != DisaggScheduleStyle.GENERATION_FIRST: - session.unpack_draft_history(req) - return session.unpack_aux(req) first_gen_tokens = req.py_first_gen_tokens # type: ignore[attr-defined] draft_tokens = req.py_draft_tokens @@ -1291,7 +1218,6 @@ def request_and_receive_sync(self, req: LlmRequest) -> None: if self._need_aux_transfer(req): self._apply_aux(session, req) self._assert_disagg_history_declared(req) - self._restore_draft_history(req) req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE else: req.state = LlmRequestState.DISAGG_TRANS_ERROR @@ -1558,7 +1484,6 @@ def check_gen_transfer_status(self, at_least_request_num: Optional[int]) -> GenT if self._need_aux_transfer(req): self._apply_aux(session, req) self._assert_disagg_history_declared(req) - self._restore_draft_history(req) self._close_session_or_raise(session, rid, "completed") req.state = LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE del self._recv_reqs[rid] diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index da9d5ff65834..1f1d8ffd1212 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -492,13 +492,13 @@ def _dspark_attention_forward_legacy( start_pos: int, kv_cache: torch.Tensor, *, - wq_a: torch.Tensor, + wq_a: Optional[torch.Tensor], q_norm_w: torch.Tensor, - wq_b: torch.Tensor, - wkv: torch.Tensor, + wq_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + wkv: Optional[torch.Tensor], kv_norm_w: torch.Tensor, - wo_a: torch.Tensor, - wo_b: torch.Tensor, + wo_a: Optional[torch.Tensor], + wo_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], attn_sink: torch.Tensor, n_heads: int, head_dim: int, @@ -511,6 +511,13 @@ def _dspark_attention_forward_legacy( freqs_cis: torch.Tensor, persist: bool = False, q_b_norm_enabled: bool = True, + qkv_proj: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + o_a_proj: Optional[torch.Tensor] = None, + o_a_scale: Optional[torch.Tensor] = None, + o_a_dequant: Optional[torch.Tensor] = None, + use_fused_output_projection: Optional[bool] = None, + use_fused_qkv_input_projection: bool = False, + use_fused_q_input_projection: bool = False, ) -> torch.Tensor: """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). @@ -522,8 +529,7 @@ def _dspark_attention_forward_legacy( block's own positions, attention-sink softmax, inverse-RoPE on the output, and a grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). - Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / - dequantized projection weights); ``wo_a`` is the raw grouped weight matrix + Projections accept tensor weights or loaded projection modules; ``wo_a`` is the raw grouped weight matrix ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is ``[b, window_size, head_dim]`` and is updated functionally (cloned). @@ -543,13 +549,24 @@ def _dspark_attention_forward_legacy( main_freqs = freqs_cis[start_pos : start_pos + 1] blk_freqs = freqs_cis[start_pos + 1 : start_pos + 1 + block] + q_latent, block_kv_input, main_kv_input = _qkv_projections( + x, + main_x, + wq_a=wq_a, + wkv=wkv, + q_norm_w=q_norm_w, + qkv_proj=qkv_proj, + use_fused_fp8_input=use_fused_qkv_input_projection, + ) + # Captured-context K/V from main_x (MQA, shared across heads). - main_kv = _rmsnorm(F.linear(main_x, wkv), kv_norm_w, eps) # [b, 1, head_dim] + main_kv = _rmsnorm(main_kv_input, kv_norm_w, eps) # [b, 1, head_dim] main_kv = _rope_last_dims(main_kv, rd, main_freqs) # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm(F.linear(x, wq_a), q_norm_w, eps) - q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [b, block, h, head_dim] + q = _q_projection(q_latent, q_norm_w, eps, wq_b, use_fused_q_input_projection).unflatten( + -1, (n_heads, head_dim) + ) # Per-head RMS in the query dtype (matches the reference inline normalization, # which is NOT the fp32 RMSNorm path). if q_b_norm_enabled: @@ -557,7 +574,7 @@ def _dspark_attention_forward_legacy( q = _rope_last_dims(q, rd, blk_freqs) # Block K/V. - kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] + kv = _rmsnorm(block_kv_input, kv_norm_w, eps) # [b, block, head_dim] kv = _rope_last_dims(kv, rd, blk_freqs) # Write the context K/V into the rolling window, then attend over @@ -572,10 +589,17 @@ def _dspark_attention_forward_legacy( o = _rope_last_dims(o, rd, blk_freqs, inverse=True) # Grouped low-rank O projection. - o = o.reshape(b, block, n_groups, -1) - wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) - o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) - return F.linear(o.flatten(2), wo_b) + return _output_projection( + o, + wo_a=wo_a, + wo_b=wo_b, + n_groups=n_groups, + o_lora_rank=o_lora_rank, + o_a_proj=o_a_proj, + o_a_scale=o_a_scale, + o_a_dequant=o_a_dequant, + use_fused_output_projection=use_fused_output_projection, + ) def dspark_attention_forward( @@ -647,6 +671,13 @@ def dspark_attention_forward( freqs_cis=freqs_cis, q_b_norm_enabled=q_b_norm_enabled, persist=persist, + qkv_proj=qkv_proj, + o_a_proj=o_a_proj, + o_a_scale=o_a_scale, + o_a_dequant=o_a_dequant, + use_fused_output_projection=use_fused_output_projection, + use_fused_qkv_input_projection=use_fused_qkv_input_projection, + use_fused_q_input_projection=use_fused_q_input_projection, ) if capacities is None or valid_len is None: raise ValueError("Paged DSpark attention requires capacities and valid lengths") @@ -800,13 +831,13 @@ def dspark_attention_forward_batched( slots: torch.Tensor, valid_len: torch.Tensor | None = None, *, - wq_a: torch.Tensor, + wq_a: Optional[torch.Tensor], q_norm_w: torch.Tensor, - wq_b: torch.Tensor, - wkv: torch.Tensor, + wq_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], + wkv: Optional[torch.Tensor], kv_norm_w: torch.Tensor, - wo_a: torch.Tensor, - wo_b: torch.Tensor, + wo_a: Optional[torch.Tensor], + wo_b: torch.Tensor | Callable[[torch.Tensor], torch.Tensor], attn_sink: torch.Tensor, n_heads: int, head_dim: int, @@ -819,6 +850,13 @@ def dspark_attention_forward_batched( freqs_cis: torch.Tensor, persist: bool = False, q_b_norm_enabled: bool = True, + qkv_proj: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + o_a_proj: Optional[torch.Tensor] = None, + o_a_scale: Optional[torch.Tensor] = None, + o_a_dequant: Optional[torch.Tensor] = None, + use_fused_output_projection: Optional[bool] = None, + use_fused_qkv_input_projection: bool = False, + use_fused_q_input_projection: bool = False, ) -> torch.Tensor: """Batched, CUDA-graph-safe captured-context DSpark draft attention. @@ -869,11 +907,25 @@ def dspark_attention_forward_batched( # Keep the two GEMM outputs live until dispatch: the fused path lets # their RMSNorm/RoPE kernels store directly into attention's physical inputs. - main_kv_input = F.linear(main_x, wkv) + q_latent, block_kv_input, main_kv_input = _qkv_projections( + x, + main_x, + wq_a=wq_a, + wkv=wkv, + q_norm_w=q_norm_w, + qkv_proj=qkv_proj, + use_fused_fp8_input=use_fused_qkv_input_projection, + ) # Query: low-rank + per-head RMS + RoPE. - q = _rmsnorm_rope_batched(F.linear(x, wq_a), q_norm_w, eps, 0, blk_freqs) - q = F.linear(q, wq_b).unflatten(-1, (n_heads, head_dim)) # [G, block, h, head_dim] + q = _q_projection( + q_latent, + q_norm_w, + eps, + wq_b, + use_fused_q_input_projection, + batched_freqs_cis=blk_freqs, + ).unflatten(-1, (n_heads, head_dim)) q = _rmsnorm_rope_batched( q, None, @@ -885,8 +937,6 @@ def dspark_attention_forward_batched( apply_rmsnorm=q_b_norm_enabled, ) - block_kv_input = F.linear(x, wkv) - # ``persist=True`` writes through to the worker-owned rolling window; # single-shot callers keep the original functional clone behavior. write_target = kv_cache if persist else kv_cache.clone() @@ -969,10 +1019,17 @@ def dspark_attention_forward_batched( ) # Grouped low-rank O projection. - o = o.reshape(g, block, n_groups, -1) - wo_a_v = wo_a.view(n_groups, o_lora_rank, -1) - o = torch.einsum("bsgd,grd->bsgr", o, wo_a_v) - return F.linear(o.flatten(2), wo_b) + return _output_projection( + o, + wo_a=wo_a, + wo_b=wo_b, + n_groups=n_groups, + o_lora_rank=o_lora_rank, + o_a_proj=o_a_proj, + o_a_scale=o_a_scale, + o_a_dequant=o_a_dequant, + use_fused_output_projection=use_fused_output_projection, + ) # ---------------------------------------------------------------------------- @@ -1812,9 +1869,20 @@ def write_context_windows( freqs = self._dspark_freqs_table(main_x.device)[positions] slots = positions % win # [M] mx = main_x.unsqueeze(0) # [1, M, hidden] for the per-position RoPE layout + projected_wkv = self._project_context_wkv(mx) + head_dim = int(self._attn_params["head_dim"]) for s, stage in enumerate(self.mtp_layers): a = stage._dspark_attn - kv = _rmsnorm(F.linear(mx, a["wkv"]), a["kv_norm_w"], eps) # [1, M, head_dim] + kv = ( + projected_wkv[..., s, :] + if projected_wkv is not None + else ( + a["qkv_proj"](mx)[..., -head_dim:] + if a.get("qkv_proj") is not None + else F.linear(mx, a["wkv"]) + ) + ) + kv = _rmsnorm(kv, a["kv_norm_w"], eps) kv = _rope_last_dims(kv, rd, freqs) # [1, M, head_dim] stage_windows[s, slots] = kv[0].to(stage_windows.dtype) @@ -1862,9 +1930,20 @@ def write_context_windows_batched( mask3 = mask.unsqueeze(-1) # [G, M, 1] stage0 = self.mtp_layers[0] main_x = stage0.main_norm(stage0.main_proj(main_hidden)) # [G, M, hidden] + projected_wkv = self._project_context_wkv(main_x) + head_dim = int(self._attn_params["head_dim"]) for s, stage in enumerate(self.mtp_layers): a = stage._dspark_attn - kv = _rmsnorm(F.linear(main_x, a["wkv"]), a["kv_norm_w"], eps) # [G, M, head_dim] + kv = ( + projected_wkv[..., s, :] + if projected_wkv is not None + else ( + a["qkv_proj"](main_x)[..., -head_dim:] + if a.get("qkv_proj") is not None + else F.linear(main_x, a["wkv"]) + ) + ) + kv = _rmsnorm(kv, a["kv_norm_w"], eps) kv = _rope_last_dims_batched(kv, rd, freqs) # [G, M, head_dim] win_s = kv_windows[:, s] # [N, win, head_dim] view onto the base buffer # Read-modify-write so masked-out (g, m) entries keep their current diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b25bd5acfc7c..d60b199ffbba 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -24,9 +24,7 @@ import tensorrt_llm.bindings.executor as trtllm from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, is_sm_100f, prefer_pinned, - str_dtype_to_binding, torch_dtype_to_binding, - torch_dtype_to_str) -from tensorrt_llm.bindings import DataType + str_dtype_to_binding, torch_dtype_to_str) from tensorrt_llm.inputs.multimodal import MultimodalParams # isort: off @@ -998,7 +996,6 @@ def __init__( self._disable_overlap_scheduler = llm_args.disable_overlap_scheduler self._draft_config = draft_config self._skip_est = skip_est - self._validate_standalone_draft_cache() # Admission cap (tokens of summed context attended-KV) that the fp8 context-MLA workspace reservation # covers, computed in configure_kv_cache_capacity and carried to the KV manager so the scheduler # reads it directly instead of re-deriving it from pool layout. None until reserved (or w == 0). @@ -1250,15 +1247,11 @@ def _get_kv_size_per_token(self, model_config = self._model_engine.model.model_config use_separate_draft_kv_cache = ( self._should_create_separate_draft_kv_cache()) - draft_kwargs = {} - if self._uses_unified_standalone_draft_cache(): - draft_kwargs = self._get_unified_draft_cache_kwargs() total = self._per_manager_cache_cost( self._kv_cache_manager_cls, model_config, kv_cache_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - **draft_kwargs) + use_separate_draft_kv_cache=use_separate_draft_kv_cache) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) draft_cost = self._get_draft_cache_cost( @@ -1290,6 +1283,11 @@ def _get_draft_cache_cost( helix_cp_scale = self._mapping.cp_size def scaled(cost: CacheCost) -> CacheCost: + # Rescale only the slope: every other cost component (e.g. the + # fixed intercept) is per-request rank-local bytes and stays + # unscaled. dataclasses.replace keeps any component the rescale + # does not touch, where rebuilding the CacheCost from named + # fields would silently drop it. return dataclasses.replace(cost, slope=cost.slope * helix_cp_scale) if use_separate_draft_kv_cache: @@ -1602,8 +1600,6 @@ def _get_token_num_for_estimation(self) -> int: if spec_cfg is not None: num_extra_tokens_per_seq += spec_cfg.tokens_per_gen_step - 1 num_extra_tokens_per_seq += get_num_extra_kv_tokens(spec_cfg) - if self._uses_unified_standalone_draft_cache(): - num_extra_tokens_per_seq += self._get_draft_scratch_tokens() if self._dummy_reqs is None: self._dummy_reqs = self._create_dummy_context_requests( @@ -1648,8 +1644,6 @@ def _get_token_num_for_estimation(self) -> int: self._model_engine.max_seq_len, self._kv_cache_config, ) - if self._uses_unified_standalone_draft_cache(): - num_pool_groups += 1 num_cache_blocks *= num_pool_groups # Dummy context requests use the configured maximum beam width. Scale @@ -2091,14 +2085,12 @@ def _create_kv_cache_manager( kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( kv_cache_config) - # The target and draft layers have independent storage properties. - # spec_config also determines the target's extra-token capacity. + # When using separate draft KV cache in one-model speculative decoding, + # use layer_mask to include only target layers. The draft layers should + # only be in the separate draft KV cache manager. + # We still pass spec_config so that num_extra_kv_tokens is calculated. spec_dec_layer_mask = None - draft_cache_kwargs = (self._get_unified_draft_cache_kwargs() - if self._uses_unified_standalone_draft_cache() - else {}) - if (self._should_create_separate_draft_kv_cache() - or draft_cache_kwargs): + if self._should_create_separate_draft_kv_cache(): num_target_layers = self._model_engine.model.model_config.pretrained_config.num_hidden_layers spec_dec_layer_mask = [True] * num_target_layers @@ -2127,7 +2119,6 @@ def _create_kv_cache_manager( self._llm_args.kv_cache_config.kv_events_config, cold_page_codec_provider=cold_page_codec_provider, joint_kv_cache_reuse=self._joint_kv_cache_reuse, - **draft_cache_kwargs, ) if not self._skip_est: @@ -2145,101 +2136,6 @@ def _create_kv_cache_manager( return kv_cache_manager - def _is_standalone_dspark(self) -> bool: - spec_config = self._speculative_config - return (spec_config is not None - and spec_config.spec_dec_mode.is_dspark() - and not spec_config.draft_is_embedded_in_target - and not spec_config._use_shared_kv_cache) - - def _is_embedded_dspark(self) -> bool: - spec_config = self._speculative_config - return (spec_config is not None - and spec_config.spec_dec_mode.is_dspark() - and spec_config.draft_is_embedded_in_target) - - def _uses_unified_standalone_draft_cache(self) -> bool: - if (not (self._is_standalone_dspark() or self._is_embedded_dspark()) - or not self._is_kv_cache_manager_v2): - return False - # Disaggregation validates unified support; aggregate may use legacy state. - return (self._is_disagg - or self._unified_draft_cache_unsupported_reason() is None) - - def _validate_standalone_draft_cache(self) -> None: - """Validate manager-owned DSpark history for disaggregated serving.""" - if (not (self._is_standalone_dspark() or self._is_embedded_dspark()) - or not self._is_disagg): - return - if not self._is_kv_cache_manager_v2: - raise ValueError("DSpark disaggregation requires " - "kv_cache_config.use_kv_cache_manager_v2=True " - "on both workers.") - reason = self._unified_draft_cache_unsupported_reason() - if reason is not None: - raise ValueError(reason) - - def _unified_draft_cache_unsupported_reason(self) -> Optional[str]: - """Shared admission requirements for unified aggregate and disagg KV.""" - if (self._is_standalone_dspark() - and is_mla(self._draft_config.pretrained_config)): - return "Unified DSpark KV cache does not yet support MLA drafters." - return None - - def _get_draft_scratch_tokens(self) -> int: - return 0 if self._is_embedded_dspark( - ) else self._speculative_config.max_draft_len + 1 - - def _get_unified_draft_cache_kwargs(self) -> dict: - """Pass draft storage properties separately from its history requirements.""" - draft = self._model_engine.model.draft_model - if self._is_embedded_dspark(): - return dict( - draft_num_layers=draft.num_stages, - draft_num_kv_heads=1, - draft_head_dim=int(draft._attn_params["head_dim"]), - draft_dtype=torch_dtype_to_binding(torch.bfloat16), - draft_kv_factor=1, - draft_scratch_tokens=0, - draft_attention_backend="DSv4", - draft_window_size=int(draft._attn_params["window_size"]), - ) - config = self._draft_config.pretrained_config - num_heads = config.num_attention_heads - num_kv_heads = getattr(config, "num_key_value_heads", num_heads) - head_dim = getattr(config, "head_dim", None) - if head_dim is None: - head_dim = config.hidden_size // num_heads - draft_mapping = self._get_draft_mapping() - attention_tp_size = (1 if draft_mapping.enable_attention_dp else - draft_mapping.tp_size) - attention_backend = self._speculative_config.attention_backend - if attention_backend == "AUTO": - attention_backend = draft.dflash_attention_backend - draft_model_config = getattr(draft, "config", None) - dtype = draft.fc.weight.dtype if hasattr(draft, - "fc") else torch.bfloat16 - if dtype not in (torch.float16, torch.bfloat16): - raise ValueError( - "Standalone draft KV supports FP16 and BF16 storage") - if attention_backend == "TRTLLM": - # The TRTLLM drafter attention reads FP8 pages with FP8 queries; - # the drafter's history is stored at half the BF16 footprint. - dtype = torch.float8_e4m3fn - return dict( - draft_num_layers=config.num_hidden_layers, - draft_num_kv_heads=(num_kv_heads + attention_tp_size - 1) // - attention_tp_size, - draft_head_dim=head_dim, - draft_dtype=torch_dtype_to_binding(dtype), - draft_total_num_kv_heads=num_kv_heads, - draft_scratch_tokens=self._get_draft_scratch_tokens(), - draft_attention_backend=attention_backend, - draft_max_position_embeddings=(getattr( - draft_model_config, "max_position_embeddings", None) - if draft_model_config else None), - ) - def _should_create_separate_draft_kv_cache(self) -> bool: """ Check if we need a separate draft KV cache manager for one-model mode. @@ -2252,8 +2148,6 @@ def _should_create_separate_draft_kv_cache(self) -> bool: if self._speculative_config is None: # No drafter at all, so there is nothing to give a manager to. return False - if self._uses_unified_standalone_draft_cache(): - return False # Narrower than is_external_drafter(): PARD and DRAFT_TARGET_ONE_MODEL # never reach the arena this carve-out exists for. spec_dec_mode = self._speculative_config.spec_dec_mode @@ -3410,16 +3304,6 @@ def _create_kv_cache_manager( disable_overlap_scheduler: bool = False, cold_page_codec_provider: Optional[object] = None, kv_events_config: Optional[KVEventsConfig] = None, - draft_num_layers: int = 0, - draft_num_kv_heads: int = 0, - draft_head_dim: int = 0, - draft_dtype: DataType = DataType.BF16, - draft_kv_factor: int = 2, - draft_total_num_kv_heads: Optional[int] = None, - draft_scratch_tokens: int = 0, - draft_window_size: Optional[int] = None, - draft_attention_backend: Optional[str] = None, - draft_max_position_embeddings: Optional[int] = None, joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None) -> KVCacheManager: """ @@ -3432,9 +3316,6 @@ def _create_kv_cache_manager( "Cold-page quantization requires the resolved KV cache manager " f"to be KVCacheManagerV2; selected {kv_cache_manager_cls.__name__}") - if draft_num_layers and not issubclass(kv_cache_manager_cls, - KVCacheManagerV2): - raise ValueError("Unified draft cache layers require KVCacheManagerV2.") if (estimating_kv_cache and issubclass(kv_cache_manager_cls, KVCacheManagerV2) and kv_cache_config.pool_ratio is None @@ -3591,18 +3472,6 @@ def _create_kv_cache_manager( "Could not resolve vocab_size from the model config; " "multimodal requests will fail when block reuse is on. " "Disable block reuse to serve them.") - manager_extra_kwargs.update( - draft_num_layers=draft_num_layers, - draft_num_kv_heads=draft_num_kv_heads, - draft_head_dim=draft_head_dim, - draft_dtype=draft_dtype, - draft_kv_factor=draft_kv_factor, - draft_total_num_kv_heads=draft_total_num_kv_heads, - draft_scratch_tokens=draft_scratch_tokens, - draft_window_size=draft_window_size, - draft_attention_backend=draft_attention_backend, - draft_max_position_embeddings=draft_max_position_embeddings, - ) elif kv_events_config is not None and kv_events_config.enable_kv_cache_events: logger.warning( "kv_cache_config.kv_events_config is set but streaming KV event " diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 28ef2728ec0a..e7300783f5c3 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -18,8 +18,6 @@ ) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.distributed import Distributed -from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin -from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import ( MoeLoadBalancer, @@ -636,22 +634,6 @@ def _forward_decoder( is_dummy=is_dummy, gather_context_logits=gather_context_logits, ) - spec_worker = None - if enable_spec_decode: - speculative_model = self._model - if isinstance(speculative_model, MultimodalModelMixin): - speculative_model = speculative_model.language_model - if not isinstance(speculative_model, SpecDecOneEngineForCausalLM): - speculative_model = getattr(speculative_model, "llm", None) - if isinstance(speculative_model, SpecDecOneEngineForCausalLM): - spec_worker = speculative_model.spec_worker - if spec_worker is not None: - spec_worker.prepare_managed_draft_cache( - speculative_model.draft_model, - spec_metadata, - attn_metadata, - resource_manager, - ) inputs = prepared.kwargs runtime_draft_len = prepared.runtime_draft_len if execution_promoted_context_ids: @@ -700,18 +682,6 @@ def postprocess_inputs(model_inputs: dict[str, Any]) -> None: if captures_model_body: return outputs - if ( - spec_worker is not None - and not is_dummy - and not getattr(self._graphs.cuda_graph_runner, "is_warmup_only", False) - ): - history_update = spec_worker.snapshot_managed_draft_history() - if history_update is not None: - # Graph dictionaries are reused; each overlapped step owns - # its own readback until sampling commits the history. - outputs = dict(outputs) - outputs["draft_history_update"] = history_update - if self.forward_pass_callable is not None: self.forward_pass_callable() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index cafc0ee08105..8984856fd55e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -78,7 +78,6 @@ BatchDesc, BeamIndex, BufferConfig, - CacheDomain, CacheLevel, CacheTier, CacheTierConfig, @@ -138,7 +137,6 @@ scan_kv_buffer_for_invalid_values, ) from ..scheduler import ScheduledRequests -from ..standalone_draft_cache import StandaloneDraftHistory if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata @@ -682,7 +680,6 @@ def _estimate_cache_size_components( scratch: bool, generation_capacity_headroom: int, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, - standalone_draft_reserve: int = 0, ) -> tuple[int, int, int]: """Return context/generation bytes per token and generation bytes per request. @@ -711,7 +708,7 @@ def _estimate_cache_size_components( return ( full_attn_size + context_swa_size, full_attn_size + generation_swa_size, - generation_swa_per_request + standalone_draft_reserve * sum(layer_sizes), + generation_swa_per_request, ) @@ -1230,8 +1227,6 @@ def _update_kv_cache_draft_token_location( run_kv_cache_relocation = True if not run_kv_cache_relocation: return - if getattr(cache_manager, "_standalone_draft_layer_ids", ()): - raise ValueError("Unified standalone draft KV does not support tree-token relocation") requests = scheduled_batch.all_requests() ( accepted_draft_token_offsets, @@ -1777,19 +1772,6 @@ class KVCacheManagerV2(BaseResourceManager): # _init_mixed_precision assigns fresh per-instance values. low_precision_dtype: Optional[DataType] = None high_layers: dict[int, int] = {} - - _draft_layer_sizes: tuple[int, ...] = () - - cp_as_tp_layer_ids: frozenset[int] = frozenset() - - draft_window_size: int | None = None - - _standalone_draft_layer_ids: tuple[int, ...] = () - - _standalone_draft_reserve: int = 0 - - _draft_registered_as_model_layers: bool = False - # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True @@ -1853,39 +1835,8 @@ def __init__( joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, - draft_num_layers: int = 0, - draft_num_kv_heads: int = 0, - draft_head_dim: int = 0, - draft_dtype: DataType = DataType.BF16, - draft_kv_factor: int = 2, - draft_total_num_kv_heads: int | None = None, - draft_scratch_tokens: int = 0, - draft_window_size: int | None = None, - draft_attention_backend: str | None = None, - draft_max_position_embeddings: int | None = None, **kwargs, ) -> None: - self._draft_layer_sizes = _get_draft_layer_sizes( - draft_num_layers=draft_num_layers, - draft_num_kv_heads=draft_num_kv_heads, - draft_head_dim=draft_head_dim, - draft_dtype=draft_dtype, - draft_kv_factor=draft_kv_factor, - ) - if draft_total_num_kv_heads is not None and draft_total_num_kv_heads < draft_num_kv_heads: - raise ValueError("Total draft KV heads must cover the rank-local heads") - self._draft_total_num_kv_heads = draft_total_num_kv_heads - if draft_scratch_tokens < 0: - raise ValueError("Standalone draft scratch capacity must be nonnegative") - if draft_window_size is not None and draft_window_size <= 0: - raise ValueError("Draft history window must be positive") - self.draft_window_size = draft_window_size - self.draft_attention_backend = draft_attention_backend - self.draft_max_position_embeddings = draft_max_position_embeddings - self._standalone_draft_layer_ids: tuple[int, ...] = () - self.draft_history: dict[int, StandaloneDraftHistory] = {} - self._draft_dummy_request_ids: set[int] = set() - self._standalone_draft_reserve = draft_scratch_tokens self.mapping = mapping self.dtype = dtype self._declared_target_layer_count = num_layers @@ -1906,8 +1857,6 @@ def __init__( self._init_mixed_precision( kv_cache_config, kwargs.get("pretrained_config"), max_seq_len, dtype ) - self._declared_target_layer_count = num_layers - self.layer_specs = tuple(layer_specs) if layer_specs is not None else None self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager @@ -2083,7 +2032,6 @@ def __init__( self._kv_reserve_draft_tokens, self._generation_kv_capacity_headroom = ( _get_generation_kv_capacity(spec_config, is_draft=self.is_draft) ) - self._generation_kv_capacity_headroom += self._standalone_draft_reserve self.event_buffer_max_size = kv_cache_config.event_buffer_max_size self.enable_stats = enable_stats @@ -2234,9 +2182,6 @@ def append_to_kv_heads_per_layer( f"unique values={set(self.head_dim_per_layer)}" ) - self.dtype_per_layer = [self.dtype] * self.num_local_layers - self.kv_factor_per_layer = [self.kv_factor] * self.num_local_layers - self.is_vswa = uses_vswa_kv_cache_layout(self.max_attention_window_vec) max_util_for_resume = kv_cache_config.max_util_for_resume @@ -2376,14 +2321,6 @@ def append_to_kv_heads_per_layer( cache_tiers=cache_tiers, ) config = self._build_cache_config(config) - config = self._append_standalone_draft_layers( - config, - num_kv_heads=draft_num_kv_heads, - head_dim=draft_head_dim, - dtype=draft_dtype, - kv_factor=draft_kv_factor, - ) - config = self._configure_unified_draft_pool_ratio(config) config = self._remove_zero_size_buffers(config) config = self._add_mixed_precision_layers(config) has_host_cache_tier = any( @@ -2653,10 +2590,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) if self._sparse_layer_group_ids: self.sparse_metadata_batch = Batch( - self.impl, - index_mapper_capacity, - self.max_blocks_per_seq, - self.max_beam_width, + self.impl, index_mapper_capacity, self.max_blocks_per_seq, self.max_beam_width ) self._log_kv_cache_pool_lifecycle_mapping() @@ -2955,13 +2889,10 @@ def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: layer_id = self._pool_layer_ids_by_role[(pool_id, Role.KEY)] roles = self._get_layer_roles_by_local_id(int(layer_id)) return Role.KEY, Role.VALUE if Role.VALUE in roles else None - layer_id = int(self.impl.layer_grouping[pool_id][0]) - role_b = Role.VALUE if self._get_layer_kv_factor_by_local_id(layer_id) == 2 else None + role_b = None if self.kv_cache_type == CacheTypeCpp.SELFKONLY else Role.VALUE return Role.KEY, role_b def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: - if self._is_standalone_draft_layer(local_layer_id): - return self.kv_factor_per_layer[local_layer_id] if self.layer_specs is not None: return self.layer_specs[local_layer_id].kv_factor return self.kv_factor @@ -2974,8 +2905,6 @@ def _get_layer_page_table_roles( return Role.KEY, (Role.VALUE if self.layer_specs[local_layer_id].kv_factor == 2 else None) def _get_layer_dtype_by_local_id(self, local_layer_id: int) -> DataType: - if self._is_standalone_draft_layer(local_layer_id): - return self.dtype_per_layer[local_layer_id] if self.layer_specs is not None: return self.layer_specs[local_layer_id].dtype return self.dtype @@ -3270,17 +3199,7 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Opti layer_sizes.append( self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=Role.ALL) ) - if self._is_standalone_draft_layer(local_layer_idx): - attention_windows.append(self.draft_retention_window_size) - else: - attention_windows.append(self.max_attention_window_vec[local_layer_idx]) - # Quota sizing precedes specialized target-layer construction. Include - # draft storage then; after append it is already in the per-layer arrays. - if self._draft_layer_sizes and not self._draft_registered_as_model_layers: - layer_sizes.extend(self._draft_layer_sizes) - attention_windows.extend( - [self.draft_retention_window_size] * len(self._draft_layer_sizes) - ) + attention_windows.append(self.max_attention_window_vec[local_layer_idx]) return layer_sizes, attention_windows def _get_generation_request_capacity(self) -> int: @@ -3299,24 +3218,19 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: return tokens def _get_max_tokens_from_quota_impl(self, quota: int) -> float: - if self._draft_layer_sizes: - (context_size_per_token, generation_size_per_token, generation_swa_size_per_request) = ( - self._get_runtime_cache_size_components() - ) - else: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - ( - context_size_per_token, - generation_size_per_token, - generation_swa_size_per_request, - ) = _estimate_cache_size_components( - layer_sizes, - attention_windows, - self.tokens_per_block, - scratch=self.enable_swa_scratch_reuse, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - layer_specs=self.layer_specs, - ) + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + ( + context_size_per_token, + generation_size_per_token, + generation_swa_size_per_request, + ) = _estimate_cache_size_components( + layer_sizes, + attention_windows, + self.tokens_per_block, + scratch=self.enable_swa_scratch_reuse, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, + ) size_per_batch = self._get_generation_request_capacity() * generation_swa_size_per_request if quota < size_per_batch: return 0 @@ -3341,24 +3255,19 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: return self._get_quota_from_max_tokens_impl(max_tokens) def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: - if self._draft_layer_sizes: - (context_size_per_token, generation_size_per_token, generation_swa_size_per_request) = ( - self._get_runtime_cache_size_components() - ) - else: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - ( - context_size_per_token, - generation_size_per_token, - generation_swa_size_per_request, - ) = _estimate_cache_size_components( - layer_sizes, - attention_windows, - self.tokens_per_block, - scratch=self.enable_swa_scratch_reuse, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - layer_specs=self.layer_specs, - ) + layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() + ( + context_size_per_token, + generation_size_per_token, + generation_swa_size_per_request, + ) = _estimate_cache_size_components( + layer_sizes, + attention_windows, + self.tokens_per_block, + scratch=self.enable_swa_scratch_reuse, + generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, + ) context_tokens = min(max_tokens, self.max_num_tokens) generation_tokens = max_tokens - context_tokens return int( @@ -3524,12 +3433,6 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: ] for role_idx, role in enumerate(roles): converter = self.impl.get_page_index_converter(layer_id, role) - if converter.expansion != 1 and self._is_standalone_draft_layer(layer_id): - pool_ids[local_layer_idx, role_idx] = pool_id - scales[local_layer_idx, role_idx] = 1 - layer_offsets[local_layer_idx, role_idx] = 0 - scratch_pages[local_layer_idx, role_idx] = 0 - continue if converter.expansion != 1: raise NotImplementedError( "SWA scratch block-table conversion does not support " @@ -3714,8 +3617,6 @@ def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: """Return the primary cache roles allocated for one local layer.""" if self.layer_specs is not None: return list(self.layer_specs[local_layer_idx].buffer_roles) - if self._is_standalone_draft_layer(local_layer_idx): - return [Role.KEY, Role.VALUE][: self.kv_factor_per_layer[local_layer_idx]] roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: roles.append(Role.VALUE) @@ -3745,11 +3646,9 @@ def _build_base_config( """ scratch_reuse_config = None if self.enable_swa_scratch_reuse: - # Unwritten speculative and draft reserves must stay outside the - # scratch range when a context chunk advances its history. - scratch_reuse_config = SwaScratchReuseConfig( - max_rewind_len=self.num_extra_kv_tokens + self._standalone_draft_reserve - ) + # Context requests allocate num_extra_kv_tokens for spec decoding. + # They should not count toward the scratch range. + scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens) # Beam search only replicates blocks from the prompt tail onward, so # KVCacheDesc splits each request into a shared prefix (prompt_length) @@ -4344,8 +4243,6 @@ def target_layer_ids(self) -> tuple[int, ...]: @property def draft_layer_ids(self) -> tuple[int, ...]: - if getattr(self, "_standalone_draft_layer_ids", ()): - return self._standalone_draft_layer_ids return tuple( layer_id for layer_id in self.pp_layers @@ -4354,7 +4251,7 @@ def draft_layer_ids(self) -> tuple[int, ...]: def get_draft_kv_cache_view(self) -> Optional["KVCacheManagerV2"]: """Return this manager when explicit draft layers share its request timeline.""" - if not self.draft_layer_ids: + if self.layer_specs is None or not self.draft_layer_ids: return None return self @@ -4642,11 +4539,7 @@ def get_index_k_buffer( return full_view[:, 0] def get_num_available_tokens( - self, - *, - token_num_upper_bound: int, - batch_size: int = 1, - max_num_draft_tokens: int = 0, + self, *, token_num_upper_bound: int, batch_size: int = 1, max_num_draft_tokens: int = 0 ) -> int: """Clamp ``token_num_upper_bound`` to the allocatable token capacity. @@ -4659,9 +4552,7 @@ def get_num_available_tokens( replicates all tokens on every rank, so both bounds constrain the same request-length variable. """ - extra_tokens = ( - self.num_extra_kv_tokens + max_num_draft_tokens + self._standalone_draft_reserve - ) + extra_tokens = self.num_extra_kv_tokens + max_num_draft_tokens # Token num upper bound is the maximum number of tokens that can be allocated in the kv cache manager. # We need to add extra tokens to the token num upper bound to account for the extra tokens. clamped = ( @@ -4703,10 +4594,10 @@ def get_num_free_blocks(self) -> int: ) # Legacy managers estimate primary-pool capacity from page indices. - max_num_blocks = max( + max_num_pages = max( [ - self._get_attention_pool_num_blocks(local_layer_idx) - for local_layer_idx in range(self.num_local_layers) + self.impl.get_page_index_upper_bound(layer_id, Role.KEY) + for layer_id in typed_range(LayerId(self.num_local_layers)) ] ) # Subtract the pages the reserved sequences hold, not the count of @@ -4714,7 +4605,7 @@ def get_num_free_blocks(self) -> int: # page (the guard is today). Summing ``num_blocks`` keeps this a page # count if a reservation ever spans more than one. reserved_pages = sum(int(self.kv_cache_map[req_id].num_blocks) for req_id in reserved) - return max_num_blocks - reserved_pages + return max_num_pages // self.kv_factor - reserved_pages def commit_scheduled_kv_cache_stats(self, scheduled_batch: ScheduledRequests) -> None: if self.is_draft or (not self.enable_stats and not self._request_stats_enabled_ids): @@ -5018,8 +4909,6 @@ def _context_reuse_tokens( """Radix-tree evidence for one context prefix. Raw prompt, so block hashes stay comparable with KV cache events, the cache-aware router and disagg peers. """ - if self._draft_context_exceeds_limit(req): - return [] all_tokens = self._reuse_token_source(req) # The target must recompute the last prompt token, but a one-model draft # also reads D prompt tokens ahead. Include those tokens as match evidence; @@ -5240,7 +5129,7 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) return None kv_cache.cuda_stream = self._cache_stream(req.py_request_id) - if not self.enable_block_reuse or self._draft_context_exceeds_limit(req): + if not self.enable_block_reuse: kv_cache.stop_committing() else: self._record_branch_snapshot_point(req, kv_cache, num_lookup_tokens) @@ -5269,19 +5158,6 @@ def prepare_context_cache(self, req: LlmRequest, reuse_limit: int | None = None) kv_cache.enable_swa_scratch_reuse = False if not self._resume_context_cache(req, kv_cache): return None - if ( - self._standalone_draft_layer_ids - and not self._draft_history_uses_kv_history() - and req.py_request_id not in self.draft_history - ): - # The reuse match covers every cache domain. Initialize once, - # after resume succeeds, so overlap/chunk retries cannot rewind - # draft history that has already advanced on the device. - reused = kv_cache.num_committed_tokens - valid_length = reused - if self.draft_window_size is not None: - valid_length = min(valid_length, self.draft_window_size) - self.set_draft_history(req.py_request_id, valid_length, reused) return kv_cache.num_committed_tokens # Subsequent chunk: cache must exist from first chunk. It may be @@ -5356,12 +5232,7 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: if kv_cache is None: return False - target = ( - req.context_current_position - + num_tokens - + self.num_extra_kv_tokens - + self._standalone_draft_reserve - ) + target = req.context_current_position + num_tokens + self.num_extra_kv_tokens capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -5438,12 +5309,7 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: else (req.total_input_len_cp if self._has_cp_helix else req.prompt_len) ) prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - target = ( - prompt_len - + get_draft_token_length(req) - + self.num_extra_kv_tokens - + self._standalone_draft_reserve - ) + target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -6961,9 +6827,6 @@ def add_dummy_requests( on this manager and the supplied draft manager when it is no longer needed, including when subsequent warmup or graph capture fails. """ - # A unified manager creates target and draft pools in one cache. - if draft_kv_cache_manager is self: - draft_kv_cache_manager = None _kv_draft = ( kv_reserve_draft_tokens if kv_reserve_draft_tokens is not None else max_num_draft_tokens ) @@ -7129,7 +6992,6 @@ def try_commit_blocks(self, request: LlmRequest) -> None: self.enable_block_reuse and self._can_publish_block_reuse and not request.is_dummy_request - and not self._draft_context_exceeds_limit(request) ) if not should_block_reuse: return @@ -7262,9 +7124,6 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._disagg_receive_ready.pop(request.py_request_id, None) self._on_allocation_stream.discard(request.py_request_id) self._allocation_ready.pop(request.py_request_id, None) - if self._standalone_draft_layer_ids: - self.draft_history.pop(request.py_request_id, None) - self._draft_dummy_request_ids.discard(request.py_request_id) kv_cache = self.kv_cache_map.pop(request.py_request_id, None) if kv_cache is None: self.impl.clear_stats_excluded(request.py_request_id) @@ -7445,9 +7304,9 @@ def get_cache_bytes_per_token(self) -> int: data_roles.append(Role.VALUE_BLOCK_SCALE) return sum( - self.get_layer_bytes_per_token(local_id, Role.ALL) - * (self._helix_cp_size if local_id in self.cp_as_tp_layer_ids else 1) - for local_id in range(len(self.dtype_per_layer)) + self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=data_role) + for local_layer_idx in range(self.num_local_layers) + for data_role in data_roles ) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): @@ -7459,37 +7318,24 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): if self._is_mixed_precision_layer(local_layer_idx): return self._mixed_precision_layer_bytes_per_token(local_layer_idx, data_role) - if self.layer_specs is not None: - layer = self.layer_specs[local_layer_idx] - if data_role == Role.ALL: - return layer.bytes_per_token() - return layer.bytes_per_token(data_role) - - if data_role not in (Role.ALL, *self._get_buffer_roles_for_layer(local_layer_idx)): - return 0 - if self._get_layer_dtype_by_local_id(local_layer_idx) not in ( + if self.dtype not in ( DataType.FP8, DataType.HALF, DataType.BF16, DataType.FLOAT, DataType.NVFP4, ): - raise ValueError(f"Cannot support {self.dtype_per_layer[local_layer_idx]} KV cache.") + raise ValueError(f"Cannot support {self.dtype} KV cache.") if data_role == Role.ALL: - kv_factor = self._get_layer_kv_factor_by_local_id(local_layer_idx) - elif data_role in [ - Role.KEY, - Role.VALUE, - Role.KEY_BLOCK_SCALE, - Role.VALUE_BLOCK_SCALE, - ]: + kv_factor = self.kv_factor + elif data_role in [Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: if data_role in [Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE]: - assert self.dtype_per_layer[local_layer_idx] == DataType.NVFP4, ( + assert self.dtype == DataType.NVFP4, ( "NVFP4 is the only supported dtype for block quant data roles" ) if data_role == Role.VALUE: - assert self._get_layer_kv_factor_by_local_id(local_layer_idx) == 2, ( + assert self.kv_cache_type != CacheTypeCpp.SELFKONLY, ( "VALUE data role is not supported for SELFKONLY cache type" ) kv_factor = 1 @@ -7502,16 +7348,14 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): * self.head_dim_per_layer[local_layer_idx] ) - cache_size_bytes_per_token = get_size_in_bytes( - cache_size_per_token, self.dtype_per_layer[local_layer_idx] - ) + cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, self.dtype) if data_role in [Role.KEY, Role.VALUE]: return cache_size_bytes_per_token quant_size_per_token = 0 - if self.dtype_per_layer[local_layer_idx] == DataType.NVFP4: + if self.dtype == DataType.NVFP4: quant_size_per_token = self.calculate_scaling_factor_size_bytes( cache_size_per_token, quant_vector_size=16, @@ -7604,9 +7448,6 @@ def shutdown(self): for kv_cache in self.kv_cache_map.values(): kv_cache.close() self.kv_cache_map.clear() - if self._standalone_draft_layer_ids: - self.draft_history.clear() - self._draft_dummy_request_ids.clear() self._disagg_receive_ready.clear() self._on_allocation_stream.clear() self._allocation_ready.clear() @@ -7661,8 +7502,6 @@ def get_cache_size_per_token( spec_config=None, is_draft: bool = False, layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, - draft_scratch_tokens: int = 0, - draft_window_size: int | None = None, **kwargs, ): if layer_specs is not None: @@ -7697,32 +7536,6 @@ def get_cache_size_per_token( backoff, max_seq_len, ) - draft_layer_sizes = _get_draft_layer_sizes(**kwargs) - if draft_layer_sizes: - _, headroom = _get_generation_kv_capacity(spec_config, is_draft=is_draft) - headroom += draft_scratch_tokens - context, generation, fixed = _estimate_cache_size_components( - layer_sizes, - attention_windows, - tokens_per_block, - scratch=bool(kv_cache_config and kv_cache_config.enable_swa_scratch_reuse), - generation_capacity_headroom=headroom, - standalone_draft_reserve=draft_scratch_tokens, - ) - draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( - draft_layer_sizes, - tokens_per_block, - scratch_tokens=draft_scratch_tokens, - window_size=draft_window_size, - generation_capacity_headroom=headroom, - helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, - ) - return ( - generation + draft_generation, - (fixed + draft_fixed) * int(max_batch_size or 0) - + (context + draft_context - generation - draft_generation) - * int(max_num_tokens or 0), - ) _, generation_capacity_headroom = _get_generation_kv_capacity( spec_config, is_draft=is_draft ) @@ -8096,7 +7909,6 @@ def update_context_resources(self, scheduled_batch: ScheduledRequests): self.enable_block_reuse and self._can_publish_block_reuse and not req.is_dummy_request - and not self._draft_context_exceeds_limit(req) ) is_all_reusable = self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE # A chunk that ends at a branch point commits there. History has @@ -8191,16 +8003,6 @@ def update_resources( if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) else kv_cache.capacity - rewind_len ) - embedded_draft = KVCacheManagerV2._draft_history_uses_kv_history(self) - if self._standalone_draft_layer_ids and not embedded_draft and new_capacity is not None: - draft_history = self.get_draft_history(req.py_request_id) - if draft_history is not None: - # Rejected target verification slots do not revoke draft - # feature history. Retain it and its next-forward scratch. - new_capacity = max( - new_capacity, - draft_history.position + self._standalone_draft_reserve, - ) history_length = ( None # Reuse (history's consumer) is disabled under helix, and @@ -8208,14 +8010,6 @@ def update_resources( if self.kv_compression_manages_history or self._has_cp_helix else req.max_beam_num_tokens - 1 ) - if embedded_draft and self._has_cp_helix: - # HELIX stores rank-local prompt tokens but generated tokens - # are replicated. The cache ledger uses global positions. - history_length = ( - req.total_input_len_cp + req.max_beam_num_tokens - req.py_prompt_len - 1 - ) - if embedded_draft and new_capacity is not None and history_length is not None: - new_capacity = max(new_capacity, history_length) success = kv_cache.resize(new_capacity, history_length) if not success: raise ValueError( @@ -8386,8 +8180,6 @@ def _create_kv_cache( kv_cache.stop_committing() self.impl.mark_stats_excluded(request_id) kv_cache.discard_pending_stats() - if self._standalone_draft_layer_ids: - self._draft_dummy_request_ids.add(request_id) index = self.index_mapper.add_new_sequence(request_id) if self.sparse_metadata_batch is not None: self.sparse_metadata_batch.add(kv_cache, index) @@ -8489,499 +8281,3 @@ def reset_reuse_state(self): self.conversation_manager.clear() if self._sharing is not None: self._sharing._on_reset() - - def _get_runtime_cache_size_components(self) -> tuple[int, int, int]: - layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() - num_target_layers = len(layer_sizes) - (len(self._draft_layer_sizes)) - context_size, generation_size, fixed_size = _estimate_cache_size_components( - layer_sizes[:num_target_layers], - attention_windows[:num_target_layers], - self.tokens_per_block, - scratch=self.enable_swa_scratch_reuse, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - standalone_draft_reserve=self._standalone_draft_reserve, - ) - if self._draft_layer_sizes: - draft_context, draft_generation, draft_fixed = _estimate_draft_cache_size_components( - self._draft_layer_sizes, - self.tokens_per_block, - scratch_tokens=self._standalone_draft_reserve, - window_size=self.draft_window_size, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - helix_cp_size=self._helix_cp_size, - ) - context_size += draft_context - generation_size += draft_generation - fixed_size += draft_fixed - return context_size, generation_size, fixed_size - - def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: - """Draft bytes within the unified manager's usable token capacity.""" - assert self._draft_layer_sizes - if self._has_cp_helix: - # Each target ledger block retains all of its draft subpages. - max_tokens = ( - math.ceil(max_tokens / self._ledger_tokens_per_block) - * self._ledger_tokens_per_block - ) - context_size, generation_size, fixed_size = _estimate_draft_cache_size_components( - self._draft_layer_sizes, - self.tokens_per_block, - scratch_tokens=self._standalone_draft_reserve, - window_size=self.draft_window_size, - generation_capacity_headroom=self._generation_kv_capacity_headroom, - ) - context_tokens = ( - max_tokens if self.max_num_tokens is None else min(max_tokens, self.max_num_tokens) - ) - return math.ceil( - context_tokens * context_size - + (max_tokens - context_tokens) * generation_size - + self._get_generation_request_capacity() * fixed_size - ) - - def _append_standalone_draft_layers( - self, - config: KVCacheManagerConfigPy, - *, - num_kv_heads: int, - head_dim: int, - dtype: DataType, - kv_factor: int, - register_model_layers: bool = True, - ) -> KVCacheManagerConfigPy: - if not self._draft_layer_sizes: - return config - first_global_id = max(self.pp_layers, default=-1) + 1 - first_local_id = len(config.layers) - # FP4 MLA adds native HP-tail layers outside the model-layer arrays. - register_model_layers = register_model_layers and first_local_id == len(self.pp_layers) - self._draft_registered_as_model_layers = register_model_layers - # Specialized managers can replace model layers with multiple native layers. - padding = first_local_id - len(self.dtype_per_layer) - assert padding >= 0 - num_draft_layers = len(self._draft_layer_sizes) - self.dtype_per_layer.extend([self.dtype] * padding + [dtype] * num_draft_layers) - self.kv_factor_per_layer.extend([self.kv_factor] * padding + [kv_factor] * num_draft_layers) - self.num_kv_heads_per_layer.extend([0] * padding + [num_kv_heads] * num_draft_layers) - self.head_dim_per_layer.extend([0] * padding + [head_dim] * num_draft_layers) - self.cp_as_tp_layer_ids = frozenset( - range(first_local_id, first_local_id + num_draft_layers) - ) - self._standalone_draft_layer_ids = tuple( - range(first_global_id, first_global_id + len(self._draft_layer_sizes)) - ) - layers = list(config.layers) - for offset, global_id in enumerate(self._standalone_draft_layer_ids): - local_id = first_local_id + offset - layers.append( - AttentionLayerConfig( - layer_id=LayerId(local_id), - buffers=[ - BufferConfig( - role=role, - size=KVCacheManagerV2.get_layer_bytes_per_token(self, local_id, role) - * self.tokens_per_block, - tokens_per_block_override=( - self.tokens_per_block if self._helix_cp_size > 1 else None - ), - ) - for role in self._get_buffer_roles_for_layer(local_id) - ], - sliding_window_size=self.draft_retention_window_size, - cache_domain=CacheDomain.STANDALONE_DRAFT, - ) - ) - self.layer_offsets[global_id] = local_id - if register_model_layers: - self.pp_layers.append(global_id) - self.total_num_kv_heads_per_layer.append( - self._draft_total_num_kv_heads or num_kv_heads - ) - if register_model_layers: - self.num_local_layers = len(self.pp_layers) - self.num_layers += len(self._draft_layer_sizes) - - def reserve_scratch(batch: BatchDesc) -> BatchDesc: - return BatchDesc( - [ - KVCacheDesc( - capacity=desc.capacity - + (self._standalone_draft_reserve if desc.capacity else 0), - history_length=desc.history_length, - ) - for desc in batch.kv_caches - ], - system_prompt_length=batch.system_prompt_length, - ) - - return replace( - config, - layers=layers, - constraints=[reserve_scratch(batch) for batch in config.constraints], - typical_step=reserve_scratch(config.typical_step) - if config.typical_step is not None - else None, - ) - - def _configure_unified_draft_pool_ratio( - self, config: KVCacheManagerConfigPy - ) -> KVCacheManagerConfigPy: - """Preserve target ratios within their share of the unified byte quota.""" - if not self._standalone_draft_layer_ids: - return config - target_ratios = config.initial_pool_ratio - if target_ratios is None: - return config - - # Native lifecycle equality includes window, rounded sink blocks, sparsity, - # and cache domain. All SSM layers share one lifecycle, represented by None. - target_life_cycles: list[AttnLifeCycle | None] = [] - for layer in config.layers: - if isinstance(layer, AttentionLayerConfig): - if layer.cache_domain == CacheDomain.STANDALONE_DRAFT: - continue - life_cycle = AttnLifeCycle.make( - layer.sliding_window_size, - layer.num_sink_tokens, - config.tokens_per_block, - cache_domain=layer.cache_domain, - ) - else: - life_cycle = None - if not any( - type(existing) is type(life_cycle) and existing == life_cycle - for existing in target_life_cycles - ): - target_life_cycles.append(life_cycle) - - if len(target_ratios) != len(target_life_cycles): - raise ValueError( - "kv_cache_config.pool_ratio must have one entry per target layer group " - f"({len(target_life_cycles)}), got {len(target_ratios)}" - ) - # Use the same model-specific capacity accounting as quota selection; - # specialized targets can include compressed KV or recurrent snapshots. - usable_quota = int(config.cache_tiers[GPU_LEVEL].quota * config.max_util_for_resume) - max_tokens = self._get_max_tokens_from_quota(usable_quota) - if math.isinf(max_tokens): - max_tokens = self.max_num_tokens or self.max_seq_len - if usable_quota <= 0 or self._get_quota_from_max_tokens(max_tokens) > usable_quota: - raise ValueError("Unified KV cache quota cannot hold the target and draft reserves") - draft_ratio = self._get_draft_quota_from_max_tokens(max_tokens) / usable_quota - if not 0 < draft_ratio < 1: - raise ValueError("Unified KV cache quota must leave room for target and draft pools") - - # Appended draft layers share one lifecycle, after all target groups. - return replace( - config, - initial_pool_ratio=[ratio * (1 - draft_ratio) for ratio in target_ratios] - + [draft_ratio], - ) - - def _get_layer_buffer_view(self, local_layer_idx: int, kv_layout: str) -> torch.Tensor: - return self._get_kv_buffer_view( - local_layer_idx, - dtype=self.dtype_per_layer[local_layer_idx], - kv_factor=self.kv_factor_per_layer[local_layer_idx], - num_kv_heads=self.num_kv_heads_per_layer[local_layer_idx], - head_dim=self.head_dim_per_layer[local_layer_idx], - kv_layout=kv_layout, - page_expansion=self._helix_cp_size if local_layer_idx in self.cp_as_tp_layer_ids else 1, - ) - - def _get_kv_buffer_view( - self, - layer_id: int, - *, - dtype: DataType | torch.dtype, - kv_factor: int, - num_kv_heads: int, - head_dim: int, - kv_layout: str, - page_expansion: int = 1, - ) -> torch.Tensor: - """View shared pool pages for a manager-local layer and its KV geometry.""" - if kv_layout not in ("NHD", "HND"): - raise ValueError(f"Unsupported KV layout: {kv_layout}") - addr_key = self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED) - if kv_factor == 2: - addr_value = self.impl.get_mem_pool_base_address( - layer_id, Role.VALUE, PageIndexMode.SHARED - ) - page_size_key = self.impl.get_page_stride(layer_id, Role.KEY) - page_size_value = self.impl.get_page_stride(layer_id, Role.VALUE) - if ( - addr_key + page_size_key * page_expansion != addr_value - or page_size_key != page_size_value - ): - raise ValueError("K/V buffers must have adjacent equal-sized pages") - - element_per_container = 1 - if dtype == DataType.NVFP4: - element_per_container = 2 - dtype = torch.int8 - - dimensions = ( - [self.tokens_per_block, num_kv_heads] - if kv_layout == "NHD" - else [num_kv_heads, self.tokens_per_block] - ) - head_dim //= element_per_container - page_upper = self.impl.get_page_index_upper_bound(layer_id, Role.KEY) - num_pages = page_upper // kv_factor - strides = None - if page_expansion > 1: - # Each ledger slot contains all of its K subpages, then all V - # subpages. Expanded KEY indices skip the V spans between slots. - num_pages = page_upper - (kv_factor - 1) * page_expansion - page_elements = self.tokens_per_block * num_kv_heads * head_dim - strides = [ - page_elements, - page_expansion * page_elements, - dimensions[1] * head_dim, - head_dim, - 1, - ] - return convert_to_torch_tensor( - TensorWrapper( - addr_key, - dtype, - [ - num_pages, - kv_factor, - *dimensions, - head_dim, - ], - strides=strides, - ) - ) - - def _is_standalone_draft_layer(self, local_layer_idx: int) -> bool: - return local_layer_idx in self.cp_as_tp_layer_ids - - def get_draft_buffers(self, local_layer_idx: int, kv_layout: str = "HND") -> torch.Tensor: - """View one draft layer using the same storage contract as every other layer.""" - if not 0 <= local_layer_idx < len(self._standalone_draft_layer_ids): - raise ValueError("No standalone draft cache layer at this index") - layer_idx = self._standalone_draft_layer_ids[local_layer_idx] - return self._get_layer_buffer_view(self.layer_offsets[layer_idx], kv_layout) - - @property - def draft_max_blocks_per_seq(self) -> int: - """Maximum physical draft pages, including every Helix ledger subpage.""" - return self.max_blocks_per_seq * self._helix_cp_size - - def get_draft_num_blocks(self, request_id: int) -> int: - """Allocated draft pages covering the request's global token positions.""" - return self.kv_cache_map[request_id].num_blocks * self._helix_cp_size - - @property - def draft_retention_window_size(self) -> int | None: - return self.draft_window_size + 1 if self.draft_window_size is not None else None - - def get_draft_transfer_identity(self) -> dict: - local_layer_idx = self.layer_offsets[self._standalone_draft_layer_ids[0]] - return { - "num_layers": len(self._standalone_draft_layer_ids), - "num_kv_heads": self._draft_total_num_kv_heads - or self.num_kv_heads_per_layer[local_layer_idx], - "head_dim": self.head_dim_per_layer[local_layer_idx], - "dtype": str(binding_to_torch_dtype(self.dtype_per_layer[local_layer_idx])), - "attention_backend": ( - "VANILLA" - if self.draft_attention_backend == "TRTLLM" - else self.draft_attention_backend - ), - "kv_factor": self.kv_factor_per_layer[local_layer_idx], - "window_size": self.draft_window_size, - } - - def _draft_history_uses_kv_history(self) -> bool: - """Embedded DSpark writes the complete live suffix of target history.""" - return getattr(self, "draft_attention_backend", None) == "DSv4" - - def get_draft_block_indices( - self, - request_ids: List[int], - histories: Optional[Sequence[StandaloneDraftHistory]] = None, - ) -> List[List[int]]: - """Validated rank-local mappings for allocated draft pages, without padding.""" - for request_id in request_ids: - cache = self.kv_cache_map.get(request_id) - if cache is None or not cache.is_active: - raise ValueError(f"Standalone draft request {request_id} has no active cache") - expanded = self._helix_cp_size > 1 - layer_idx = self._standalone_draft_layer_ids[0] - batch_indices = self.get_batch_cache_indices(request_ids, layer_idx, raw_indices=expanded) - if expanded: - converter = self.impl.get_page_index_converter(self.layer_offsets[layer_idx], Role.KEY) - batch_indices = [converter(indices, PageIndexMode.SHARED) for indices in batch_indices] - for row, indices in enumerate(batch_indices): - cache = self.kv_cache_map[request_ids[row]] - if ( - len(indices) != self.get_draft_num_blocks(request_ids[row]) - or len(indices) > self.draft_max_blocks_per_seq - ): - raise ValueError("Standalone draft cache has an incomplete or oversized page table") - first_required_block = 0 - if self.draft_window_size is not None: - history = None - if not self._draft_history_uses_kv_history(): - history = ( - histories[row] - if histories is not None - else self.get_draft_history(request_ids[row]) - ) - history_start = ( - history.position - history.valid_length - if history is not None - else max(0, cache.history_length - self.draft_window_size) - ) - first_required_block = history_start // self.tokens_per_block - if any(index == BAD_PAGE_INDEX for index in indices[first_required_block:]): - raise ValueError( - "Draft cache contains missing pages in its required history or scratch" - ) - return batch_indices - - def get_draft_block_table( - self, - request_ids: List[int], - histories: Optional[Sequence[StandaloneDraftHistory]] = None, - ) -> torch.Tensor: - """Current rank-local page mappings; unused tail entries point to page zero.""" - batch_indices = self.get_draft_block_indices(request_ids, histories) - table = torch.zeros((len(request_ids), self.draft_max_blocks_per_seq), dtype=torch.int32) - for row, indices in enumerate(batch_indices): - table[row, : len(indices)] = torch.tensor(indices, dtype=torch.int32) - return table - - def get_draft_history(self, request_id: int) -> Optional[StandaloneDraftHistory]: - return self.draft_history.get(request_id) - - def set_draft_history(self, request_id: int, valid_length: int, position: int) -> None: - cache = self.kv_cache_map.get(request_id) - if cache is None or not cache.is_active: - raise ValueError( - f"Standalone draft request {request_id} has no active cache allocation" - ) - history = StandaloneDraftHistory(valid_length, position) - if history.position > cache.capacity: - raise ValueError("Standalone draft history exceeds allocated capacity") - if self.draft_window_size is not None and history.valid_length > self.draft_window_size: - raise ValueError("Draft history exceeds its retention window") - self.draft_history[request_id] = history - - def export_draft_history(self, request_id: int) -> Optional[dict]: - if not self._standalone_draft_layer_ids: - return None - if self._draft_history_uses_kv_history(): - cache = self.kv_cache_map[request_id] - if not cache.is_active: - raise ValueError("Embedded DSpark transfer requires an active cache") - return { - "valid_length": min(cache.history_length, self.draft_window_size), - "position": cache.history_length, - "layout": self.get_draft_transfer_identity(), - } - history = self.get_draft_history(request_id) - if history is None: - raise ValueError( - f"Standalone draft request {request_id} has no valid history to transfer" - ) - return { - "valid_length": history.valid_length, - "position": history.position, - "layout": self.get_draft_transfer_identity(), - } - - def restore_draft_history(self, request_id: int, metadata: dict) -> None: - if ( - not self._standalone_draft_layer_ids - or metadata.get("layout") != self.get_draft_transfer_identity() - ): - raise ValueError("Standalone draft transfer layout does not match the receiving worker") - valid_length = metadata.get("valid_length") - position = metadata.get("position") - if self._draft_history_uses_kv_history(): - cache = self.kv_cache_map[request_id] - # Receive admission establishes the global prompt watermark; the - # transfer completion path validates that the received KV covers it. - if ( - type(position) is not int - or type(valid_length) is not int - or position != cache.history_length - or valid_length != min(position, self.draft_window_size) - ): - raise ValueError("Embedded DSpark transfer does not match the admitted KV history") - self.get_draft_block_table([request_id]) - return - history = StandaloneDraftHistory(valid_length, position) - # Validate receiver-local allocation before publishing history. - self.get_draft_block_table([request_id], [history]) - self.set_draft_history(request_id, valid_length, position) - - def _get_attention_pool_num_blocks(self, local_layer_idx: int) -> int: - """Count ledger blocks without counting expanded physical subpages.""" - layer_id = LayerId(local_layer_idx) - expansion = self.impl.get_page_index_converter(layer_id, Role.KEY).expansion - return self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // ( - self.get_layer_kv_factor(self.pp_layers[local_layer_idx]) * expansion - ) - - def _draft_context_exceeds_limit(self, req: LlmRequest) -> bool: - """Whether draft prefill skips this prompt, preventing unified block reuse.""" - if ( - not self._standalone_draft_layer_ids - or self.draft_window_size is not None - or self.draft_max_position_embeddings is None - ): - return False - prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len - return prompt_len > min(self.max_seq_len, self.draft_max_position_embeddings) - - -def _get_draft_layer_sizes( - *, - draft_num_layers: int = 0, - draft_num_kv_heads: int = 0, - draft_head_dim: int = 0, - draft_dtype: DataType = DataType.BF16, - draft_kv_factor: int = 2, - **kwargs, -) -> tuple[int, ...]: - if draft_num_layers < 0: - raise ValueError("Draft layer count must be nonnegative") - if draft_num_layers: - if min(draft_num_kv_heads, draft_head_dim) <= 0: - raise ValueError("Standalone draft cache dimensions must be positive") - if draft_dtype not in (DataType.HALF, DataType.BF16, DataType.FP8): - raise ValueError("Standalone draft KV supports FP16, BF16 and FP8 storage") - if draft_kv_factor not in (1, 2): - raise ValueError("KV storage requires one or two planes") - size = get_size_in_bytes(draft_num_kv_heads * draft_head_dim * draft_kv_factor, draft_dtype) - return (size,) * draft_num_layers - - -def _estimate_draft_cache_size_components( - layer_sizes: Sequence[int], - tokens_per_block: int, - *, - generation_capacity_headroom: int, - scratch_tokens: int = 0, - window_size: int | None = None, - helix_cp_size: int = 1, -) -> tuple[int, int, int]: - """Draft storage cost, with history and scratch requirements supplied separately.""" - retention_window = window_size + 1 if window_size is not None else None - context, generation, per_request = _estimate_cache_size_components( - layer_sizes, - [retention_window] * len(layer_sizes), - tokens_per_block, - scratch=False, - generation_capacity_headroom=generation_capacity_headroom, - standalone_draft_reserve=scratch_tokens, - ) - extra_bytes_per_token = (helix_cp_size - 1) * sum(layer_sizes) - return context + extra_bytes_per_token, generation + extra_bytes_per_token, per_request diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py index 50202e8e1dab..fac049c226c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -24,11 +24,6 @@ import torch -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - _estimate_draft_cache_size_components, - _get_draft_layer_sizes, -) - if sys.version_info[:2] >= (3, 12): from typing import override else: @@ -598,7 +593,6 @@ def _estimate_mamba_hybrid_cache_cost( include_explicit_snapshots: bool, cap_partial_attention_snapshots: bool, extra_state_bytes_per_rank: int = 0, - draft_scratch_tokens: int = 0, is_draft: bool = False, use_separate_draft_kv_cache: bool = False, **kwargs, @@ -657,22 +651,6 @@ def _estimate_mamba_hybrid_cache_cost( regular_slope = math.ceil(state_bytes_per_rank / interval) if has_unaligned_periodic_snapshot and not cap_partial_attention_snapshots: regular_slope += math.ceil(attention_block_bytes / interval) - draft_layer_sizes = _get_draft_layer_sizes(**kwargs) - if draft_layer_sizes: - # Hybrid target attention is full attention. Reserve capture/noise - # capacity only in its attention pools and the distinct draft pools; - # recurrent state retains the fixed/snapshot accounting above. - # Hybrid profiling bounds draft storage by full-history attention. - _, draft_slope, draft_fixed = _estimate_draft_cache_size_components( - draft_layer_sizes, - tokens_per_block, - generation_capacity_headroom=1, - scratch_tokens=draft_scratch_tokens, - helix_cp_size=mapping.cp_size if mapping.has_cp_helix() else 1, - ) - draft_fixed += draft_scratch_tokens * attention_slope - attention_slope += draft_slope - intercept += max_batch_size * draft_fixed return attention_slope + regular_slope, intercept diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index e978406cc15c..5019ed8b8a33 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -424,16 +424,11 @@ def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConf ] def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: - if self._is_standalone_draft_layer(local_layer_idx): - return False - layer_id = self.pp_layers[local_layer_idx] - return layer_id < len(self._mamba_layer_mask) and self._mamba_layer_mask[layer_id] + return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] @override def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: layer_id = int(self.impl.layer_grouping[pool_id][0]) - if self._is_standalone_draft_layer(layer_id): - return KVCacheManagerV2._get_pool_roles(self, pool_id) if self._is_local_mamba_layer(layer_id): return MambaRole.SSM_STATE, None return super()._get_pool_roles(pool_id) @@ -579,19 +574,6 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: ) return attention_quota + state_quota + extra_attention_quota - @override - def _get_draft_quota_from_max_tokens(self, max_tokens: float) -> int: - draft_quota = super()._get_draft_quota_from_max_tokens(max_tokens) - snapshot_slots = self._num_ssm_snapshots_for_capacity(max_tokens, self.kv_cache_config) - if snapshot_slots > 0: - # The snapshot plan's partial-page reserve also covers draft KV. - draft_quota += ( - self._max_resident_sequences() - * sum(self._draft_layer_sizes) - * self.tokens_per_block - ) - return draft_quota - @override def _get_max_tokens_from_quota(self, quota: int) -> float: if self._get_quota_from_max_tokens(0) > quota: @@ -865,7 +847,9 @@ def get_num_free_blocks(self) -> int: // self._ssm_page_index_scale ) else: - attention_pages.append(self._get_attention_pool_num_blocks(local_layer_idx)) + attention_pages.append( + self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // self.kv_factor + ) if attention_pages: return max(attention_pages) return max(ssm_pages) if ssm_pages else 0 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 818d9f0cbaca..a0e323196a53 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -9116,9 +9116,6 @@ def _sample_async(self, scheduled_batch, sample_state = self.sampler.sample_async( scheduled_batch, batch_outputs, num_context_logits_prefix_sum) - if sample_state is not None: - sample_state.draft_history_update = batch_outputs.get( - 'draft_history_update') self._maybe_record_hang_diagnostic_phase( "sampling_returned", scheduled_batch) return sample_state @@ -9145,15 +9142,10 @@ def _setup_sampler_step(self, requests: ScheduledRequests): @nvtx_range("_update_requests") def _update_requests(self, - sample_state: SampleState | None, + sample_state: SampleState, resource_manager: Optional[ResourceManager] = None): - if sample_state is None: - return try: self.sampler.update_requests(sample_state, resource_manager) - if sample_state.draft_history_update is not None: - sample_state.draft_history_update.publish() - sample_state.draft_history_update = None self._accumulate_spec_dec_stats(sample_state) except Exception as e: traceback.print_exc() diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 397d67a33c41..b1f256e560c7 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -42,7 +42,6 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager from ..scheduler import ScheduledRequests -from ..standalone_draft_cache import DraftHistoryUpdate from .beam_search import ( BeamHistory, BeamHistoryBuilder, @@ -115,7 +114,6 @@ class SampleState(Generic[GenericSampleStateTensorsHost, GenericSampleStateTenso host: Optional[GenericSampleStateTensorsHost] = None sampler_event: Optional[SamplerEvent] = None runtime_draft_len: Optional[int] = None - draft_history_update: DraftHistoryUpdate | None = None # Generic bounds not supported, https://github.com/python/typing/issues/548 diff --git a/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py b/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py deleted file mode 100644 index 3e049ef8eb23..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py +++ /dev/null @@ -1,66 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Storage and history contracts for manager-owned DSpark/DFlash drafters.""" - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -import torch - -if TYPE_CHECKING: - from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 - - -@dataclass(frozen=True) -class StandaloneDraftHistory: - """Committed length and next absolute position, independent of target/scratch state.""" - - valid_length: int - position: int - - def __post_init__(self) -> None: - if type(self.valid_length) is not int or type(self.position) is not int: - raise ValueError("Standalone draft history requires integer length and position") - if self.valid_length < 0 or self.position < self.valid_length: - raise ValueError("Invalid standalone draft history length or position") - - -@dataclass(frozen=True) -class DraftHistoryUpdate: - """One execution's draft history, published after its sampling event completes.""" - - manager: "KVCacheManagerV2" - request_ids: tuple[int, ...] - cache_instances: tuple[object, ...] - values_host: torch.Tensor - - @classmethod - def capture( - cls, - manager: "KVCacheManagerV2", - request_ids: tuple[int, ...], - values: torch.Tensor, - ) -> "DraftHistoryUpdate": - """Queue an independent length/position readback on the execution stream.""" - if values.shape != (len(request_ids), 2): - raise ValueError("Draft history update requires one length/position pair per request") - values_host = torch.empty( - values.shape, dtype=values.dtype, device="cpu", pin_memory=values.is_cuda - ) - values_host.copy_(values, non_blocking=True) - return cls( - manager, - tuple(request_ids), - tuple(manager.kv_cache_map[request_id] for request_id in request_ids), - values_host, - ) - - def publish(self) -> None: - """Publish completed writes without reviving released or replaced requests.""" - for request_id, cache, (length, position) in zip( - self.request_ids, self.cache_instances, self.values_host.tolist() - ): - current = self.manager.kv_cache_map.get(request_id) - if current is cache and current.is_active: - self.manager.set_draft_history(request_id, length, position) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 8a8ce92796b8..6b7e6b517876 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -31,11 +31,9 @@ from ..attention.backends import AttentionMetadata from ..modules.logits_processor import LogitsProcessor -from ..pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID -from ..pyexecutor.resource_manager import BaseResourceManager, ResourceManager, ResourceManagerType -from ..pyexecutor.standalone_draft_cache import DraftHistoryUpdate +from ..pyexecutor.resource_manager import BaseResourceManager from .accept_stats import maybe_create_recorder from .dflash_attention import ( get_dflash_paged_append, @@ -601,7 +599,7 @@ def prepare(self): # Update slot mapping for DFlash context buffers worker = getattr(self, "_dflash_worker", None) - if worker is not None and worker._ctx_buf_inited and not worker._has_unified_draft_cache(): + if worker is not None and worker._ctx_buf_inited: current = set(self.request_ids) evicted = {} for rid in list(worker._req_to_slot.keys()): @@ -616,41 +614,24 @@ def prepare(self): # prefilling, so _store_prefill_context -- the only other assigner -- # never runs and every request would land on the dummy slot. num_contexts = max(0, len(self.request_ids) - self.num_generations) - assigned = {} - for request_idx, rid in enumerate(self.request_ids[num_contexts:], num_contexts): + for rid in self.request_ids[num_contexts:]: if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - slot = worker._assign_slot(rid, updates=assigned) - if slot is None: + if worker._assign_slot(rid) is None: # Not fatal -- spec dec is lossless, so the dummy-slot # fallback only costs acceptance. Warn so it is not silent. logger.warning( f"DFlash: no free drafter slot for generation request {rid}; " "it shares the dummy slot and will draft poorly." ) - elif worker._ctx_unified_cache: - # New requests have no in-flight predecessor. Subsequent - # accepted tokens advance the device length during replay. - token_offset = sum(self.seq_lens[:request_idx]) - assigned[slot] = min( - int(self.host_position_ids[token_offset]), worker._max_ctx - ) - worker._write_ctx_len(assigned) # Route unknown request IDs (cuda graph padding or warmup dummies) - # to dummy slot to avoid corrupting real request's context. A - # generation row with the idle id is a dummy even when a context - # warmup stored a prompt under that id: its length is not this row's. + # to dummy slot to avoid corrupting real request's context mapping = torch.tensor( - [ - worker._dummy_slot - if i >= num_contexts and rid == ATTENTION_DP_DUMMY_REQUEST_ID - else worker._req_to_slot.get(rid, worker._dummy_slot) - for i, rid in enumerate(self.request_ids) - ], + [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], dtype=torch.long, ) worker._batch_to_slot_buffer.update(mapping) @@ -771,21 +752,12 @@ def __init__( # graph compatible. self._ctx_buf_inited = False self._ctx_len = None - self._ctx_position_offset = None - self._managed_cache_bindings = {} - self._managed_request_ids = () - self._managed_snapshot_slots: Optional[torch.Tensor] = None # Host shadows of _ctx_len and of each request's prompt progress. self._ctx_len_host = None self._req_ctx_pos = {} - # Decode the manager draft pool's device block offsets instead of - # building the table on the host (see _refresh_unified_block_tables). - self._device_draft_tables = False - self._draft_table_checks_left = 0 # Snapshot for rolling back in-place _ctx_len updates when a forward # fails (or after warmup). See _ensure_spec_dec_state_restored. self._saved_ctx_len = None - self._saved_ctx_position_offset = None self._saved_ctx_len_host = None self._saved_req_ctx_pos = None self._ctx_len_restore_pending = False @@ -821,9 +793,6 @@ def __init__( # the table width on the default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None self._ctx_pool_idx = 0 - self._ctx_layer_ids = None - self._ctx_unified_cache = False - self._ctx_block_offset = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index self._ctx_kv_indptr = None self._ctx_kv_last_page_len = None @@ -887,7 +856,7 @@ def _validate_draft_attention_backend(self, draft_model) -> None: draft_model, "dflash_attention_backend", None ) if draft_model_dflash_attention_backend is None: - raise ValueError("DFlash draft model is missing dflash_attention.backends.") + raise ValueError("DFlash draft model is missing dflash_attention_backend.") if self._dflash_attention_backend == "AUTO": # The draft model resolved AUTO against its own family and this # build; adopt that, because _lazy_init_ctx_buffers keys paging and @@ -1090,21 +1059,13 @@ def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype, kv_facto f"scales it down linearly too." ) - def _managed_ctx_pool( - self, - draft_kv_cache_manager: BaseResourceManager | None, - L: int, - nkv: int, - hd: int, - dtype: torch.dtype, - kv_factor: int = 2, - ) -> list[torch.Tensor] | None: + def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor=2): """Per-layer views of the draft KV cache manager's pool, or None. - The manager's buffer accessor supplies the same per-layer HND shape - as the private arena. Unified storage uses ``get_draft_buffers`` and - requires the pool to match the drafter and the manager's page size; - legacy storage uses ``get_buffers`` and may fall back to the arena. + The manager already funds a pool sized from the drafter's own config + (K3 at TP8: 5 layers x 2 x 2 kv heads x 64 x 2B = 2560 B/token), and + ``get_buffers(..., "HND")`` hands back exactly the per-layer shape the + private arena uses, so nothing but the ownership of the memory changes. A list, not one stacked tensor: only the two consumers that index it by layer need it, and the two managers disagree on what a stack would mean. @@ -1114,45 +1075,19 @@ def _managed_ctx_pool( Note the page count is not the drafter's to interpret: a V2 view spans the whole interleaved pool, not this layer's slice, so only shape[1:] is - checked here. Legacy offsets are settled in _init_ctx_block_offsets; - unified storage supplies draft block indices directly. + checked here and the index space is settled in _init_ctx_block_tables. """ if draft_kv_cache_manager is None: # No separate draft KV cache. Attention DP alone does NOT disable # it: _util.py skips that bail for an external drafter, which is # what DSpark is, so this drafter does get a manager under DP. return None - draft_layer_ids = tuple(getattr(draft_kv_cache_manager, "draft_layer_ids", ())) - layer_ids = draft_layer_ids if draft_layer_ids else tuple(range(L)) - if len(layer_ids) != L: - logger.warning( - f"DFlash: KV cache manager exposes {len(layer_ids)} draft layers " - f"but the drafter has {L}; using the private ctx arena." - ) - return None - self._ctx_layer_ids = layer_ids - unified = bool(getattr(draft_kv_cache_manager, "_draft_layer_sizes", ())) - if unified: - layers = [ - draft_kv_cache_manager.get_draft_buffers(i, kv_layout="HND") for i in range(L) - ] - else: - layers = [ - draft_kv_cache_manager.get_buffers(layer_id, kv_layout="HND") - for layer_id in layer_ids - ] + layers = [draft_kv_cache_manager.get_buffers(i, kv_layout="HND") for i in range(L)] base = layers[0] - page_size = draft_kv_cache_manager.tokens_per_block if unified else base.size(-2) # kv_factor 1 is the MLA drafter's SELFKONLY pool: one latent, no V. - expected = (kv_factor, nkv, page_size, hd) - # A unified pool may store the drafter's K/V in FP8; writes convert. - allowed_dtypes = (dtype, torch.float8_e4m3fn) if unified else (dtype,) + expected = (kv_factor, nkv, base.size(-2), hd) for i, layer in enumerate(layers): - if tuple(layer.shape[1:]) != expected or layer.dtype not in allowed_dtypes: - if unified: - raise ValueError( - "Unified DSpark draft pool does not match the drafter KV layout" - ) + if tuple(layer.shape[1:]) != expected or layer.dtype != dtype: # Fall back rather than fail: the private arena is what every # DFlash model shipped with, so an unrecognized pool costs the # old memory footprint, not the run. @@ -1180,42 +1115,10 @@ def _init_ctx_block_tables( """ mgr = draft_kv_cache_manager pool_mapping = mgr.kv_cache_pool_mapping - first_layer_id = self._ctx_layer_ids[0] if self._ctx_layer_ids is not None else 0 - layer_offset = mgr.layer_offsets[first_layer_id] + layer_offset = mgr.layer_offsets[0] self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) - per_layer_tables = bool(getattr(mgr, "_use_per_layer_page_tables", False)) num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) - self._ctx_block_offset = 0 - if isinstance(mgr, KVCacheManagerV2): - converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY) - layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset] - unsupported = None - if mgr.enable_swa_scratch_reuse: - unsupported = "DFlash managed context does not support SWA scratch pages" - for layer_id in self._ctx_layer_ids: - local_id = mgr.layer_offsets[layer_id] - if ( - mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id - or mgr.get_layer_page_index_scale(layer_id) != converter.scale - ): - unsupported = ( - "DFlash's shared context block table requires draft layers " - "with the same lifecycle and page-index scale" - ) - break - if unsupported is not None: - # A unified manager was selected for this layout, so the drafter - # cannot move to private storage without breaking allocation, - # reuse and transfer. A separate manager keeps the private arena. - if self._ctx_unified_cache: - raise NotImplementedError(unsupported) - logger.warning(f"{unsupported}; using the private ctx arena.") - return False - if per_layer_tables: - # Attention tables address the physical pool origin, whereas - # get_buffers() starts at this layer's own K buffer. - self._ctx_block_offset = int(converter.layer_offset) - elif num_pool_layers != L: + if num_pool_layers != L: logger.warning( f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " f"the drafter has {L}; using the private ctx arena." @@ -1224,23 +1127,12 @@ def _init_ctx_block_tables( # V1 slices a page-major pool, so stride(0) spans every layer; V2 wraps # the layer base in a dense view over interleaved pages, so stride(0) is # one layer slot. Guessing wrong is silent: acceptance quietly drops. - manager_kv_factor = ( - mgr.get_layer_kv_factor(first_layer_id) - if hasattr(mgr, "get_layer_kv_factor") - else mgr.kv_factor - ) - if manager_kv_factor != getattr(self, "_ctx_kv_factor", manager_kv_factor): - logger.warning( - "DFlash: draft KV layer factor does not match the drafter; " - "using the private ctx arena." - ) - return False - layer_block_elems = manager_kv_factor * nkv * page_size * hd + layer_block_elems = mgr.kv_factor * nkv * page_size * hd stride0 = pool[0].stride(0) if stride0 == layer_block_elems: - self._ctx_block_divisor = manager_kv_factor + self._ctx_block_divisor = mgr.kv_factor elif stride0 == num_pool_layers * layer_block_elems: - self._ctx_block_divisor = num_pool_layers * manager_kv_factor + self._ctx_block_divisor = num_pool_layers * mgr.kv_factor else: logger.warning( f"DFlash: draft KV pool layer stride {stride0} matches neither V1's " @@ -1264,36 +1156,7 @@ def _init_ctx_block_tables( ) return True - def _init_managed_ctx_block_tables(self, num_slots: int, max_blocks: int) -> None: - self._ctx_block_tables = torch.zeros( - (num_slots, max_blocks), dtype=torch.int32, device="cuda" - ) - self._ctx_block_indptr = torch.arange( - 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" - ) - self._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") - logger.info(f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}") - - @staticmethod - def _check_block_within_allocation( - num_ctx: torch.Tensor, allocated_ctx_limit: torch.Tensor, block_size: int - ) -> None: - """Raise when a request's drafter block would leave its allocated pages. - - The shared manager sizes the request, so a shortfall there would have the - block's scratch write land on the last context position, or on another - request's page. Asynchronous, so it adds no host sync and is safe to - capture in a CUDA graph; on the device it traps instead of raising. - """ - torch._assert_async( - (num_ctx <= allocated_ctx_limit).all(), - f"DFlash: the drafter's block of {block_size} positions falls outside the " - "pages the shared KV cache manager allocated for the request", - ) - - def _refresh_ctx_block_tables( - self, attn_metadata, num_seqs: int, request_ids: list[int] | None = None - ) -> bool: + def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: """Decode this iteration's draft block table into the persistent buffer. In place and fixed-shape so a captured graph keeps reading the same @@ -1302,48 +1165,22 @@ def _refresh_ctx_block_tables( """ if self._ctx_block_tables is None or num_seqs <= 0: return False - if self._has_unified_draft_cache(): - if self._device_draft_tables and self._refresh_unified_block_tables( - attn_metadata, num_seqs, request_ids - ): - return True - table = self._ctx_kv_manager.get_draft_block_table(request_ids[:num_seqs]) - self._ctx_block_tables[:num_seqs].copy_(table, non_blocking=True) - self._ctx_block_counts[:num_seqs].copy_( - torch.tensor( - [ - self._ctx_kv_manager.get_draft_num_blocks(rid) - for rid in request_ids[:num_seqs] - ], - dtype=torch.long, - device="cpu", - pin_memory=prefer_pinned(), - ), - non_blocking=True, - ) - return True - offsets_name = ( - "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" - ) - src = getattr(attn_metadata, offsets_name, None) + src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None) if src is None: # The table is bound but unfillable. Leaving it at zeros would send # every request's every layer to pool block 0 -- silently. raise RuntimeError( "DFlash: bound to the draft KV pool but " f"{type(attn_metadata).__name__} carries no " - f"{offsets_name} this iteration." + "draft_kv_cache_block_offsets this iteration." ) # Field 0 of the K/V pair: both name the same block for a kv_factor=2 # pool, and the drafter addresses blocks, not planes. The block count # is taken before the clamp erases the placeholders; it is informative # only on the python V2 backend, the cpp one already mapped # BAD_PAGE_INDEX to 0. - encoded = src[self._ctx_pool_idx, :num_seqs, 0] - if self._ctx_block_offset: - encoded = encoded.to(torch.int64) - self._ctx_block_offset decode_draft_block_table( - encoded, + src[self._ctx_pool_idx, :num_seqs, 0], self._ctx_block_divisor, self._ctx_block_tables, self._ctx_block_counts, @@ -1354,25 +1191,18 @@ def _lazy_init_ctx_buffers( self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None ): self._validate_draft_attention_backend(draft_model) - self._ctx_unified_cache = ( - draft_kv_cache_manager is not None - and draft_kv_cache_manager is getattr(attn_metadata, "kv_cache_manager", None) - ) - offsets_name = ( - "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets" - ) - if draft_kv_cache_manager is not None and not hasattr(attn_metadata, offsets_name): - if self._ctx_unified_cache: - raise NotImplementedError( - f"DFlash unified KV requires {offsets_name} in attention metadata" - ) + # Only TrtllmAttentionMetadata builds draft_kv_cache_block_offsets. + # Without it the table stays all-zero and every request would read and + # write pool block 0 -- silently, at acceptance 1.0. + if draft_kv_cache_manager is not None and not hasattr( + attn_metadata, "draft_kv_cache_block_offsets" + ): logger.warning( f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; " "using the private ctx arena." ) draft_kv_cache_manager = None - unified = bool(getattr(draft_kv_cache_manager, "_draft_layer_sizes", ())) if self._ctx_buf_inited: if self._ctx_kv_manager is draft_kv_cache_manager: return @@ -1440,21 +1270,15 @@ def _lazy_init_ctx_buffers( self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") - self._ctx_position_offset = torch.zeros_like(self._ctx_len) self._ctx_len_host = [0] * num_slots self._batch_to_slot_buffer = DeviceCopyBuffer( torch.zeros(max_batch, dtype=torch.long, device="cuda"), pin_memory=prefer_pinned(), ) - self._managed_snapshot_slots = ( - torch.empty_like(self._batch_to_slot_buffer.tensor) if unified else None - ) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} self._req_ctx_pos = {} - self._managed_cache_bindings = {} - self._managed_request_ids = () # checkpoint's trained block width self._resolved_block_size = getattr(draft_model, "block_size", None) or ( @@ -1481,7 +1305,6 @@ def _lazy_init_ctx_buffers( nkv = draft_model._num_kv_heads hd = draft_model._head_dim kv_factor = getattr(draft_model, "_kv_factor", 2) - self._ctx_kv_factor = kv_factor capacity = self._max_ctx + self._compute_block_size # Published before any drafter table is built, so a drafter may read it # instead of its config cap. _max_ctx is attn_metadata.max_seq_len, which @@ -1497,9 +1320,14 @@ def _lazy_init_ctx_buffers( use_paged = self._dflash_attention_backend in _PAGED_ATTENTION_BACKENDS or getattr( draft_model, "_paged_ctx_cache", False ) - self._ctx_paged = unified or use_paged - if self._ctx_paged: - pool = self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor) + self._ctx_paged = use_paged + if use_paged: + # FA4 pages, but on the private arena with its own kernel. + pool = ( + None + if self._dflash_attention_backend == "FA4" + else self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor) + ) # The manager's page size wins when bound to it: its pool is already # carved, so the drafter adopts the geometry rather than imposing one. page_size = self._ctx_page_size if pool is None else pool[0].size(-2) @@ -1527,17 +1355,11 @@ def _lazy_init_ctx_buffers( # Settle the block table before committing to the pool: it is the # last thing that can rule the pool out, and falling back after # taking the pool branch would leave no buffer allocated at all. - if ( - not unified - and pool is not None - and not self._init_ctx_block_offsets( - draft_kv_cache_manager, pool, L, nkv, hd, page_size - ) + if pool is not None and not self._init_ctx_block_tables( + draft_kv_cache_manager, pool, num_slots, L, nkv, hd, page_size ): pool = None if pool is None: - if self._ctx_unified_cache: - raise ValueError("DFlash could not bind the unified manager's draft KV layers") # Per-slot reservation is an arena-only notion; the pool path # takes its pages from the manager one iteration at a time. self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size @@ -1572,31 +1394,18 @@ def _lazy_init_ctx_buffers( # projected from the captured target features through fc and # hidden_norm (which every loaded DFlash drafter owns), as the # eager path does; LiLiCorr drafters keep the eager path. - # Unified caches also track retained-history positions and - # validate per-request allocation in the eager preparation path. - self._paged_latent_prepare = ( - kv_factor == 1 and not draft_model._uses_lilicorr and not unified - ) - max_blocks = ( - draft_kv_cache_manager.draft_max_blocks_per_seq - if unified - else draft_kv_cache_manager.max_blocks_per_seq - ) - self._init_managed_ctx_block_tables(num_slots, max_blocks) - if unified: - self._init_unified_draft_tables(draft_kv_cache_manager) + self._paged_latent_prepare = kv_factor == 1 and not draft_model._uses_lilicorr # Only a manager that publishes and matches its own blocks can # hand back a reused prefix's drafter K/V; the private arena and # an unpaired draft pool both start every request at 0. - self._ctx_reuse_addressable = self._ctx_unified_cache or bool( + self._ctx_reuse_addressable = bool( getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) ) self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) - if not use_paged: - if not unified: - self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) + else: # dense private arena (VANILLA, unpaged) + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype, kv_factor) kv_shape = (num_slots, L, capacity, nkv, hd) self._ctx_k_buf = torch.zeros(kv_shape, dtype=dtype, device="cuda") # An MLA drafter stores one latent per token; leaving _ctx_v_buf as @@ -1665,14 +1474,9 @@ def clear(slot: int) -> None: self._free_slots.append(old_slot) if req_id not in self._req_to_slot: if not self._free_slots: - if self._has_unified_draft_cache(): - raise RuntimeError("Unified DSpark has no free request staging slot") return None slot = self._free_slots.popleft() self._req_to_slot[req_id] = slot - if self._has_unified_draft_cache(): - self._managed_cache_bindings.pop(req_id, None) - self._ctx_position_offset[slot] = 0 clear(slot) return self._req_to_slot[req_id] @@ -1689,9 +1493,10 @@ def _store_context_kv_paged( rows: torch.Tensor, positions: torch.Tensor, ) -> None: - fp8_cache = self._ctx_kv_buf[0].dtype == torch.float8_e4m3fn - if v is None or self._dflash_attention_backend == "VANILLA" or fp8_cache: - # VANILLA, single-latent (MLA) and FP8 caches use direct page writes. + if v is None: + # Single-latent (MLA) cache: one vector per token per layer, so + # there is no K/V interleave for append_paged_kv_cache to do and the + # write is a plain scatter into the page the block table names. table = ( self._ctx_block_tables if self._ctx_block_tables is not None @@ -1700,23 +1505,8 @@ def _store_context_kv_paged( page_size = self._ctx_page_size pages = table[rows.long(), positions.long() // page_size].long() offsets = positions.long() % page_size - for layer_idx in range(k.size(1)): - pool = self._ctx_kv_buf[layer_idx] - if fp8_cache: - # Byte views: FP8 index_put support varies across builds. - pool_bytes = pool.view(torch.uint8) - pool_bytes[pages, 0, :, offsets] = ( - k[:, layer_idx].to(pool.dtype).view(torch.uint8) - ) - if v is not None: - pool_bytes[pages, 1, :, offsets] = ( - v[:, layer_idx].to(pool.dtype).view(torch.uint8) - ) - continue - pool[pages, 0, :, offsets] = k[:, layer_idx].to(pool.dtype) - if v is not None: - pool[pages, 1, :, offsets] = v[:, layer_idx].to(pool.dtype) + self._ctx_kv_buf[layer_idx][pages, 0, 0, offsets] = k[:, layer_idx, 0] return append_paged_kv_cache = self._get_ctx_paged_append() @@ -1727,13 +1517,12 @@ def _store_context_kv_paged( positions_i32 = positions.to(torch.int32) kv_indices, kv_indptr = self._ctx_paged_index_args() for layer_idx in range(k.size(1)): - pool = self._ctx_kv_buf[layer_idx] append_paged_kv_cache( - append_key=k[:, layer_idx].to(pool.dtype).contiguous(), - append_value=v[:, layer_idx].to(pool.dtype).contiguous(), + append_key=k[:, layer_idx].contiguous(), + append_value=v[:, layer_idx].contiguous(), batch_indices=rows_i32, positions=positions_i32, - paged_kv_cache=pool, + paged_kv_cache=self._ctx_kv_buf[layer_idx], kv_indices=kv_indices, kv_indptr=kv_indptr, kv_last_page_len=self._ctx_kv_last_page_len, @@ -1823,8 +1612,6 @@ def _restore_ctx_len_host(self) -> None: self._ctx_len_host = list(self._saved_ctx_len_host) if self._saved_req_ctx_pos is not None: self._req_ctx_pos = dict(self._saved_req_ctx_pos) - if self._has_unified_draft_cache() and self._saved_ctx_position_offset is not None: - self._ctx_position_offset.copy_(self._saved_ctx_position_offset) def _store_prefill_context( self, @@ -1876,20 +1663,13 @@ def _store_prefill_context( if self._ctx_block_tables is not None else None ) - # Host copies of each chunk's start avoid a device sync per request. - cached_tokens = None - kv_cache_params = getattr(attn_metadata, "kv_cache_params", None) - if self._has_unified_draft_cache() and kv_cache_params is not None: - cached_tokens = getattr(kv_cache_params, "num_cached_tokens_per_seq", None) for i in range(num_contexts): req_id = spec_metadata.request_ids[i] slen = int(attn_metadata._seq_lens[i]) chunk_proj = ctx_proj[offset : offset + slen].detach() chunk_pos = position_ids[offset : offset + slen].long().detach() - if cached_tokens is not None and i < len(cached_tokens): - first_pos = int(cached_tokens[i]) - elif slen == 0: + if slen == 0: first_pos = 0 elif host_pos is not None: first_pos = int(host_pos[offset]) @@ -1900,9 +1680,7 @@ def _store_prefill_context( # where the previous one ended. Everything else -- a fresh request, a # request id reused after completion, a prefill restarted after # preemption -- has to start from a clean slot. - reset = ( - not self._has_unified_draft_cache() and self._req_ctx_pos.get(req_id) != first_pos - ) + reset = self._req_ctx_pos.get(req_id) != first_pos if self._assign_slot(req_id, reset=reset, updates=ctx_len_updates) is None: logger.warning("DFlash: no free slots, skipping context store") self._req_ctx_pos.pop(req_id, None) @@ -1918,14 +1696,6 @@ def _store_prefill_context( # paired pool: precompute_context_kv is per-token in (hidden, p). cur = first_pos if cur + slen > cap: - if self._has_unified_draft_cache(): - self._ctx_len[slot] = 0 - self._ctx_len_host[slot] = 0 - self._req_ctx_pos.pop(req_id, None) - self._ctx_position_offset[slot] = first_pos + slen - self._managed_cache_bindings[req_id] = self._ctx_kv_manager.kv_cache_map[req_id] - offset += slen - continue # Request-level, like the no-free-slots path above: truncating # would silently draft from a stale prefix, but killing the # forward would take every other in-flight request with it. @@ -1950,10 +1720,6 @@ def _store_prefill_context( cache_dtype = ( self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype ) - if cache_dtype == torch.float8_e4m3fn: - # The K/V projection runs in the drafter's weight dtype; - # the store converts its outputs to the cache dtype. - cache_dtype = chunk_proj.dtype chunk_proj_cast = chunk_proj[:actual].to(cache_dtype) ctx_len_updates[slot] = end # Precompute post-norm/post-RoPE K,V for this prefill chunk @@ -1979,16 +1745,12 @@ def _store_prefill_context( chunk_k, chunk_v, torch.full((actual,), row, dtype=torch.long, device="cuda"), - torch.arange(cur, end, dtype=torch.long, device="cuda"), + local_pos, ) - if self._ctx_k_buf is not None: + else: # VANILLA DFlash backend (FlashAttention) self._ctx_k_buf[slot, :, cur:end] = chunk_k.permute(1, 0, 2, 3) if chunk_v is not None: self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) - self._req_ctx_pos[req_id] = first_pos + slen - if self._has_unified_draft_cache(): - self._ctx_position_offset[slot] = self._req_ctx_pos[req_id] - end - self._managed_cache_bindings[req_id] = self._ctx_kv_manager.kv_cache_map[req_id] offset += slen self._write_ctx_len(ctx_len_updates) @@ -2054,27 +1816,16 @@ def _forward_impl( draft_model, ) + # Lazy init buffers and attach worker reference for prepare() draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - if bool(getattr(draft_kv_cache_manager, "_draft_layer_sizes", ())): - if not self._ctx_buf_inited or self._ctx_kv_manager is not draft_kv_cache_manager: - raise RuntimeError("Unified DSpark draft inputs must be prepared before forward") - else: - previous_ctx_len = self._ctx_len - self._lazy_init_ctx_buffers( - draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager - ) - spec_metadata._dflash_worker = self - if ( - self._ctx_unified_cache - and self._ctx_len is not previous_ctx_len - and spec_metadata.request_ids - ): - # New buffers clear request slots; seed them before generation. - spec_metadata.prepare() - # Before any store: prefill and decode both address pages through it. - # Returning False here means an empty batch -- the missing-offsets case - # raises inside, with the metadata type in the message. - self._refresh_ctx_block_tables(attn_metadata, batch_size, spec_metadata.request_ids) + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) + spec_metadata._dflash_worker = self + # Before any store: prefill and decode both address pages through it. + # Returning False here means an empty batch -- the missing-offsets case + # raises inside, with the metadata type in the message. + self._refresh_ctx_block_tables(attn_metadata, batch_size) # Save context lengths so both warmup and a failed forward can roll # back the in-place _ctx_len updates made during drafting. @@ -2087,8 +1838,6 @@ def _forward_impl( # during capture aborts the graph itself, and captured ops do not # mutate _ctx_len until replay. self._saved_ctx_len = self._ctx_len.clone() - if self._has_unified_draft_cache(): - self._saved_ctx_position_offset = self._ctx_position_offset.clone() self._saved_ctx_len_host = list(self._ctx_len_host) self._saved_req_ctx_pos = dict(self._req_ctx_pos) self._ctx_len_restore_pending = True @@ -2160,10 +1909,7 @@ def _forward_impl( ) if num_gens > 0: - with self.draft_kv_cache_context( - attn_metadata, - None if self._has_unified_draft_cache() else draft_kv_cache_manager, - ): + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): hidden_states_out = draft_model.dflash_forward( noise_embedding=inputs["noise_embedding"], query_positions=inputs["query_positions"], @@ -2320,12 +2066,7 @@ def _draft_block_width(self, draft_model) -> int: this alongside :meth:`_draft_slot_ids` — the two must agree, since the width bounds the slots the gather is allowed to name. """ - return self.block_width_for(self.spec_config) - - @staticmethod - def block_width_for(spec_config) -> int: - """Block slots a worker of this family computes for ``spec_config``.""" - return spec_config.max_draft_len + 1 + return self.max_draft_len + 1 def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int @@ -2573,16 +2314,17 @@ def prepare_1st_drafter_inputs( j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") - query_position = (ctx_len_gen + gen_num_accepted.long()).clamp(max=self._max_ctx) - query_position_ids = query_position.unsqueeze(1) + j_block.unsqueeze(0) + # _ctx_len is clamped to _max_ctx only AFTER this step's accepted + # tokens are folded in (see the update below), so the running length + # used here has to be clamped on its own -- otherwise a request that + # already sits at the ceiling indexes num_accepted positions past + # any position the sequence can legitimately reach. + ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) + query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) noise_embed_2d = dflash_noise_block_embedding( - embed_tokens, - bonus, - mask_token_id, - query_tokens_per_req, - trained_mask_embedding, + embed_tokens, bonus, mask_token_id, query_tokens_per_req, trained_mask_embedding ) # Accumulate new accepted features into context buffers @@ -2599,16 +2341,15 @@ def prepare_1st_drafter_inputs( projected_rows = projected_to_store.reshape(num_gens, K_plus_1, -1) lilicorr_anchor = last_accepted_hidden(projected_rows, gen_num_accepted) gen_num_accepted_long = gen_num_accepted.long() - col_idx = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) + col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) - if self._ctx_block_tables is not None: # Bound by what the manager handed THIS request: the masked # entries below are still written (fixed-size, for graph # safety) and a placeholder page resolves to another's. gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) - if self._ctx_block_tables is None or self._ctx_k_buf is not None: + else: col_idx = col_idx.clamp(max=self._max_ctx - 1) # Fixed-size writes for CUDA graph compatibility: @@ -2626,9 +2367,6 @@ def prepare_1st_drafter_inputs( cache_dtype = ( self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype ) - if cache_dtype == torch.float8_e4m3fn: - # Project in the drafter's weight dtype; the store converts. - cache_dtype = proj_flat.dtype k_new, v_new = draft_model.precompute_context_kv( proj_flat.to(cache_dtype), pos_flat ) @@ -2646,15 +2384,11 @@ def prepare_1st_drafter_inputs( else: rows_long = slot_long self._store_context_kv_paged(k_new, v_new, rows_long, col_long) - if self._ctx_k_buf is not None: + else: # VANILLA DFlash backend (FlashAttention) self._ctx_k_buf[slot_long, :, col_long] = k_new if v_new is not None: self._ctx_v_buf[slot_long, :, col_long] = v_new - if self._has_unified_draft_cache(): - self._ctx_position_offset[slots] += ( - self._ctx_len[slots] + gen_num_accepted_long - self._max_ctx - ).clamp_min(0) self._advance_ctx_len(slots, gen_num_accepted_long) num_ctx_per_req_t = self._ctx_len[slots] @@ -2666,10 +2400,6 @@ def prepare_1st_drafter_inputs( allocated = dflash_allocated_ctx_limit( self._ctx_block_counts, self._ctx_page_size, block_size ) - if self._ctx_unified_cache: - self._check_block_within_allocation( - num_ctx_per_req_t, allocated[gen_rows_out], block_size - ) num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) noise_embedding = noise_embed_2d query_positions = query_position_ids.long() @@ -2697,11 +2427,7 @@ def prepare_1st_drafter_inputs( "ctx_v_cache": self._ctx_v_buf, # Slots index the private arena; the manager's block table is keyed # by batch position, so the drafter reads its own rows there. - "ctx_cache_batch_idx": ( - gen_rows_out - if self._ctx_block_tables is not None and self._ctx_k_buf is None - else slots - ), + "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, # Anchor token per gen request (block slot 0): last accepted # token. The dspark Markov chain conditions its first step on it. "first_prev_tokens": bonus, @@ -2804,269 +2530,3 @@ def _prepare_paged_latent_gen_inputs( prep.blk_pages, prep.blk_offsets, prep.blk_seq_lens, prep.blk_kv_bounds ), } - - def get_draft_kv_cache_manager( - self, resource_manager: Optional["ResourceManager"] - ) -> Optional[BaseResourceManager]: - if resource_manager is not None: - manager = resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) - if bool(getattr(manager, "draft_layer_ids", ())): - return manager - return super().get_draft_kv_cache_manager(resource_manager) - - def _has_unified_draft_cache(self) -> bool: - return bool(getattr(self._ctx_kv_manager, "_draft_layer_sizes", ())) - - def _prepare_managed_history(self, request_ids: list[int]) -> list[int]: - """Restore newly resident histories without rewinding overlapping iterations.""" - updates = {self._dummy_slot: 0} - offsets = {self._dummy_slot: 0} - restore_rows = [] - real_request_ids = { - request_id - for request_id in request_ids - if request_id != ATTENTION_DP_DUMMY_REQUEST_ID - and request_id < self._graph_dummy_id_floor - and request_id not in self._ctx_kv_manager._draft_dummy_request_ids - } - for request_id in list(self._req_to_slot): - if request_id not in real_request_ids: - slot = self._req_to_slot.pop(request_id) - updates[slot] = 0 - offsets[slot] = 0 - self._req_ctx_pos.pop(request_id, None) - self._managed_cache_bindings.pop(request_id, None) - self._free_slots.append(slot) - for row, request_id in enumerate(request_ids): - if request_id in real_request_ids: - self._assign_slot(request_id) - slot = self._req_to_slot.get(request_id) - if slot is not None: - cache = self._ctx_kv_manager.kv_cache_map[request_id] - if self._managed_cache_bindings.get(request_id) is cache: - continue - history = self._ctx_kv_manager.get_draft_history(request_id) - updates[slot] = history.valid_length if history is not None else 0 - self._req_ctx_pos[request_id] = history.position if history is not None else 0 - offsets[slot] = self._req_ctx_pos[request_id] - updates[slot] - self._managed_cache_bindings[request_id] = cache - restore_rows.append(row) - self._write_ctx_len(updates) - self._write_slot_values(self._ctx_position_offset, offsets) - self._batch_to_slot_buffer.update( - torch.tensor( - [self._req_to_slot.get(request_id, self._dummy_slot) for request_id in request_ids], - dtype=torch.long, - ) - ) - self._managed_request_ids = tuple( - request_id for request_id in request_ids if request_id in real_request_ids - ) - # Each upload owns fresh host storage. PyTorch's pinned allocator defers - # recycling it until the asynchronous copy completes. - self._managed_snapshot_slots[: len(self._managed_request_ids)].copy_( - torch.tensor( - [self._req_to_slot[request_id] for request_id in self._managed_request_ids], - dtype=torch.long, - device="cpu", - pin_memory=prefer_pinned(), - ), - non_blocking=True, - ) - return restore_rows - - def _gather_managed_context(self, request_ids: list[int], restore_rows: list[int]) -> None: - """Refresh VANILLA's dense staging; manager pages remain authoritative.""" - if self._dflash_attention_backend != "VANILLA": - return - for row in restore_rows: - request_id = request_ids[row] - slot = self._req_to_slot.get(request_id) - if slot is None: - continue - length = self._ctx_len_host[slot] - positions = torch.arange(length, device=self._ctx_k_buf.device) - pages = self._ctx_block_tables[row, positions // self._ctx_page_size].long() - offsets = positions % self._ctx_page_size - for layer_idx, pool in enumerate(self._ctx_kv_buf): - self._ctx_k_buf[slot, layer_idx, :length] = pool[pages, 0, :, offsets, :] - self._ctx_v_buf[slot, layer_idx, :length] = pool[pages, 1, :, offsets, :] - - def prepare_managed_draft_cache( - self, draft_model, spec_metadata, attn_metadata, resource_manager - ) -> None: - """Refresh managed inputs before eager execution, capture, or graph replay.""" - manager = self.get_draft_kv_cache_manager(resource_manager) - if not getattr(manager, "draft_layer_ids", ()): - return - self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata, manager) - spec_metadata._dflash_worker = self - restore_rows = self._prepare_managed_history(spec_metadata.request_ids) - self._refresh_ctx_block_tables( - attn_metadata, attn_metadata.num_seqs, spec_metadata.request_ids - ) - self._gather_managed_context(spec_metadata.request_ids, restore_rows) - - def _write_slot_values(self, target: torch.Tensor, values: dict[int, int]) -> None: - """Scatter a slot-to-value mapping into a device tensor with one upload.""" - if not values: - return - slots, vals = zip(*values.items()) - pinned = prefer_pinned() - target.index_copy_( - 0, - torch.tensor(slots, dtype=torch.long, pin_memory=pinned).to( - target.device, non_blocking=True - ), - torch.tensor(vals, dtype=target.dtype, pin_memory=pinned).to( - target.device, non_blocking=True - ), - ) - - def snapshot_managed_draft_history(self) -> DraftHistoryUpdate | None: - """Queue iteration-owned readback; the executor publishes after completion.""" - if not self._has_unified_draft_cache() or not self._managed_request_ids: - return None - # Preparation, forward, snapshot gathers and the next preparation run - # on the execution stream, so its next upload cannot overtake this read. - slots = self._managed_snapshot_slots[: len(self._managed_request_ids)] - lengths = self._ctx_len[slots] - positions = lengths + self._ctx_position_offset[slots] - return DraftHistoryUpdate.capture( - self._ctx_kv_manager, - self._managed_request_ids, - torch.stack((lengths, positions), dim=1), - ) - - def _init_ctx_block_offsets(self, draft_kv_cache_manager, pool, L, nkv, hd, page_size) -> bool: - """Validate the legacy pool layout and resolve its offset decode constant. - - ``kv_cache_block_offsets`` entries are ``pool_block_index * - num_pool_layers * kv_factor`` plus the K/V field index. What that has to - be divided by to index a ``get_buffers`` view depends on the manager -- - see the stride check below; attention_backend/utils.py hardcodes the V1 - divisor and is not a general contract. - """ - mgr = draft_kv_cache_manager - pool_mapping = mgr.kv_cache_pool_mapping - layer_offset = mgr.layer_offsets[0] - self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) - num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) - if num_pool_layers != L: - logger.warning( - f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " - f"the drafter has {L}; using the private ctx arena." - ) - return False - # The two managers hand back different index spaces and only the - # tensor's own stride distinguishes them. V1 slices one page-major - # allocation, so index i steps a whole block (stride = num_pool_layers x - # a layer's block) and the raw block index is what's wanted. V2 wraps - # the layer's base address in a dense view over the *interleaved* page - # space -- layer 0's view spans the entire pool -- so index i steps a - # single layer slot and the same block is num_pool_layers apart. - # Getting this backwards does not fail: the drafter's layers simply - # overwrite each other and acceptance quietly drops. - layer_block_elems = mgr.kv_factor * nkv * page_size * hd - stride0 = pool[0].stride(0) - if stride0 == layer_block_elems: - self._ctx_block_divisor = mgr.kv_factor - elif stride0 == num_pool_layers * layer_block_elems: - self._ctx_block_divisor = num_pool_layers * mgr.kv_factor - else: - logger.warning( - f"DFlash: draft KV pool layer stride {stride0} matches neither the per-layer " - f"({num_pool_layers * layer_block_elems}) nor the interleaved " - f"({layer_block_elems}) layout; using the private ctx arena rather than " - f"guessing block indices." - ) - return False - logger.info( - f"DFlash: ctx block offsets pool_idx={self._ctx_pool_idx}, " - f"divisor={self._ctx_block_divisor}, view_pages={pool[0].size(0)}" - ) - return True - - def _init_unified_draft_tables(self, manager) -> None: - """Choose how the unified pool's block table is built each iteration.""" - self._device_draft_tables = False - layer = manager.draft_layer_ids[0] - pool = manager.layer_to_pool_mapping_dict[manager.layer_offsets[layer]] - if ( - getattr(manager, "_helix_cp_size", 1) != 1 - or getattr(manager, "enable_swa_scratch_reuse", False) - or manager._index_scale_ints[pool] != manager.get_layer_page_index_scale(layer) - ): - logger.warning( - "DFlash: the draft pool's block offsets cannot be decoded on the device; " - "building draft block tables on the host." - ) - return - self._ctx_pool_idx = pool - self._ctx_block_divisor = manager.get_layer_kv_factor(layer) - self._device_draft_tables = True - self._draft_table_checks_left = 8 - logger.info( - f"DFlash: draft block tables decoded on the device from pool {pool} " - f"(divisor {self._ctx_block_divisor})" - ) - - def _refresh_unified_block_tables( - self, attn_metadata, num_seqs: int, request_ids: list[int] - ) -> bool: - """Decode the draft pool's rows of the target's device block offsets. - - The target metadata already uploads every pool's offsets for this - batch, so the draft rows need one device divide plus the per-request - block counts, rather than a host table rebuilt from page lists. - Returns False when the host path must build this iteration's table. - """ - src = getattr(attn_metadata, "kv_cache_block_offsets", None) - if src is None or src.size(0) <= self._ctx_pool_idx: - return False - width = min(src.size(-1), self._ctx_block_tables.size(1)) - encoded = src[self._ctx_pool_idx, :num_seqs, 0, :width] - decoded = torch.div(encoded.clamp(min=0), self._ctx_block_divisor, rounding_mode="floor") - self._ctx_block_tables[:num_seqs, :width].copy_(decoded) - cache_map = self._ctx_kv_manager.kv_cache_map - counts = [] - for request_id in request_ids[:num_seqs]: - cache = cache_map.get(request_id) - counts.append(cache.num_blocks if cache is not None else 0) - self._ctx_block_counts[:num_seqs].copy_( - torch.tensor(counts, dtype=torch.long, pin_memory=prefer_pinned()), - non_blocking=True, - ) - if self._draft_table_checks_left > 0 and not torch.cuda.is_current_stream_capturing(): - self._draft_table_checks_left -= 1 - if not self._check_device_block_tables(num_seqs, request_ids, counts): - return False - return True - - def _check_device_block_tables( - self, num_seqs: int, request_ids: list[int], counts: list[int] - ) -> bool: - """Compare device-decoded rows with the manager's page lists (startup only).""" - manager = self._ctx_kv_manager - rows = [] - for row, request_id in enumerate(request_ids[:num_seqs]): - cache = manager.kv_cache_map.get(request_id) - if counts[row] > 0 and cache is not None and cache.is_active: - rows.append(row) - if not rows: - return True - expected = manager.get_batch_cache_indices( - [request_ids[row] for row in rows], manager.draft_layer_ids[0] - ) - device_rows = self._ctx_block_tables[rows].cpu() - for k, row in enumerate(rows): - got = device_rows[k, : counts[row]].tolist() - want = expected[k][: counts[row]] - if any(w >= 0 and g != w for g, w in zip(got, want)): - logger.error( - "DFlash: device-decoded draft block table disagrees with the manager " - f"for request {request_ids[row]}; building draft tables on the host." - ) - self._device_draft_tables = False - return False - return True diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 7ee7203b8ec4..94668e965ba7 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -33,7 +33,6 @@ from ..distributed import AllReduceStrategy from ..modules.linear import Linear, TensorParallelMode from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID -from ..pyexecutor.resource_manager import ResourceManagerType from .dflash import DFlashWorker, VocabParallelDraftTokens, dflash_draft_slot_ids from .dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT from .interface import SpecMetadata, SpecWorkerBase @@ -80,7 +79,7 @@ def _dspark_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) """Return the number of RoPE entries needed by the DSv4 block drafter. ``start_pos`` is a FRAME index, one above the absolute token position: the - prompt token at position p occupies frame p+1 (``_seed_context_pages``), so a + prompt token at position p occupies frame p+1 (``_seed_context_windows``), so a request served to ``max_ctx`` bootstraps at ``max_ctx + 1``. A verification then accepts up to ``max_draft_len + 1`` tokens and the block drafter indexes ``block_size`` further positions, making the largest index @@ -227,6 +226,62 @@ def prepare(self): ) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) + # CUDA-graph-safe path: maintain the request->slot mapping on the host + # (outside the captured region) and mirror it into ``_batch_to_slot`` so the + # captured gen forward can index the rolling windows by tensor. Mirrors + # ``DFlashSpecMetadata.prepare`` (dflash.py:96-113). + worker = getattr(self, "_dspark_worker", None) + if worker is not None and worker._win_inited: + current = set(self.request_ids) + for rid in list(worker._req_to_slot.keys()): + if rid not in current: + slot = worker._req_to_slot.pop(rid) + worker._ctx_len[slot].zero_() + worker._valid_len[slot].zero_() + worker._position_initialized[slot].zero_() + worker._kv_windows[slot].zero_() + worker._free_slots.append(slot) + # Assign a persistent rolling-window slot to every real generation + # request that never ran a context/seed forward on this worker. In + # disaggregated serving the prompt is prefilled (and the window + # seeded) on the *context* server, so ``_seed_context_windows`` never + # runs on the generation server and ``_req_to_slot`` stays empty; + # without this, all concurrent gen requests fall through to the shared + # scratch row below and corrupt each other's draft window at batch + # size > 1 (GitHub #16767). Context-prefix entries are left to + # ``_seed_context_windows``. CUDA-graph capture dummies and real + # generation requests acquire a persistent slot only once; prepare() + # runs before every replay, so it must never reset a live slot. + # ADP-idle (id 0) and high-ID CUDA-graph padding dummies use scratch. + num_contexts = max(0, len(self.request_ids) - self.num_generations) + for rid in self.request_ids[num_contexts:]: + if ( + rid != ATTENTION_DP_DUMMY_REQUEST_ID + and rid < worker._graph_dummy_id_floor + and rid not in worker._req_to_slot + ): + worker._assign_slot(rid, reset=False) + # Unknown request IDs (synthetic warmup / CUDA-graph padding, ADP idle + # requests, or disagg seed forwards without a real id) map to the + # dedicated throwaway scratch row so they cannot overwrite a live + # request's rolling window (they previously aliased to slot 0). + scratch = worker._scratch_slot + # Reset every iteration. The scratch row has no prefill to seed it + # and no completion to free it, so _advance_generation_state's + # unconditional advance would leave its ABSOLUTE position climbing + # across graph shapes and iterations until it indexes past the RoPE + # table. A real slot is bounded instead by its request's lifetime. + worker._ctx_len[scratch].zero_() + worker._valid_len[scratch].zero_() + worker._position_initialized[scratch].zero_() + mapping = torch.tensor( + [worker._req_to_slot.get(rid, scratch) for rid in self.request_ids], + dtype=torch.long, + device="cpu", + pin_memory=prefer_pinned(), + ) + worker._batch_to_slot[:num_seqs].copy_(mapping, non_blocking=True) + def is_layer_capture(self, layer_id: int) -> bool: return layer_id in self._capture_layer_set @@ -330,16 +385,11 @@ def __init__( self._win_inited = False self._output_verify_lens: Optional[torch.Tensor] = None self._attention_warmup_attempted = False + self._kv_windows: Optional[torch.Tensor] = None # [max_batch, num_stages, win, hd] self._ctx_len: Optional[torch.Tensor] = None # [max_batch] abs decode position self._valid_len: Optional[torch.Tensor] = None # [max_batch] written window entries + self._position_initialized: Optional[torch.Tensor] = None # [max_batch] bool self._win = 0 - self._draft_kv_manager = None - self._draft_kv_buffers = () - self._draft_block_tables = None - self._draft_capacities = None - self._draft_page_indices: dict[int, tuple[int, ...]] = {} - self._draft_metadata_host_buffers: list[dict] = [] - self._managed_residency = {} # Set in _lazy_init from the RoPE table the drafter will build; None # leaves positions unbounded (direct construction in tests). self._position_cap: Optional[int] = None @@ -348,11 +398,11 @@ def __init__( # source of truth, updated in prepare()/forward(); ``_batch_to_slot`` is the # CUDA mirror (request-order -> slot) read by the CUDA-graph-safe batched # gen path (set on the host in prepare(), so the captured forward indexes - # request progress through a tensor instead of a python dict lookup). + # the rolling windows through a tensor instead of a python dict lookup). self._req_to_slot = {} # request_id -> slot index self._free_slots = deque() # available slot indices self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda - # Index of the throwaway "scratch" progress row that absorbs padded / + # Index of the throwaway "scratch" window row that absorbs padded / # unknown request IDs (set in ``_lazy_init`` to ``max_batch``); it is # never handed out through ``_free_slots``. self._scratch_slot = 0 @@ -373,8 +423,8 @@ def __init__( self.verify_planner = None # The generation draft path is the batched, host-sync-free - # ``_draft_gen_block_batched`` + ``DSv4DSparkDraftModel.forward`` + - # ``dspark_attention_forward``: it is correct in eager mode AND safe + # ``_draft_gen_block_batched`` + ``DSv4DSparkDraftModel.forward_batched`` + + # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe # to capture into the target's CUDA graph (DSpark is a one-engine drafter — # its worker forward runs inside that graph, so the draft path MUST be # capture-safe whenever ``cuda_graph_config`` is set). @@ -479,7 +529,7 @@ def _publish_position_ceiling(self, draft_model, attn_metadata, block_size: int) draft_model._runtime_position_ceiling = ceiling if inner_model is not draft_model: inner_model._runtime_position_ceiling = ceiling - # Largest absolute position the block drafter may hold: forward + # Largest absolute position the block drafter may hold: forward_batched # gathers ``freqs[start_pos + block_size]`` and the interim back-fill # ``freqs[old + block_size]``, so keep one block of headroom. table_len = int(getattr(inner_model, "_runtime_position_ceiling", 0) or config_cap) @@ -495,15 +545,20 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: self._publish_position_ceiling(draft_model, attn_metadata, block_size) - if self._ctx_len is None: + if not self._win_inited: max_batch = spec_metadata.max_num_requests self._output_verify_lens = torch.empty(max_batch, dtype=torch.int32, device="cuda") - draft_model.configure_generation_autotuner(max_batch) - self._num_stages = draft_model.num_stages + num_stages = draft_model.num_stages self._win = int(draft_model._attn_params["window_size"]) - self._head_dim = int(draft_model._attn_params["head_dim"]) - - # Padding requests share an extra progress slot, never a real request slot. + head_dim = int(draft_model._attn_params["head_dim"]) + + # Real requests occupy slots ``[0, max_batch)``; one extra "scratch" row + # at index ``max_batch`` absorbs padded / unknown request IDs (CUDA-graph + # padding, ADP idle requests, or disagg seed forwards that arrive without + # a real request id) so they can never overwrite a live request's rolling + # window. Previously such IDs aliased to slot 0 and corrupted whichever + # real request occupied it. The scratch row is never handed out through + # ``_free_slots`` and its contents are throwaway. self._scratch_slot = max_batch num_rows = max_batch + 1 @@ -518,14 +573,17 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len + self._kv_windows = torch.zeros( + (num_rows, num_stages, self._win, head_dim), + dtype=torch.bfloat16, + device="cuda", + ) self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._valid_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") + self._position_initialized = torch.zeros(num_rows, dtype=torch.bool, device="cuda") self._batch_to_slot = torch.full( (max_batch,), self._scratch_slot, dtype=torch.long, device="cuda" ) - self._free_slots = deque(range(max_batch)) - self._req_to_slot = {} - self._position_initialized = torch.zeros(num_rows, dtype=torch.bool, device="cuda") if self.return_confidence: # Neutral fill = large positive logit (sigmoid ~ 1.0): # unscored slots fail safe by verifying the full block. The @@ -547,8 +605,44 @@ def _lazy_init(self, draft_model, spec_metadata, attn_metadata=None) -> None: with torch.inference_mode(False): self._draft_seq_cuda = torch.zeros((), dtype=torch.int32, device="cuda") self._build_verify_planner(draft_model, block_size) + self._free_slots = deque(range(max_batch)) + self._req_to_slot = {} + logger.info( + f"DSpark: allocated rolling KV windows " + f"[{num_rows}, {num_stages}, {self._win}, {head_dim}] " + f"({max_batch} request slots + 1 scratch row)" + ) + # Buffer state is complete independently of CuTe DSL prewarming. + # A failed prewarm must not cause the next forward to recreate the + # windows or reset the live slot maps. self._win_inited = True + if self._attention_warmup_attempted: + return + + # Prewarm the same self-JIT ops used by production before CUDA graph + # capture. This is best-effort inside the named warmup entry; the ops + # remain able to compile themselves on a later eager first use. + from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + if IS_CUTLASS_DSL_AVAILABLE: + from ..custom_ops.dspark_attention_custom_op import ( + is_dsv4_dspark_attention_config_supported, + warmup_fused_dsv4_dspark_attention, + ) + + if is_dsv4_dspark_attention_config_supported( + block_size, + int(draft_model._attn_params["n_heads"]), + int(draft_model._attn_params["head_dim"]), + int(draft_model._attn_params["window_size"]), + ): + warmup_fused_dsv4_dspark_attention( + block_size, + eps=float(draft_model._attn_params["eps"]), + ) + self._attention_warmup_attempted = True + def _build_verify_planner(self, draft_model, block_size: int) -> None: """Construct the host-side verify planner (once, at lazy init). @@ -602,16 +696,19 @@ def _build_verify_planner(self, draft_model, block_size: int) -> None: f"device_windows={self.verify_planner.device_windows}" ) - def _assign_slot(self, req_id: int, reset: bool = False) -> int: - """Get the persistent progress slot for a request.""" + def _assign_slot(self, req_id: int, reset: bool) -> int: + """Get (or refresh) the slot for a request; reset clears its window.""" if reset and req_id in self._req_to_slot: old = self._req_to_slot.pop(req_id) - self._managed_residency.pop(req_id, None) + self._ctx_len[old].zero_() + self._valid_len[old].zero_() + self._position_initialized[old].zero_() + self._kv_windows[old].zero_() self._free_slots.append(old) if req_id not in self._req_to_slot: if not self._free_slots: raise RuntimeError( - "DSpark has no free progress slots for request " + "DSpark has no free rolling-window slots for request " f"{req_id}; increase max_num_requests" ) slot = self._free_slots.popleft() @@ -619,6 +716,7 @@ def _assign_slot(self, req_id: int, reset: bool = False) -> int: self._ctx_len[slot].zero_() self._valid_len[slot].zero_() self._position_initialized[slot].zero_() + self._kv_windows[slot].zero_() if self._confidence_logits is not None: # A recycled slot still holds the previous occupant's scores; # reset to neutral so the first scheduling decision verifies @@ -626,6 +724,54 @@ def _assign_slot(self, req_id: int, reset: bool = False) -> int: self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] + def _seed_context_windows( + self, + draft_model, + spec_metadata: "DSparkSpecMetadata", + attn_metadata, + position_ids: torch.Tensor, + total_target_tokens: int, + ) -> None: + """Seed context chunks using their absolute positions. + + A request can arrive in multiple prefill chunks. Only its first chunk + starts at position zero and resets the persistent rolling window; + continuation chunks append to the same request slot. + """ + captured = spec_metadata.get_hidden_states(total_target_tokens) + capture_lens = getattr(spec_metadata, "context_capture_lens", None) + if capture_lens is not None and len(capture_lens) != attn_metadata.num_contexts: + raise ValueError("DSpark capture lengths must match the context requests") + flat_position_ids = position_ids.reshape(-1) + context_offset = 0 + for i in range(attn_metadata.num_contexts): + chunk_len = int(attn_metadata._seq_lens[i]) + capture_len = chunk_len if capture_lens is None else capture_lens[i] + if not 0 <= capture_len <= chunk_len: + raise ValueError("DSpark capture length must be a suffix of its context chunk") + chunk_positions = flat_position_ids[context_offset : context_offset + chunk_len].long() + if chunk_len == 0: + context_offset += chunk_len + continue + + req_id = spec_metadata.request_ids[i] + first_position = int(chunk_positions[0].item()) + slot = self._assign_slot(req_id, reset=first_position == 0) + self._ctx_len[slot] = chunk_positions[-1] + 1 + self._position_initialized[slot] = True + + if captured is not None and capture_len: + self._valid_len[slot] = torch.clamp( + self._valid_len[slot] + capture_len, max=self._win + ) + keep = min(self._win, capture_len) + hidden = captured[context_offset + chunk_len - keep : context_offset + chunk_len] + # A prompt token at absolute position p is stored in frame p+1, + # matching the generation path's start_pos convention. + window_positions = chunk_positions[-keep:] + 1 + draft_model.write_context_windows(hidden, window_positions, self._kv_windows[slot]) + context_offset += chunk_len + def _advance_generation_state( self, slots: torch.Tensor, @@ -652,10 +798,10 @@ def _advance_generation_state( old = torch.clamp(old, max=self._position_cap) start_pos = torch.clamp(start_pos, max=self._position_cap) self._ctx_len[slots] = start_pos - self._position_initialized[slots] = True self._valid_len[slots] = torch.clamp( self._valid_len[slots] + num_accepted_tokens, max=self._win ) + self._position_initialized[slots] = torch.ones_like(slots, dtype=torch.bool) return old, start_pos def _draft_gen_block_batched( @@ -677,7 +823,7 @@ def _draft_gen_block_batched( (``nacc``, the bonus, ``main_hidden``, ``start_pos``, the multi-accept back-fill) are gathered as tensors, slots come from the host-built ``_batch_to_slot`` mirror, and the backbone runs once via - ``DSv4DSparkDraftModel.forward``. Returns the per-position corrected + ``DSv4DSparkDraftModel.forward_batched``. Returns the per-position corrected block logits ``[num_gens, K, vocab]`` (or ``None`` when there is nothing to draft); the worker feeds them to ``SpecWorkerBase.sample_draft_tokens``. Confidence truncation stays disabled — the full block is proposed. When @@ -752,27 +898,19 @@ def _draft_gen_block_batched( min=0, max=captured.shape[0] - 1 ) # [G, K] (clamped; invalid entries are masked out anyway) interim_hidden = captured[interim_base] # [G, K, ncap*hidden] - tables = self._draft_block_tables[num_contexts:batch_size] - capacities = self._draft_capacities[num_contexts:batch_size] - draft_model.write_context_pages( - interim_hidden, - interim_pos, - interim_valid, - self._draft_kv_buffers, - tables, - capacities, + draft_model.write_context_windows_batched( + interim_hidden, interim_pos, slots, interim_valid, self._kv_windows ) # Surface the per-position corrected block logits ([num_gens, K, vocab]) # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) # sampling + TP gather + draft_probs scatter, rather than argmaxing here. - _toks, confidence, block_logits = draft_model.forward( + _toks, confidence, block_logits = draft_model.forward_batched( main_hidden, bonus, start_pos, - kv_pages=self._draft_kv_buffers, - block_tables=tables, - capacities=capacities, + kv_windows=self._kv_windows, + slots=slots, valid_len=self._valid_len[slots], temperature=0.0, return_confidence=self.return_confidence, @@ -846,9 +984,10 @@ def _forward_impl( raw_logits = logits K = self.max_draft_len - if self._draft_kv_manager is None: - raise RuntimeError("DSpark KV pages must be prepared before worker forward") - + self._lazy_init(draft_model, spec_metadata, attn_metadata) + # Backref so DSparkSpecMetadata.prepare() can maintain the host slot map + # and mirror it into _batch_to_slot for the CUDA-graph-safe gen path. + spec_metadata._dspark_worker = self self._execute_guided_decoder_if_present(logits) # Target-verify acceptance via the unified SpecWorkerBase entry: it @@ -878,6 +1017,7 @@ def _forward_impl( saved_ctx_len = self._ctx_len.clone() saved_valid_len = self._valid_len.clone() saved_position_initialized = self._position_initialized.clone() + saved_windows = self._kv_windows.clone() # The confidence buffer is persistent state too; warmup scores # must not reach the verification planner. saved_confidence = ( @@ -890,7 +1030,7 @@ def _forward_impl( # first generation step drafts against real context instead of an all-zero # window (acceptance-rate only; verified decoding keeps output correct). if num_contexts > 0: - self._seed_context_pages( + self._seed_context_windows( draft_model, spec_metadata, attn_metadata, @@ -997,6 +1137,7 @@ def _forward_impl( self._ctx_len.copy_(saved_ctx_len) self._valid_len.copy_(saved_valid_len) self._position_initialized.copy_(saved_position_initialized) + self._kv_windows.copy_(saved_windows) if saved_confidence is not None: self._confidence_logits.copy_(saved_confidence) if saved_stamp is not None: @@ -1018,210 +1159,6 @@ def _forward_impl( ) return outputs - def _release_inactive_slots(self, request_ids: list[int]) -> None: - current = set(request_ids) - for request_id in list(self._req_to_slot): - if request_id not in current: - slot = self._req_to_slot.pop(request_id) - self._managed_residency.pop(request_id, None) - self._ctx_len[slot] = 0 - self._valid_len[slot] = 0 - self._position_initialized[slot] = False - self._free_slots.append(slot) - - def _is_managed_request(self, request_id: int) -> bool: - return ( - request_id != ATTENTION_DP_DUMMY_REQUEST_ID - and request_id < self._graph_dummy_id_floor - and request_id not in self._draft_kv_manager._draft_dummy_request_ids - ) - - def _bind_managed_history(self, resource_manager) -> None: - manager = ( - resource_manager.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER) - if resource_manager is not None - else None - ) - if not getattr(manager, "draft_layer_ids", ()): - raise ValueError("Embedded DSpark requires manager-owned draft KV pages") - if manager is self._draft_kv_manager: - return - if ( - manager.draft_window_size != self._win - or len(manager.draft_layer_ids) != self._num_stages - or any( - manager.kv_factor_per_layer[manager.layer_offsets[layer_id]] != 1 - or manager.head_dim_per_layer[manager.layer_offsets[layer_id]] != self._head_dim - for layer_id in manager.draft_layer_ids - ) - ): - raise ValueError("Embedded DSpark draft cache does not match its attention dimensions") - buffers = tuple( - manager.get_draft_buffers(stage) for stage in range(len(manager.draft_layer_ids)) - ) - self._draft_kv_manager = manager - self._draft_kv_buffers = buffers - self._managed_residency.clear() - self._draft_page_indices.clear() - self._draft_block_tables = torch.zeros( - (self._batch_to_slot.shape[0], manager.draft_max_blocks_per_seq), - dtype=torch.int32, - device=self._ctx_len.device, - ) - self._draft_capacities = torch.zeros( - self._batch_to_slot.shape[0], dtype=torch.long, device=self._ctx_len.device - ) - - def prepare_managed_draft_cache( - self, draft_model, spec_metadata, attn_metadata, resource_manager - ) -> None: - """Refresh persistent draft inputs before eager execution or graph replay.""" - self._lazy_init(draft_model, spec_metadata, attn_metadata) - self._bind_managed_history(resource_manager) - self._prepare_managed_history(spec_metadata.request_ids, attn_metadata.num_contexts) - - def _prepare_managed_history(self, request_ids: list[int], num_contexts: int) -> None: - """Bind live page tables and initialize progress for newly resident histories.""" - # Startup probes and graph/ADP padding never publish synthetic history. - real_rows = [ - row - for row, request_id in enumerate(request_ids) - if self._is_managed_request(request_id) - ] - real_ids = [request_ids[row] for row in real_rows] - self._release_inactive_slots(real_ids) - batch_indices = self._draft_kv_manager.get_draft_block_indices(real_ids) - mappings = {row: tuple(indices) for row, indices in zip(real_rows, batch_indices)} - updates = [] - for row in sorted(self._draft_page_indices.keys() | mappings.keys()): - previous = self._draft_page_indices.get(row, ()) - current = mappings.get(row, ()) - if current == previous: - continue - # Clear removed mappings as well as uploading new/reassigned pages. - width = max(len(previous), len(current)) - old = previous + (0,) * (width - len(previous)) - new = current + (0,) * (width - len(current)) - first = next((i for i in range(width) if old[i] != new[i]), width) - if first == width: - continue - last = width - while old[last - 1] == new[last - 1]: - last -= 1 - updates.append((row, first, new[first:last])) - self._draft_page_indices = mappings - self._ctx_len[self._scratch_slot] = 0 - self._valid_len[self._scratch_slot] = 0 - self._position_initialized[self._scratch_slot] = False - batch_slots = [self._scratch_slot] * len(request_ids) - for row in real_rows: - request_id = request_ids[row] - cache = self._draft_kv_manager.kv_cache_map[request_id] - slot = self._assign_slot(request_id) - batch_slots[row] = slot - if self._managed_residency.get(request_id) is cache: - continue - position = cache.history_length - if row >= num_contexts and position == 0: - raise ValueError( - f"Embedded DSpark generation request {request_id} has no committed KV history" - ) - self._ctx_len[slot] = position - self._valid_len[slot] = min(self._win, position) - self._position_initialized[slot] = True - self._managed_residency[request_id] = cache - # Staging may still be read by an earlier asynchronous H2D. Reuse it - # only after its copy event, independently of overlapped model execution. - count = sum(len(indices) for _, _, indices in updates) - host = next( - (buffers for buffers in self._draft_metadata_host_buffers if buffers["event"].query()), - None, - ) - if host is None: - host = { - "tables": torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()), - "capacities": torch.empty( - self._draft_capacities.shape, dtype=torch.long, pin_memory=prefer_pinned() - ), - "slots": torch.empty( - self._batch_to_slot.shape, dtype=torch.long, pin_memory=prefer_pinned() - ), - "event": torch.cuda.Event(), - } - self._draft_metadata_host_buffers.append(host) - elif host["tables"].numel() < count: - host["tables"] = torch.empty(count, dtype=torch.int32, pin_memory=prefer_pinned()) - offset = 0 - for row, first, indices in updates: - size = len(indices) - source = host["tables"][offset : offset + size] - source.copy_(torch.tensor(indices, dtype=torch.int32)) - self._draft_block_tables[row, first : first + size].copy_(source, non_blocking=True) - offset += size - host["capacities"].zero_() - for row, request_id in zip(real_rows, real_ids): - host["capacities"][row] = self._draft_kv_manager.kv_cache_map[request_id].capacity - self._draft_capacities.copy_(host["capacities"], non_blocking=True) - host["slots"].fill_(self._scratch_slot) - host["slots"][: len(batch_slots)] = torch.tensor(batch_slots, dtype=torch.long) - self._batch_to_slot.copy_(host["slots"], non_blocking=True) - host["event"].record() - - def _seed_context_pages( - self, - draft_model, - spec_metadata: "DSparkSpecMetadata", - attn_metadata, - position_ids: torch.Tensor, - total_target_tokens: int, - ) -> None: - """Seed context chunks using their absolute positions. - - A request can arrive in multiple prefill chunks; continuation chunks - append to the same request's pages. - """ - captured = spec_metadata.get_hidden_states(total_target_tokens) - capture_lens = getattr(spec_metadata, "context_capture_lens", None) - if capture_lens is not None and len(capture_lens) != attn_metadata.num_contexts: - raise ValueError("DSpark capture lengths must match the context requests") - flat_position_ids = position_ids.reshape(-1) - context_offset = 0 - for i in range(attn_metadata.num_contexts): - chunk_len = int(attn_metadata._seq_lens[i]) - capture_len = chunk_len if capture_lens is None else capture_lens[i] - if not 0 <= capture_len <= chunk_len: - raise ValueError("DSpark capture length must be a suffix of its context chunk") - chunk_positions = flat_position_ids[context_offset : context_offset + chunk_len].long() - if chunk_len == 0: - context_offset += chunk_len - continue - - req_id = spec_metadata.request_ids[i] - slot = self._req_to_slot.get(req_id, self._scratch_slot) - if captured is None and slot != self._scratch_slot: - raise RuntimeError( - "Embedded DSpark requires captured context to initialize KV history" - ) - if captured is not None and capture_len: - keep = min(self._win, capture_len) - # Prefix reuse can publish any complete page in this chunk. - count = capture_len if self._draft_kv_manager.enable_block_reuse else keep - positions = (chunk_positions[-count:] + 1).unsqueeze(0) - draft_model.write_context_pages( - captured[ - context_offset + chunk_len - count : context_offset + chunk_len - ].unsqueeze(0), - positions, - torch.ones_like(positions, dtype=torch.bool), - self._draft_kv_buffers, - self._draft_block_tables[i : i + 1], - self._draft_capacities[i : i + 1], - ) - self._valid_len[slot] = torch.clamp(self._valid_len[slot] + keep, max=self._win) - self._ctx_len[slot] = chunk_positions[-1] + 1 - self._position_initialized[slot] = True - context_offset += chunk_len - class DSparkWorker(DFlashWorker): """Worker for a *standalone* DSpark drafter (DFlash lineage). @@ -1300,16 +1237,6 @@ def bind_draft_lm_head(self, draft_model) -> None: ): self._markov_tp_exchange = tp_pair_exchange(mapping, AllReduceStrategy.NCCL) - @staticmethod - def block_width_for(spec_config) -> int: - """Block slots a DSpark worker computes for ``spec_config``. - - The ``shift_label`` convention, which the DSpark drafter defaults to, - reads slots 0..K-1; a checkpoint that turns it off needs K+1 and is - caught by the block-within-allocation check. - """ - return spec_config.max_draft_len - def _draft_block_width(self, draft_model) -> int: """Block width under the dspark ``shift_label`` convention. @@ -1319,8 +1246,8 @@ def _draft_block_width(self, draft_model) -> int: drafters are meant to run. """ if getattr(draft_model, "_dspark_shift_label", False): - return self.block_width_for(self.spec_config) - return DFlashWorker.block_width_for(self.spec_config) + return self.max_draft_len + return super()._draft_block_width(draft_model) def _draft_slot_ids( self, draft_model, num_gens: int, block_size: int, num_draft_tokens: int diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 0f012b000373..3947fe41f020 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -37,8 +37,6 @@ if TYPE_CHECKING: from ..pyexecutor.guided_decoder import CapturableGuidedDecoder from ..pyexecutor.llm_request import LlmRequest - from ..pyexecutor.resource_manager import ResourceManager - from ..pyexecutor.standalone_draft_cache import DraftHistoryUpdate if IS_FLASHINFER_AVAILABLE: import flashinfer @@ -1693,19 +1691,6 @@ def commit_auxiliary_speculative_states( handler.commit_speculative_states(num_accepted_tokens, state_indices, num_contexts) - def prepare_managed_draft_cache( - self, - draft_model: nn.Module | None, - spec_metadata: "SpecMetadata", - attn_metadata: AttentionMetadata, - resource_manager: "ResourceManager", - ) -> None: - """Stage manager-owned draft history before eager execution or replay.""" - - def snapshot_managed_draft_history(self) -> Optional["DraftHistoryUpdate"]: - """Snapshot this execution's history for publication at completion.""" - return None - def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) if "forward" in cls.__dict__: diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index feb34389991b..f8e1a7340145 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -754,9 +754,9 @@ def get_spec_worker(spec_config, if spec_dec_mode.is_dflash(): return DFlashWorker(spec_config, mapping, use_separate_draft_kv_cache) # DSpark splits by deployment form, mirroring the draft-model side. The - # embedded DeepSeek-V4-Pro draft needs DSv4DSparkWorker, whose captured-context - # plumbing reads V4-draft-only attributes (num_stages, write_context_pages, - # forward). A standalone drafter is DFlash lineage and is served by + # embedded DeepSeek-V4-Pro draft needs DSv4DSparkWorker, whose rolling-window + # plumbing reads V4-draft-only attributes (num_stages, write_context_windows, + # forward_batched). A standalone drafter is DFlash lineage and is served by # DSparkWorker, which adds only the Markov bias and the shift_label # slot convention on top of DFlashWorker. if spec_dec_mode.is_dspark(): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 97f1b146017a..3eb903af9973 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -117,8 +117,6 @@ class _BatchDescFieldSpec: # field), matching the binding's constructor. import dataclasses as _dataclasses # noqa: E402 -CacheDomain = _cpp.CacheDomain - @_dataclasses.dataclass class _KVCacheManagerConfigFieldSpec: @@ -309,5 +307,3 @@ def typed_range(*args: int) -> range: "take_poison", "typed_range", ] - -__all__.append("CacheDomain") diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index ed7828864a56..e418a261853e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -75,10 +75,6 @@ def take_poison() -> str | None: def num_live_managers() -> int: """Number of constructed, not-yet-destroyed managers.""" -class CacheDomain(enum.Enum): - STANDALONE_DRAFT = 0 - TARGET = 1 - class CacheTier(enum.IntEnum): GPU_MEM = 0 HOST_MEM = 1 @@ -105,7 +101,6 @@ class AttnLifeCycle: num_sink_blocks: int, is_sparse: bool = False, reuse_policy: AttentionReusePolicy = ..., - cache_domain: CacheDomain = CacheDomain.TARGET, ) -> None: ... @staticmethod def make( @@ -114,7 +109,6 @@ class AttnLifeCycle: tokens_per_block: int, is_sparse: bool = False, reuse_policy: AttentionReusePolicy = ..., - cache_domain: CacheDomain = CacheDomain.TARGET, ) -> "AttnLifeCycle": ... @property def window_size(self) -> int | None: ... @@ -124,8 +118,6 @@ class AttnLifeCycle: def is_sparse(self) -> bool: ... @property def reuse_policy(self) -> AttentionReusePolicy: ... - @property - def cache_domain(self) -> CacheDomain: ... def get_stale_range(self, history_length: int, tokens_per_block: int) -> HalfOpenRange: ... CacheLevel = NewType("CacheLevel", int) @@ -285,7 +277,6 @@ class AttentionLayerConfig: sliding_window_size: int | None = None num_sink_tokens: int | None = None reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED - cache_domain: CacheDomain = CacheDomain.TARGET @property def window_size(self) -> int | None: ... diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index 6e26a946765b..8ae3b7a00c1f 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -160,11 +160,7 @@ def test_indexer_runtime_cache_cost_matches_registered_buffers(fp8, snapshot_int manager.head_dim_per_layer = [512, 512] manager.kv_cache_type = CacheTypeCpp.SELFKONLY manager.kv_factor = 1 - manager.kv_factor_per_layer = [1, 1] - manager.cp_as_tp_layer_ids = set() - manager._helix_cp_size = 1 manager.dtype = DataType.FP8 if fp8 else DataType.BF16 - manager.dtype_per_layer = [manager.dtype, manager.dtype] manager.tokens_per_block = 32 manager.max_batch_size = 4 manager.max_num_tokens = 16 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 43f36d08c3dd..e1992e41feee 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -92,7 +92,6 @@ def _make_creator( ) ) c._should_create_separate_draft_kv_cache = Mock(return_value=True) - c._uses_unified_standalone_draft_cache = Mock(return_value=False) c._get_draft_cache_cost = Mock( return_value=CacheCost( slope=total_kv_per_token - target_kv_per_token, @@ -256,7 +255,6 @@ def get_cache_size_per_token(model_config, *args, **kwargs): creator._kv_cache_manager_cls = DraftCostKVCacheManager creator._is_disagg = False creator._should_create_separate_draft_kv_cache = Mock(return_value=True) - creator._uses_unified_standalone_draft_cache = Mock(return_value=False) creator._get_effective_draft_config = Mock(return_value=draft_model_config) creator._get_num_draft_layers = Mock(return_value=1) get_manager_cls = mocker.patch( @@ -376,7 +374,6 @@ def get_cache_size_per_token( creator._is_disagg = False creator._cache_transceiver_config = None creator._should_create_separate_draft_kv_cache = Mock(return_value=True) - creator._uses_unified_standalone_draft_cache = Mock(return_value=False) creator._get_effective_draft_config = Mock(return_value=draft_model_config) creator._get_num_draft_layers = Mock(return_value=1) @@ -505,7 +502,6 @@ def test_v1_mixed_draft_build_uses_original_max_seq_len(self, mocker): c._enable_kv_cache_stats = Mock(return_value=False) c._fp8_ctx_mla_kv_len_cap = None c._should_create_separate_draft_kv_cache = Mock(return_value=True) - c._uses_unified_standalone_draft_cache = Mock(return_value=False) c._speculative_config = Mock() c._speculative_config.spec_dec_mode.is_external_drafter.return_value = True @@ -723,7 +719,6 @@ def get_cache_size_per_token(model_config, *args, **kwargs): c._is_disagg = True c._is_encoder_decoder = Mock(return_value=False) c._should_create_separate_draft_kv_cache = Mock(return_value=True) - c._uses_unified_standalone_draft_cache = Mock(return_value=False) c._get_num_draft_layers = Mock(return_value=5) mocker.patch( "tensorrt_llm._torch.pyexecutor._util.get_kv_cache_manager_cls", @@ -1046,7 +1041,6 @@ def _make_build_creator( c._fp8_ctx_mla_kv_len_cap = None c._model_engine.model.model_config.is_encoder_decoder = False c._should_create_separate_draft_kv_cache = Mock(return_value=True) - c._uses_unified_standalone_draft_cache = Mock(return_value=False) c._needs_gpu_kv_cache_budget_split = Mock(return_value=True) c._create_kv_cache_manager = Mock(return_value=Mock()) c._create_one_model_draft_kv_cache_manager = Mock(return_value=Mock()) @@ -1197,7 +1191,6 @@ def get_cache_size_per_token(model_config, *args, **kwargs): c._cache_transceiver_config = None c._is_encoder_decoder = Mock(return_value=False) c._should_create_separate_draft_kv_cache = Mock(return_value=True) - c._uses_unified_standalone_draft_cache = Mock(return_value=False) c._get_num_draft_layers = Mock(return_value=self.DRAFT_LAYERS) mocker.patch( "tensorrt_llm._torch.pyexecutor._util.get_kv_cache_manager_cls", diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index b8ad36ffd1b0..ef1f2d68b17b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -37,13 +37,11 @@ KVCacheLayerSpec, KVCacheManagerV2, Role, - _get_draft_layer_sizes, _get_static_cache_size_layer_components, ) from tensorrt_llm._torch.pyexecutor.model_engine import ModelEngine from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode -from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings import DataType from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( @@ -2061,47 +2059,3 @@ def test_create_kv_cache_manager_layer_specs_preserve_window_guards(guard: str) assert manager_config is kv_cache_config assert manager_config.max_attention_window == ([1024] if guard == "explicit_window" else None) - - -@pytest.mark.parametrize( - "attention_backend,dtype,bytes_per_layer", - [("TRTLLM", torch.float8_e4m3fn, 2048), ("VANILLA", torch.bfloat16, 4096)], -) -def test_standalone_draft_kv_dtype_follows_the_drafter_attention( - attention_backend, dtype, bytes_per_layer -) -> None: - """The TRTLLM drafter attention stores the drafter's history in FP8; VANILLA keeps BF16.""" - creator = object.__new__(KvCacheCreator) - draft = SimpleNamespace( - fc=SimpleNamespace(weight=torch.empty(1, dtype=torch.bfloat16)), - dflash_attention_backend=attention_backend, - config=SimpleNamespace(max_position_embeddings=4096), - ) - creator._model_engine = SimpleNamespace(model=SimpleNamespace(draft_model=draft)) - creator._speculative_config = SimpleNamespace( - spec_dec_mode=SimpleNamespace(is_dspark=lambda: True), - draft_is_embedded_in_target=False, - _use_shared_kv_cache=False, - attention_backend="AUTO", - max_draft_len=3, - ) - # Drafter shapes remain unsharded under attention DP. - creator._draft_config = SimpleNamespace( - pretrained_config=SimpleNamespace( - num_attention_heads=64, - num_key_value_heads=16, - head_dim=64, - hidden_size=4096, - num_hidden_layers=5, - ) - ) - creator._mapping = SimpleNamespace( - has_cp_helix=lambda: False, enable_attention_dp=True, tp_size=16 - ) - - creator._draft_config.mapping = creator._mapping - - kwargs = creator._get_unified_draft_cache_kwargs() - - assert kwargs["draft_dtype"] == torch_dtype_to_binding(dtype) - assert _get_draft_layer_sizes(**kwargs) == (bytes_per_layer,) * 5 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index 47c0df353fd9..7f9d7ef36e06 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -229,7 +229,6 @@ def test_update_resources_leaves_history_untouched_under_helix() -> None: # rewind falls back to the reserve width. _allocated_draft_lens={}, _inflight_draft_lens={}, - _standalone_draft_layer_ids=(), kv_cache_map={7: kv}, kv_compression_manages_history=False, _has_cp_helix=True, @@ -388,7 +387,6 @@ def _context_resize_mgr(is_disagg: bool, resizes: list) -> SimpleNamespace: is_disagg=is_disagg, kv_cache_map={7: kv}, num_extra_kv_tokens=0, - _standalone_draft_reserve=0, _connector_reservations_enabled=lambda: False, _fill_fresh_kv_pages=lambda req_id: None, _log_window_crossing=lambda *args: None, diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 58259fbf57ca..51321fbb476b 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2745,15 +2745,10 @@ def allocated_memory(pool_ratio): _install_default_v2_state(mgr) mgr._num_reserved_dummy_slots = 1 mgr.dtype = DataType.HALF - mgr.dtype_per_layer = [mgr.dtype] * mgr.num_local_layers - mgr.cp_as_tp_layer_ids = set() - mgr._helix_cp_size = 1 - mgr._standalone_draft_layer_ids = () - mgr._draft_layer_sizes = [] mgr.enable_swa_scratch_reuse = False mgr.enable_stats = False mgr.num_extra_kv_tokens = 0 - mgr.get_layer_bytes_per_token = lambda local_layer_idx, data_role: 8 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 kv_cache_config = KvCacheConfig( pool_ratio=pool_ratio, diff --git a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py index 5913fc3089df..89d9e727ff11 100644 --- a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py +++ b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py @@ -64,7 +64,6 @@ def _manager() -> Mock: manager._has_cp_helix = False manager.num_extra_kv_tokens = 0 manager.kv_cache_map = {} - manager._standalone_draft_layer_ids = () manager.prepare_context_cache.return_value = 0 manager.prepare_disagg_gen_init = MethodType(KVCacheManagerV2.prepare_disagg_gen_init, manager) manager.prepare_resources = MethodType(KVCacheManagerV2.prepare_resources, manager) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py index 6c08d6f38322..f88ba9687763 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py @@ -17,7 +17,6 @@ def _worker(req_to_slot): worker = SimpleNamespace( _ctx_buf_inited=True, - _has_unified_draft_cache=lambda: False, _req_to_slot=dict(req_to_slot), _req_ctx_pos={}, _free_slots=[], diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py index 991bebfd89ce..269e1de42197 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -40,7 +40,6 @@ def _worker(lengths): w = SimpleNamespace( _ctx_buf_inited=True, - _has_unified_draft_cache=lambda: False, _req_to_slot={11: 0, 12: 1, 13: 2}, _req_ctx_pos={}, _free_slots=[3], diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py index fc71e7deae1a..e77b14140dd5 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py @@ -164,124 +164,3 @@ def precompute_context_kv(projected, positions): real_len = 5 + 2 * (step + 1) assert worker._ctx_len.tolist() == [real_len, 0, 0] assert inputs["num_ctx_per_req"].tolist() == [real_len, 0] - - -def _make_worker(monkeypatch): - monkeypatch.setattr( - "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", - lambda **kwargs: None, - ) - config = DFlashDecodingConfig(max_draft_len=7, attention_backend="TRTLLM") - return DFlashWorker(config, Mapping()) - - -class _FakeDraftCacheManager: - """Unified manager whose draft layer 61 lives in pool 1 with two KV planes.""" - - draft_layer_ids = [61] - layer_offsets = {61: 7} - layer_to_pool_mapping_dict = {7: 1} - _index_scale_ints = [1, 1] - enable_swa_scratch_reuse = False - - def __init__(self, pages, helix_cp_size=1): - self.pages = pages - self._helix_cp_size = helix_cp_size - self.kv_cache_map = { - request_id: SimpleNamespace(num_blocks=len(request_pages), is_active=True) - for request_id, request_pages in pages.items() - } - - def get_layer_page_index_scale(self, layer_idx): - return 1 - - def get_layer_kv_factor(self, layer_idx): - return 2 - - def get_batch_cache_indices(self, request_ids, layer_idx): - return [list(self.pages[request_id]) for request_id in request_ids] - - def block_offsets(self, request_ids, width): - """The target's [pool, batch, plane, block] offsets: K at page * 2, V at page * 2 + 1.""" - offsets = torch.full((2, len(request_ids), 2, width), -1, dtype=torch.int32) - for row, request_id in enumerate(request_ids): - pages = torch.tensor(self.pages[request_id], dtype=torch.int32) - offsets[0, row, :, : len(pages)] = 1000 # the target pool's own pages - offsets[1, row, 0, : len(pages)] = pages * 2 - offsets[1, row, 1, : len(pages)] = pages * 2 + 1 - return offsets.cuda() - - -def _bind_device_tables(worker, manager): - worker._ctx_kv_manager = manager - worker._ctx_layer_ids = [61] - manager.kv_cache_pool_mapping = torch.tensor([[0, 0]] * 7 + [[1, 0]]) - manager.max_blocks_per_seq = 4 - pool = [torch.empty(16, 2, 1, 32, 64, device="cuda")] - assert worker._init_ctx_block_tables(manager, pool, 3, 1, 1, 64, 32) - worker._init_unified_draft_tables(manager) - - -def test_unified_draft_tables_are_decoded_from_target_block_offsets(monkeypatch): - worker = _make_worker(monkeypatch) - manager = _FakeDraftCacheManager({11: [5, 9, 2], 12: [7]}) - _bind_device_tables(worker, manager) - assert worker._device_draft_tables - assert (worker._ctx_pool_idx, worker._ctx_block_divisor) == (1, 2) - - attn_metadata = SimpleNamespace(kv_cache_block_offsets=manager.block_offsets([11, 12], 4)) - assert worker._refresh_unified_block_tables(attn_metadata, 2, [11, 12]) - - tables = worker._ctx_block_tables.cpu() - assert tables[0, :3].tolist() == [5, 9, 2] - assert tables[1, :1].tolist() == [7] - assert worker._ctx_block_counts[:2].tolist() == [3, 1] - # The startup check compared the rows with the manager's page lists. - assert worker._device_draft_tables - assert worker._draft_table_checks_left == 7 - - -def test_unified_draft_tables_fall_back_to_host_on_mismatch(monkeypatch): - worker = _make_worker(monkeypatch) - manager = _FakeDraftCacheManager({11: [5, 9, 2], 12: [7]}) - _bind_device_tables(worker, manager) - attn_metadata = SimpleNamespace(kv_cache_block_offsets=manager.block_offsets([11, 12], 4)) - manager.pages[12] = [8] - - assert not worker._refresh_unified_block_tables(attn_metadata, 2, [11, 12]) - assert not worker._device_draft_tables - - -def test_unified_draft_tables_stay_on_host_for_helix_managers(monkeypatch): - worker = _make_worker(monkeypatch) - _bind_device_tables(worker, _FakeDraftCacheManager({}, helix_cp_size=2)) - - assert not worker._device_draft_tables - - -def test_fp8_draft_cache_stores_context_kv_by_direct_page_writes(monkeypatch): - worker = _make_worker(monkeypatch) - num_layers, nkv, page_size, hd = 2, 2, 32, 64 - worker._ctx_kv_buf = [ - torch.zeros(4, 2, nkv, page_size, hd, dtype=torch.float8_e4m3fn, device="cuda") - for _ in range(num_layers) - ] - worker._ctx_block_tables = torch.tensor([[3, 1], [0, 2]], dtype=torch.int32, device="cuda") - worker._ctx_page_size = page_size - rows = torch.tensor([0, 0, 1], device="cuda") - positions = torch.tensor([5, 40, 33], device="cuda") - k = torch.randn(3, num_layers, nkv, hd, dtype=torch.bfloat16, device="cuda") - v = torch.randn_like(k) - - worker._store_context_kv_paged(k, v, rows, positions) - - pages = torch.tensor([3, 1, 2], device="cuda") - offsets = positions % page_size - for layer, pool in enumerate(worker._ctx_kv_buf): - stored = pool.view(torch.uint8) - expected_k = k[:, layer].to(pool.dtype).view(torch.uint8) - expected_v = v[:, layer].to(pool.dtype).view(torch.uint8) - assert torch.equal(stored[pages, 0, :, offsets], expected_k) - assert torch.equal(stored[pages, 1, :, offsets], expected_v) - stored[pages, :, :, offsets] = 0 - assert not stored.any(), "the store wrote outside the addressed slots" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index 392ac59cc4b5..d8d8cd962939 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -1364,3 +1364,111 @@ def transform_weights(self): assert draft.main_proj.transforms == 1 assert stage.self_attn[0].transforms == 0 + + +@pytest.mark.parametrize("batched", [False, True]) +def test_rolling_attention_uses_projection_modules(monkeypatch, batched): + monkeypatch.setattr(modeling_dspark, "IS_CUTLASS_DSL_AVAILABLE", False) + g = _make_attn_inputs() + weights = dict( + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + persist=True, + ) + expected_cache = g["kv_cache0"].clone() + actual_cache = g["kv_cache0"].clone() + if batched: + fn = modeling_dspark.dspark_attention_forward_batched + args = (torch.full((g["b"],), g["start_pos"]),) + tail = (torch.arange(g["b"]),) + else: + fn = dspark_attention_forward + args = (g["start_pos"],) + tail = () + expected = fn(g["x"], g["main_x"], *args, expected_cache, *tail, **weights) + qkv_weight = torch.cat((g["wq_a"], g["wkv"]), dim=0) + weights.update( + wq_a=None, + wkv=None, + wo_a=None, + qkv_proj=lambda x: F.linear(x, qkv_weight), + wq_b=lambda x: F.linear(x, g["wq_b"]), + wo_b=lambda x: F.linear(x, g["wo_b"]), + o_a_proj=g["wo_a"].view(g["n_groups"], g["o_lora"], -1), + o_a_scale=None, + o_a_dequant=None, + use_fused_output_projection=False, + use_fused_qkv_input_projection=False, + use_fused_q_input_projection=False, + ) + actual = fn(g["x"], g["main_x"], *args, actual_cache, *tail, **weights) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual_cache, expected_cache) + + +@pytest.mark.parametrize("batched", [False, True]) +@pytest.mark.parametrize("stacked", [False, True]) +def test_context_window_writes_use_projection_modules(batched, stacked): + g = _make_attn_inputs() + + def make_model(use_modules): + model = object.__new__(DSv4DSparkDraftModel) + torch.nn.Module.__init__(model) + model._attn_params = dict( + window_size=g["window"], + rope_head_dim=g["rope_head_dim"], + head_dim=g["head_dim"], + eps=g["eps"], + ) + model._context_wkv_fp8 = None + model._dspark_freqs_table = lambda device: g["freqs"].to(device) + qkv_weight = torch.cat((g["wq_a"], g["wkv"]), dim=0) + model.mtp_layers = [ + types.SimpleNamespace( + main_proj=torch.nn.Identity(), + main_norm=torch.nn.Identity(), + _dspark_attn=dict( + wkv=None if use_modules else g["wkv"], + kv_norm_w=g["kv_norm"], + qkv_proj=(lambda x: F.linear(x, qkv_weight)) if use_modules else None, + ), + ) + for _ in range(2) + ] + if use_modules and stacked: + model._project_context_wkv = lambda x: torch.stack([F.linear(x, g["wkv"])] * 2, dim=-2) + return model + + reference, model = make_model(False), make_model(True) + if batched: + hidden = g["x"] + positions = torch.tensor([[1, 2, 3], [4, 5, 6]]) + slots = torch.tensor([1, 0]) + mask = torch.tensor([[True, False, True], [False, True, True]]) + expected = torch.randn(3, 2, g["window"], g["head_dim"], dtype=torch.bfloat16) + actual = expected.clone() + reference.write_context_windows_batched(hidden, positions, slots, mask, expected) + model.write_context_windows_batched(hidden, positions, slots, mask, actual) + else: + hidden = g["x"][0] + positions = torch.tensor([1, 4, 6]) + expected = torch.randn(2, g["window"], g["head_dim"], dtype=torch.bfloat16) + actual = expected.clone() + reference.write_context_windows(hidden, positions, expected) + model.write_context_windows(hidden, positions, actual) + torch.testing.assert_close(actual, expected) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index 5d77d8cff968..95fa963affed 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -214,12 +214,14 @@ def test_batched_topk_respects_partial_window_valid_len(): @pytest.mark.parametrize("ndim", [3, 4]) def test_batched_rotary_matches_scalar_per_row(ndim): + """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row.""" torch.manual_seed(0) G, s, h, rd = 3, 4, 2, 8 x = torch.randn(G, s, h, rd) if ndim == 4 else torch.randn(G, s, rd) table = precompute_dspark_freqs_cis(rd, 64) + # Per-row absolute start positions -> per-row freq windows. starts = [1, 9, 30] - per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) + per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) # [G, s, rd/2] got = apply_dspark_rotary_batched(x, per_row) for i, sp in enumerate(starts): ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py index 113b4be2591b..8fd59c23d4f6 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py @@ -370,10 +370,6 @@ def test_v41_stage_collapses_with_previous_not_current_mix(monkeypatch): 1, torch.empty(0), torch.zeros(2), - kv_cache=torch.empty(0), - block_tables=torch.empty(0, dtype=torch.int32), - capacities=torch.empty(0, dtype=torch.long), - valid_len=torch.empty(0, dtype=torch.long), pre_mix=previous_pre, ) assert attention_mixer.pre_mapping_lagged.call_args.args[1] is previous_pre diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 866918e842ed..40786627b200 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -22,7 +22,6 @@ """ import types -from unittest.mock import Mock import pytest import torch @@ -123,7 +122,6 @@ def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): "softmax_scale": 1.0, "eps": 1e-6, }, - configure_generation_autotuner=Mock(), ) @@ -132,13 +130,11 @@ def test_worker_lazy_init_window_buffers(): dm = _fake_draft_model(num_stages=3, window_size=128, head_dim=64) meta = _make_metadata(max_num_requests=8) worker._lazy_init(dm, meta) - manager = _bind_manager(worker) # max_batch (8) request slots + 1 scratch row for padded / unknown IDs. - assert worker._draft_block_tables.shape == (8, manager.draft_max_blocks_per_seq) - assert worker._draft_kv_buffers[0] is manager.get_draft_buffers(0) + assert worker._kv_windows.shape == (9, 3, 128, 64) assert worker._ctx_len.shape == (9,) assert worker._valid_len.shape == (9,) - assert worker._draft_capacities.shape == (8,) + assert worker._position_initialized.shape == (9,) assert worker._scratch_slot == 8 # Dummy-id floor separates real request ids from CUDA-graph padding ids. assert worker._graph_dummy_id_floor == (1 << 64) - 1 - worker.max_draft_len @@ -149,9 +145,9 @@ def test_worker_lazy_init_window_buffers(): assert worker._batch_to_slot.device.type == "cuda" assert worker._batch_to_slot.tolist() == [worker._scratch_slot] * 8 # idempotent - buf_id = worker._draft_block_tables.data_ptr() + buf_id = id(worker._kv_windows) worker._lazy_init(dm, meta) - assert worker._draft_block_tables.data_ptr() == buf_id + assert id(worker._kv_windows) == buf_id def test_worker_rejects_mismatched_block_size(): @@ -167,36 +163,94 @@ def test_worker_slot_assignment_and_reset(): worker = _make_worker() worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=4)) - s0 = worker._assign_slot(100) - s1 = worker._assign_slot(101) + s0 = worker._assign_slot(100, reset=False) + s1 = worker._assign_slot(101, reset=False) assert s0 != s1 # same request id -> same slot (no reset) - assert worker._assign_slot(100) == s0 + assert worker._assign_slot(100, reset=False) == s0 # mark a position, then reset -> slot freed + window/pos cleared worker._ctx_len[s0] = 42 worker._valid_len[s0] = 8 - worker._release_inactive_slots([101]) - s0b = worker._assign_slot(100) + worker._position_initialized[s0] = True + worker._kv_windows[s0].fill_(1.0) + s0b = worker._assign_slot(100, reset=True) assert int(worker._ctx_len[s0b]) == 0 assert int(worker._valid_len[s0b]) == 0 - assert worker._req_to_slot[101] == s1 + assert not bool(worker._position_initialized[s0b]) + assert float(worker._kv_windows[s0b].abs().sum()) == 0.0 def test_worker_slot_exhaustion_preserves_live_request(): worker = _make_worker() worker._lazy_init(_fake_draft_model(), _make_metadata(max_num_requests=1)) - slot = worker._assign_slot(100) + slot = worker._assign_slot(100, reset=False) worker._ctx_len[slot] = 42 - worker._valid_len[slot] = 1 + worker._kv_windows[slot].fill_(1.0) - with pytest.raises(RuntimeError, match="no free progress slots"): - worker._assign_slot(101) + with pytest.raises(RuntimeError, match="no free rolling-window slots"): + worker._assign_slot(101, reset=False) assert worker._req_to_slot == {100: slot} assert int(worker._ctx_len[slot]) == 42 - assert int(worker._valid_len[slot]) == 1 + assert torch.all(worker._kv_windows[slot] == 1.0) + + +def test_seed_context_windows_preserves_state_across_prefill_chunks(): + class DraftModel: + num_stages = 1 + block_size = 5 + _attn_params = { + "window_size": 8, + "head_dim": 4, + "n_heads": 128, + "softmax_scale": 1.0, + "eps": 1e-6, + } + + def __init__(self): + self.written_positions = [] + + def write_context_windows(self, hidden, positions, windows): + self.written_positions.append(positions.clone()) + windows.add_(1) + + worker = _make_worker() + draft_model = DraftModel() + metadata = types.SimpleNamespace( + max_num_requests=1, + request_ids=[100], + get_hidden_states=lambda _num_tokens: torch.zeros( + 3, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 + ), + ) + worker._lazy_init(draft_model, metadata) + + first_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[3]) + worker._seed_context_windows( + draft_model, metadata, first_chunk, torch.tensor([[0, 1, 2]], device="cuda"), 3 + ) + slot = worker._req_to_slot[100] + assert int(worker._ctx_len[slot]) == 3 + assert int(worker._valid_len[slot]) == 3 + assert bool(worker._position_initialized[slot]) + + metadata.get_hidden_states = lambda _num_tokens: torch.zeros( + 2, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 + ) + second_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[2]) + worker._seed_context_windows( + draft_model, metadata, second_chunk, torch.tensor([[3, 4]], device="cuda"), 2 + ) + + assert int(worker._ctx_len[slot]) == 5 + assert int(worker._valid_len[slot]) == 5 + assert [positions.tolist() for positions in draft_model.written_positions] == [ + [1, 2, 3], + [4, 5], + ] + assert torch.all(worker._kv_windows[slot] == 2.0) @pytest.mark.parametrize( @@ -211,25 +265,21 @@ def test_worker_slot_exhaustion_preserves_live_request(): ) @pytest.mark.parametrize("window", [64, 128, 256]) @pytest.mark.parametrize("final_window", [False, True]) -def test_seed_context_pages_reads_only_each_chunk_tail(chunks, window, final_window): +def test_seed_context_windows_reads_only_each_chunk_tail(chunks, window, final_window): """Never consume dropped captures, including an entirely skipped chunk.""" worker = _make_worker() draft = _fake_draft_model() draft._attn_params = dict(draft._attn_params, window_size=window) written = [] - def write(hidden, positions, valid, pages, tables, capacities): + def write(hidden, positions, windows): assert torch.isfinite(hidden).all() - torch.testing.assert_close(hidden[..., 0].float(), positions.float() - 1) - written.append(positions.squeeze(0).tolist()) + torch.testing.assert_close(hidden[:, 0].float(), positions.float() - 1) + written.append(positions.tolist()) - draft.write_context_pages = write + draft.write_context_windows = write meta = types.SimpleNamespace(max_num_requests=2, request_ids=[100, 101]) worker._lazy_init(draft, meta) - manager = _bind_manager(worker, meta.request_ids) - for request in manager.kv_cache_map.values(): - request.history_length = 0 - worker._prepare_managed_history(meta.request_ids, 2) expected = [] total_kept = 0 prompt_end = chunks[-1][0] + chunks[-1][1] @@ -255,7 +305,7 @@ def write(hidden, positions, valid, pages, tables, capacities): packed_positions = torch.cat( (positions, positions.new_tensor([0, 1, 2, 50, 51, 52, 53, 54, 55])) ) - worker._seed_context_pages(draft, meta, attn, packed_positions, len(captured)) + worker._seed_context_windows(draft, meta, attn, packed_positions, len(captured)) end = start + length if kept: expected.append(list(range(end - min(window, kept) + 1, end + 1))) @@ -274,11 +324,11 @@ def test_prepare_builds_batch_to_slot_on_batched_path(): meta._dspark_worker = worker # Assign slots for two requests (as the prefill path would). - sa = worker._assign_slot(100) - sb = worker._assign_slot(101) + sa = worker._assign_slot(100, reset=True) + sb = worker._assign_slot(101, reset=True) meta.request_ids = [101, 100] - _prepare(worker, meta) + meta.prepare() # Mirror reflects request-order -> slot. assert worker._batch_to_slot[:2].tolist() == [sb, sa] @@ -290,19 +340,20 @@ def test_prepare_frees_stale_slots_on_batched_path(): worker._lazy_init(_fake_draft_model(), meta) meta._dspark_worker = worker - sa = worker._assign_slot(100) - worker._assign_slot(101) + sa = worker._assign_slot(100, reset=True) + worker._assign_slot(101, reset=True) worker._ctx_len[sa] = 17 worker._valid_len[sa] = 8 + worker._position_initialized[sa] = True # Only request 101 survives; 100's slot must be freed + cleared. meta.request_ids = [101] - _prepare(worker, meta) + meta.prepare() assert 100 not in worker._req_to_slot assert sa in worker._free_slots assert int(worker._ctx_len[sa]) == 0 assert int(worker._valid_len[sa]) == 0 - assert 100 not in worker._managed_residency + assert not bool(worker._position_initialized[sa]) def test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(): @@ -318,16 +369,14 @@ def test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(): meta._dspark_worker = worker # A live request takes the first free slot (0) and populates its window. - s_real = worker._assign_slot(100) + s_real = worker._assign_slot(100, reset=True) assert s_real == 0 worker._ctx_len[s_real] = 17 - worker._valid_len[s_real] = 5 - manager = _bind_manager(worker) - manager._draft_dummy_request_ids.add(999) + worker._kv_windows[s_real].fill_(1.0) # Batch contains the live request plus an unknown id (e.g. graph padding). meta.request_ids = [100, 999] - _prepare(worker, meta) + meta.prepare() # The unknown id maps to the scratch row, not to slot 0. assert worker._batch_to_slot[:2].tolist() == [s_real, worker._scratch_slot] @@ -337,7 +386,7 @@ def test_prepare_maps_unknown_request_to_scratch_row_not_slot_zero(): assert list(worker._free_slots) == [1, 2, 3] # The live request's rolling window and position are untouched. assert int(worker._ctx_len[s_real]) == 17 - assert int(worker._valid_len[s_real]) == 5 + assert torch.all(worker._kv_windows[s_real] == 1.0) def test_prepare_assigns_slots_to_disagg_generation_requests(): @@ -345,7 +394,7 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): Regression for GitHub #16767: on the disaggregated generation server the prompt is prefilled (and the DSpark window seeded) on the *context* server, - so ``_seed_context_pages`` never runs here and ``_req_to_slot`` stays empty. + so ``_seed_context_windows`` never runs here and ``_req_to_slot`` stays empty. ``prepare()`` must therefore assign each real generation request its own rolling-window slot instead of collapsing them all onto the shared scratch row (which corrupts drafts and collapses accept length at batch size > 1). @@ -358,7 +407,7 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): # All-generation batch (num_contexts == 0), no prior seeding. meta.request_ids = [1000, 1001] meta.num_generations = 2 - _prepare(worker, meta) + meta.prepare() s0 = worker._req_to_slot[1000] s1 = worker._req_to_slot[1001] @@ -367,7 +416,7 @@ def test_prepare_assigns_slots_to_disagg_generation_requests(): assert worker._batch_to_slot[:2].tolist() == [s0, s1] # Stable across steps: the same ids keep their slots (no churn / reassignment). - _prepare(worker, meta) + meta.prepare() assert worker._req_to_slot[1000] == s0 assert worker._req_to_slot[1001] == s1 assert worker._batch_to_slot[:2].tolist() == [s0, s1] @@ -386,7 +435,7 @@ def test_prepare_keeps_dummy_generation_requests_on_scratch_row(): graph_dummy = CUDA_GRAPH_DUMMY_REQUEST_ID - worker.max_draft_len meta.request_ids = [1000, ATTENTION_DP_DUMMY_REQUEST_ID, graph_dummy] meta.num_generations = 3 - _prepare(worker, meta) + meta.prepare() s_real = worker._req_to_slot[1000] assert s_real != worker._scratch_slot @@ -410,15 +459,17 @@ def test_prepare_keeps_small_real_request_slots_across_steps(): meta._dspark_worker = worker meta.request_ids = [1, 2] meta.num_generations = 2 - _prepare(worker, meta) + meta.prepare() slots = worker._batch_to_slot[:2] worker._ctx_len[slots] = torch.tensor([7, 8], device="cuda") worker._valid_len[slots] = torch.tensor([1, 2], device="cuda") + worker._position_initialized[slots] = True - _prepare(worker, meta) + meta.prepare() assert worker._ctx_len[slots].tolist() == [7, 8] assert worker._valid_len[slots].tolist() == [1, 2] + assert worker._position_initialized[slots].tolist() == [True, True] def test_prepare_keeps_cuda_graph_slots_across_replays(): @@ -427,21 +478,25 @@ def test_prepare_keeps_cuda_graph_slots_across_replays(): meta = _make_metadata(max_num_requests=4) worker._lazy_init(_fake_draft_model(), meta) meta._dspark_worker = worker - meta.request_ids = [100, 101, 102] + # Not 0: that is ATTENTION_DP_DUMMY_REQUEST_ID, which prepare() routes to + # the scratch row by design, so it would never take a persistent slot. + meta.request_ids = [11, 12, 13] meta.num_generations = 3 meta.is_cuda_graph = True - _prepare(worker, meta) + meta.prepare() slots = worker._batch_to_slot[:3] assert len(set(slots.tolist())) == 3 assert all(slot != worker._scratch_slot for slot in slots.tolist()) worker._ctx_len[slots] = torch.tensor([7, 8, 9], device="cuda") worker._valid_len[slots] = torch.tensor([1, 2, 3], device="cuda") + worker._position_initialized[slots] = True - _prepare(worker, meta) + meta.prepare() slots = worker._batch_to_slot[:3] assert worker._ctx_len[slots].tolist() == [7, 8, 9] assert worker._valid_len[slots].tolist() == [1, 2, 3] + assert worker._position_initialized[slots].tolist() == [True, True, True] class _RecordingDraftModel: @@ -458,13 +513,10 @@ class _RecordingDraftModel: def __init__(self): self.forward_calls = [] - def configure_generation_autotuner(self, _max_batch): + def write_context_windows_batched(self, *args): pass - def write_context_pages(self, *args): - pass - - def forward(self, main_hidden, bonus, start_pos, **kwargs): + def forward_batched(self, main_hidden, bonus, start_pos, **kwargs): self.forward_calls.append( { "main_hidden": main_hidden.clone(), @@ -489,11 +541,13 @@ def test_generation_state_cuda_graph_bootstrap_and_replay(): with torch.cuda.graph(graph): worker._advance_generation_state(slots, num_accepted, input_positions) - worker._ctx_len[slots] = input_positions + worker._ctx_len[slots] = 0 worker._valid_len[slots] = 0 + worker._position_initialized[slots] = False graph.replay() assert worker._ctx_len[slots].tolist() == [4017, 89] assert worker._valid_len[slots].tolist() == [1, 2] + assert worker._position_initialized[slots].tolist() == [True, True] num_accepted.copy_(torch.tensor([2, 3], device="cuda")) graph.replay() @@ -540,13 +594,8 @@ def test_disagg_position_bootstrap_uses_actual_positions_and_target_width(): verify_lens=None, ) worker._lazy_init(draft_model, metadata) - manager = _bind_manager(worker, [1000, 1001]) - manager.kv_cache_map[1000].history_length = 4016 - manager.kv_cache_map[1001].history_length = 87 - worker._prepare_managed_history([1000, 1001], 0) - worker._valid_len.zero_() - slots = [worker._assign_slot(1000), worker._assign_slot(1001)] + slots = [worker._assign_slot(1000, reset=False), worker._assign_slot(1001, reset=False)] worker._batch_to_slot[:2] = torch.tensor(slots, device="cuda") captured = torch.stack( @@ -579,6 +628,7 @@ def test_disagg_position_bootstrap_uses_actual_positions_and_target_width(): assert worker._ctx_len[slots].tolist() == [4017, 88] assert first_call["valid_len"].tolist() == [1, 1] assert worker._valid_len[slots].tolist() == [1, 1] + assert worker._position_initialized[slots].tolist() == [True, True] # Existing slots retain their state. The normal K+1 target layout must still # select each accepted bonus hidden from its own packed request row. @@ -660,7 +710,7 @@ def fake_accept(logits, am, sm): monkeypatch.setattr(worker, "sample_and_accept_draft_tokens", fake_accept) # Context-window seeding is covered by its own test; stub it out here. - monkeypatch.setattr(worker, "_seed_context_pages", lambda *a, **k: None) + monkeypatch.setattr(worker, "_seed_context_windows", lambda *a, **k: None) # The gen-block helper now returns the corrected block logits [num_gens,K,vocab]. gen_logits = torch.randn(num_gens, K, vocab, device="cuda") @@ -694,8 +744,6 @@ def fake_sample_draft_tokens(gl, sm, bs, *, num_contexts): hidden = torch.zeros(batch_size, HIDDEN, device="cuda", dtype=torch.bfloat16) logits = torch.zeros(batch_size, vocab, device="cuda") - worker._lazy_init(dm, meta) - _prepare(worker, meta, num_contexts) out = worker.forward(input_ids, position_ids, hidden, logits, attn_metadata, meta, dm) # Acceptance went through the unified entry with the right metadata objects. @@ -762,7 +810,7 @@ def test_forward_guided_batch_masks_and_advances_matcher_per_step(monkeypatch): monkeypatch.setattr( worker, "sample_and_accept_draft_tokens", lambda *a, **k: (accepted, num_accepted) ) - monkeypatch.setattr(worker, "_seed_context_pages", lambda *a, **k: None) + monkeypatch.setattr(worker, "_seed_context_windows", lambda *a, **k: None) gen_logits = torch.randn(num_gens, K, vocab, device="cuda") monkeypatch.setattr(worker, "_draft_gen_block_batched", lambda *a, **k: gen_logits) @@ -825,8 +873,6 @@ def fake_sample_draft_tokens(gl, sm, bs, *, draft_step): hidden = torch.zeros(batch_size, HIDDEN, device="cuda", dtype=torch.bfloat16) logits = torch.zeros(batch_size, vocab, device="cuda") - worker._lazy_init(dm, meta) - _prepare(worker, meta, num_contexts) out = worker.forward(input_ids, position_ids, hidden, logits, attn_metadata, meta, dm) # Matcher advanced once per draft position, in order 0..K-1. @@ -937,10 +983,10 @@ def test_lazy_init_publishes_the_runtime_position_ceiling(): worker._lazy_init(dm, _make_metadata(max_num_requests=2), attn_metadata) - # A full target verification accepts K+1 tokens, then DSpark indexes K - # block positions from that new start. The table length is max index + 1. - # Managed cache positions include the next RoPE frame in addition to the logical index. - expected = 4096 + 5 + 1 + 5 + 2 + # start_pos is a frame index (position + 1); a full verification then accepts + # K+1 tokens and DSpark indexes K block positions from that start. Length is + # max index + 1. + expected = 4096 + 1 + 5 + 1 + 5 + 1 assert dm._runtime_position_ceiling == expected assert dm.dspark_model._runtime_position_ceiling == expected @@ -990,7 +1036,8 @@ def test_position_bound_holds_across_cuda_graph_replays(): with torch.cuda.graph(graph): worker._advance_generation_state(slots, num_accepted, input_positions) - worker._ctx_len[slots] = input_positions + worker._ctx_len[slots] = 0 + worker._position_initialized[slots] = False for _ in range(5): graph.replay() assert worker._ctx_len[0].item() <= cap @@ -1049,21 +1096,23 @@ def test_prepare_resets_the_scratch_row_every_iteration(): # State a previous padding / ADP-idle forward left behind. worker._ctx_len[scratch] = 4096 worker._valid_len[scratch] = 8 + worker._position_initialized[scratch] = True - live = worker._assign_slot(100) + live = worker._assign_slot(100, reset=True) worker._ctx_len[live] = 17 worker._valid_len[live] = 5 + worker._position_initialized[live] = True meta.request_ids = [100, CUDA_GRAPH_DUMMY_REQUEST_ID] - _prepare(worker, meta) + meta.prepare() assert int(worker._ctx_len[scratch]) == 0 assert int(worker._valid_len[scratch]) == 0 - assert worker._draft_capacities[1].item() == 0 + assert not bool(worker._position_initialized[scratch]) # A live request's own state is bounded by its lifetime and must survive. assert int(worker._ctx_len[live]) == 17 assert int(worker._valid_len[live]) == 5 - assert worker._managed_residency[100] is worker._draft_kv_manager.kv_cache_map[100] + assert bool(worker._position_initialized[live]) def test_ragged_draft_gathers_each_requests_hidden_state(): @@ -1075,8 +1124,6 @@ def test_ragged_draft_gathers_each_requests_hidden_state(): verify_lens=torch.tensor([4, 2], dtype=torch.int32, device="cuda"), ) worker._lazy_init(draft_model, metadata) - _bind_manager(worker, [1000, 1001]) - worker._prepare_managed_history([1000, 1001], 0) worker._batch_to_slot[:2] = torch.tensor([0, 1], device="cuda") captured = torch.stack( [torch.full((HIDDEN * NCAP,), float(index), device="cuda") for index in range(6)] @@ -1135,139 +1182,3 @@ def test_confidence_state_is_absent_when_scheduling_is_disabled(): assert worker.staged_confidence_buffer() is None assert worker.confidence_stamp_buffer() is None assert worker.verify_planner is None - - -def _bind_manager(worker, request_ids=()): - from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType - - manager = types.SimpleNamespace( - draft_layer_ids=tuple(range(worker._num_stages)), - draft_window_size=worker._win, - kv_factor_per_layer=[1] * worker._num_stages, - head_dim_per_layer=[worker._head_dim] * worker._num_stages, - layer_offsets=list(range(worker._num_stages)), - draft_max_blocks_per_seq=8192, - tokens_per_block=128, - _draft_dummy_request_ids=set(), - enable_block_reuse=False, - kv_cache_map={}, - ) - buffers = tuple( - torch.zeros(65, 1, 1, 128, worker._head_dim, device="cuda", dtype=torch.bfloat16) - for _ in range(worker._num_stages) - ) - manager.get_draft_buffers = lambda stage: buffers[stage] - manager.get_draft_block_indices = lambda ids: [manager.kv_cache_map[rid].indices for rid in ids] - resource_manager = types.SimpleNamespace( - get_resource_manager=lambda kind: manager - if kind == ResourceManagerType.KV_CACHE_MANAGER - else None - ) - worker._bind_managed_history(resource_manager) - for i, rid in enumerate(request_ids): - manager.kv_cache_map[rid] = types.SimpleNamespace( - history_length=1, capacity=4096, indices=[i + 1] * 32 - ) - return manager - - -def _prepare(worker, meta, num_contexts=0): - if worker._draft_kv_manager is None: - _bind_manager(worker) - manager = worker._draft_kv_manager - for rid in meta.request_ids: - if worker._is_managed_request(rid) and rid not in manager.kv_cache_map: - manager.kv_cache_map[rid] = types.SimpleNamespace( - history_length=1, capacity=4096, indices=[len(manager.kv_cache_map) + 1] * 32 - ) - # A previously assigned slot represents already initialized resident history. - for rid in worker._req_to_slot: - if rid in manager.kv_cache_map: - worker._managed_residency.setdefault(rid, manager.kv_cache_map[rid]) - meta.prepare() - worker._prepare_managed_history(meta.request_ids, num_contexts) - - -def test_worker_lazy_init_configures_generation_autotuner_once(): - worker = _make_worker() - draft_model = _fake_draft_model() - draft_model.configure_generation_autotuner = Mock() - metadata = _make_metadata(max_num_requests=8) - - worker._lazy_init(draft_model, metadata) - worker._lazy_init(draft_model, metadata) - - draft_model.configure_generation_autotuner.assert_called_once_with(8) - - -def test_seed_context_pages_preserves_state_across_prefill_chunks(): - class DraftModel: - num_stages = 1 - block_size = 5 - _attn_params = { - "window_size": 8, - "head_dim": 4, - "n_heads": 128, - "softmax_scale": 1.0, - "eps": 1e-6, - } - - def __init__(self): - self.written_positions = [] - - def configure_generation_autotuner(self, _max_batch): - pass - - def write_context_pages(self, hidden, positions, mask, pools, tables, capacities): - from tensorrt_llm._torch.attention.backends.dspark import write_dspark_context - - self.written_positions.append(positions.clone()) - write_dspark_context( - torch.ones(*positions.shape, 4, device="cuda", dtype=torch.bfloat16), - positions, - mask, - pools[0][:, 0, 0], - tables, - capacities, - ) - - worker = _make_worker() - draft_model = DraftModel() - metadata = types.SimpleNamespace( - max_num_requests=1, - request_ids=[100], - get_hidden_states=lambda _num_tokens: torch.zeros( - 3, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 - ), - ) - worker._lazy_init(draft_model, metadata) - manager = _bind_manager(worker, [100]) - manager.kv_cache_map[100].history_length = 0 - worker._prepare_managed_history([100], 1) - - first_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[3]) - worker._seed_context_pages( - draft_model, metadata, first_chunk, torch.tensor([[0, 1, 2]], device="cuda"), 3 - ) - slot = worker._req_to_slot[100] - assert int(worker._ctx_len[slot]) == 3 - assert int(worker._valid_len[slot]) == 3 - assert worker._managed_residency[100] is manager.kv_cache_map[100] - assert bool(worker._position_initialized[slot]) - - metadata.get_hidden_states = lambda _num_tokens: torch.zeros( - 2, HIDDEN * NCAP, device="cuda", dtype=torch.bfloat16 - ) - second_chunk = types.SimpleNamespace(num_contexts=1, _seq_lens=[2]) - worker._seed_context_pages( - draft_model, metadata, second_chunk, torch.tensor([[3, 4]], device="cuda"), 2 - ) - - assert int(worker._ctx_len[slot]) == 5 - assert int(worker._valid_len[slot]) == 5 - assert [positions.tolist() for positions in draft_model.written_positions] == [ - [[1, 2, 3]], - [[4, 5]], - ] - assert torch.all(worker._draft_kv_buffers[0][1, 0, 0, :5] == 1.0) - assert torch.count_nonzero(worker._draft_kv_buffers[0][1, 0, 0, 5:]) == 0 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py index 842d46025f46..d94e995df727 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py @@ -38,10 +38,7 @@ def _worker(num_rows=2): worker._ctx_len = torch.zeros(num_rows + 1, dtype=torch.long) worker._valid_len = torch.zeros(num_rows + 1, dtype=torch.long) worker._position_initialized = torch.zeros(num_rows + 1, dtype=torch.bool) - worker._draft_kv_manager = object() - worker._draft_kv_buffers = () - worker._draft_block_tables = torch.zeros(num_rows + 1, 1, dtype=torch.int32) - worker._draft_capacities = torch.full((num_rows + 1,), 4, dtype=torch.long) + worker._kv_windows = torch.zeros(num_rows + 1, 1, 4, 2) worker._win = 4 worker._position_cap = None worker._confidence_logits = torch.full((num_rows + 2, 3), 30.0) @@ -60,8 +57,8 @@ def forward(main_hidden, bonus, start_pos, **kwargs): return SimpleNamespace( block_size=3, - forward=Mock(side_effect=forward), - write_context_pages=Mock(), + forward_batched=Mock(side_effect=forward), + write_context_windows_batched=Mock(), ) @@ -120,10 +117,10 @@ def test_mixed_ragged_draft_uses_generation_offsets_for_hidden_and_positions(ret positions=positions, ) - args = model.forward.call_args.args + args = model.forward_batched.call_args.args assert args[0][:, 0].tolist() == [5.0, 8.0] assert args[2].tolist() == [103, 202] - backfill = model.write_context_pages.call_args.args + backfill = model.write_context_windows_batched.call_args.args assert backfill[0][:, :, 0][backfill[3]].tolist() == [3.0, 4.0, 7.0] assert backfill[3].tolist() == [[True, True, False], [True, False, False]] @@ -156,7 +153,7 @@ def test_padding_offsets_do_not_index_past_captured_hidden_or_positions(): accepted_counts=[2, 2], ) - assert model.forward.call_args.args[0][0].tolist() == [2.0, 2.0] + assert model.forward_batched.call_args.args[0][0].tolist() == [2.0, 2.0] assert worker._confidence_logits[0].tolist() == [-3.0, -3.0, -3.0] assert worker._confidence_logits[worker._scratch_slot + 1].tolist() == [30.0, 30.0, 30.0] @@ -267,7 +264,6 @@ def test_embedded_checkpoint_load_preserves_optional_confidence_bias(monkeypatch stage.confidence_head = DSparkConfidenceHead(hidden_size=2, block_size=2) draft.mtp_layers = nn.ModuleList([stage]) draft.num_stages = 1 - draft._use_bf16_linear = False draft.post_load_weights = Mock() draft.cache_attn_weights_from_state_dict = Mock() wrapper = modeling_dspark.DSv4DSparkForCausalLM.__new__(modeling_dspark.DSv4DSparkForCausalLM) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index d324bb1bcc1e..947ee792d725 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -990,7 +990,6 @@ def test_dspark_forward_publishes_execution_not_confidence_flag(return_confidenc from tensorrt_llm._torch.speculative.dspark import DSv4DSparkWorker worker = DSv4DSparkWorker.__new__(DSv4DSparkWorker) - worker._draft_kv_manager = object() worker.spec_config = types.SimpleNamespace(max_draft_len=1) worker.return_confidence = return_confidence worker.guided_decoder = None diff --git a/tests/unittest/disaggregated/region/test_aux.py b/tests/unittest/disaggregated/region/test_aux.py index 2571b5869084..7c2528fbf81c 100644 --- a/tests/unittest/disaggregated/region/test_aux.py +++ b/tests/unittest/disaggregated/region/test_aux.py @@ -7,8 +7,6 @@ AuxBuffer, AuxBufferMeta, AuxSlot, - _decode_draft_history, - _encode_draft_history, build_aux_transfer_layout, ) @@ -187,23 +185,3 @@ def test_aux_transfer_layout_ctx_no_spec_gen_sa(): np.testing.assert_array_equal(layout.dst_item_sizes, gen.meta.item_sizes[expected_keep]) np.testing.assert_array_equal(layout.src_base_ptrs, ctx.meta.ptrs[expected_keep]) np.testing.assert_array_equal(layout.dst_base_ptrs, gen.meta.ptrs[expected_keep]) - - -@pytest.mark.parametrize("dtype", ["torch.bfloat16", "torch.float8_e4m3fn"]) -def test_draft_history_record_round_trips_the_storage_dtype(dtype): - """Prefill and generation servers compare the decoded layout, dtype included.""" - history = { - "valid_length": 89, - "position": 89, - "layout": { - "num_layers": 5, - "num_kv_heads": 16, - "head_dim": 64, - "dtype": dtype, - "attention_backend": "VANILLA", - "kv_factor": 2, - "window_size": None, - }, - } - - assert _decode_draft_history(_encode_draft_history(history)) == history diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index a6db3d7a83cf..ec6ef42ad0f7 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -313,9 +313,7 @@ def _make_fake_v2_manager(attrs, role_mapper_kinds, *, num_pools=1, slot_bytes_l get_disagg_ignored_roles=lambda: frozenset(), impl=impl, pp_layers=[0, 1], - draft_layer_ids=(), num_kv_heads_per_layer=[1, 1], - cp_as_tp_layer_ids=frozenset(), get_disagg_role_mapper_kinds=lambda: role_mapper_kinds, get_disagg_role_layouts=lambda: {}, ) From 1d08c8c8dc65ade5ea407d86e772ad7457fecf95 Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Fri, 30 Oct 2026 02:25:07 +0000 Subject: [PATCH 14/16] Fix DSpark test fixtures after mergeback Initialize the selected linear backend in the confidence-loader fixture, exercise the rolling-window stage entry point in V4.1 tests, and expect the retained SM107 allocator support. Preserve the existing behavioral assertions. Signed-off-by: Mingyang Hao --- .../_torch/speculative/hw_agnostic/test_dspark_v41.py | 6 +++--- .../hw_agnostic/test_dspark_worker_confidence_state.py | 2 +- .../_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py | 3 +-- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py index 8fd59c23d4f6..20077d7c4e9e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py @@ -311,7 +311,7 @@ def fake_stage(self, stage, hidden, *args, pre_mix, **kwargs): seen.append(pre_mix.clone()) return hidden + 1, mixes[index] - model._forward_stage = MethodType(fake_stage, model) + model._forward_stage_legacy = MethodType(fake_stage, model) model.forward_embed = Mock(return_value=(initial, torch.zeros(1, 3), torch.zeros(1, 2))) model._dspark_freqs_table = Mock(return_value=torch.empty(0)) call = model.forward_batched if batched else model.forward @@ -360,9 +360,9 @@ def test_v41_stage_collapses_with_previous_not_current_mix(monkeypatch): model_config=SimpleNamespace(mapping=SimpleNamespace(enable_attention_dp=False, tp_size=1)), ) monkeypatch.setattr( - modeling_dspark, "dspark_attention_forward", Mock(return_value=attention_output) + modeling_dspark, "_dspark_attention_forward_legacy", Mock(return_value=attention_output) ) - result, next_pre = modeling_dspark.DSv4DSparkDraftModel._forward_stage( + result, next_pre = modeling_dspark.DSv4DSparkDraftModel._forward_stage_legacy( model, stage, hidden, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py index d94e995df727..e355eaa83179 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker_confidence_state.py @@ -265,7 +265,7 @@ def test_embedded_checkpoint_load_preserves_optional_confidence_bias(monkeypatch draft.mtp_layers = nn.ModuleList([stage]) draft.num_stages = 1 draft.post_load_weights = Mock() - draft.cache_attn_weights_from_state_dict = Mock() + draft._use_bf16_linear = False wrapper = modeling_dspark.DSv4DSparkForCausalLM.__new__(modeling_dspark.DSv4DSparkForCausalLM) nn.Module.__init__(wrapper) wrapper.dspark_model = draft diff --git a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py index 68acc4846bc8..c0d6a6566fc3 100644 --- a/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py +++ b/tests/unittest/_torch/speculative/test_dspark_cute_dsl_rmsnorm_rope.py @@ -76,8 +76,7 @@ def test_fused_dspark_rmsnorm_rope_supported_architectures(): assert all(_is_dspark_rmsnorm_rope_sm_supported(sm) for sm in (100, 103, 107)) assert not any(_is_dspark_rmsnorm_rope_sm_supported(sm) for sm in (90, 101, 110, 120)) assert _is_dspark_rmsnorm_rope_sm_supported() - # The DSV4 DSpark attention kernel keeps its own, narrower gate. - assert _get_dspark_arch_str(107) is None + assert _get_dspark_arch_str(107) == "sm_107" @pytest.mark.parametrize( From 0671031c9a8f93ab54d2d5c0090f20d6ee1a7466 Mon Sep 17 00:00:00 2001 From: Mingyang Hao Date: Fri, 30 Oct 2026 02:37:18 +0000 Subject: [PATCH 15/16] Preserve FP8 cache precision in windowed DFlash attention Keep the converted query when dispatching through the new windowed helper, and round the future-block K/V to the cache precision before merging attention. This preserves the retained FP8 cache path after main refactored windowed decode. Cover BF16 and FP8 causal/noncausal calls at the CPU kernel boundary and compare against a dense oracle over the rounded cache values. Signed-off-by: Mingyang Hao --- tensorrt_llm/_torch/models/modeling_dflash.py | 12 ++++- .../test_dspark_noncausal_window.py | 49 +++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 280dc94495c2..af1c9f93cc99 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -160,6 +160,8 @@ def dflash_trtllm_windowed_block_attention( out: [B * blk, Hq, D] output. """ B, block_size, num_heads, head_dim = query.shape + # The decode and future-block merge must see the precision stored in the cache. + query = query.to(layer_cache.dtype) window_left, window_right = window_size merge_future = not causal and window_left >= 0 decoded = trtllm_gen_ops.batch_decode_with_kv_cache( @@ -190,7 +192,13 @@ def dflash_trtllm_windowed_block_attention( lse = decoded[1] * TRTLLM_GEN_LSE_LN2 out.copy_( dflash_merge_future_block_attention( - out, lse, query, block_key, block_value, head_dim**-0.5, window_right + out, + lse, + query, + block_key.to(layer_cache.dtype), + block_value.to(layer_cache.dtype), + head_dim**-0.5, + window_right, ) ) @@ -1952,7 +1960,7 @@ def dflash_forward( if causal or window_left >= 0: dflash_trtllm_windowed_block_attention( trtllm_gen_ops, - Q_bshd, + q_flat.view_as(Q_bshd), k_noise_bshd, v_noise_bshd, layer_cache, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py index 63926ee8abfb..0b3f854957e0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py @@ -106,6 +106,55 @@ def test_merge_matches_the_non_causal_oracle(block_size, window, heads): # --------------------------------------------------------------------------- +@pytest.mark.cpu_only +@pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("causal", [False, True]) +def test_windowed_decode_and_merge_use_cache_precision(cache_dtype, causal): + """Check the decode boundary and future-block merge against rounded cache values.""" + torch.manual_seed(17) + batch, block_size, num_heads, num_kv_heads, head_dim = 2, 3, 4, 2, 8 + query = torch.randn(batch, block_size, num_heads, head_dim, dtype=torch.bfloat16) + block_key = torch.randn(batch, block_size, num_kv_heads, head_dim, dtype=torch.bfloat16) + block_value = torch.randn_like(block_key) + rounded_query = query.to(cache_dtype) + keys = block_key.to(cache_dtype).float().unbind() + values = block_value.to(cache_dtype).float().unbind() + context_lengths = [0] * batch + causal_out, causal_lse = _oracle( + rounded_query, keys, values, context_lengths, 1, 0, causal=True + ) + expected, _ = _oracle(rounded_query, keys, values, context_lengths, 1, 1, causal=causal) + cache = torch.empty(1, 2, num_kv_heads, 4, head_dim, dtype=cache_dtype) + out = torch.empty_like(causal_out) + + def decode(**kwargs): + assert kwargs["query"].dtype == kwargs["kv_cache"][0].dtype == cache_dtype + torch.testing.assert_close( + kwargs["query"].float(), rounded_query.float().reshape_as(kwargs["query"]) + ) + kwargs["out"].copy_(causal_out) + if kwargs["return_lse"]: + return kwargs["out"], causal_lse / math.log(2.0) + return kwargs["out"] + + dflash_trtllm_windowed_block_attention( + SimpleNamespace(batch_decode_with_kv_cache=decode), + query, + block_key, + block_value, + cache, + torch.zeros(batch, 1, dtype=torch.int32), + torch.full((batch,), block_size, dtype=torch.int32), + 4, + causal, + (1, 0 if causal else 1), + torch.empty(0), + torch.empty(0), + out, + ) + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + def _wrapper(cls, config, layer_windows=(), backend="TRTLLM"): wrapper = cls.__new__(cls) torch.nn.Module.__init__(wrapper) From 72e6cc75cadde16482b6e56e0c9070005c818987 Mon Sep 17 00:00:00 2001 From: Jonas Li Date: Fri, 30 Oct 2026 07:41:31 +0000 Subject: [PATCH 16/16] Fix FC12 input-ready event joins Backport cc6b46096defce52488a27f06c9be25a43ad55f4. Define the join on the fused FC12 implementation, which does not inherit the grouped-GEMM implementation, and consume the input-ready event on empty batches too. Cover repeated dispatch within a forward, fresh events on the next forward, empty inputs, absent events, and output allocation/reuse with CPU tests. Signed-off-by: Jonas Li Signed-off-by: Mingyang Hao --- .../moe/fused_moe/cutedsl/fused_fc12.py | 13 +++- tests/unittest/_torch/moe/test_moe_backend.py | 63 +++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py b/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py index ad39a7429cb3..4262f23cf2c0 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py +++ b/tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py @@ -274,6 +274,13 @@ def _build_local_weight_view(self) -> NvFp4WeightView: # ------------------------------------------------------------------ quantize_input = quantize_input_nvfp4 + def _join_pending_x_ready(self) -> None: + """Wait for the input-ready event handed to ``run_moe_nvfp4`` (once per forward).""" + event = getattr(self, "_pending_x_ready_event", None) + if event is not None: + event.wait() + self._pending_x_ready_event = None + def run_moe(self, ctx: MoERunContext, *, workspace: dict | None = None) -> torch.Tensor: del workspace # the fused kernel allocates its own intermediates plan = require_comm_plan(self, ctx) @@ -313,13 +320,15 @@ def run_moe_nvfp4( assert moe_output.size() == (token_final_scales.size(0), self.hidden_size) assert moe_output.dtype == output_dtype + # Join before the fused kernel reads x, or before returning an empty batch. + self._pending_x_ready_event = x_ready_event + # Empty micro-batches: skip autotuning (synthetic grouped-GEMM inputs # require at least one output row). if token_selected_experts.size(0) == 0: + self._join_pending_x_ready() return moe_output - self._pending_x_ready_event = x_ready_event - effective_top_k = token_selected_experts.size(-1) tuner = AutoTuner.get() runner = CuteDslFc12FusedMoENvfp4Runner( diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 3285b5caa1d7..f27238a14504 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -5452,6 +5452,69 @@ def _runner(identity): assert situ(inputs, tactic=-1, do_preparation=True) is inputs[4] +@pytest.mark.cpu_only +@pytest.mark.parametrize("num_tokens", [0, 4]) +@pytest.mark.parametrize("with_event", [False, True]) +@pytest.mark.parametrize("provided_output", [False, True]) +def test_cutedsl_fc12_input_ready_event(monkeypatch, num_tokens, with_event, provided_output): + """Each forward joins once before input consumption, including empty batches.""" + backend = CuteDslFc12FusedMoE.__new__(CuteDslFc12FusedMoE) + torch.nn.Module.__init__(backend) + backend._weights_created = True + backend.quant_config = SimpleNamespace(layer_quant_mode=SimpleNamespace(has_nvfp4=lambda: True)) + backend.hidden_size = 16 + backend.num_slots = 8 + backend.use_fused_finalize = True + backend.activation_type = ActivationType.Swiglu + backend.act_clamp = float("inf") + backend.fc2_input_scale = torch.ones(1) + backend.mapping = SimpleNamespace(moe_ep_size=1) + weights = MagicMock(expert_size_per_partition=8, slot_start=0) + calls = [] + + def fused_op(**kwargs): + calls.append("kernel") + + monkeypatch.setattr(torch.ops.trtllm, "moe_sort", lambda **kwargs: (None,) * 6, raising=False) + monkeypatch.setattr( + torch.ops.trtllm, "cute_dsl_nvfp4_fc12_fused_rubin", fused_op, raising=False + ) + monkeypatch.setattr(CuteDslFc12FusedMoENvfp4Runner, "get_tuning_config", lambda self: None) + + def choose_one(name, runners, config, inputs): + assert num_tokens > 0, "Empty batches must not enter autotuning" + # Exercise repeated dispatch within one forward, as during autotuning. + runners[0](inputs, tactic=128) + return runners[0], 128 + + tuner = SimpleNamespace(choose_one=Mock(side_effect=choose_one)) + monkeypatch.setattr(AutoTuner, "get", staticmethod(lambda: tuner)) + output = torch.empty((num_tokens, backend.hidden_size), dtype=torch.bfloat16) + for _ in range(2): + calls.clear() + event = Mock() if with_event else None + if event is not None: + event.wait.side_effect = lambda: calls.append("wait") + result = backend.run_moe_nvfp4( + x=torch.zeros((num_tokens, backend.hidden_size // 2), dtype=torch.uint8), + token_selected_experts=torch.zeros((num_tokens, 2), dtype=torch.int32), + token_final_scales=torch.ones((num_tokens, 2), dtype=torch.float32), + x_sf=torch.ones(1, dtype=torch.uint8), + moe_output=output if provided_output else None, + weight_view=weights, + x_ready_event=event, + ) + assert calls == (["wait"] if with_event else []) + (["kernel"] * 2 if num_tokens else []) + if event is not None: + event.wait.assert_called_once_with() + assert getattr(backend, "_pending_x_ready_event", None) is None + assert result.shape == output.shape + assert result.dtype == torch.bfloat16 + if provided_output: + assert result is output + assert tuner.choose_one.call_count == (2 if num_tokens else 0) + + def test_unresolvable_layer_error_carries_rejection_details(): """describe() prints reason codes only, so impl_class_for has to add the details -- without them a shape rejection reaches the operator as a bare