From 6a5aa93c690643c988bb7d8fd414d31bf65e8d17 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Tue, 7 Jul 2026 02:17:45 -0700
Subject: [PATCH 01/12] [None][feat] add DSA KV cache manager V2 support

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../kernels/IndexerKCacheGather.h             |   4 +-
 .../kernels/IndexerKCacheScatter.h            |   4 +-
 .../kernels/indexerKCacheGather.cu            |  26 +-
 .../kernels/indexerKCacheScatter.cu           |  27 +-
 .../thop/IndexerKCacheGatherOp.cpp            |   9 +-
 .../thop/IndexerKCacheScatterOp.cpp           |   8 +-
 .../attention_backend/sparse/dsa/__init__.py  |   3 +-
 .../sparse/dsa/cache_manager.py               | 354 +++++++++++++++++-
 .../attention_backend/sparse/dsa/indexer.py   | 145 ++++---
 .../attention_backend/sparse/dsa/metadata.py  | 122 +++---
 .../attention_backend/sparse/registry.py      |   9 +-
 .../_torch/custom_ops/cpp_custom_ops.py       |  10 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |  22 +-
 .../defs/accuracy/test_llm_api_pytorch.py     |  56 +++
 .../test_lists/qa/llm_function_core.txt       |   4 +
 .../test_lists/test-db/l0_dgx_b200.yml        |   4 +
 .../attention/sparse/dsa/test_dsa_indexer.py  | 157 +++++++-
 .../attention/sparse/test_cpp_custom_ops.py   | 118 ++++++
 18 files changed, 940 insertions(+), 142 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
index 60422dea5910..7519d17bf1f0 100644
--- a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
+++ b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,7 +28,7 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
     int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2,
     int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
-    cudaStream_t stream = 0);
+    int32_t page_index_scale, cudaStream_t stream = 0);
 
 }
 
diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h b/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
index 9e316c0b4ebe..cd8123873eef 100644
--- a/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
+++ b/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,7 +28,7 @@ void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_sca
     int64_t const* slot_mapping_fp8, int64_t const* slot_mapping_scale, int32_t num_tokens, int32_t head_dim,
     int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
     int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
-    cudaStream_t stream = 0);
+    int32_t page_index_scale, cudaStream_t stream = 0);
 
 }
 
diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
index cbcb31f07111..81f3ca21fae7 100644
--- a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
+++ b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
@@ -49,6 +49,18 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
     return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3;
 }
 
+__device__ __forceinline__ int64_t scalePageIndex(int64_t flatIndex, int64_t blockStride, int32_t pageIndexScale)
+{
+    if (pageIndexScale == 1)
+    {
+        return flatIndex;
+    }
+
+    int64_t const blockIndex = flatIndex / blockStride;
+    int64_t const inBlockOffset = flatIndex - blockIndex * blockStride;
+    return blockIndex * pageIndexScale * blockStride + inBlockOffset;
+}
+
 } // anonymous namespace
 
 /**
@@ -63,7 +75,7 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
  * @param out_scale           Output scale data [num_tokens, scale_size] contiguous
  * @param k_token_start       Start offset into slot_mapping arrays
  * @param num_tokens          Number of tokens to gather
- * @param head_dim            Head dimension (must be 128)
+ * @param head_dim            Payload byte width (128 for FP8, 64 for packed FP4)
  * @param scale_size          Scale size in bytes (must be 4)
  * @param cache_stride_0      Stride for k_cache dimension 0 (in bytes)
  * @param cache_stride_1      Stride for k_cache dimension 1 (in bytes)
@@ -73,12 +85,14 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
  * @param cache_dim_1         Size of k_cache dimension 1
  * @param cache_dim_2         Size of k_cache dimension 2
  * @param cache_dim_3         Size of k_cache dimension 3
+ * @param page_index_scale    Scale from physical block ID to V2 shared-page index
  */
 __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_cache,
     int64_t const* __restrict__ slot_mapping_fp8, int64_t const* __restrict__ slot_mapping_scale,
     uint8_t* __restrict__ out_fp8, uint8_t* __restrict__ out_scale, int32_t k_token_start, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2,
-    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3)
+    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
+    int32_t page_index_scale)
 {
     // For head_dim=128, each thread handles 4 bytes/elements per read/write instruction
     constexpr int VEC_SIZE = 4;
@@ -104,6 +118,9 @@ __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_c
 
     int32_t head_dim_idx = threadIdx.x * VEC_SIZE;
     int64_t flat_idx = flat_idx_fp8_base + head_dim_idx;
+    int64_t const blockStride = static_cast<int64_t>(cache_dim_1) * cache_dim_2 * cache_dim_3;
+    flat_idx = scalePageIndex(flat_idx, blockStride, page_index_scale);
+    flat_idx_scale_base = scalePageIndex(flat_idx_scale_base, blockStride, page_index_scale);
 
     // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)
     int64_t src_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
@@ -130,7 +147,7 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
     int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2,
     int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
-    cudaStream_t stream)
+    int32_t page_index_scale, cudaStream_t stream)
 {
     if (num_tokens == 0)
     {
@@ -145,6 +162,7 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
     TLLM_CHECK_WITH_INFO(head_dim % VEC_SIZE == 0, "head_dim (%d) must be a multiple of %d", head_dim, VEC_SIZE);
     TLLM_CHECK_WITH_INFO(scale_size == 4,
         "scale_size must equal 4 bytes (packed UE8M0 x4 for FP4, 1 float32 for FP8, got %d)", scale_size);
+    TLLM_CHECK_WITH_INFO(page_index_scale >= 1, "page_index_scale must be at least 1 (got %d)", page_index_scale);
 
     int32_t const threads_per_block = head_dim / VEC_SIZE;
 
@@ -153,7 +171,7 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
 
     indexerKCacheGatherUnifiedKernel<<<grid, block, 0, stream>>>(k_cache, slot_mapping_fp8, slot_mapping_scale, out_fp8,
         out_scale, k_token_start, num_tokens, head_dim, scale_size, cache_stride_0, cache_stride_1, cache_stride_2,
-        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3);
+        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, page_index_scale);
 
     // Check for kernel launch errors
     TLLM_CUDA_CHECK(cudaGetLastError());
diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu
index f0a41a5009fc..332fdd2eb030 100644
--- a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu
+++ b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu
@@ -49,6 +49,18 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
     return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3;
 }
 
+__device__ __forceinline__ int64_t scalePageIndex(int64_t flatIndex, int64_t blockStride, int32_t pageIndexScale)
+{
+    if (pageIndexScale == 1)
+    {
+        return flatIndex;
+    }
+
+    int64_t const blockIndex = flatIndex / blockStride;
+    int64_t const inBlockOffset = flatIndex - blockIndex * blockStride;
+    return blockIndex * pageIndexScale * blockStride + inBlockOffset;
+}
+
 } // anonymous namespace
 
 /**
@@ -61,7 +73,7 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
  * @param slot_mapping_fp8  Flat element index for FP8 data start position [num_tokens]
  * @param slot_mapping_scale Flat element index for scale data start position [num_tokens]
  * @param num_tokens        Number of tokens
- * @param head_dim          Head dimension (must be 128)
+ * @param head_dim          Payload byte width (128 for FP8, 64 for packed FP4)
  * @param scale_size        Scale size in bytes (must be 4)
  * @param cache_stride_0    Stride for k_cache dimension 0 (in bytes)
  * @param cache_stride_1    Stride for k_cache dimension 1 (in bytes)
@@ -71,12 +83,14 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
  * @param cache_dim_1       Size of k_cache dimension 1
  * @param cache_dim_2       Size of k_cache dimension 2
  * @param cache_dim_3       Size of k_cache dimension 3
+ * @param page_index_scale  Scale from physical block ID to V2 shared-page index
  */
 __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_fp8_bytes,
     uint8_t const* __restrict__ k_scale_bytes, uint8_t* __restrict__ k_cache,
     int64_t const* __restrict__ slot_mapping_fp8, int64_t const* __restrict__ slot_mapping_scale, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2,
-    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3)
+    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
+    int32_t page_index_scale)
 {
     // For head_dim=128, each thread handles 4 bytes/elements per read/write instruction
     constexpr int VEC_SIZE = 4;
@@ -99,6 +113,9 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_
 
     int32_t head_dim_idx = threadIdx.x * VEC_SIZE;
     int64_t flat_idx = flat_idx_fp8_base + head_dim_idx;
+    int64_t const blockStride = static_cast<int64_t>(cache_dim_1) * cache_dim_2 * cache_dim_3;
+    flat_idx = scalePageIndex(flat_idx, blockStride, page_index_scale);
+    flat_idx_scale_base = scalePageIndex(flat_idx_scale_base, blockStride, page_index_scale);
 
     // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)
     int64_t dst_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
@@ -124,7 +141,8 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_
 void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_scale_bytes, uint8_t* k_cache,
     int64_t const* slot_mapping_fp8, int64_t const* slot_mapping_scale, int32_t num_tokens, int32_t head_dim,
     int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
-    int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3, cudaStream_t stream)
+    int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
+    int32_t page_index_scale, cudaStream_t stream)
 {
     if (num_tokens == 0)
     {
@@ -140,6 +158,7 @@ void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_sca
     TLLM_CHECK_WITH_INFO(head_dim % VEC_SIZE == 0, "head_dim (%d) must be a multiple of %d", head_dim, VEC_SIZE);
     TLLM_CHECK_WITH_INFO(scale_size == 4,
         "scale_size must equal 4 bytes (packed UE8M0 x4 for FP4, 1 float32 for FP8, got %d)", scale_size);
+    TLLM_CHECK_WITH_INFO(page_index_scale >= 1, "page_index_scale must be at least 1 (got %d)", page_index_scale);
 
     int32_t const threads_per_block = head_dim / VEC_SIZE;
 
@@ -148,7 +167,7 @@ void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_sca
 
     indexerKCacheScatterUnifiedKernel<<<grid, block, 0, stream>>>(k_fp8_bytes, k_scale_bytes, k_cache, slot_mapping_fp8,
         slot_mapping_scale, num_tokens, head_dim, scale_size, cache_stride_0, cache_stride_1, cache_stride_2,
-        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3);
+        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, page_index_scale);
 
     // Check for kernel launch errors
     TLLM_CUDA_CHECK(cudaGetLastError());
diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
index bd094dac58fc..501263a5cae3 100644
--- a/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
+++ b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
@@ -30,7 +30,7 @@ namespace torch_ext
 
 std::tuple<th::Tensor, th::Tensor> indexer_k_cache_gather_op(th::Tensor const& k_cache,
     th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t k_token_start, int64_t num_tokens,
-    int64_t head_dim)
+    int64_t head_dim, int64_t page_index_scale)
 {
     // head_dim is the number of payload bytes per token: 128 for FP8 (one byte
     // per float8 value) or 64 for FP4 (two packed E2M1 codes per byte).
@@ -86,6 +86,7 @@ std::tuple<th::Tensor, th::Tensor> indexer_k_cache_gather_op(th::Tensor const& k
 
     // Validation for indexer k cache pool for DeepSeek-V3.2 constraints
     TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1 for DeepSeek-V3.2, got %d", cache_dim_2);
+    TORCH_CHECK(page_index_scale >= 1, "page_index_scale must be at least 1, got %ld", page_index_scale);
 
     int64_t cache_stride_0 = static_cast<int64_t>(k_cache.stride(0));
     int64_t cache_stride_1 = static_cast<int64_t>(k_cache.stride(1));
@@ -102,7 +103,8 @@ std::tuple<th::Tensor, th::Tensor> indexer_k_cache_gather_op(th::Tensor const& k
     tk::invokeIndexerKCacheGather(k_cache.data_ptr<uint8_t>(), slot_mapping_fp8.data_ptr<int64_t>(),
         slot_mapping_scale.data_ptr<int64_t>(), out_fp8.data_ptr<uint8_t>(), out_scale.data_ptr<uint8_t>(),
         static_cast<int32_t>(k_token_start), num_tokens_i32, head_dim_i32, SCALE_SIZE, cache_dim_0, cache_dim_1,
-        cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3, stream);
+        cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3,
+        static_cast<int32_t>(page_index_scale), stream);
 
     // View-cast to final dtypes (no copy). The FP8 call site keeps its existing
     // float8_e4m3fn semantics; the FP4 call site reshapes the raw bytes into
@@ -121,7 +123,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def(
         "indexer_k_cache_gather_op(Tensor k_cache, Tensor slot_mapping_fp8, "
-        "Tensor slot_mapping_scale, int k_token_start, int num_tokens, int head_dim) -> (Tensor, Tensor)");
+        "Tensor slot_mapping_scale, int k_token_start, int num_tokens, int head_dim, "
+        "int page_index_scale=1) -> (Tensor, Tensor)");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
index 59c77177be5c..1adfed4fd252 100644
--- a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
+++ b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
@@ -29,7 +29,8 @@ namespace torch_ext
 {
 
 void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_scale, th::Tensor& k_cache,
-    th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t num_tokens)
+    th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t num_tokens,
+    int64_t page_index_scale)
 {
     // k_fp8: [>=num_tokens, head_dim] in FP8 (1 byte/element) — reinterpreted as uint8
     // k_scale: [>=num_tokens, head_dim // quant_block_size] in float32 — reinterpreted as uint8 bytes
@@ -78,6 +79,7 @@ void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_sca
     TORCH_CHECK(
         head_dim == 128 || head_dim == 64, "k_fp8 head_dim must be 128 (FP8) or 64 (FP4 packed), got %d", head_dim);
     TORCH_CHECK(scale_size == 4, "k_scale scale_size must be 4 bytes, got %d", scale_size);
+    TORCH_CHECK(page_index_scale >= 1, "page_index_scale must be at least 1, got %ld", page_index_scale);
 
     int64_t const cache_stride_0 = static_cast<int64_t>(k_cache.stride(0));
     int64_t const cache_stride_1 = static_cast<int64_t>(k_cache.stride(1));
@@ -92,7 +94,7 @@ void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_sca
         reinterpret_cast<uint8_t const*>(k_scale.data_ptr()), k_cache.data_ptr<uint8_t>(),
         slot_mapping_fp8.data_ptr<int64_t>(), slot_mapping_scale.data_ptr<int64_t>(), static_cast<int32_t>(num_tokens),
         head_dim, scale_size, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1,
-        cache_stride_2, cache_stride_3, stream);
+        cache_stride_2, cache_stride_3, static_cast<int32_t>(page_index_scale), stream);
 }
 
 } // namespace torch_ext
@@ -103,7 +105,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def(
         "indexer_k_cache_scatter_op(Tensor k_fp8, Tensor k_scale, Tensor(a!) k_cache, "
-        "Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int num_tokens) -> ()");
+        "Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int num_tokens, int page_index_scale=1) -> ()");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
index 6cf3aec0a6d1..685c2e35bfda 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
@@ -4,7 +4,7 @@
 """DeepSeek Sparse Attention backend package."""
 
 from .backend import DSATrtllmAttention
-from .cache_manager import DSACacheManager
+from .cache_manager import DSACacheManager, DSACacheManagerV2
 from .indexer import (
     _DG_SCHEDULE_BLOCK_KV,
     HAS_FAST_HADAMARD,
@@ -29,6 +29,7 @@
     "HAS_FAST_HADAMARD",
     "DSABackendForwardArgs",
     "DSACacheManager",
+    "DSACacheManagerV2",
     "DSAMetadataParams",
     "DSAParams",
     "DSATrtllmAttention",
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
index e6ff3d8b8b80..620e5452656b 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
@@ -5,17 +5,21 @@
 from __future__ import annotations
 
 import math
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
+import torch
 
 import tensorrt_llm
 import tensorrt_llm.bindings
-from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
-from tensorrt_llm._utils import get_size_in_bytes
+from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role
+from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager, get_pp_layers
+from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.bindings.executor import KvCacheConfig
 from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, DataRole, PageIndexMode
 
 from .params import DSAParams
 
@@ -25,6 +29,41 @@
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, SparseAttentionConfig
 
 
+def _get_indexer_k_cache_bytes_per_token(
+    index_head_dim: int, quant_block_size: int, use_fp4: bool
+) -> int:
+    """Return the raw indexer K-cache footprint for one token."""
+    data_bytes = index_head_dim // 2 if use_fp4 else index_head_dim
+    scale_bytes = index_head_dim // quant_block_size * 4
+    return data_bytes + scale_bytes
+
+
+def _get_indexer_k_cache_size_per_token(
+    model_config: ModelConfig,
+    mapping: Mapping,
+    num_layers: Optional[int] = None,
+) -> int:
+    """Estimate the indexer-only cache cost across local attention layers."""
+    sparse_attention_config = model_config.sparse_attention_config
+    if sparse_attention_config is None:
+        raise ValueError("sparse_attention_config is required for DSA cache")
+    sparse_params = sparse_attention_config.to_sparse_params(
+        pretrained_config=model_config.pretrained_config
+    )
+    if not isinstance(sparse_params, DSAParams):
+        raise ValueError("DSA cache requires DSA sparse parameters")
+
+    num_attention_layers = KVCacheManager._resolve_num_attention_layers(
+        model_config, mapping, num_layers
+    )
+    bytes_per_layer = _get_indexer_k_cache_bytes_per_token(
+        sparse_params.index_head_dim,
+        128,
+        sparse_params.indexer_k_dtype == "fp4",
+    )
+    return num_attention_layers * bytes_per_layer
+
+
 def derive_indexer_k_cache_layer_mask(
     sparse_attention_config: "SparseAttentionConfig",
     pretrained_config,
@@ -123,6 +162,8 @@ def __init__(
             **kwargs,
         )
         self.num_blocks = self.blocks_in_primary_pool
+        # V1 stores one INDEX_KEY page per physical pool slot.
+        self.indexer_k_cache_page_scale = 1
 
         # Indexer K cache pool for DSA attention
         # Shape: [num_blocks, self.tokens_per_block * (index_head_dim + scale_size)]
@@ -141,11 +182,16 @@ def __init__(
                 f"{self.num_local_layers} local layers own an indexer k-cache."
             )
 
+    def get_primary_pool_page_index_params(self, local_layer_idx: int) -> Tuple[int, int]:
+        """Return V1 page scale and layer offset for sparse MLA indices."""
+        return self.num_local_layers, local_layer_idx
+
     def get_indexer_k_cache_buffers(self, layer_idx: int):
         """Get indexer k cache buffer from a specific layer pool."""
         block_size = self.tokens_per_block
-        data_bytes = self.index_head_dim // 2 if self.use_fp4 else self.index_head_dim
-        per_token_size = data_bytes + self.index_head_dim // self.quant_block_size * 4
+        per_token_size = _get_indexer_k_cache_bytes_per_token(
+            self.index_head_dim, self.quant_block_size, self.use_fp4
+        )
         layer_offset = self.layer_offsets[layer_idx]
         pool = self.indexer_k_cache_pool_per_layer[layer_offset]
         assert pool is not None, (
@@ -154,6 +200,20 @@ def get_indexer_k_cache_buffers(self, layer_idx: int):
         )
         return pool.view(self.num_blocks, block_size, 1, per_token_size)
 
+    def get_pool_block_indices(
+        self,
+        num_seqs: int,
+        *,
+        request_ids: Optional[List[int]] = None,
+        num_contexts: int = 0,
+        beam_width: int = 1,
+    ) -> torch.Tensor:
+        """Decode V1 block offsets into physical memory-pool block indices."""
+        del request_ids, num_contexts, beam_width
+        encoded = self.host_kv_cache_block_offsets[0, :num_seqs, 0, :]
+        max_pool_idx = self.blocks_in_primary_pool - 1
+        return (encoded // self.num_local_layers).clamp(min=0, max=max_pool_idx).to(torch.int32)
+
     def get_batch_indexer_k_cache_indices(self, request_ids: List[int]) -> List[List[int]]:
         """
         Get the indices for the indexer k cache for a specific batch of requests.
@@ -277,3 +337,287 @@ def get_cache_bytes_per_token(self):
         cache_size_bytes_per_token += indexer_bytes_per_token
 
         return cache_size_bytes_per_token
+
+
+class DSACacheManagerV2(KVCacheManagerV2):
+    """KVCacheManagerV2-backed cache manager with a DSA indexer K-cache."""
+
+    def __init__(
+        self,
+        kv_cache_config: KvCacheConfig,
+        kv_cache_type: CacheTypeCpp,
+        *,
+        num_layers: int,
+        num_kv_heads: Union[int, List[Optional[int]]],
+        head_dim: int,
+        tokens_per_block: int,
+        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,
+        max_num_tokens: int = 8192,
+        model_config: Optional[ModelConfig] = None,
+        max_beam_width: int = 1,
+        sparse_attention_config: Optional["SparseAttentionConfig"] = None,
+        pretrained_config=None,
+        **kwargs,
+    ) -> None:
+        if sparse_attention_config is None:
+            sparse_attention_config = kwargs.pop("sparse_attn_config", None)
+        if sparse_attention_config is None and model_config is not None:
+            sparse_attention_config = model_config.sparse_attention_config
+        if sparse_attention_config is None:
+            raise ValueError("sparse_attention_config is required for DSA cache")
+        sparse_params = sparse_attention_config.to_sparse_params(
+            pretrained_config=pretrained_config
+        )
+        if not isinstance(sparse_params, DSAParams):
+            raise ValueError("DSA cache requires DSA sparse parameters")
+
+        self.quant_block_size = 128
+        self.index_head_dim = sparse_params.index_head_dim
+        self.use_fp4 = sparse_params.indexer_k_dtype == "fp4"
+        self._unique_primary_pool: Optional[torch.Tensor] = None
+
+        from tensorrt_llm._torch.speculative import get_num_spec_layers
+
+        total_num_layers = len(layer_mask) if layer_mask is not None else num_layers
+        if spec_config is not None and layer_mask is None:
+            total_num_layers += get_num_spec_layers(spec_config)
+        indexer_k_cache_layer_mask = derive_indexer_k_cache_layer_mask(
+            sparse_attention_config, pretrained_config, total_num_layers
+        )
+        pp_layers, _ = get_pp_layers(
+            num_layers, mapping, spec_config=spec_config, layer_mask=layer_mask
+        )
+        self.indexer_k_cache_local_layer_mask = [
+            indexer_k_cache_layer_mask[layer_idx] for layer_idx in pp_layers
+        ]
+
+        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=layer_mask,
+            max_num_tokens=max_num_tokens,
+            model_config=model_config,
+            max_beam_width=max_beam_width,
+            **kwargs,
+        )
+        if self.num_local_layers == 0:
+            raise ValueError("DSA requires at least one local attention layer")
+        primary_converters = [
+            self.impl.get_page_index_converter(local_layer_idx, Role.KEY)
+            for local_layer_idx in range(self.num_local_layers)
+        ]
+        self._primary_pool_page_index_params = [
+            (int(converter.scale), int(converter.layer_offset))
+            for converter in primary_converters
+        ]
+        full_indexer_local_layers = [
+            local_layer_idx
+            for local_layer_idx, has_indexer in enumerate(self.indexer_k_cache_local_layer_mask)
+            if has_indexer
+        ]
+        indexer_converters = [
+            self.impl.get_page_index_converter(local_layer_idx, Role.INDEX_KEY)
+            for local_layer_idx in full_indexer_local_layers
+        ]
+        if any(int(converter.expansion) != 1 for converter in indexer_converters):
+            raise ValueError("DSA indexer K-cache does not support page expansion")
+        # Each cache view starts at its SHARED base address, so the converter's
+        # layer offset is already represented by the base pointer.
+        self.indexer_k_cache_page_scale = (
+            int(indexer_converters[0].scale) if indexer_converters else 1
+        )
+        if any(
+            int(converter.scale) != self.indexer_k_cache_page_scale
+            for converter in indexer_converters[1:]
+        ):
+            raise ValueError("DSA requires a uniform shared INDEX_KEY page mapping across layers")
+        self.num_blocks = self.blocks_in_primary_pool
+        self.indexer_k_cache_pool_per_layer = [
+            self._get_indexer_k_cache_pool_data(local_layer_idx)
+            if self.indexer_k_cache_local_layer_mask[local_layer_idx]
+            else None
+            for local_layer_idx in range(self.num_local_layers)
+        ]
+        num_full = sum(self.indexer_k_cache_local_layer_mask)
+        if num_full < self.num_local_layers:
+            logger.info(
+                f"[DSACacheManagerV2] Indexer k-cache: {num_full} of "
+                f"{self.num_local_layers} local layers own an indexer k-cache."
+            )
+
+    def get_primary_pool_page_index_params(self, local_layer_idx: int) -> Tuple[int, int]:
+        """Return the formal V2 page scale and layer offset for sparse MLA."""
+        return self._primary_pool_page_index_params[local_layer_idx]
+
+    def _extra_buffers_per_layer(
+        self, *, tokens_per_block: int
+    ) -> dict[int, List[BufferConfig]]:
+        return {
+            local_layer_idx: [
+                BufferConfig(
+                    role=Role.INDEX_KEY,
+                    size=self.get_layer_bytes_per_token(local_layer_idx, Role.INDEX_KEY)
+                    * tokens_per_block,
+                )
+            ]
+            for local_layer_idx in range(self.num_local_layers)
+            if self.indexer_k_cache_local_layer_mask[local_layer_idx]
+        }
+
+    @property
+    def blocks_in_primary_pool(self) -> int:
+        """Return the physical slot count rather than the V2 page bound."""
+        converter = self.impl.get_page_index_converter(0, Role.KEY)
+        page_upper = self.impl.get_page_index_upper_bound(0, Role.KEY)
+        expansion = int(converter.expansion)
+        assert page_upper % expansion == 0
+        num_pages_with_offset = page_upper // expansion + int(converter.layer_offset)
+        scale = int(converter.scale)
+        assert num_pages_with_offset % scale == 0
+        return num_pages_with_offset // scale
+
+    def _get_indexer_k_cache_pool_data(self, local_layer_idx: int) -> torch.Tensor:
+        """Return a contiguous shared-page view for one indexer K-cache."""
+        address = self.impl.get_mem_pool_base_address(
+            local_layer_idx, Role.INDEX_KEY, PageIndexMode.SHARED
+        )
+        page_upper = self.impl.get_page_index_upper_bound(local_layer_idx, Role.INDEX_KEY)
+        flat_page_size = self.tokens_per_block * self.get_layer_bytes_per_token(
+            local_layer_idx, Role.INDEX_KEY
+        )
+        return convert_to_torch_tensor(
+            TensorWrapper(address, torch.uint8, [page_upper, flat_page_size])
+        )
+
+    def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor:
+        """Return the page-indexed indexer K-cache view for a global layer."""
+        layer_offset = self.layer_offsets[layer_idx]
+        pool = self.indexer_k_cache_pool_per_layer[layer_offset]
+        assert pool is not None, (
+            f"Layer {layer_idx} is a shared-indexer layer and owns no indexer "
+            f"k-cache; only full-indexer layers may access it."
+        )
+        per_token_size = self.get_layer_bytes_per_token(layer_offset, Role.INDEX_KEY)
+        return pool.view(pool.shape[0], self.tokens_per_block, 1, per_token_size)
+
+    def get_pool_block_indices(
+        self,
+        num_seqs: int,
+        *,
+        request_ids: Optional[List[int]] = None,
+        num_contexts: int = 0,
+        beam_width: int = 1,
+    ) -> torch.Tensor:
+        """Read V2 stable slots in current-batch order as physical block IDs."""
+        if request_ids is None:
+            raise ValueError("DSACacheManagerV2 requires request_ids to map stable slots")
+        if len(request_ids) != num_seqs:
+            raise ValueError(f"Expected {num_seqs} request IDs, got {len(request_ids)}")
+        copy_idx = self.index_mapper.get_copy_index(list(request_ids), num_contexts, beam_width)
+        copy_idx = copy_idx.to(device="cpu", dtype=torch.long)
+        block_indices = self.host_kv_cache_block_offsets[0, copy_idx, 0, :]
+        return block_indices.clamp(min=0, max=self.num_blocks - 1).to(torch.int32)
+
+    def get_unique_primary_pool(self) -> torch.Tensor:
+        """Return the uniform MLA K pool in the V1-compatible layout."""
+        if self._unique_primary_pool is not None:
+            return self._unique_primary_pool
+        if self.num_local_layers == 0:
+            raise ValueError("DSA requires at least one local attention layer")
+        if self.kv_factor != 1:
+            raise ValueError("DSA requires a SELFKONLY KV cache")
+
+        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):
+            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):
+            raise ValueError("DSA requires a uniform KV head count")
+
+        first_converter = self.impl.get_page_index_converter(0, Role.KEY)
+        if int(first_converter.expansion) != 1:
+            raise ValueError("DSA MLA K-cache does not support page expansion")
+        if int(first_converter.layer_offset) != 0:
+            raise ValueError("The first DSA layer must start at pool offset 0")
+        if int(first_converter.scale) != self.num_local_layers:
+            raise ValueError("DSA requires one uniformly coalesced K page per local layer")
+
+        page_stride = self.impl.get_page_stride(0, Role.KEY)
+        base_address = self.impl.get_mem_pool_base_address(0, Role.KEY, PageIndexMode.SHARED)
+        for local_layer_idx in range(1, self.num_local_layers):
+            converter = self.impl.get_page_index_converter(local_layer_idx, Role.KEY)
+            if int(converter.scale) != int(first_converter.scale) or int(converter.expansion) != 1:
+                raise ValueError("DSA requires a uniform page-index mapping across layers")
+            if int(converter.layer_offset) != local_layer_idx:
+                raise ValueError("DSA requires K pages to follow local-layer order")
+            address = self.impl.get_mem_pool_base_address(
+                local_layer_idx, Role.KEY, PageIndexMode.SHARED
+            )
+            expected_address = base_address + int(converter.layer_offset) * page_stride
+            if int(address) != int(expected_address):
+                raise ValueError("DSA requires contiguous per-layer K pages in each slot")
+
+        element_per_container = 2 if self.dtype == DataType.NVFP4 else 1
+        dtype = torch.int8 if self.dtype == DataType.NVFP4 else self.dtype
+        elements_per_layer = (
+            self.tokens_per_block
+            * first_num_heads
+            * first_head_dim
+            // element_per_container
+        )
+        shape = [
+            self.blocks_in_primary_pool,
+            self.num_local_layers,
+            1,
+            elements_per_layer,
+        ]
+        self._unique_primary_pool = convert_to_torch_tensor(
+            TensorWrapper(base_address, dtype, shape)
+        )
+        return self._unique_primary_pool
+
+    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 data_role == Role.ALL 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:
+        return sum(
+            self.get_layer_bytes_per_token(local_layer_idx, Role.ALL)
+            for local_layer_idx in range(self.num_local_layers)
+        )
+
+    @staticmethod
+    def get_cache_size_per_token(
+        model_config: ModelConfig,
+        mapping: Mapping,
+        num_layers: Optional[int] = None,
+        **kwargs,
+    ):
+        return DSACacheManager.get_cache_size_per_token(
+            model_config, mapping, num_layers=num_layers, **kwargs
+        )
+
+    def shutdown(self) -> None:
+        self.indexer_k_cache_pool_per_layer = []
+        self._unique_primary_pool = None
+        super().shutdown()
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
index ded1bda7694f..eae55ac98ae7 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py
@@ -294,6 +294,9 @@ def transform_local_topk_and_prepare_pool_view(
     assert topk_indices.dtype == torch.int32
 
     attn_metadata._ensure_pool_view_cached()
+    page_index_scale, layer_offset = (
+        attn_metadata.kv_cache_manager.get_primary_pool_page_index_params(layer_idx)
+    )
 
     if is_generation:
         block_table = attn_metadata._cached_block_table_gen
@@ -308,8 +311,8 @@ def transform_local_topk_and_prepare_pool_view(
         topk_indices,
         attn_metadata._cached_tokens_per_block,
         topk_indices.shape[1],
-        attn_metadata._cached_stride_factor,
-        layer_idx,
+        page_index_scale * attn_metadata._cached_tokens_per_block,
+        layer_offset,
     )
 
     return global_indices, attn_metadata._cached_pool_view
@@ -932,6 +935,86 @@ def recompute_slot_mappings(
             metadata.host_slot_mapping_scale[:total_new_kv_tokens], non_blocking=True
         )
 
+    @staticmethod
+    def recompute_context_kv_gather_mappings(
+        metadata: DSAtrtllmAttentionMetadata,
+        indexer_params: Optional[IndexerParams] = None,
+    ) -> None:
+        """Recompute context KV-cache gather mappings for the active cache manager."""
+        if not metadata.enable_context_mla_with_cached_kv:
+            metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8
+            metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale
+            return
+
+        if indexer_params is None:
+            indexer_params = Indexer.build_indexer_params(metadata)
+            if indexer_params is None:
+                return
+
+        num_contexts = indexer_params.num_contexts
+        if num_contexts == 0:
+            # Generation does not consume full-KV gather mappings. In
+            # particular, avoid allocating temporary zero-length tensors
+            # while capturing generation CUDA graphs.
+            return
+
+        if metadata.skip_indexer_for_ctx_reqs:
+            # The dense short-context path does not gather from the indexer K
+            # cache, so full-KV mappings would be allocated but never consumed.
+            metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8
+            metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale
+            return
+
+        cached_kv_tokens = indexer_params.cached_kv_tokens[:num_contexts]
+        if indexer_params.compress_ratio == 1 and cached_kv_tokens.count_nonzero().item() == 0:
+            # Without compression or cached tokens, the regular mappings cover
+            # every context KV token in the same order and remain equivalent
+            # when on_update_kv_lens() refreshes them in place.
+            metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8
+            metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale
+            return
+
+        total_kv_per_request = indexer_params.kv_lens[:num_contexts]
+        total_kv_len = total_kv_per_request.sum().item()
+        if total_kv_len == 0:
+            # No chunk can gather indexer KV, so avoid zero-length allocation
+            # and H2D work for short compressed contexts.
+            metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8
+            metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale
+            return
+
+        host_slot_mapping_fp8_fullkv = torch.empty(
+            total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()
+        )
+        host_slot_mapping_scale_fullkv = torch.empty(
+            total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()
+        )
+
+        req_indices = torch.repeat_interleave(
+            torch.arange(num_contexts, dtype=torch.int64, device="cpu"), total_kv_per_request
+        )
+
+        cu_kv = torch.zeros(num_contexts + 1, dtype=torch.int64, device="cpu")
+        cu_kv[1:] = total_kv_per_request.to(torch.int64).cumsum(0)
+        kv_positions = torch.arange(total_kv_len, dtype=torch.int64, device="cpu") - cu_kv[
+            :-1
+        ].repeat_interleave(total_kv_per_request)
+
+        fp8_flat_indices, scale_flat_indices = _compute_slot_mappings(
+            kv_positions,
+            metadata.host_indexer_k_cache_block_offsets,
+            req_indices,
+            indexer_params.head_dim,
+            indexer_params.tokens_per_block,
+            indexer_params.quant_block_size,
+            data_bytes_per_token=indexer_params.data_bytes_per_token,
+        )
+
+        host_slot_mapping_fp8_fullkv.copy_(fp8_flat_indices)
+        host_slot_mapping_scale_fullkv.copy_(scale_flat_indices)
+        metadata.slot_mapping_fp8_fullkv = host_slot_mapping_fp8_fullkv.cuda(non_blocking=True)
+        metadata.slot_mapping_scale_fullkv = host_slot_mapping_scale_fullkv.cuda(non_blocking=True)
+
     @staticmethod
     def prepare_for_update_k_cache(
         metadata: DSAtrtllmAttentionMetadata, indexer_params: IndexerParams
@@ -953,8 +1036,6 @@ def prepare_for_chunked_prefill(
         """
         num_contexts = indexer_params.num_contexts
         seq_lens = indexer_params.seq_lens
-        tokens_per_block = indexer_params.tokens_per_block
-        head_dim = indexer_params.head_dim
 
         # When MLA chunked prefill is active, it already handles chunking
         # Indexer should just process the current MLA chunk as a single chunk
@@ -1001,57 +1082,9 @@ def prepare_for_chunked_prefill(
             else:
                 metadata.indexer_prefill_chunks = None
 
-        # Chunked prefill and KV-cache reuse require the full KV for indexer
-        # logits. The indexer's own chunking gathers only the current chunk.
-        if metadata.enable_context_mla_with_cached_kv:
-            # Use kv_lens which correctly computes (raw_past + seq_lens) // compress_ratio.
-            total_kv_per_request = indexer_params.kv_lens[:num_contexts]
-            total_kv_len = total_kv_per_request.sum().item()
-            host_slot_mapping_fp8_fullkv = torch.empty(
-                total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()
-            )
-            host_slot_mapping_scale_fullkv = torch.empty(
-                total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()
-            )
-
-            req_indices = torch.repeat_interleave(
-                torch.arange(num_contexts, dtype=torch.int64, device="cpu"), total_kv_per_request
-            )
-
-            cu_kv = torch.zeros(num_contexts + 1, dtype=torch.int64, device="cpu")
-            cu_kv[1:] = total_kv_per_request.to(torch.int64).cumsum(0)
-            kv_positions = torch.arange(total_kv_len, dtype=torch.int64, device="cpu") - cu_kv[
-                :-1
-            ].repeat_interleave(total_kv_per_request)
-
-            fp8_flat_indices, scale_flat_indices = _compute_slot_mappings(
-                kv_positions,
-                metadata.host_indexer_k_cache_block_offsets,
-                req_indices,
-                head_dim,
-                tokens_per_block,
-                indexer_params.quant_block_size,
-                data_bytes_per_token=head_dim // 2
-                if metadata.kv_cache_manager.use_fp4
-                else head_dim,
-            )
-
-            host_slot_mapping_fp8_fullkv[:total_kv_len] = fp8_flat_indices
-            host_slot_mapping_scale_fullkv[:total_kv_len] = scale_flat_indices
-
-            assert len(fp8_flat_indices) == total_kv_len, (
-                "host_slot_mapping_fp8_fullkv/host_slot_mapping_scale_fullkv "
-                f"length mismatch: {len(fp8_flat_indices)} != total_kv_len={total_kv_len}"
-            )
-
-            # Store extended mappings for indexer full KV gathering
-            metadata.slot_mapping_fp8_fullkv = host_slot_mapping_fp8_fullkv.cuda(non_blocking=True)
-            metadata.slot_mapping_scale_fullkv = host_slot_mapping_scale_fullkv.cuda(
-                non_blocking=True
-            )
-        else:
-            metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8
-            metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale
+        # Chunked prefill and KV-cache reuse require gather mappings that cover
+        # all cached and newly appended indexer tokens.
+        Indexer.recompute_context_kv_gather_mappings(metadata, indexer_params)
 
     @staticmethod
     def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata):
@@ -1216,7 +1249,7 @@ def _gather_k_cache_for_chunk(
 
         k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx)
 
-        head_dim = self.head_dim
+        head_dim = self.head_dim // 2 if self.use_fp4 else self.head_dim
         scale_size = 4  # float32 = 4 bytes
 
         # Extract slot mappings using chunk's k_token_start/end
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
index c7b0dcafae38..10513814e39d 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
@@ -124,7 +124,6 @@ def __init__(self, *args, **kwargs):
         self._pool_cache_valid = False
         self._cached_kv_mgr_id = 0
         self._cached_pool_view = None
-        self._cached_stride_factor = 0
         self._cached_tokens_per_block = 0
         self._cached_block_table_ctx = None
         self._cached_block_table_gen = None
@@ -138,10 +137,10 @@ def __init__(self, *args, **kwargs):
 
     def __post_init__(self):
         """Allocate indexer K-cache buffers and heuristic TopK metadata."""
-        from .cache_manager import DSACacheManager
+        from .cache_manager import DSACacheManager, DSACacheManagerV2
 
         super().__post_init__()
-        if not isinstance(self.kv_cache_manager, DSACacheManager):
+        if not isinstance(self.kv_cache_manager, (DSACacheManager, DSACacheManagerV2)):
             has_deepseek_v4_cache_interface = all(
                 hasattr(self.kv_cache_manager, attr)
                 for attr in ("compressed_block_sizes", "get_cache_indices")
@@ -570,6 +569,50 @@ def create_buffers_for_indexer(self, capture_graph=False):
             device="cpu",
             pin_memory=prefer_pinned(),
         )
+        # Allocate separate indexer block-offset and slot-mapping buffers for
+        # the draft KV cache manager, mirroring draft_kv_cache_block_offsets:
+        # the draft-replay context swaps these in by rebinding, so CUDA graph
+        # capture bakes distinct addresses for the target and draft segments
+        # and both sides can be refreshed eagerly outside the graph.
+        if self.draft_kv_cache_manager is not None and hasattr(
+            self.draft_kv_cache_manager, "index_head_dim"
+        ):
+            self.draft_indexer_k_cache_block_offsets = self.get_empty(
+                self.cuda_graph_buffers,
+                [self.max_num_sequences, self.draft_kv_cache_manager.max_blocks_per_seq],
+                cache_name="draft_indexer_k_cache_block_offsets",
+                dtype=torch.int32,
+                capture_graph=capture_graph,
+            )
+            self.host_draft_indexer_k_cache_block_offsets = torch.zeros_like(
+                self.draft_indexer_k_cache_block_offsets,
+                device="cpu",
+                pin_memory=prefer_pinned(),
+            )
+            self.draft_slot_mapping_fp8 = self.get_empty(
+                self.cuda_graph_buffers,
+                (self.max_num_tokens,),
+                cache_name="draft_slot_mapping_fp8",
+                dtype=torch.int64,
+                capture_graph=capture_graph,
+            )
+            self.host_draft_slot_mapping_fp8 = torch.zeros_like(
+                self.draft_slot_mapping_fp8,
+                device="cpu",
+                pin_memory=prefer_pinned(),
+            )
+            self.draft_slot_mapping_scale = self.get_empty(
+                self.cuda_graph_buffers,
+                (self.max_num_tokens,),
+                cache_name="draft_slot_mapping_scale",
+                dtype=torch.int64,
+                capture_graph=capture_graph,
+            )
+            self.host_draft_slot_mapping_scale = torch.zeros_like(
+                self.draft_slot_mapping_scale,
+                device="cpu",
+                pin_memory=prefer_pinned(),
+            )
         # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation.
         # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory.
         if self.enable_context_mla_with_cached_kv:
@@ -805,31 +848,26 @@ def update_spec_dec_param(
             # elements").
             self._create_radix_aux_buffers(capture_graph=capture_graph)
 
-    def _get_pool_block_indices(self) -> torch.Tensor:
-        """Extract memory pool block indices from host_kv_cache_block_offsets.
-
-        The C++ setOffsets() encodes offsets as:
-            encoded = memPoolBlockIndex * numLayers * kvFactor
-        For SELFKONLY (MLA/DSA), kvFactor=1, so:
-            memPoolBlockIndex = encoded // num_local_layers
-
-        Returns a (num_seqs, max_blocks_per_seq) int32 CPU tensor with valid
-        pool indices clamped to [0, blocks_in_primary_pool - 1].
-        """
-        num_local_layers = self.kv_cache_manager.num_local_layers
-        max_pool_idx = self.kv_cache_manager.blocks_in_primary_pool - 1
-        # DSA uses SELFKONLY mode where only key cache is stored (kv_factor=1).
-        # host_kv_cache_block_offsets shape: (num_pools, max_batch*beam, 2, max_blocks_per_seq)
-        # Note: dim=2 is always 2 in the tensor layout (K and V slots), but for
-        # SELFKONLY only the K slot (index 0) contains valid data.
-        assert self.kv_cache_manager.kv_factor == 1, (
-            f"DSA requires SELFKONLY mode (kv_factor=1), got kv_factor={self.kv_cache_manager.kv_factor}"
+    def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor:
+        """Refresh INDEX_KEY offsets and return their physical pool slots."""
+        cache_manager = self.kv_cache_manager
+        # Raw block IDs can exceed the primary pool after host-cache offload;
+        # the manager resolves their current physical slots.
+        pool_indices = cache_manager.get_pool_block_indices(
+            self.num_seqs,
+            request_ids=self.request_ids,
+            num_contexts=self.num_contexts,
+            beam_width=self.beam_width,
         )
-        # Pool 0, first num_seqs entries, field 0 (key offsets)
-        encoded = self.kv_cache_manager.host_kv_cache_block_offsets[0, : self.num_seqs, 0, :]
-        pool_indices = encoded // num_local_layers
-        # Clamp for safety: handles garbage padding from torch.empty in uninitialized slots
-        pool_indices = pool_indices.clamp(min=0, max=max_pool_idx).to(torch.int32)
+        page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale
+        num_blocks = page_indices.shape[1]
+        self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices)
+        self.indexer_k_cache_block_offsets[: self.num_seqs].copy_(
+            self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True
+        )
+        # Sanitize graph-padding entries that may be stale after cache
+        # eviction or host-cache onboarding.
+        self.indexer_k_cache_block_offsets.clamp_(min=0)
         return pool_indices
 
     def set_skip_topk(self, skip: bool) -> None:
@@ -853,10 +891,10 @@ def _ensure_pool_view_cached(self):
         """Compute and cache values used by
         transform_local_topk_and_prepare_pool_view().
 
-        These values (pool view, stride factor, block table slices, request
-        index slices) are constant across all layers sharing the same KV pool
-        and batch dimensions within a forward pass. Caching them avoids
-        redundant Python/CUDA overhead per layer.
+        These values (pool view, block table slices, and request index slices)
+        are constant across all layers sharing the same KV pool and batch
+        dimensions within a forward pass. Caching them avoids redundant
+        Python/CUDA overhead per layer.
 
         Safety: _invalidate_pool_view_cache() is called unconditionally at the
         start of every step (prepare() and on_update_kv_lens()), so the boolean
@@ -867,11 +905,9 @@ def _ensure_pool_view_cached(self):
 
         pool = self.kv_cache_manager.get_unique_primary_pool()
         kv_cache_manager = self.kv_cache_manager
-        num_blocks, num_layers, _, _ = pool.shape
         self._cached_tokens_per_block = kv_cache_manager.tokens_per_block
         head_dim = kv_cache_manager.head_dim
         self._cached_pool_view = pool.squeeze(2).view(-1, 1, head_dim)
-        self._cached_stride_factor = num_layers * self._cached_tokens_per_block
         self._cached_block_table_ctx = self.block_table[: self.num_contexts]
         self._cached_block_table_gen = self.block_table[self.num_contexts : self.num_seqs]
         self._cached_req_idx_ctx = self.req_idx_per_token[: self.num_ctx_tokens]
@@ -970,8 +1006,8 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
             )
 
             # Expand indexer_k_cache_block_offsets (only generation)
-            # host_indexer_k_cache_block_offsets already contains correct pool
-            # indices from _get_pool_block_indices() in prepare_for_indexer_k_cache().
+            # host_indexer_k_cache_block_offsets already contains INDEX_KEY
+            # page indices from prepare_for_indexer_k_cache().
             if self.kv_cache_manager is not None and self.num_generations > 0:
                 max_len = self.host_indexer_k_cache_block_offsets.shape[1]
                 gen_block_tensor = self.host_indexer_k_cache_block_offsets[
@@ -1029,22 +1065,10 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
             self.dsl_atom = 1 + self.max_draft_tokens
 
     def prepare_for_indexer_k_cache(self):
-        # Build indexer_k_cache_block_offsets using pool block indices derived
-        # from host_kv_cache_block_offsets (populated by super().prepare()).
-        # This correctly resolves block IDs to memory pool indices, which is
-        # required when host cache offload is enabled (block IDs != pool indices
-        # for onboarded secondary blocks).
         if self.kv_cache_manager is None:
             return
-        pool_indices = self._get_pool_block_indices()
-        self.host_indexer_k_cache_block_offsets[: self.num_seqs].copy_(pool_indices)
-        self.indexer_k_cache_block_offsets[: self.num_seqs].copy_(
-            self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True
-        )
-        # Safety clamp: prevent OOB from CUDA graph padding entries which
-        # may contain stale negative or out-of-range values after block
-        # eviction/onboarding with host cache offload.
-        self.indexer_k_cache_block_offsets.clamp_(min=0)
+        # Keep physical slots for the primary-pool TopK conversion below.
+        pool_indices = self._update_indexer_k_cache_block_offsets()
 
         # Build block_table for topk_indices conversion (actual block allocation)
         cached_token_lens = torch.tensor(
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/registry.py b/tensorrt_llm/_torch/attention_backend/sparse/registry.py
index cd2cea5d2db2..298462b0f409 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/registry.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/registry.py
@@ -18,18 +18,21 @@
 # loaded.
 
 
-def get_sparse_attn_kv_cache_manager(sparse_attention_config: "SparseAttentionConfig") -> type:
+def get_sparse_attn_kv_cache_manager(
+    sparse_attention_config: "SparseAttentionConfig",
+    use_kv_cache_manager_v2: bool = False,
+) -> type:
     from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 
     from .deepseek_v4 import DeepseekV4CacheManager
-    from .dsa import DSACacheManager
+    from .dsa import DSACacheManager, DSACacheManagerV2
     from .minimax_m3 import MiniMaxM3KVCacheManagerV2
     from .rocket import RocketKVCacheManager
 
     if sparse_attention_config.algorithm == "rocket":
         return RocketKVCacheManager
     elif sparse_attention_config.algorithm == "dsa":
-        return DSACacheManager
+        return DSACacheManagerV2 if use_kv_cache_manager_v2 else DSACacheManager
     elif sparse_attention_config.algorithm == "deepseek_v4":
         return DeepseekV4CacheManager
     elif sparse_attention_config.algorithm == "skip_softmax":
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index 9ebbd8df355e..a73111374f75 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -1527,9 +1527,13 @@ def _(req_id: torch.Tensor, block_table: torch.Tensor,
         return torch.empty_like(token_indices)
 
     @torch.library.register_fake("trtllm::indexer_k_cache_gather_op")
-    def _(k_cache: torch.Tensor, slot_mapping_fp8: torch.Tensor,
-          slot_mapping_scale: torch.Tensor, k_token_start: int, num_tokens: int,
-          head_dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    def _(k_cache: torch.Tensor,
+          slot_mapping_fp8: torch.Tensor,
+          slot_mapping_scale: torch.Tensor,
+          k_token_start: int,
+          num_tokens: int,
+          head_dim: int,
+          page_index_scale: int = 1) -> Tuple[torch.Tensor, torch.Tensor]:
         # head_dim is payload bytes per token: 128 for FP8 or 64 for FP4
         # (two packed E2M1 codes per byte). k_fp8 holds raw gathered bytes
         # view-cast as float8_e4m3fn; the FP4 caller reinterprets as int8
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index c4fb115111fe..5c2d6ab8a340 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -136,7 +136,10 @@ def get_kv_cache_manager_cls(
             logger.info("Hybrid linear model has 0 mamba layers; using "
                         "KV cache manager without mamba caching")
             if sparse_attn_config is not None:
-                return get_sparse_attn_kv_cache_manager(sparse_attn_config)
+                return get_sparse_attn_kv_cache_manager(
+                    sparse_attn_config,
+                    use_kv_cache_manager_v2=kv_cache_config.
+                    use_kv_cache_manager_v2)
             return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config)
 
         if (sparse_attn_config is not None
@@ -262,7 +265,9 @@ def get_kv_cache_manager_cls(
                 "yet model retained recurrent-state snapshots.")
         return MambaHybridCacheManagerV2
     elif sparse_attn_config is not None:
-        return get_sparse_attn_kv_cache_manager(sparse_attn_config)
+        return get_sparse_attn_kv_cache_manager(
+            sparse_attn_config,
+            use_kv_cache_manager_v2=kv_cache_config.use_kv_cache_manager_v2)
     else:
         return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config)
 
@@ -619,6 +624,11 @@ def _fallback_if_unsupported_kv_cache_manager_v2(
                 incompat.append("kv_connector_manager")
             if self._max_beam_width is not None and self._max_beam_width > 1:
                 incompat.append("max_beam_width > 1")
+            sparse_attn_config = model_config.sparse_attention_config
+            if (sparse_attn_config is not None
+                    and sparse_attn_config.algorithm == "dsa"
+                    and self._mapping.cp_config.get("cp_type") == CpType.STAR):
+                incompat.append("STAR context parallelism")
             if incompat:
                 incompat_str = ", ".join(incompat)
                 # Some models are structurally bound to V2 and cannot fall
@@ -630,8 +640,14 @@ def _fallback_if_unsupported_kv_cache_manager_v2(
                 #   * Gemma4 hybrid uses per-layer head_dim that V1 would
                 #     coerce to ``max(head_dim)``, changing per-layer KV
                 #     byte sizes — correctness bug, not just efficiency.
-                sparse_attn_config = model_config.sparse_attention_config
                 if sparse_attn_config is not None:
+                    if sparse_attn_config.algorithm == "dsa":
+                        logger.warning(
+                            "DSA with KVCacheManagerV2 is not supported with %s. "
+                            "Falling back to DSACacheManager (V1).",
+                            incompat_str)
+                        return get_sparse_attn_kv_cache_manager(
+                            sparse_attn_config, use_kv_cache_manager_v2=False)
                     raise NotImplementedError(
                         f"Sparse-attention models "
                         f"(algorithm={sparse_attn_config.algorithm!r}) require "
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 8b00d737c793..6e3762648f50 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -3529,6 +3529,62 @@ def test_dsa_host_cache_offload(self, tp_size, pp_size, ep_size, mtp_nextn,
             task = MMLU(self.MODEL_NAME)
             task.evaluate(llm)
 
+    @pytest.mark.skip_less_mpi_world_size(4)
+    @skip_pre_blackwell
+    @pytest.mark.parametrize(
+        "indexer_k_dtype,mtp_nextn,use_cute_dsl,host_cache_size_gb,chunked_prefill,max_batch_size",
+        [
+            ("fp8", 0, False, 0, False, 24),
+            ("fp4", 1, False, 0, True, 24),
+            ("fp4", 3, True, 0, False, 24),
+            ("fp8", 0, False, 10, False, 24),
+        ],
+        ids=[
+            "baseline",
+            "fp4_mtp1_chunked",
+            "fp4_dsl_mtp3",
+            "host_cache_offload",
+        ],
+    )
+    def test_kv_cache_manager_v2(self, indexer_k_dtype, mtp_nextn, use_cute_dsl,
+                                 host_cache_size_gb, chunked_prefill,
+                                 max_batch_size):
+        """Exercise DSA indexer cache paths backed by KVCacheManagerV2."""
+        kv_cache_config = KvCacheConfig(
+            free_gpu_memory_fraction=0.4 if host_cache_size_gb else 0.7,
+            host_cache_size=host_cache_size_gb * (1 << 30),
+            use_kv_cache_manager_v2=True,
+        )
+        sparse_attention_config = DeepSeekSparseAttentionConfig(
+            indexer_k_dtype=indexer_k_dtype,
+            skip_indexer_for_short_seqs=False,
+            use_cute_dsl_paged_mqa_logits=use_cute_dsl,
+        )
+        speculative_config = (MTPDecodingConfig(
+            max_draft_len=mtp_nextn) if mtp_nextn else None)
+        llm_kwargs = {}
+        if chunked_prefill:
+            llm_kwargs.update(enable_chunked_prefill=True, max_num_tokens=512)
+
+        with LLM(
+                f"{llm_models_root()}/DeepSeek-V3.2-Exp-FP4-v2",
+                max_batch_size=max_batch_size,
+                tensor_parallel_size=4,
+                pipeline_parallel_size=1,
+                moe_expert_parallel_size=4,
+                kv_cache_config=kv_cache_config,
+                disable_overlap_scheduler=False,
+                cuda_graph_config=CudaGraphConfig(
+                    enable_padding=True, max_batch_size=max_batch_size),
+                moe_config=MoeConfig(backend="CUTLASS", max_num_tokens=16384),
+                enable_attention_dp=True,
+                speculative_config=speculative_config,
+                sparse_attention_config=sparse_attention_config,
+                **llm_kwargs,
+        ) as llm:
+            task = MMLU(self.MODEL_NAME)
+            task.evaluate(llm)
+
     @pytest.mark.skip_less_mpi_world_size(8)
     @skip_pre_blackwell
     @pytest.mark.parametrize(
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index 765ad10586f3..efce87353b52 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -175,6 +175,10 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus_corner_c
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp]
+accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[baseline]
+accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_mtp1_chunked]
+accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_dsl_mtp3]
+accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[host_cache_offload]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline_fp8kv]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline_mtp1]
diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 74d7936cab15..7153ad6bd542 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -179,6 +179,10 @@ l0_dgx_b200:
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] TIMEOUT (60)
+  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[baseline] TIMEOUT (60)
+  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_mtp1_chunked] TIMEOUT (60)
+  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_dsl_mtp3] TIMEOUT (60)
+  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[host_cache_offload] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4_mtp_index_share[tp_size=8-ep_size=8] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240)
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index 3cd06c02f863..b8fd3efb1d6e 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -43,6 +43,7 @@
 from tensorrt_llm._torch.attention_backend.sparse.dsa import (
     DSABackendForwardArgs,
     DSACacheManager,
+    DSACacheManagerV2,
     DSATrtllmAttention,
     DSAtrtllmAttentionMetadata,
     Indexer,
@@ -51,19 +52,22 @@
     compute_cu_seqlen_kv_bounds_with_cache,
     split_prefill_chunks,
 )
+from tensorrt_llm._torch.attention_backend.sparse.registry import get_sparse_attn_kv_cache_manager
+from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
 from tensorrt_llm._torch.speculative.interface import (
     prepare_attn_metadata_for_draft_replay,
     restore_attn_metadata_after_draft_replay,
 )
 from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.bindings import DataType
-from tensorrt_llm.bindings.executor import KvCacheConfig
+from tensorrt_llm.bindings.executor import KvCacheConfig as BindingKvCacheConfig
 from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp
 from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits
 from tensorrt_llm.functional import PositionEmbeddingType
 from tensorrt_llm.llmapi.llm_args import (
     DeepSeekSparseAttentionConfig,
     DeepSeekV4SparseAttentionConfig,
+    KvCacheConfig,
 )
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.quantization.utils import fp8_utils
@@ -250,8 +254,9 @@ def create_dsa_cache_manager(
     num_layers: int = 1,
     indexer_k_dtype: str = "fp8",
     index_topk: int = 2048,
+    use_kv_cache_manager_v2: bool = False,
 ):
-    """Helper to create a DSACacheManager for testing."""
+    """Helper to create a DSA cache manager for testing."""
 
     sparse_attn_config = DeepSeekSparseAttentionConfig(
         index_head_dim=head_dim,
@@ -261,7 +266,8 @@ def create_dsa_cache_manager(
     )
 
     # Create KV cache config
-    kv_cache_config = KvCacheConfig(
+    kv_cache_config_cls = KvCacheConfig if use_kv_cache_manager_v2 else BindingKvCacheConfig
+    kv_cache_config = kv_cache_config_cls(
         enable_block_reuse=False,
         max_tokens=max_seq_len * batch_size,
     )
@@ -271,7 +277,8 @@ def create_dsa_cache_manager(
 
     # Create cache manager
     # Use SELFKONLY for DSA (similar to MLA usage in _util.py)
-    cache_manager = DSACacheManager(
+    cache_manager_cls = DSACacheManagerV2 if use_kv_cache_manager_v2 else DSACacheManager
+    cache_manager = cache_manager_cls(
         kv_cache_config=kv_cache_config,
         kv_cache_type=CacheTypeCpp.SELFKONLY,
         num_layers=num_layers,
@@ -288,6 +295,148 @@ def create_dsa_cache_manager(
     return cache_manager, sparse_attn_config
 
 
+def test_dsa_cache_manager_selection_honors_v2_config():
+    sparse_attn_config = DeepSeekSparseAttentionConfig()
+
+    assert (
+        get_sparse_attn_kv_cache_manager(sparse_attn_config, use_kv_cache_manager_v2=False)
+        is DSACacheManager
+    )
+    assert (
+        get_sparse_attn_kv_cache_manager(sparse_attn_config, use_kv_cache_manager_v2=True)
+        is DSACacheManagerV2
+    )
+
+
+@pytest.mark.parametrize(
+    "indexer_k_dtype,indexer_bytes_per_token",
+    [
+        ("fp8", 132),
+        pytest.param("fp4", 68, marks=skip_pre_blackwell),
+    ],
+)
+def test_dsa_cache_manager_v2_cache_size_estimate_includes_indexer(
+    indexer_k_dtype, indexer_bytes_per_token
+):
+    """Warmup estimation includes main MLA K and indexer K for every layer."""
+    num_layers = 3
+    head_dim = 128
+    model_config = SimpleNamespace(
+        pretrained_config=SimpleNamespace(
+            num_attention_heads=1,
+            num_key_value_heads=1,
+            hidden_size=head_dim,
+            kv_lora_rank=96,
+            qk_rope_head_dim=32,
+        ),
+        quant_config=None,
+        sparse_attention_config=DeepSeekSparseAttentionConfig(
+            index_head_dim=head_dim,
+            indexer_k_dtype=indexer_k_dtype,
+        ),
+        get_num_attention_layers=lambda: num_layers,
+    )
+    mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1)
+
+    estimated_bytes_per_token = DSACacheManagerV2.get_cache_size_per_token(
+        model_config,
+        mapping,
+    )
+
+    mla_k_bytes_per_layer = head_dim * 2
+    assert estimated_bytes_per_token == num_layers * (
+        mla_k_bytes_per_layer + indexer_bytes_per_token
+    )
+
+
+@pytest.mark.parametrize(
+    "indexer_k_dtype",
+    ["fp8", pytest.param("fp4", marks=skip_pre_blackwell)],
+)
+def test_dsa_cache_manager_v2_indexer_pool_layout(indexer_k_dtype):
+    """V2 exposes shared-page indexer buffers and physical block metadata."""
+    num_layers = 3
+    tokens_per_block = 16
+    head_dim = 128
+    cache_manager, _ = create_dsa_cache_manager(
+        batch_size=2,
+        head_dim=head_dim,
+        tokens_per_block=tokens_per_block,
+        max_seq_len=64,
+        num_layers=num_layers,
+        indexer_k_dtype=indexer_k_dtype,
+        use_kv_cache_manager_v2=True,
+    )
+
+    try:
+        assert isinstance(cache_manager, DSACacheManagerV2)
+        page_index_scale = cache_manager.get_indexer_k_cache_page_index_scale(0)
+        assert page_index_scale == num_layers
+        payload_bytes = head_dim // 2 if indexer_k_dtype == "fp4" else head_dim
+        indexer_bytes_per_token = payload_bytes + 4
+        key_bytes_per_token = cache_manager.get_layer_bytes_per_token(0, Role.KEY)
+        all_bytes_per_token = cache_manager.get_layer_bytes_per_token(0, Role.ALL)
+        assert cache_manager.get_layer_bytes_per_token(0, Role.INDEX_KEY) == indexer_bytes_per_token
+        assert all_bytes_per_token == key_bytes_per_token + indexer_bytes_per_token
+        assert cache_manager.get_cache_bytes_per_token() == num_layers * all_bytes_per_token
+
+        layer_sizes, attention_windows = cache_manager._get_runtime_cache_size_layer_components()
+        assert layer_sizes == [all_bytes_per_token] * num_layers
+        assert attention_windows == [None] * num_layers
+        assert cache_manager._get_quota_from_max_tokens(64) == (
+            64 * num_layers * all_bytes_per_token
+        )
+
+        for layer_config in cache_manager.kv_cache_manager_py_config.layers:
+            indexer_buffer_config = next(
+                buffer for buffer in layer_config.buffers if buffer.role == Role.INDEX_KEY
+            )
+            assert indexer_buffer_config.size == (indexer_bytes_per_token * tokens_per_block)
+
+        physical_blocks = torch.tensor([[0, 1, 3]], device="cuda", dtype=torch.int32)
+        scaled_blocks = cache_manager.scale_indexer_k_cache_block_table(physical_blocks, 0)
+        assert torch.equal(scaled_blocks, physical_blocks * page_index_scale)
+
+        indexer_buffer = cache_manager.get_indexer_k_cache_buffers(0)
+        assert indexer_buffer.shape == (
+            cache_manager.blocks_in_primary_pool * page_index_scale,
+            tokens_per_block,
+            1,
+            payload_bytes + 4,
+        )
+
+        primary_pool = cache_manager.get_unique_primary_pool()
+        assert primary_pool.shape[:3] == (
+            cache_manager.blocks_in_primary_pool,
+            num_layers,
+            1,
+        )
+
+        request_ids = [11, 22]
+        cache_manager.add_dummy_requests(
+            request_ids,
+            token_nums=[16, 32],
+            prepare_resource=True,
+        )
+        pool_indices = cache_manager.get_pool_block_indices(
+            len(request_ids),
+            request_ids=request_ids,
+            num_contexts=len(request_ids),
+        )
+        reversed_pool_indices = cache_manager.get_pool_block_indices(
+            len(request_ids),
+            request_ids=list(reversed(request_ids)),
+            num_contexts=len(request_ids),
+        )
+        assert pool_indices.shape[0] == len(request_ids)
+        assert pool_indices.min() >= 0
+        assert pool_indices.max() < cache_manager.blocks_in_primary_pool
+        assert torch.equal(reversed_pool_indices[0], pool_indices[1])
+        assert torch.equal(reversed_pool_indices[1], pool_indices[0])
+    finally:
+        cache_manager.shutdown()
+
+
 def create_indexer(sparse_attn_config, layer_idx=0):
     """Helper to create an Indexer for testing."""
     # Create RopeParams
diff --git a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
index 09f16a5e7666..d27f7a1664e4 100644
--- a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
+++ b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
@@ -15,6 +15,7 @@
 """Unit tests for DSA C++ custom ops:
 
 - ``torch.ops.trtllm.indexer_k_cache_gather_op``
+- ``torch.ops.trtllm.indexer_k_cache_scatter_op``
 - ``torch.ops.trtllm.convert_req_index_to_global``
 - ``torch.ops.trtllm.fused_cat_fp4``
 - ``torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell``
@@ -126,6 +127,17 @@ def _create_4d_cache_and_mappings(
     return k_cache, slot_mapping_fp8, slot_mapping_scale
 
 
+def _scale_flat_slot_mapping(
+    slot_mapping: torch.Tensor,
+    block_stride: int,
+    page_index_scale: int,
+) -> torch.Tensor:
+    """Convert physical-block flat offsets to shared-page flat offsets."""
+    block_idx = torch.div(slot_mapping, block_stride, rounding_mode="floor")
+    in_block_offset = slot_mapping - block_idx * block_stride
+    return block_idx * page_index_scale * block_stride + in_block_offset
+
+
 @pytest.mark.parametrize(
     "total_kv_len,num_blocks,block_size,k_token_start,num_tokens",
     [
@@ -247,6 +259,112 @@ def test_indexer_k_cache_gather_noncontiguous(
     )
 
 
+@pytest.mark.parametrize("payload_bytes", [128, 64], ids=["fp8", "fp4"])
+def test_indexer_k_cache_gather_page_index_scale(payload_bytes):
+    """Gather converts physical block IDs to V2 shared-page indices."""
+    device = torch.device("cuda")
+    page_index_scale = 4
+    num_blocks = 8
+    block_size = 16
+    num_tokens = 64
+    per_token_size = payload_bytes + SCALE_BYTES
+    block_stride = block_size * per_token_size
+
+    physical_cache = torch.randint(
+        0,
+        256,
+        (num_blocks, block_size, 1, per_token_size),
+        dtype=torch.uint8,
+        device=device,
+    )
+    k_cache = torch.zeros(
+        (num_blocks * page_index_scale, block_size, 1, per_token_size),
+        dtype=torch.uint8,
+        device=device,
+    )
+    k_cache[::page_index_scale] = physical_cache
+
+    perm = torch.randperm(num_blocks * block_size, device=device)[:num_tokens]
+    slot_data = perm.to(torch.int64) * per_token_size
+    slot_scale = slot_data + payload_bytes
+    gathered_data, gathered_scale = torch.ops.trtllm.indexer_k_cache_gather_op(
+        k_cache,
+        slot_data,
+        slot_scale,
+        0,
+        num_tokens,
+        payload_bytes,
+        page_index_scale,
+    )
+
+    scaled_data = _scale_flat_slot_mapping(slot_data, block_stride, page_index_scale)
+    scaled_scale = _scale_flat_slot_mapping(slot_scale, block_stride, page_index_scale)
+    flat_cache = k_cache.reshape(-1)
+    data_offsets = torch.arange(payload_bytes, device=device, dtype=torch.int64)
+    scale_offsets = torch.arange(SCALE_BYTES, device=device, dtype=torch.int64)
+    expected_data = flat_cache[scaled_data[:, None] + data_offsets[None, :]]
+    expected_scale = flat_cache[scaled_scale[:, None] + scale_offsets[None, :]]
+
+    assert torch.equal(gathered_data.view(torch.uint8), expected_data)
+    assert torch.equal(gathered_scale.view(torch.uint8), expected_scale)
+
+
+@pytest.mark.parametrize("payload_bytes", [128, 64], ids=["fp8", "fp4"])
+def test_indexer_k_cache_scatter_page_index_scale(payload_bytes):
+    """Scatter converts physical block IDs to V2 shared-page indices."""
+    device = torch.device("cuda")
+    page_index_scale = 4
+    num_blocks = 8
+    block_size = 16
+    num_tokens = 64
+    per_token_size = payload_bytes + SCALE_BYTES
+    block_stride = block_size * per_token_size
+
+    k_cache = torch.zeros(
+        (num_blocks * page_index_scale, block_size, 1, per_token_size),
+        dtype=torch.uint8,
+        device=device,
+    )
+    expected = torch.zeros_like(k_cache)
+    perm = torch.randperm(num_blocks * block_size, device=device)[:num_tokens]
+    slot_data = perm.to(torch.int64) * per_token_size
+    slot_scale = slot_data + payload_bytes
+    data_bytes = torch.randint(
+        0,
+        256,
+        (num_tokens, payload_bytes),
+        dtype=torch.uint8,
+        device=device,
+    )
+    scale_bytes = torch.randint(
+        0,
+        256,
+        (num_tokens, SCALE_BYTES),
+        dtype=torch.uint8,
+        device=device,
+    )
+
+    torch.ops.trtllm.indexer_k_cache_scatter_op(
+        data_bytes.view(torch.float8_e4m3fn),
+        scale_bytes.view(torch.float32).view(num_tokens, 1),
+        k_cache,
+        slot_data,
+        slot_scale,
+        num_tokens,
+        page_index_scale,
+    )
+
+    scaled_data = _scale_flat_slot_mapping(slot_data, block_stride, page_index_scale)
+    scaled_scale = _scale_flat_slot_mapping(slot_scale, block_stride, page_index_scale)
+    flat_expected = expected.reshape(-1)
+    data_offsets = torch.arange(payload_bytes, device=device, dtype=torch.int64)
+    scale_offsets = torch.arange(SCALE_BYTES, device=device, dtype=torch.int64)
+    flat_expected[scaled_data[:, None] + data_offsets[None, :]] = data_bytes
+    flat_expected[scaled_scale[:, None] + scale_offsets[None, :]] = scale_bytes
+
+    assert torch.equal(k_cache, expected)
+
+
 @pytest.mark.parametrize(
     "total_kv_len,num_blocks,block_size,k_token_start,num_tokens",
     [

From 032e7c8be931fdedfa1d02767e8fad999ed93713 Mon Sep 17 00:00:00 2001
From: yizhang-nv <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 8 Jul 2026 01:29:55 -0700
Subject: [PATCH 02/12] [None][fix] resolve DSA indexer cache page indices

Signed-off-by: yizhang-nv <187001205+yizhang-nv@users.noreply.github.com>
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../kernels/IndexerKCacheGather.h             |   2 +-
 .../kernels/IndexerKCacheScatter.h            |   2 +-
 .../kernels/indexerKCacheGather.cu            |  24 +---
 .../kernels/indexerKCacheScatter.cu           |  25 +---
 .../thop/IndexerKCacheGatherOp.cpp            |   9 +-
 .../thop/IndexerKCacheScatterOp.cpp           |   8 +-
 .../_torch/custom_ops/cpp_custom_ops.py       |  10 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |  32 ++---
 tensorrt_llm/_torch/speculative/interface.py  |  22 +---
 .../attention/sparse/dsa/test_dsa_indexer.py  |  43 ++++++-
 .../attention/sparse/test_cpp_custom_ops.py   | 117 ------------------
 11 files changed, 67 insertions(+), 227 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
index 7519d17bf1f0..c0a461e3c72d 100644
--- a/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
+++ b/cpp/tensorrt_llm/kernels/IndexerKCacheGather.h
@@ -28,7 +28,7 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
     int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2,
     int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
-    int32_t page_index_scale, cudaStream_t stream = 0);
+    cudaStream_t stream = 0);
 
 }
 
diff --git a/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h b/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
index cd8123873eef..a62da181cbb5 100644
--- a/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
+++ b/cpp/tensorrt_llm/kernels/IndexerKCacheScatter.h
@@ -28,7 +28,7 @@ void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_sca
     int64_t const* slot_mapping_fp8, int64_t const* slot_mapping_scale, int32_t num_tokens, int32_t head_dim,
     int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
     int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
-    int32_t page_index_scale, cudaStream_t stream = 0);
+    cudaStream_t stream = 0);
 
 }
 
diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
index 81f3ca21fae7..cd33862afb4c 100644
--- a/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
+++ b/cpp/tensorrt_llm/kernels/indexerKCacheGather.cu
@@ -49,18 +49,6 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
     return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3;
 }
 
-__device__ __forceinline__ int64_t scalePageIndex(int64_t flatIndex, int64_t blockStride, int32_t pageIndexScale)
-{
-    if (pageIndexScale == 1)
-    {
-        return flatIndex;
-    }
-
-    int64_t const blockIndex = flatIndex / blockStride;
-    int64_t const inBlockOffset = flatIndex - blockIndex * blockStride;
-    return blockIndex * pageIndexScale * blockStride + inBlockOffset;
-}
-
 } // anonymous namespace
 
 /**
@@ -85,14 +73,12 @@ __device__ __forceinline__ int64_t scalePageIndex(int64_t flatIndex, int64_t blo
  * @param cache_dim_1         Size of k_cache dimension 1
  * @param cache_dim_2         Size of k_cache dimension 2
  * @param cache_dim_3         Size of k_cache dimension 3
- * @param page_index_scale    Scale from physical block ID to V2 shared-page index
  */
 __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_cache,
     int64_t const* __restrict__ slot_mapping_fp8, int64_t const* __restrict__ slot_mapping_scale,
     uint8_t* __restrict__ out_fp8, uint8_t* __restrict__ out_scale, int32_t k_token_start, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2,
-    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
-    int32_t page_index_scale)
+    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3)
 {
     // For head_dim=128, each thread handles 4 bytes/elements per read/write instruction
     constexpr int VEC_SIZE = 4;
@@ -118,9 +104,6 @@ __global__ void indexerKCacheGatherUnifiedKernel(uint8_t const* __restrict__ k_c
 
     int32_t head_dim_idx = threadIdx.x * VEC_SIZE;
     int64_t flat_idx = flat_idx_fp8_base + head_dim_idx;
-    int64_t const blockStride = static_cast<int64_t>(cache_dim_1) * cache_dim_2 * cache_dim_3;
-    flat_idx = scalePageIndex(flat_idx, blockStride, page_index_scale);
-    flat_idx_scale_base = scalePageIndex(flat_idx_scale_base, blockStride, page_index_scale);
 
     // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)
     int64_t src_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
@@ -147,7 +130,7 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
     int64_t const* slot_mapping_scale, uint8_t* out_fp8, uint8_t* out_scale, int32_t k_token_start, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2,
     int32_t cache_dim_3, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
-    int32_t page_index_scale, cudaStream_t stream)
+    cudaStream_t stream)
 {
     if (num_tokens == 0)
     {
@@ -162,7 +145,6 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
     TLLM_CHECK_WITH_INFO(head_dim % VEC_SIZE == 0, "head_dim (%d) must be a multiple of %d", head_dim, VEC_SIZE);
     TLLM_CHECK_WITH_INFO(scale_size == 4,
         "scale_size must equal 4 bytes (packed UE8M0 x4 for FP4, 1 float32 for FP8, got %d)", scale_size);
-    TLLM_CHECK_WITH_INFO(page_index_scale >= 1, "page_index_scale must be at least 1 (got %d)", page_index_scale);
 
     int32_t const threads_per_block = head_dim / VEC_SIZE;
 
@@ -171,7 +153,7 @@ void invokeIndexerKCacheGather(uint8_t const* k_cache, int64_t const* slot_mappi
 
     indexerKCacheGatherUnifiedKernel<<<grid, block, 0, stream>>>(k_cache, slot_mapping_fp8, slot_mapping_scale, out_fp8,
         out_scale, k_token_start, num_tokens, head_dim, scale_size, cache_stride_0, cache_stride_1, cache_stride_2,
-        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, page_index_scale);
+        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3);
 
     // Check for kernel launch errors
     TLLM_CUDA_CHECK(cudaGetLastError());
diff --git a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu
index 332fdd2eb030..f0c7c27a235c 100644
--- a/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu
+++ b/cpp/tensorrt_llm/kernels/indexerKCacheScatter.cu
@@ -49,18 +49,6 @@ __device__ __forceinline__ int64_t flatIndexToMemoryOffset(
     return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3;
 }
 
-__device__ __forceinline__ int64_t scalePageIndex(int64_t flatIndex, int64_t blockStride, int32_t pageIndexScale)
-{
-    if (pageIndexScale == 1)
-    {
-        return flatIndex;
-    }
-
-    int64_t const blockIndex = flatIndex / blockStride;
-    int64_t const inBlockOffset = flatIndex - blockIndex * blockStride;
-    return blockIndex * pageIndexScale * blockStride + inBlockOffset;
-}
-
 } // anonymous namespace
 
 /**
@@ -83,14 +71,12 @@ __device__ __forceinline__ int64_t scalePageIndex(int64_t flatIndex, int64_t blo
  * @param cache_dim_1       Size of k_cache dimension 1
  * @param cache_dim_2       Size of k_cache dimension 2
  * @param cache_dim_3       Size of k_cache dimension 3
- * @param page_index_scale  Scale from physical block ID to V2 shared-page index
  */
 __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_fp8_bytes,
     uint8_t const* __restrict__ k_scale_bytes, uint8_t* __restrict__ k_cache,
     int64_t const* __restrict__ slot_mapping_fp8, int64_t const* __restrict__ slot_mapping_scale, int32_t num_tokens,
     int32_t head_dim, int32_t scale_size, int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2,
-    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
-    int32_t page_index_scale)
+    int64_t cache_stride_3, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3)
 {
     // For head_dim=128, each thread handles 4 bytes/elements per read/write instruction
     constexpr int VEC_SIZE = 4;
@@ -113,9 +99,6 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_
 
     int32_t head_dim_idx = threadIdx.x * VEC_SIZE;
     int64_t flat_idx = flat_idx_fp8_base + head_dim_idx;
-    int64_t const blockStride = static_cast<int64_t>(cache_dim_1) * cache_dim_2 * cache_dim_3;
-    flat_idx = scalePageIndex(flat_idx, blockStride, page_index_scale);
-    flat_idx_scale_base = scalePageIndex(flat_idx_scale_base, blockStride, page_index_scale);
 
     // Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)
     int64_t dst_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
@@ -141,8 +124,7 @@ __global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_
 void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_scale_bytes, uint8_t* k_cache,
     int64_t const* slot_mapping_fp8, int64_t const* slot_mapping_scale, int32_t num_tokens, int32_t head_dim,
     int32_t scale_size, int32_t cache_dim_0, int32_t cache_dim_1, int32_t cache_dim_2, int32_t cache_dim_3,
-    int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3,
-    int32_t page_index_scale, cudaStream_t stream)
+    int64_t cache_stride_0, int64_t cache_stride_1, int64_t cache_stride_2, int64_t cache_stride_3, cudaStream_t stream)
 {
     if (num_tokens == 0)
     {
@@ -158,7 +140,6 @@ void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_sca
     TLLM_CHECK_WITH_INFO(head_dim % VEC_SIZE == 0, "head_dim (%d) must be a multiple of %d", head_dim, VEC_SIZE);
     TLLM_CHECK_WITH_INFO(scale_size == 4,
         "scale_size must equal 4 bytes (packed UE8M0 x4 for FP4, 1 float32 for FP8, got %d)", scale_size);
-    TLLM_CHECK_WITH_INFO(page_index_scale >= 1, "page_index_scale must be at least 1 (got %d)", page_index_scale);
 
     int32_t const threads_per_block = head_dim / VEC_SIZE;
 
@@ -167,7 +148,7 @@ void invokeIndexerKCacheScatter(uint8_t const* k_fp8_bytes, uint8_t const* k_sca
 
     indexerKCacheScatterUnifiedKernel<<<grid, block, 0, stream>>>(k_fp8_bytes, k_scale_bytes, k_cache, slot_mapping_fp8,
         slot_mapping_scale, num_tokens, head_dim, scale_size, cache_stride_0, cache_stride_1, cache_stride_2,
-        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, page_index_scale);
+        cache_stride_3, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3);
 
     // Check for kernel launch errors
     TLLM_CUDA_CHECK(cudaGetLastError());
diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
index 501263a5cae3..bd094dac58fc 100644
--- a/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
+++ b/cpp/tensorrt_llm/thop/IndexerKCacheGatherOp.cpp
@@ -30,7 +30,7 @@ namespace torch_ext
 
 std::tuple<th::Tensor, th::Tensor> indexer_k_cache_gather_op(th::Tensor const& k_cache,
     th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t k_token_start, int64_t num_tokens,
-    int64_t head_dim, int64_t page_index_scale)
+    int64_t head_dim)
 {
     // head_dim is the number of payload bytes per token: 128 for FP8 (one byte
     // per float8 value) or 64 for FP4 (two packed E2M1 codes per byte).
@@ -86,7 +86,6 @@ std::tuple<th::Tensor, th::Tensor> indexer_k_cache_gather_op(th::Tensor const& k
 
     // Validation for indexer k cache pool for DeepSeek-V3.2 constraints
     TORCH_CHECK(cache_dim_2 == 1, "k_cache dimension 2 must be 1 for DeepSeek-V3.2, got %d", cache_dim_2);
-    TORCH_CHECK(page_index_scale >= 1, "page_index_scale must be at least 1, got %ld", page_index_scale);
 
     int64_t cache_stride_0 = static_cast<int64_t>(k_cache.stride(0));
     int64_t cache_stride_1 = static_cast<int64_t>(k_cache.stride(1));
@@ -103,8 +102,7 @@ std::tuple<th::Tensor, th::Tensor> indexer_k_cache_gather_op(th::Tensor const& k
     tk::invokeIndexerKCacheGather(k_cache.data_ptr<uint8_t>(), slot_mapping_fp8.data_ptr<int64_t>(),
         slot_mapping_scale.data_ptr<int64_t>(), out_fp8.data_ptr<uint8_t>(), out_scale.data_ptr<uint8_t>(),
         static_cast<int32_t>(k_token_start), num_tokens_i32, head_dim_i32, SCALE_SIZE, cache_dim_0, cache_dim_1,
-        cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3,
-        static_cast<int32_t>(page_index_scale), stream);
+        cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3, stream);
 
     // View-cast to final dtypes (no copy). The FP8 call site keeps its existing
     // float8_e4m3fn semantics; the FP4 call site reshapes the raw bytes into
@@ -123,8 +121,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def(
         "indexer_k_cache_gather_op(Tensor k_cache, Tensor slot_mapping_fp8, "
-        "Tensor slot_mapping_scale, int k_token_start, int num_tokens, int head_dim, "
-        "int page_index_scale=1) -> (Tensor, Tensor)");
+        "Tensor slot_mapping_scale, int k_token_start, int num_tokens, int head_dim) -> (Tensor, Tensor)");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
diff --git a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
index 1adfed4fd252..59c77177be5c 100644
--- a/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
+++ b/cpp/tensorrt_llm/thop/IndexerKCacheScatterOp.cpp
@@ -29,8 +29,7 @@ namespace torch_ext
 {
 
 void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_scale, th::Tensor& k_cache,
-    th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t num_tokens,
-    int64_t page_index_scale)
+    th::Tensor const& slot_mapping_fp8, th::Tensor const& slot_mapping_scale, int64_t num_tokens)
 {
     // k_fp8: [>=num_tokens, head_dim] in FP8 (1 byte/element) — reinterpreted as uint8
     // k_scale: [>=num_tokens, head_dim // quant_block_size] in float32 — reinterpreted as uint8 bytes
@@ -79,7 +78,6 @@ void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_sca
     TORCH_CHECK(
         head_dim == 128 || head_dim == 64, "k_fp8 head_dim must be 128 (FP8) or 64 (FP4 packed), got %d", head_dim);
     TORCH_CHECK(scale_size == 4, "k_scale scale_size must be 4 bytes, got %d", scale_size);
-    TORCH_CHECK(page_index_scale >= 1, "page_index_scale must be at least 1, got %ld", page_index_scale);
 
     int64_t const cache_stride_0 = static_cast<int64_t>(k_cache.stride(0));
     int64_t const cache_stride_1 = static_cast<int64_t>(k_cache.stride(1));
@@ -94,7 +92,7 @@ void indexer_k_cache_scatter_op(th::Tensor const& k_fp8, th::Tensor const& k_sca
         reinterpret_cast<uint8_t const*>(k_scale.data_ptr()), k_cache.data_ptr<uint8_t>(),
         slot_mapping_fp8.data_ptr<int64_t>(), slot_mapping_scale.data_ptr<int64_t>(), static_cast<int32_t>(num_tokens),
         head_dim, scale_size, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3, cache_stride_0, cache_stride_1,
-        cache_stride_2, cache_stride_3, static_cast<int32_t>(page_index_scale), stream);
+        cache_stride_2, cache_stride_3, stream);
 }
 
 } // namespace torch_ext
@@ -105,7 +103,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
     m.def(
         "indexer_k_cache_scatter_op(Tensor k_fp8, Tensor k_scale, Tensor(a!) k_cache, "
-        "Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int num_tokens, int page_index_scale=1) -> ()");
+        "Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int num_tokens) -> ()");
 }
 
 TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
index a73111374f75..9ebbd8df355e 100644
--- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
@@ -1527,13 +1527,9 @@ def _(req_id: torch.Tensor, block_table: torch.Tensor,
         return torch.empty_like(token_indices)
 
     @torch.library.register_fake("trtllm::indexer_k_cache_gather_op")
-    def _(k_cache: torch.Tensor,
-          slot_mapping_fp8: torch.Tensor,
-          slot_mapping_scale: torch.Tensor,
-          k_token_start: int,
-          num_tokens: int,
-          head_dim: int,
-          page_index_scale: int = 1) -> Tuple[torch.Tensor, torch.Tensor]:
+    def _(k_cache: torch.Tensor, slot_mapping_fp8: torch.Tensor,
+          slot_mapping_scale: torch.Tensor, k_token_start: int, num_tokens: int,
+          head_dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
         # head_dim is payload bytes per token: 128 for FP8 or 64 for FP4
         # (two packed E2M1 codes per byte). k_fp8 holds raw gathered bytes
         # view-cast as float8_e4m3fn; the FP4 caller reinterprets as int8
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 5c2d6ab8a340..2a42c0e9c9d3 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -631,29 +631,19 @@ def _fallback_if_unsupported_kv_cache_manager_v2(
                 incompat.append("STAR context parallelism")
             if incompat:
                 incompat_str = ", ".join(incompat)
-                # Some models are structurally bound to V2 and cannot fall
-                # back to V1 without producing wrong outputs:
-                #   * Sparse-attention models (e.g. MiniMax-M3) need V2's
-                #     per-layer split-pool to allocate the per-sparse-layer
-                #     INDEX_KEY pool with a different stride than the main
-                #     K/V pool. V1's unified pool cannot represent that.
-                #   * Gemma4 hybrid uses per-layer head_dim that V1 would
-                #     coerce to ``max(head_dim)``, changing per-layer KV
-                #     byte sizes — correctness bug, not just efficiency.
+                # Never silently replace a sparse V2 manager with V1. Some
+                # sparse models require V2 structurally; for models such as DSA
+                # that support both managers, fallback would ignore the user's
+                # explicit manager selection.
                 if sparse_attn_config is not None:
-                    if sparse_attn_config.algorithm == "dsa":
-                        logger.warning(
-                            "DSA with KVCacheManagerV2 is not supported with %s. "
-                            "Falling back to DSACacheManager (V1).",
-                            incompat_str)
-                        return get_sparse_attn_kv_cache_manager(
-                            sparse_attn_config, use_kv_cache_manager_v2=False)
                     raise NotImplementedError(
-                        f"Sparse-attention models "
-                        f"(algorithm={sparse_attn_config.algorithm!r}) require "
-                        f"KVCacheManagerV2, which is not yet supported with "
-                        f"{incompat_str}. Disable these KvCacheConfig features "
-                        f"to run sparse-attention models.")
+                        f"KVCacheManagerV2 for sparse-attention models "
+                        f"(algorithm={sparse_attn_config.algorithm!r}) is not "
+                        f"supported with "
+                        f"{incompat_str}. Disable the incompatible features to "
+                        f"run sparse-attention models.")
+                # Gemma4 hybrid uses per-layer head_dim that V1 would coerce to
+                # ``max(head_dim)``, changing per-layer KV byte sizes.
                 if is_gemma4_hybrid(config):
                     raise NotImplementedError(
                         f"Gemma4 hybrid attention requires KVCacheManagerV2, "
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index bdb122ac4457..9c49354b2ecb 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -175,26 +175,8 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
             'slot_mapping_scale':
             m.slot_mapping_scale.clone(),
         }
-        # Derive pool indices from the draft manager's encoded block
-        # offsets (via _get_pool_block_indices) instead of using raw block
-        # IDs.  With host cache offload, block IDs can exceed
-        # blocks_in_primary_pool after offload swaps (the block keeps its
-        # original high ID even though its memory now lives in the primary
-        # GPU pool).  Using raw block IDs as pool indices causes OOB access
-        # in the indexer k-cache buffers.  _get_pool_block_indices correctly
-        # decodes memPoolBlockIndex from the C++ encoded offsets.
-        # Note: kv_cache_manager was already swapped to draft above (line 67).
-        pool_indices = m._get_pool_block_indices()
-        num_blocks = pool_indices.shape[1]
-        m.host_indexer_k_cache_block_offsets[:m.num_seqs, :num_blocks].copy_(
-            pool_indices)
-        m.indexer_k_cache_block_offsets[:m.num_seqs].copy_(
-            m.host_indexer_k_cache_block_offsets[:m.num_seqs],
-            non_blocking=True)
-        # Safety clamp: sanitize stale padding entries beyond num_seqs
-        # that may contain negative or out-of-range values, matching the
-        # regular DSA prepare() flow.
-        m.indexer_k_cache_block_offsets.clamp_(min=0)
+        # kv_cache_manager was already swapped to the draft manager above.
+        m._update_indexer_k_cache_block_offsets()
         Indexer.recompute_slot_mappings(m)
     return saved
 
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index b8fd3efb1d6e..2349765a09a4 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -44,6 +44,7 @@
     DSABackendForwardArgs,
     DSACacheManager,
     DSACacheManagerV2,
+    DSAMetadataParams,
     DSATrtllmAttention,
     DSAtrtllmAttentionMetadata,
     Indexer,
@@ -53,6 +54,7 @@
     split_prefill_chunks,
 )
 from tensorrt_llm._torch.attention_backend.sparse.registry import get_sparse_attn_kv_cache_manager
+from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
 from tensorrt_llm._torch.speculative.interface import (
     prepare_attn_metadata_for_draft_replay,
@@ -308,6 +310,32 @@ def test_dsa_cache_manager_selection_honors_v2_config():
     )
 
 
+def test_dsa_metadata_accepts_kv_cache_manager_v2():
+    cache_manager = object.__new__(DSACacheManagerV2)
+    cache_manager.tokens_per_block = 32
+    metadata = object.__new__(DSAtrtllmAttentionMetadata)
+    metadata.kv_cache_manager = cache_manager
+    metadata.sparse_metadata_params = DSAMetadataParams(
+        indexer_max_chunk_size=8192,
+        max_sparse_topk=2048,
+        index_head_dim=128,
+        enable_indexer_skip=False,
+        enable_heuristic_topk=False,
+        use_cute_dsl_paged_mqa_logits=False,
+        q_split_threshold=8192,
+    )
+    metadata.is_cuda_graph = False
+
+    with (
+        patch.object(TrtllmAttentionMetadata, "__post_init__"),
+        patch.object(DSAtrtllmAttentionMetadata, "create_buffers_for_mla_rope_append"),
+        patch.object(DSAtrtllmAttentionMetadata, "create_buffers_for_indexer"),
+    ):
+        metadata.__post_init__()
+
+    assert metadata._tokens_per_block == cache_manager.tokens_per_block
+
+
 @pytest.mark.parametrize(
     "indexer_k_dtype,indexer_bytes_per_token",
     [
@@ -354,7 +382,7 @@ def test_dsa_cache_manager_v2_cache_size_estimate_includes_indexer(
     ["fp8", pytest.param("fp4", marks=skip_pre_blackwell)],
 )
 def test_dsa_cache_manager_v2_indexer_pool_layout(indexer_k_dtype):
-    """V2 exposes shared-page indexer buffers and physical block metadata."""
+    """V2 exposes shared-page indexer buffers and role-specific page metadata."""
     num_layers = 3
     tokens_per_block = 16
     head_dim = 128
@@ -370,7 +398,7 @@ def test_dsa_cache_manager_v2_indexer_pool_layout(indexer_k_dtype):
 
     try:
         assert isinstance(cache_manager, DSACacheManagerV2)
-        page_index_scale = cache_manager.get_indexer_k_cache_page_index_scale(0)
+        page_index_scale = cache_manager.indexer_k_cache_page_scale
         assert page_index_scale == num_layers
         payload_bytes = head_dim // 2 if indexer_k_dtype == "fp4" else head_dim
         indexer_bytes_per_token = payload_bytes + 4
@@ -393,10 +421,6 @@ def test_dsa_cache_manager_v2_indexer_pool_layout(indexer_k_dtype):
             )
             assert indexer_buffer_config.size == (indexer_bytes_per_token * tokens_per_block)
 
-        physical_blocks = torch.tensor([[0, 1, 3]], device="cuda", dtype=torch.int32)
-        scaled_blocks = cache_manager.scale_indexer_k_cache_block_table(physical_blocks, 0)
-        assert torch.equal(scaled_blocks, physical_blocks * page_index_scale)
-
         indexer_buffer = cache_manager.get_indexer_k_cache_buffers(0)
         assert indexer_buffer.shape == (
             cache_manager.blocks_in_primary_pool * page_index_scale,
@@ -428,11 +452,18 @@ def test_dsa_cache_manager_v2_indexer_pool_layout(indexer_k_dtype):
             request_ids=list(reversed(request_ids)),
             num_contexts=len(request_ids),
         )
+        indexer_page_indices = pool_indices * page_index_scale
+        reversed_indexer_page_indices = reversed_pool_indices * page_index_scale
         assert pool_indices.shape[0] == len(request_ids)
         assert pool_indices.min() >= 0
         assert pool_indices.max() < cache_manager.blocks_in_primary_pool
         assert torch.equal(reversed_pool_indices[0], pool_indices[1])
         assert torch.equal(reversed_pool_indices[1], pool_indices[0])
+        assert torch.equal(indexer_page_indices, pool_indices * page_index_scale)
+        assert indexer_page_indices.min() >= 0
+        assert indexer_page_indices.max() < indexer_buffer.shape[0]
+        assert torch.equal(reversed_indexer_page_indices[0], indexer_page_indices[1])
+        assert torch.equal(reversed_indexer_page_indices[1], indexer_page_indices[0])
     finally:
         cache_manager.shutdown()
 
diff --git a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
index d27f7a1664e4..5bc02ede7379 100644
--- a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
+++ b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
@@ -127,17 +127,6 @@ def _create_4d_cache_and_mappings(
     return k_cache, slot_mapping_fp8, slot_mapping_scale
 
 
-def _scale_flat_slot_mapping(
-    slot_mapping: torch.Tensor,
-    block_stride: int,
-    page_index_scale: int,
-) -> torch.Tensor:
-    """Convert physical-block flat offsets to shared-page flat offsets."""
-    block_idx = torch.div(slot_mapping, block_stride, rounding_mode="floor")
-    in_block_offset = slot_mapping - block_idx * block_stride
-    return block_idx * page_index_scale * block_stride + in_block_offset
-
-
 @pytest.mark.parametrize(
     "total_kv_len,num_blocks,block_size,k_token_start,num_tokens",
     [
@@ -259,112 +248,6 @@ def test_indexer_k_cache_gather_noncontiguous(
     )
 
 
-@pytest.mark.parametrize("payload_bytes", [128, 64], ids=["fp8", "fp4"])
-def test_indexer_k_cache_gather_page_index_scale(payload_bytes):
-    """Gather converts physical block IDs to V2 shared-page indices."""
-    device = torch.device("cuda")
-    page_index_scale = 4
-    num_blocks = 8
-    block_size = 16
-    num_tokens = 64
-    per_token_size = payload_bytes + SCALE_BYTES
-    block_stride = block_size * per_token_size
-
-    physical_cache = torch.randint(
-        0,
-        256,
-        (num_blocks, block_size, 1, per_token_size),
-        dtype=torch.uint8,
-        device=device,
-    )
-    k_cache = torch.zeros(
-        (num_blocks * page_index_scale, block_size, 1, per_token_size),
-        dtype=torch.uint8,
-        device=device,
-    )
-    k_cache[::page_index_scale] = physical_cache
-
-    perm = torch.randperm(num_blocks * block_size, device=device)[:num_tokens]
-    slot_data = perm.to(torch.int64) * per_token_size
-    slot_scale = slot_data + payload_bytes
-    gathered_data, gathered_scale = torch.ops.trtllm.indexer_k_cache_gather_op(
-        k_cache,
-        slot_data,
-        slot_scale,
-        0,
-        num_tokens,
-        payload_bytes,
-        page_index_scale,
-    )
-
-    scaled_data = _scale_flat_slot_mapping(slot_data, block_stride, page_index_scale)
-    scaled_scale = _scale_flat_slot_mapping(slot_scale, block_stride, page_index_scale)
-    flat_cache = k_cache.reshape(-1)
-    data_offsets = torch.arange(payload_bytes, device=device, dtype=torch.int64)
-    scale_offsets = torch.arange(SCALE_BYTES, device=device, dtype=torch.int64)
-    expected_data = flat_cache[scaled_data[:, None] + data_offsets[None, :]]
-    expected_scale = flat_cache[scaled_scale[:, None] + scale_offsets[None, :]]
-
-    assert torch.equal(gathered_data.view(torch.uint8), expected_data)
-    assert torch.equal(gathered_scale.view(torch.uint8), expected_scale)
-
-
-@pytest.mark.parametrize("payload_bytes", [128, 64], ids=["fp8", "fp4"])
-def test_indexer_k_cache_scatter_page_index_scale(payload_bytes):
-    """Scatter converts physical block IDs to V2 shared-page indices."""
-    device = torch.device("cuda")
-    page_index_scale = 4
-    num_blocks = 8
-    block_size = 16
-    num_tokens = 64
-    per_token_size = payload_bytes + SCALE_BYTES
-    block_stride = block_size * per_token_size
-
-    k_cache = torch.zeros(
-        (num_blocks * page_index_scale, block_size, 1, per_token_size),
-        dtype=torch.uint8,
-        device=device,
-    )
-    expected = torch.zeros_like(k_cache)
-    perm = torch.randperm(num_blocks * block_size, device=device)[:num_tokens]
-    slot_data = perm.to(torch.int64) * per_token_size
-    slot_scale = slot_data + payload_bytes
-    data_bytes = torch.randint(
-        0,
-        256,
-        (num_tokens, payload_bytes),
-        dtype=torch.uint8,
-        device=device,
-    )
-    scale_bytes = torch.randint(
-        0,
-        256,
-        (num_tokens, SCALE_BYTES),
-        dtype=torch.uint8,
-        device=device,
-    )
-
-    torch.ops.trtllm.indexer_k_cache_scatter_op(
-        data_bytes.view(torch.float8_e4m3fn),
-        scale_bytes.view(torch.float32).view(num_tokens, 1),
-        k_cache,
-        slot_data,
-        slot_scale,
-        num_tokens,
-        page_index_scale,
-    )
-
-    scaled_data = _scale_flat_slot_mapping(slot_data, block_stride, page_index_scale)
-    scaled_scale = _scale_flat_slot_mapping(slot_scale, block_stride, page_index_scale)
-    flat_expected = expected.reshape(-1)
-    data_offsets = torch.arange(payload_bytes, device=device, dtype=torch.int64)
-    scale_offsets = torch.arange(SCALE_BYTES, device=device, dtype=torch.int64)
-    flat_expected[scaled_data[:, None] + data_offsets[None, :]] = data_bytes
-    flat_expected[scaled_scale[:, None] + scale_offsets[None, :]] = scale_bytes
-
-    assert torch.equal(k_cache, expected)
-
-
 @pytest.mark.parametrize(
     "total_kv_len,num_blocks,block_size,k_token_start,num_tokens",
     [

From 24be8e5747de405524ed689864216a22159ebe7b Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 15 Jul 2026 04:06:17 -0700
Subject: [PATCH 03/12] [None][fix] stabilize DSA draft cache replay with V2

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 tensorrt_llm/_torch/speculative/interface.py  | 71 +++++++++---------
 .../defs/accuracy/test_llm_api_pytorch.py     | 60 +--------------
 .../test_lists/qa/llm_function_core.txt       |  4 -
 .../test_lists/test-db/l0_dgx_b200.yml        |  4 -
 .../attention/sparse/dsa/test_dsa_indexer.py  | 75 +++++++++++++++++++
 5 files changed, 112 insertions(+), 102 deletions(-)

diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 9c49354b2ecb..1edc11d149da 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -163,21 +163,34 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
         m = attn_metadata
         saved['saved_dsa_state'] = {
             'host_indexer_k_cache_block_offsets':
-            m.host_indexer_k_cache_block_offsets.clone(),
+            m.host_indexer_k_cache_block_offsets,
             'indexer_k_cache_block_offsets':
             m.indexer_k_cache_block_offsets.clone(),
-            'host_slot_mapping_fp8':
-            m.host_slot_mapping_fp8.clone(),
-            'host_slot_mapping_scale':
-            m.host_slot_mapping_scale.clone(),
-            'slot_mapping_fp8':
-            m.slot_mapping_fp8.clone(),
-            'slot_mapping_scale':
-            m.slot_mapping_scale.clone(),
+            'host_slot_mapping_fp8': m.host_slot_mapping_fp8,
+            'host_slot_mapping_scale': m.host_slot_mapping_scale,
+            'slot_mapping_fp8': m.slot_mapping_fp8.clone(),
+            'slot_mapping_scale': m.slot_mapping_scale.clone(),
+            'slot_mapping_fp8_fullkv': m.slot_mapping_fp8_fullkv,
+            'slot_mapping_scale_fullkv': m.slot_mapping_scale_fullkv,
         }
+        # Keep the target's pinned host buffers untouched while asynchronous
+        # H2D copies from the draft mappings are in flight.  Restoring their
+        # contents in place can race those copies because CPU writes are not
+        # ordered with the CUDA stream.  The fresh host buffers stay attached
+        # to the metadata until after the draft forward has been enqueued.
+        m.host_indexer_k_cache_block_offsets = torch.empty_like(
+            m.host_indexer_k_cache_block_offsets,
+            pin_memory=m.host_indexer_k_cache_block_offsets.is_pinned())
+        m.host_slot_mapping_fp8 = torch.empty_like(
+            m.host_slot_mapping_fp8,
+            pin_memory=m.host_slot_mapping_fp8.is_pinned())
+        m.host_slot_mapping_scale = torch.empty_like(
+            m.host_slot_mapping_scale,
+            pin_memory=m.host_slot_mapping_scale.is_pinned())
         # kv_cache_manager was already swapped to the draft manager above.
         m._update_indexer_k_cache_block_offsets()
         Indexer.recompute_slot_mappings(m)
+        Indexer.recompute_context_kv_gather_mappings(m)
     return saved
 
 
@@ -195,14 +208,16 @@ def restore_attn_metadata_after_draft_replay(attn_metadata, saved_state):
     saved_dsa = saved_state.get('saved_dsa_state')
     if saved_dsa is not None:
         m = attn_metadata
-        m.host_indexer_k_cache_block_offsets.copy_(
-            saved_dsa['host_indexer_k_cache_block_offsets'], non_blocking=True)
+        m.host_indexer_k_cache_block_offsets = saved_dsa[
+            'host_indexer_k_cache_block_offsets']
         m.indexer_k_cache_block_offsets.copy_(
             saved_dsa['indexer_k_cache_block_offsets'], non_blocking=True)
-        m.host_slot_mapping_fp8.copy_(saved_dsa['host_slot_mapping_fp8'])
-        m.host_slot_mapping_scale.copy_(saved_dsa['host_slot_mapping_scale'])
+        m.host_slot_mapping_fp8 = saved_dsa['host_slot_mapping_fp8']
+        m.host_slot_mapping_scale = saved_dsa['host_slot_mapping_scale']
         m.slot_mapping_fp8.copy_(saved_dsa['slot_mapping_fp8'])
         m.slot_mapping_scale.copy_(saved_dsa['slot_mapping_scale'])
+        m.slot_mapping_fp8_fullkv = saved_dsa['slot_mapping_fp8_fullkv']
+        m.slot_mapping_scale_fullkv = saved_dsa['slot_mapping_scale_fullkv']
 
 
 def get_force_num_accepted_tokens() -> int:
@@ -2121,7 +2136,8 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager):
         """
         Select draft attention metadata for one-engine speculative decoding.
 
-        TRTLLM metadata temporarily swaps its manager and block offsets.
+        TRTLLM metadata temporarily swaps its manager and cache-layout-dependent
+        buffers, including DSA indexer offsets and slot mappings.
         FlashInfer uses an independently planned metadata view because its page
         tables and kernel wrappers are manager-specific.
         """
@@ -2140,35 +2156,16 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager):
             yield attn_metadata
             return
 
-        # Check if draft KV cache block offsets are allocated
-        draft_block_offsets = getattr(attn_metadata,
-                                      'draft_kv_cache_block_offsets', None)
-        if draft_block_offsets is None:
-            # Draft KV cache block offsets not allocated, skip switching
+        saved_state = prepare_attn_metadata_for_draft_replay(
+            attn_metadata, draft_kv_cache_manager)
+        if saved_state is None:
             yield attn_metadata
             return
 
-        # Save main KV cache manager and block offsets
-        target_kv_cache_manager = attn_metadata.kv_cache_manager
-        target_kv_cache_block_offsets = attn_metadata.kv_cache_block_offsets
-        target_host_kv_cache_block_offsets = attn_metadata.host_kv_cache_block_offsets
-
-        # Switch to draft KV cache manager and its block offsets
-        attn_metadata.kv_cache_manager = draft_kv_cache_manager
-        attn_metadata.kv_cache_block_offsets = attn_metadata.draft_kv_cache_block_offsets
-        attn_metadata.host_kv_cache_block_offsets = draft_kv_cache_manager.host_kv_cache_block_offsets
-        if attn_metadata.enable_flash_mla:
-            attn_metadata.prepare_flash_mla()
-
         try:
             yield attn_metadata
         finally:
-            # Restore main KV cache manager and block offsets
-            attn_metadata.kv_cache_manager = target_kv_cache_manager
-            attn_metadata.kv_cache_block_offsets = target_kv_cache_block_offsets
-            attn_metadata.host_kv_cache_block_offsets = target_host_kv_cache_block_offsets
-            if attn_metadata.enable_flash_mla:
-                attn_metadata.prepare_flash_mla()
+            restore_attn_metadata_after_draft_replay(attn_metadata, saved_state)
 
     def _sample_tokens_for_batch(
         self,
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 6e3762648f50..d7aa17fda956 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -3529,62 +3529,6 @@ def test_dsa_host_cache_offload(self, tp_size, pp_size, ep_size, mtp_nextn,
             task = MMLU(self.MODEL_NAME)
             task.evaluate(llm)
 
-    @pytest.mark.skip_less_mpi_world_size(4)
-    @skip_pre_blackwell
-    @pytest.mark.parametrize(
-        "indexer_k_dtype,mtp_nextn,use_cute_dsl,host_cache_size_gb,chunked_prefill,max_batch_size",
-        [
-            ("fp8", 0, False, 0, False, 24),
-            ("fp4", 1, False, 0, True, 24),
-            ("fp4", 3, True, 0, False, 24),
-            ("fp8", 0, False, 10, False, 24),
-        ],
-        ids=[
-            "baseline",
-            "fp4_mtp1_chunked",
-            "fp4_dsl_mtp3",
-            "host_cache_offload",
-        ],
-    )
-    def test_kv_cache_manager_v2(self, indexer_k_dtype, mtp_nextn, use_cute_dsl,
-                                 host_cache_size_gb, chunked_prefill,
-                                 max_batch_size):
-        """Exercise DSA indexer cache paths backed by KVCacheManagerV2."""
-        kv_cache_config = KvCacheConfig(
-            free_gpu_memory_fraction=0.4 if host_cache_size_gb else 0.7,
-            host_cache_size=host_cache_size_gb * (1 << 30),
-            use_kv_cache_manager_v2=True,
-        )
-        sparse_attention_config = DeepSeekSparseAttentionConfig(
-            indexer_k_dtype=indexer_k_dtype,
-            skip_indexer_for_short_seqs=False,
-            use_cute_dsl_paged_mqa_logits=use_cute_dsl,
-        )
-        speculative_config = (MTPDecodingConfig(
-            max_draft_len=mtp_nextn) if mtp_nextn else None)
-        llm_kwargs = {}
-        if chunked_prefill:
-            llm_kwargs.update(enable_chunked_prefill=True, max_num_tokens=512)
-
-        with LLM(
-                f"{llm_models_root()}/DeepSeek-V3.2-Exp-FP4-v2",
-                max_batch_size=max_batch_size,
-                tensor_parallel_size=4,
-                pipeline_parallel_size=1,
-                moe_expert_parallel_size=4,
-                kv_cache_config=kv_cache_config,
-                disable_overlap_scheduler=False,
-                cuda_graph_config=CudaGraphConfig(
-                    enable_padding=True, max_batch_size=max_batch_size),
-                moe_config=MoeConfig(backend="CUTLASS", max_num_tokens=16384),
-                enable_attention_dp=True,
-                speculative_config=speculative_config,
-                sparse_attention_config=sparse_attention_config,
-                **llm_kwargs,
-        ) as llm:
-            task = MMLU(self.MODEL_NAME)
-            task.evaluate(llm)
-
     @pytest.mark.skip_less_mpi_world_size(8)
     @skip_pre_blackwell
     @pytest.mark.parametrize(
@@ -8015,7 +7959,8 @@ def test_nvfp4(self, tp_size, ep_size):
         # precision.
         model_name = "zai-org/GLM-5.2"
         model_path = f"{llm_models_root()}/GLM-5.2-NVFP4"
-        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7)
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7,
+                                        use_kv_cache_manager_v2=True)
 
         pytorch_config = dict(
             disable_overlap_scheduler=False,
@@ -8031,6 +7976,7 @@ def test_nvfp4(self, tp_size, ep_size):
                  pipeline_parallel_size=1,
                  moe_expert_parallel_size=ep_size,
                  kv_cache_config=kv_cache_config,
+                 max_batch_size=128,
                  max_seq_len=8192,
                  **pytorch_config) as llm:
             assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index efce87353b52..765ad10586f3 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -175,10 +175,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus_corner_c
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp]
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[baseline]
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_mtp1_chunked]
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_dsl_mtp3]
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[host_cache_offload]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline_fp8kv]
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_fp8_blockscale[baseline_mtp1]
diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 7153ad6bd542..74d7936cab15 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -179,10 +179,6 @@ l0_dgx_b200:
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp1] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_dsa_host_cache_offload[host_cache_offload_mtp3_no_adp] TIMEOUT (60)
-  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[baseline] TIMEOUT (60)
-  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_mtp1_chunked] TIMEOUT (60)
-  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[fp4_dsl_mtp3] TIMEOUT (60)
-  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_kv_cache_manager_v2[host_cache_offload] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4_mtp_index_share[tp_size=8-ep_size=8] TIMEOUT (60)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Pro::test_gsm8k_full_accuracy TIMEOUT (240)
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index 2349765a09a4..e99587062f8e 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -57,6 +57,7 @@
 from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
 from tensorrt_llm._torch.speculative.interface import (
+    SpecWorkerBase,
     prepare_attn_metadata_for_draft_replay,
     restore_attn_metadata_after_draft_replay,
 )
@@ -3620,6 +3621,80 @@ def test_prepare_swaps_and_restore_recovers(self):
         torch.testing.assert_close(meta.kv_cache_block_offsets, original_offsets)
         torch.testing.assert_close(meta.host_kv_cache_block_offsets, original_host_offsets)
 
+    def test_draft_context_recomputes_and_restores_dsa_metadata(self):
+        """The eager draft context must switch DSA indexer cache metadata."""
+        meta = self._make_mock_metadata()
+        mgr = self._make_mock_draft_manager()
+        mgr.index_head_dim = 128
+        meta.enable_flash_mla = False
+        meta.host_indexer_k_cache_block_offsets = torch.tensor([1, 2, 3])
+        meta.indexer_k_cache_block_offsets = torch.tensor([4, 5, 6])
+        meta.host_slot_mapping_fp8 = torch.tensor([7, 8, 9])
+        meta.host_slot_mapping_scale = torch.tensor([10, 11, 12])
+        meta.slot_mapping_fp8 = torch.tensor([13, 14, 15])
+        meta.slot_mapping_scale = torch.tensor([16, 17, 18])
+        meta.slot_mapping_fp8_fullkv = torch.tensor([19, 20, 21])
+        meta.slot_mapping_scale_fullkv = torch.tensor([22, 23, 24])
+        meta._update_indexer_k_cache_block_offsets = Mock()
+        original_manager = meta.kv_cache_manager
+        original_host_indexer_offsets = meta.host_indexer_k_cache_block_offsets
+        original_host_slot_mapping_fp8 = meta.host_slot_mapping_fp8
+        original_host_slot_mapping_scale = meta.host_slot_mapping_scale
+        original_slot_mapping_fp8_fullkv = meta.slot_mapping_fp8_fullkv
+        original_slot_mapping_scale_fullkv = meta.slot_mapping_scale_fullkv
+
+        def is_attn_metadata(obj, cls):
+            if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata):
+                return obj is meta
+            return builtins.isinstance(obj, cls)
+
+        with (
+            patch(
+                "tensorrt_llm._torch.speculative.interface.isinstance",
+                side_effect=is_attn_metadata,
+            ),
+            patch.object(Indexer, "recompute_slot_mappings") as recompute,
+            patch.object(
+                Indexer,
+                "recompute_context_kv_gather_mappings",
+            ) as recompute_context_gather,
+        ):
+            with SpecWorkerBase.draft_kv_cache_context(Mock(), meta, mgr):
+                assert meta.kv_cache_manager is mgr
+                assert meta.host_indexer_k_cache_block_offsets is not original_host_indexer_offsets
+                assert meta.host_slot_mapping_fp8 is not original_host_slot_mapping_fp8
+                assert meta.host_slot_mapping_scale is not original_host_slot_mapping_scale
+                meta._update_indexer_k_cache_block_offsets.assert_called_once_with()
+                recompute.assert_called_once_with(meta)
+                recompute_context_gather.assert_called_once_with(meta)
+                meta.host_indexer_k_cache_block_offsets.fill_(99)
+                meta.slot_mapping_fp8_fullkv = torch.tensor([100, 200, 300])
+                meta.slot_mapping_scale_fullkv = torch.tensor([400, 500, 600])
+
+        assert meta.kv_cache_manager is original_manager
+        assert meta.host_indexer_k_cache_block_offsets is original_host_indexer_offsets
+        assert meta.host_slot_mapping_fp8 is original_host_slot_mapping_fp8
+        assert meta.host_slot_mapping_scale is original_host_slot_mapping_scale
+        assert meta.slot_mapping_fp8_fullkv is original_slot_mapping_fp8_fullkv
+        assert meta.slot_mapping_scale_fullkv is original_slot_mapping_scale_fullkv
+
+
+def test_recompute_context_kv_gather_mappings_skips_generation():
+    """Generation graphs must not allocate unused full-KV mapping tensors."""
+    fp8_mapping = torch.tensor([1, 2, 3])
+    scale_mapping = torch.tensor([4, 5, 6])
+    metadata = SimpleNamespace(
+        enable_context_mla_with_cached_kv=True,
+        slot_mapping_fp8_fullkv=fp8_mapping,
+        slot_mapping_scale_fullkv=scale_mapping,
+    )
+    indexer_params = SimpleNamespace(num_contexts=0)
+
+    Indexer.recompute_context_kv_gather_mappings(metadata, indexer_params)
+
+    assert metadata.slot_mapping_fp8_fullkv is fp8_mapping
+    assert metadata.slot_mapping_scale_fullkv is scale_mapping
+
 
 @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
 @skip_pre_blackwell

From b89cabab11081954dbc51b9612867de6e08dcd37 Mon Sep 17 00:00:00 2001
From: Eric Tsai <ertsai@nvidia.com>
Date: Sun, 26 Jul 2026 18:42:51 -0700
Subject: [PATCH 04/12] [None][fix] pre-allocate DSA draft indexer buffers for
 CUDA graph replay

Signed-off-by: Eric Tsai <ertsai@nvidia.com>
---
 tensorrt_llm/_torch/speculative/interface.py | 51 +++++++++++---------
 1 file changed, 27 insertions(+), 24 deletions(-)

diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 1edc11d149da..a4660030a60a 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -164,32 +164,35 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
         saved['saved_dsa_state'] = {
             'host_indexer_k_cache_block_offsets':
             m.host_indexer_k_cache_block_offsets,
-            'indexer_k_cache_block_offsets':
-            m.indexer_k_cache_block_offsets.clone(),
+            'indexer_k_cache_block_offsets': m.indexer_k_cache_block_offsets,
             'host_slot_mapping_fp8': m.host_slot_mapping_fp8,
             'host_slot_mapping_scale': m.host_slot_mapping_scale,
-            'slot_mapping_fp8': m.slot_mapping_fp8.clone(),
-            'slot_mapping_scale': m.slot_mapping_scale.clone(),
+            'slot_mapping_fp8': m.slot_mapping_fp8,
+            'slot_mapping_scale': m.slot_mapping_scale,
             'slot_mapping_fp8_fullkv': m.slot_mapping_fp8_fullkv,
             'slot_mapping_scale_fullkv': m.slot_mapping_scale_fullkv,
         }
-        # Keep the target's pinned host buffers untouched while asynchronous
-        # H2D copies from the draft mappings are in flight.  Restoring their
-        # contents in place can race those copies because CPU writes are not
-        # ordered with the CUDA stream.  The fresh host buffers stay attached
-        # to the metadata until after the draft forward has been enqueued.
-        m.host_indexer_k_cache_block_offsets = torch.empty_like(
-            m.host_indexer_k_cache_block_offsets,
-            pin_memory=m.host_indexer_k_cache_block_offsets.is_pinned())
-        m.host_slot_mapping_fp8 = torch.empty_like(
-            m.host_slot_mapping_fp8,
-            pin_memory=m.host_slot_mapping_fp8.is_pinned())
-        m.host_slot_mapping_scale = torch.empty_like(
-            m.host_slot_mapping_scale,
-            pin_memory=m.host_slot_mapping_scale.is_pinned())
+        # Rebind to the draft manager's dedicated buffers instead of
+        # overwriting the target tensors in place. Rebinding is invisible to
+        # CUDA graph capture, so the target and draft segments of the graph
+        # bake distinct addresses (like draft_kv_cache_block_offsets) and no
+        # graph-recorded copy from a transient host buffer is needed.
+        m.host_indexer_k_cache_block_offsets = (
+            m.host_draft_indexer_k_cache_block_offsets)
+        m.indexer_k_cache_block_offsets = m.draft_indexer_k_cache_block_offsets
+        m.host_slot_mapping_fp8 = m.host_draft_slot_mapping_fp8
+        m.slot_mapping_fp8 = m.draft_slot_mapping_fp8
+        m.host_slot_mapping_scale = m.host_draft_slot_mapping_scale
+        m.slot_mapping_scale = m.draft_slot_mapping_scale
+        # Recording a capture executes no kernels, so the draft mappings only
+        # need refreshing when the transfers actually run: eager forwards
+        # (warmup) and the pre-replay call from model_engine. The per-step
+        # advance inside the captured graph re-derives slot mappings on
+        # device from the rebound block-offset buffer.
         # kv_cache_manager was already swapped to the draft manager above.
-        m._update_indexer_k_cache_block_offsets()
-        Indexer.recompute_slot_mappings(m)
+        if not torch.cuda.is_current_stream_capturing():
+            m._update_indexer_k_cache_block_offsets()
+            Indexer.recompute_slot_mappings(m)
         Indexer.recompute_context_kv_gather_mappings(m)
     return saved
 
@@ -210,12 +213,12 @@ def restore_attn_metadata_after_draft_replay(attn_metadata, saved_state):
         m = attn_metadata
         m.host_indexer_k_cache_block_offsets = saved_dsa[
             'host_indexer_k_cache_block_offsets']
-        m.indexer_k_cache_block_offsets.copy_(
-            saved_dsa['indexer_k_cache_block_offsets'], non_blocking=True)
+        m.indexer_k_cache_block_offsets = saved_dsa[
+            'indexer_k_cache_block_offsets']
         m.host_slot_mapping_fp8 = saved_dsa['host_slot_mapping_fp8']
         m.host_slot_mapping_scale = saved_dsa['host_slot_mapping_scale']
-        m.slot_mapping_fp8.copy_(saved_dsa['slot_mapping_fp8'])
-        m.slot_mapping_scale.copy_(saved_dsa['slot_mapping_scale'])
+        m.slot_mapping_fp8 = saved_dsa['slot_mapping_fp8']
+        m.slot_mapping_scale = saved_dsa['slot_mapping_scale']
         m.slot_mapping_fp8_fullkv = saved_dsa['slot_mapping_fp8_fullkv']
         m.slot_mapping_scale_fullkv = saved_dsa['slot_mapping_scale_fullkv']
 

From aa9d946b8b97ef3c08867ed60aebf15e6f0f54d5 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Thu, 6 Aug 2026 22:17:13 -0700
Subject: [PATCH 05/12] [None][fix] harden DSA KV cache manager V2 as default

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../_torch/attention_backend/trtllm.py        |  43 ++-
 .../_torch/models/modeling_deepseekv3.py      |   5 +
 tensorrt_llm/_torch/pyexecutor/_util.py       |   6 +-
 .../_torch/pyexecutor/kv_cache_manager_v2.py  |  29 +-
 .../_torch/pyexecutor/model_engine.py         |   3 +-
 tensorrt_llm/_torch/speculative/interface.py  |  51 ++-
 .../attention/sparse/dsa/test_dsa_indexer.py  | 355 +++++++++---------
 .../test_kv_cache_v2_extra_buffers.py         |  61 ++-
 .../executor/test_mamba_cache_manager.py      |  31 +-
 .../_torch/speculative/test_eagle3.py         |  47 +++
 tests/unittest/llmapi/test_config_database.py |  27 +-
 tests/unittest/llmapi/test_llm_args.py        |  37 ++
 12 files changed, 492 insertions(+), 203 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index dd756d561849..d421e25ff5d6 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -158,11 +158,12 @@ def effective_beam_width(self) -> int:
     kv_cache_block_offsets: Optional[torch.Tensor] = None
     host_kv_cache_block_offsets: Optional[torch.Tensor] = None
     draft_kv_cache_block_offsets: Optional[torch.Tensor] = None
-    # Block IDs per sequence; populated in __post_init__ when a KV cache
-    # manager is present. Declared here so encoder-only metadata (no KV cache)
-    # still exposes the attribute.
+    # Active block-ID buffers, defaulting to the target cache. Separate draft
+    # storage lets CUDA graphs bind stable target and draft addresses.
     block_ids_per_seq: Optional[torch.Tensor] = None
     kv_block_ids_per_seq: Optional[torch.Tensor] = None
+    draft_block_ids_per_seq: Optional[torch.Tensor] = None
+    draft_kv_block_ids_per_seq: Optional[torch.Tensor] = None
 
     # Pre-computed FlashMLA tile-scheduler metadata and num_splits.
     # Computed once per forward pass in TrtllmAttention.forward() and reused across layers.
@@ -335,6 +336,8 @@ def _post_init_with_buffers(self, buffers) -> None:
             self.host_kv_cache_block_offsets = self.kv_cache_manager.host_kv_cache_block_offsets
             self.block_ids_per_seq = None
             self.kv_block_ids_per_seq = None
+            self.draft_block_ids_per_seq = None
+            self.draft_kv_block_ids_per_seq = None
 
             # Allocate separate block offset tensors for draft KV cache manager
             # Used in one-model speculative decoding with different KV cache layouts
@@ -374,6 +377,27 @@ def _post_init_with_buffers(self, buffers) -> None:
                     dtype=torch.int32,
                     capture_graph=capture_graph,
                 )
+                if self.draft_kv_cache_manager is not None:
+                    self.draft_block_ids_per_seq = self.get_empty(
+                        buffers,
+                        [
+                            self.draft_kv_cache_manager.max_batch_size,
+                            self.draft_kv_cache_manager.max_blocks_per_seq
+                        ],
+                        cache_name="draft_block_ids_per_seq",
+                        dtype=torch.int32,
+                        capture_graph=capture_graph,
+                    )
+                    self.draft_kv_block_ids_per_seq = self.get_empty(
+                        buffers,
+                        [
+                            self.draft_kv_cache_manager.max_batch_size,
+                            self.draft_kv_cache_manager.max_blocks_per_seq
+                        ],
+                        cache_name="draft_kv_block_ids_per_seq",
+                        dtype=torch.int32,
+                        capture_graph=capture_graph,
+                    )
                 # Allocate fixed-size buffers for pre-computed FlashMLA metadata.
                 # These are pre-allocated so their GPU addresses are stable across CUDA graph captures.
                 sm_count = torch.cuda.get_device_properties(
@@ -753,9 +777,12 @@ def prepare_encoder_cuda_graph_replay(self, seq_lens: List[int],
         self.host_total_kv_lens[0] = padded_num_tokens
 
     def prepare_flash_mla(self) -> None:
-        # Invalidate the pre-computed metadata so that forward() recomputes it
-        # for this forward pass before the first attention layer runs.
         self._flash_mla_metadata_valid = False
+        # Request-specific fills and H2D copies must happen before replay, not
+        # become fixed operations in the captured graph.
+        if torch.cuda.is_current_stream_capturing():
+            return
+
         block_ids_per_seq = maybe_pin_memory(
             self.kv_cache_manager.get_block_ids_per_seq(self.request_ids))
         num_blocks = block_ids_per_seq.shape[1]
@@ -1697,9 +1724,9 @@ def forward(
             self, q, k, metadata, forward_args)
 
         # Compute FlashMLA tile-scheduler metadata once per forward pass.
-        # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger
-        # recomputation when cache_seq_lens change. The metadata must always match the
-        # compacted generation sub-batch, which is also the layout used by block_ids_per_seq.
+        # The flag is invalidated whenever FlashMLA inputs change. The metadata
+        # must always match the compacted generation sub-batch, which is also
+        # the layout used by block_ids_per_seq.
         if (metadata.enable_flash_mla and forward_args.attention_input_type
                 != AttentionInputType.context_only
                 and metadata.num_generations > 0
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index 3569426f5849..c1fb856b0021 100755
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -1907,6 +1907,11 @@ def forward(
 class DeepseekV3ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV3Model,
                                                         PretrainedConfig]):
 
+    @classmethod
+    def get_model_defaults(cls, llm_args) -> dict:
+        """Use KV cache manager V2 for models backed by this implementation."""
+        return {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+
     @classmethod
     def get_preferred_transceiver_runtime(
             cls,
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 2a42c0e9c9d3..a24a4bc73387 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -591,7 +591,7 @@ def _get_model_kv_cache_manager_cls(
             kv_cache_config,
             is_disagg=self._is_disagg,
             cache_transceiver_config=self._cache_transceiver_config)
-        cls = self._fallback_if_unsupported_kv_cache_manager_v2(
+        cls = self._validate_or_fallback_kv_cache_manager_v2(
             cls, model_config, kv_cache_config)
         if is_hybrid_linear(model_config.pretrained_config):
             logger.info_once(
@@ -609,7 +609,7 @@ def _get_model_kv_cache_manager_cls(
                     f"when using non-V2 Mamba cache manager {cls.__name__}")
         return cls
 
-    def _fallback_if_unsupported_kv_cache_manager_v2(
+    def _validate_or_fallback_kv_cache_manager_v2(
             self,
             kv_cache_manager_cls,
             model_config: ModelConfig,
@@ -1446,7 +1446,7 @@ def _create_one_model_draft_kv_cache_manager(
         # Get the appropriate KV cache manager class for the draft model
         draft_kv_cache_manager_cls = get_kv_cache_manager_cls(
             effective_draft_config, draft_kv_config, is_disagg=self._is_disagg)
-        draft_kv_cache_manager_cls = self._fallback_if_unsupported_kv_cache_manager_v2(
+        draft_kv_cache_manager_cls = self._validate_or_fallback_kv_cache_manager_v2(
             draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config)
 
         estimating_kv_cache = estimating_kv_cache and not self._skip_est
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index 2c567824a8cb..e55c8dddc5a7 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -1080,11 +1080,19 @@ def append_to_kv_heads_per_layer(
                 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
+        # config-order representative for each role in each logical group.
         self.num_pools = len(self.impl.layer_grouping)
-        # num_pools is the physical pool count owned by the KV cache manager.
-        # 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.
+        self._pool_layer_ids_by_role: Dict[Tuple[int, DataRole], LayerId] = {}
+        for layer in config.layers:
+            pool_id = int(self.impl.get_layer_group_id(layer.layer_id))
+            for buffer in layer.buffers:
+                self._pool_layer_ids_by_role.setdefault((pool_id, buffer.role), layer.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:
@@ -1192,6 +1200,15 @@ def _get_block_scale_role(self, role_a: DataRole) -> Optional[DataRole]:
             return None
         return Role.KEY_BLOCK_SCALE
 
+    def _get_pool_layer_id(self, pool_id: int, role: DataRole) -> LayerId:
+        """Return a stable config-order representative for a pool role."""
+        try:
+            return self._pool_layer_ids_by_role[(pool_id, role)]
+        except KeyError as error:
+            raise ValueError(
+                f"Logical layer group {pool_id} has no layer with role {role}"
+            ) from error
+
     def _build_pool_mapping_tensors(self):
         """Build the (kv_cache_pool_pointers, kv_cache_pool_mapping) tensors.
 
@@ -1228,8 +1245,8 @@ def _build_pool_mapping_tensors(self):
                 kv_cache_pool_mapping_list.append([int(layer_id), 0])
         else:
             for pool_id in range(self.num_pools):
-                layer_id = self.impl.layer_grouping[pool_id][0]
                 role_a, _ = self._get_pool_roles(pool_id)
+                layer_id = self._get_pool_layer_id(pool_id, role_a)
                 key_base_addr = self.impl.get_mem_pool_base_address(
                     layer_id, role_a, PageIndexMode.SHARED
                 )
@@ -1342,8 +1359,8 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
             self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu"
         )
         for pool_id in range(self.num_pools):
-            layer_id = self.impl.layer_grouping[pool_id][0]
             role_a, role_b = self._get_pool_roles(pool_id)
+            layer_id = self._get_pool_layer_id(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(
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index bd832625d042..e4cf3fc129ae 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -7189,7 +7189,8 @@ def capture_postprocess_fn(inputs: Dict[str, Any]):
                     if self.cuda_graph_runner.is_warmup_only:
                         outputs = capture_outputs
                     elif needs_capture:
-                        # Pre-replay: set DSA slot mappings for current batch's draft cache (fixes 2nd warmup)
+                        # 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:
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index a4660030a60a..a061e3b0821e 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -127,9 +127,11 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool:
 def prepare_attn_metadata_for_draft_replay(attn_metadata,
                                            draft_kv_cache_manager):
     """
-    Prepare attention metadata for CUDA graph replay when using separate draft KV cache.
-    Swaps to draft manager and (for DSA) re-prepares indexer slot mappings for the current
-    batch. Call restore_attn_metadata_after_draft_replay after replay in a finally block.
+    Prepare attention metadata for a draft forward or CUDA graph replay when using a
+    separate draft KV cache. Swaps cache-layout-dependent buffers, refreshes FlashMLA
+    block IDs outside capture, and (for DSA) re-prepares indexer slot mappings
+    for the current batch.
+    Call restore_attn_metadata_after_draft_replay in a finally block.
     Returns saved state or None if no-op.
     """
     if draft_kv_cache_manager is None:
@@ -149,6 +151,18 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
         'target_host_kv_cache_block_offsets':
         attn_metadata.host_kv_cache_block_offsets,
     }
+    if attn_metadata.enable_flash_mla:
+        if (attn_metadata.draft_block_ids_per_seq is None
+                or attn_metadata.draft_kv_block_ids_per_seq is None):
+            raise RuntimeError(
+                "FlashMLA separate draft KV cache requires dedicated draft block-ID buffers"
+            )
+        saved['target_block_ids_per_seq'] = attn_metadata.block_ids_per_seq
+        saved[
+            'target_kv_block_ids_per_seq'] = attn_metadata.kv_block_ids_per_seq
+        attn_metadata.block_ids_per_seq = attn_metadata.draft_block_ids_per_seq
+        attn_metadata.kv_block_ids_per_seq = (
+            attn_metadata.draft_kv_block_ids_per_seq)
     attn_metadata.kv_cache_manager = draft_kv_cache_manager
     attn_metadata.kv_cache_block_offsets = attn_metadata.draft_kv_cache_block_offsets
     attn_metadata.host_kv_cache_block_offsets = (
@@ -169,9 +183,17 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
             'host_slot_mapping_scale': m.host_slot_mapping_scale,
             'slot_mapping_fp8': m.slot_mapping_fp8,
             'slot_mapping_scale': m.slot_mapping_scale,
-            'slot_mapping_fp8_fullkv': m.slot_mapping_fp8_fullkv,
-            'slot_mapping_scale_fullkv': m.slot_mapping_scale_fullkv,
         }
+        # The cached-KV feature owns these references even when an optimized
+        # path aliases them to slot_mapping_*. With the feature disabled, the
+        # aliases are lazy and may not exist on the first generation replay.
+        if m.enable_context_mla_with_cached_kv:
+            saved['saved_dsa_state'].update({
+                'slot_mapping_fp8_fullkv':
+                m.slot_mapping_fp8_fullkv,
+                'slot_mapping_scale_fullkv':
+                m.slot_mapping_scale_fullkv,
+            })
         # Rebind to the draft manager's dedicated buffers instead of
         # overwriting the target tensors in place. Rebinding is invisible to
         # CUDA graph capture, so the target and draft segments of the graph
@@ -207,7 +229,14 @@ def restore_attn_metadata_after_draft_replay(attn_metadata, saved_state):
     attn_metadata.host_kv_cache_block_offsets = (
         saved_state['target_host_kv_cache_block_offsets'])
     if attn_metadata.enable_flash_mla:
-        attn_metadata.prepare_flash_mla()
+        attn_metadata.block_ids_per_seq = saved_state[
+            'target_block_ids_per_seq']
+        attn_metadata.kv_block_ids_per_seq = saved_state[
+            'target_kv_block_ids_per_seq']
+        # Target and draft block-ID buffers are independent. Restoring only
+        # needs to invalidate the scheduler metadata; refreshing the unchanged
+        # target buffers would repeat request-specific H2D work.
+        attn_metadata._flash_mla_metadata_valid = False
     saved_dsa = saved_state.get('saved_dsa_state')
     if saved_dsa is not None:
         m = attn_metadata
@@ -219,8 +248,14 @@ def restore_attn_metadata_after_draft_replay(attn_metadata, saved_state):
         m.host_slot_mapping_scale = saved_dsa['host_slot_mapping_scale']
         m.slot_mapping_fp8 = saved_dsa['slot_mapping_fp8']
         m.slot_mapping_scale = saved_dsa['slot_mapping_scale']
-        m.slot_mapping_fp8_fullkv = saved_dsa['slot_mapping_fp8_fullkv']
-        m.slot_mapping_scale_fullkv = saved_dsa['slot_mapping_scale_fullkv']
+        if 'slot_mapping_fp8_fullkv' in saved_dsa:
+            m.slot_mapping_fp8_fullkv = saved_dsa['slot_mapping_fp8_fullkv']
+            m.slot_mapping_scale_fullkv = saved_dsa['slot_mapping_scale_fullkv']
+        else:
+            # The draft recomputation rebound the aliases to the draft tensors;
+            # point them back at the restored target tensors.
+            m.slot_mapping_fp8_fullkv = m.slot_mapping_fp8
+            m.slot_mapping_scale_fullkv = m.slot_mapping_scale
 
 
 def get_force_num_accepted_tokens() -> int:
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index e99587062f8e..c3e35107db48 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -26,7 +26,7 @@
 import random
 import sys
 from pathlib import Path
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
 from unittest.mock import Mock, patch
 
 import pytest
@@ -44,7 +44,6 @@
     DSABackendForwardArgs,
     DSACacheManager,
     DSACacheManagerV2,
-    DSAMetadataParams,
     DSATrtllmAttention,
     DSAtrtllmAttentionMetadata,
     Indexer,
@@ -52,9 +51,10 @@
     _select_indexer_compress_ratio,
     compute_cu_seqlen_kv_bounds_with_cache,
     split_prefill_chunks,
+    transform_local_topk_and_prepare_pool_view,
 )
 from tensorrt_llm._torch.attention_backend.sparse.registry import get_sparse_attn_kv_cache_manager
-from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata
+from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
 from tensorrt_llm._torch.speculative.interface import (
     SpecWorkerBase,
@@ -258,6 +258,7 @@ def create_dsa_cache_manager(
     indexer_k_dtype: str = "fp8",
     index_topk: int = 2048,
     use_kv_cache_manager_v2: bool = False,
+    pretrained_config=None,
 ):
     """Helper to create a DSA cache manager for testing."""
 
@@ -293,6 +294,7 @@ def create_dsa_cache_manager(
         mapping=mapping,
         dtype=DataType.HALF,
         sparse_attention_config=sparse_attn_config,
+        pretrained_config=pretrained_config,
     )
 
     return cache_manager, sparse_attn_config
@@ -311,160 +313,126 @@ def test_dsa_cache_manager_selection_honors_v2_config():
     )
 
 
-def test_dsa_metadata_accepts_kv_cache_manager_v2():
-    cache_manager = object.__new__(DSACacheManagerV2)
-    cache_manager.tokens_per_block = 32
-    metadata = object.__new__(DSAtrtllmAttentionMetadata)
-    metadata.kv_cache_manager = cache_manager
-    metadata.sparse_metadata_params = DSAMetadataParams(
-        indexer_max_chunk_size=8192,
-        max_sparse_topk=2048,
-        index_head_dim=128,
-        enable_indexer_skip=False,
-        enable_heuristic_topk=False,
-        use_cute_dsl_paged_mqa_logits=False,
-        q_split_threshold=8192,
-    )
-    metadata.is_cuda_graph = False
-
-    with (
-        patch.object(TrtllmAttentionMetadata, "__post_init__"),
-        patch.object(DSAtrtllmAttentionMetadata, "create_buffers_for_mla_rope_append"),
-        patch.object(DSAtrtllmAttentionMetadata, "create_buffers_for_indexer"),
-    ):
-        metadata.__post_init__()
-
-    assert metadata._tokens_per_block == cache_manager.tokens_per_block
-
-
 @pytest.mark.parametrize(
-    "indexer_k_dtype,indexer_bytes_per_token",
+    "cache_manager_cls",
     [
-        ("fp8", 132),
-        pytest.param("fp4", 68, marks=skip_pre_blackwell),
+        pytest.param(DSACacheManager, id="v1"),
+        pytest.param(DSACacheManagerV2, id="v2"),
     ],
 )
-def test_dsa_cache_manager_v2_cache_size_estimate_includes_indexer(
-    indexer_k_dtype, indexer_bytes_per_token
-):
-    """Warmup estimation includes main MLA K and indexer K for every layer."""
-    num_layers = 3
-    head_dim = 128
-    model_config = SimpleNamespace(
-        pretrained_config=SimpleNamespace(
-            num_attention_heads=1,
-            num_key_value_heads=1,
-            hidden_size=head_dim,
-            kv_lora_rank=96,
-            qk_rope_head_dim=32,
-        ),
-        quant_config=None,
-        sparse_attention_config=DeepSeekSparseAttentionConfig(
-            index_head_dim=head_dim,
-            indexer_k_dtype=indexer_k_dtype,
-        ),
-        get_num_attention_layers=lambda: num_layers,
+def test_transform_local_topk_uses_page_mapping_for_pp_local_layer(cache_manager_cls):
+    """Flatten sparse indices with the active manager's PP-local page mapping."""
+    num_local_layers = 3
+    global_layer_idx = 5
+    local_layer_idx = 1
+    tokens_per_block = 4
+    page_index_scale = 3
+    layer_offset = 1
+
+    if cache_manager_cls is DSACacheManager:
+        cache_manager = SimpleNamespace(num_local_layers=num_local_layers)
+    else:
+        cache_manager = SimpleNamespace(_primary_pool_page_index_params=[(3, 0), (3, 1), (3, 2)])
+    cache_manager.layer_offsets = {global_layer_idx: local_layer_idx}
+    get_page_index_params = MethodType(
+        cache_manager_cls.get_primary_pool_page_index_params,
+        cache_manager,
     )
-    mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1)
+    cache_manager.get_primary_pool_page_index_params = Mock(wraps=get_page_index_params)
 
-    estimated_bytes_per_token = DSACacheManagerV2.get_cache_size_per_token(
-        model_config,
-        mapping,
+    req_idx = torch.tensor([0], dtype=torch.int32)
+    block_table = torch.tensor([[3, 7]], dtype=torch.int32)
+    topk_indices = torch.tensor([[6]], dtype=torch.int32)
+    pool_view = torch.empty(0)
+    metadata = SimpleNamespace(
+        kv_cache_manager=cache_manager,
+        _ensure_pool_view_cached=Mock(),
+        _cached_block_table_ctx=block_table,
+        _cached_req_idx_ctx=req_idx,
+        _cached_tokens_per_block=tokens_per_block,
+        _cached_pool_view=pool_view,
     )
+    attention = SimpleNamespace(layer_idx=global_layer_idx, local_layer_idx=None)
+    pp_local_layer_idx = TrtllmAttention.get_local_layer_idx(attention, metadata)
+    expected_flat_index = (
+        7 * (page_index_scale * tokens_per_block) + (layer_offset * tokens_per_block) + 2
+    )
+    expected_indices = torch.tensor([[expected_flat_index]], dtype=torch.int32)
 
-    mla_k_bytes_per_layer = head_dim * 2
-    assert estimated_bytes_per_token == num_layers * (
-        mla_k_bytes_per_layer + indexer_bytes_per_token
+    with patch(
+        "torch.ops.trtllm.convert_req_index_to_global",
+        return_value=expected_indices,
+    ) as convert_req_index_to_global:
+        global_indices, returned_pool_view = transform_local_topk_and_prepare_pool_view(
+            topk_indices,
+            metadata,
+            layer_idx=pp_local_layer_idx,
+        )
+
+    assert pp_local_layer_idx == local_layer_idx
+    cache_manager.get_primary_pool_page_index_params.assert_called_once_with(local_layer_idx)
+    convert_req_index_to_global.assert_called_once_with(
+        req_idx,
+        block_table,
+        topk_indices,
+        tokens_per_block,
+        1,
+        page_index_scale * tokens_per_block,
+        layer_offset,
     )
+    assert torch.equal(global_indices, expected_indices)
+    assert returned_pool_view is pool_view
+    metadata._ensure_pool_view_cached.assert_called_once_with()
 
 
-@pytest.mark.parametrize(
-    "indexer_k_dtype",
-    ["fp8", pytest.param("fp4", marks=skip_pre_blackwell)],
-)
-def test_dsa_cache_manager_v2_indexer_pool_layout(indexer_k_dtype):
-    """V2 exposes shared-page indexer buffers and role-specific page metadata."""
+def test_dsa_cache_manager_v2_respects_shared_indexer_layer_mask():
+    """V2 registers INDEX_KEY storage only for layers that own an indexer."""
     num_layers = 3
-    tokens_per_block = 16
     head_dim = 128
+    tokens_per_block = 16
+    pretrained_config = SimpleNamespace(
+        num_hidden_layers=num_layers,
+        index_topk_pattern=["F", "S", "F"],
+    )
     cache_manager, _ = create_dsa_cache_manager(
         batch_size=2,
         head_dim=head_dim,
         tokens_per_block=tokens_per_block,
         max_seq_len=64,
         num_layers=num_layers,
-        indexer_k_dtype=indexer_k_dtype,
         use_kv_cache_manager_v2=True,
+        pretrained_config=pretrained_config,
     )
 
     try:
-        assert isinstance(cache_manager, DSACacheManagerV2)
-        page_index_scale = cache_manager.indexer_k_cache_page_scale
-        assert page_index_scale == num_layers
-        payload_bytes = head_dim // 2 if indexer_k_dtype == "fp4" else head_dim
-        indexer_bytes_per_token = payload_bytes + 4
-        key_bytes_per_token = cache_manager.get_layer_bytes_per_token(0, Role.KEY)
-        all_bytes_per_token = cache_manager.get_layer_bytes_per_token(0, Role.ALL)
-        assert cache_manager.get_layer_bytes_per_token(0, Role.INDEX_KEY) == indexer_bytes_per_token
-        assert all_bytes_per_token == key_bytes_per_token + indexer_bytes_per_token
-        assert cache_manager.get_cache_bytes_per_token() == num_layers * all_bytes_per_token
-
-        layer_sizes, attention_windows = cache_manager._get_runtime_cache_size_layer_components()
-        assert layer_sizes == [all_bytes_per_token] * num_layers
-        assert attention_windows == [None] * num_layers
-        assert cache_manager._get_quota_from_max_tokens(64) == (
-            64 * num_layers * all_bytes_per_token
-        )
-
-        for layer_config in cache_manager.kv_cache_manager_py_config.layers:
-            indexer_buffer_config = next(
-                buffer for buffer in layer_config.buffers if buffer.role == Role.INDEX_KEY
-            )
-            assert indexer_buffer_config.size == (indexer_bytes_per_token * tokens_per_block)
-
-        indexer_buffer = cache_manager.get_indexer_k_cache_buffers(0)
-        assert indexer_buffer.shape == (
-            cache_manager.blocks_in_primary_pool * page_index_scale,
+        assert cache_manager.indexer_k_cache_local_layer_mask == [True, False, True]
+        assert cache_manager.indexer_k_cache_page_scale == 2
+        assert cache_manager.indexer_k_cache_pool_per_layer[1] is None
+        assert cache_manager.get_indexer_k_cache_buffers(0).shape == (
+            cache_manager.blocks_in_primary_pool * cache_manager.indexer_k_cache_page_scale,
             tokens_per_block,
             1,
-            payload_bytes + 4,
+            head_dim + 4,
         )
 
-        primary_pool = cache_manager.get_unique_primary_pool()
-        assert primary_pool.shape[:3] == (
-            cache_manager.blocks_in_primary_pool,
-            num_layers,
-            1,
-        )
+        indexer_bytes_per_token = head_dim + 4
+        expected_cache_bytes = 0
+        for local_layer_idx, owns_indexer in enumerate([True, False, True]):
+            layer_config = cache_manager.kv_cache_manager_py_config.layers[local_layer_idx]
+            has_indexer_buffer = any(
+                buffer.role == Role.INDEX_KEY for buffer in layer_config.buffers
+            )
+            assert has_indexer_buffer is owns_indexer
 
-        request_ids = [11, 22]
-        cache_manager.add_dummy_requests(
-            request_ids,
-            token_nums=[16, 32],
-            prepare_resource=True,
-        )
-        pool_indices = cache_manager.get_pool_block_indices(
-            len(request_ids),
-            request_ids=request_ids,
-            num_contexts=len(request_ids),
-        )
-        reversed_pool_indices = cache_manager.get_pool_block_indices(
-            len(request_ids),
-            request_ids=list(reversed(request_ids)),
-            num_contexts=len(request_ids),
-        )
-        indexer_page_indices = pool_indices * page_index_scale
-        reversed_indexer_page_indices = reversed_pool_indices * page_index_scale
-        assert pool_indices.shape[0] == len(request_ids)
-        assert pool_indices.min() >= 0
-        assert pool_indices.max() < cache_manager.blocks_in_primary_pool
-        assert torch.equal(reversed_pool_indices[0], pool_indices[1])
-        assert torch.equal(reversed_pool_indices[1], pool_indices[0])
-        assert torch.equal(indexer_page_indices, pool_indices * page_index_scale)
-        assert indexer_page_indices.min() >= 0
-        assert indexer_page_indices.max() < indexer_buffer.shape[0]
-        assert torch.equal(reversed_indexer_page_indices[0], indexer_page_indices[1])
-        assert torch.equal(reversed_indexer_page_indices[1], indexer_page_indices[0])
+            key_bytes = cache_manager.get_layer_bytes_per_token(local_layer_idx, Role.KEY)
+            all_bytes = cache_manager.get_layer_bytes_per_token(local_layer_idx, Role.ALL)
+            expected_all_bytes = key_bytes + indexer_bytes_per_token if owns_indexer else key_bytes
+            assert all_bytes == expected_all_bytes
+            expected_cache_bytes += expected_all_bytes
+
+        assert cache_manager.get_cache_bytes_per_token() == expected_cache_bytes
+        with pytest.raises(AssertionError, match="shared-indexer layer"):
+            cache_manager.get_indexer_k_cache_buffers(1)
     finally:
         cache_manager.shutdown()
 
@@ -3575,6 +3543,7 @@ class TestPrepareRestoreAttnMetadataForDraftReplay:
     def _make_mock_metadata():
         """Create a mock attention metadata object with KV cache block offsets."""
         meta = Mock()
+        meta.enable_flash_mla = False
         meta.kv_cache_manager = Mock(name="target_kv_cache_manager")
         meta.kv_cache_block_offsets = torch.tensor([10, 20, 30])
         meta.host_kv_cache_block_offsets = torch.tensor([10, 20, 30])
@@ -3622,26 +3591,44 @@ def test_prepare_swaps_and_restore_recovers(self):
         torch.testing.assert_close(meta.host_kv_cache_block_offsets, original_host_offsets)
 
     def test_draft_context_recomputes_and_restores_dsa_metadata(self):
-        """The eager draft context must switch DSA indexer cache metadata."""
+        """Draft context rebinds preallocated DSA buffers and restores lazy aliases."""
         meta = self._make_mock_metadata()
         mgr = self._make_mock_draft_manager()
         mgr.index_head_dim = 128
         meta.enable_flash_mla = False
+        meta.enable_context_mla_with_cached_kv = False
         meta.host_indexer_k_cache_block_offsets = torch.tensor([1, 2, 3])
         meta.indexer_k_cache_block_offsets = torch.tensor([4, 5, 6])
         meta.host_slot_mapping_fp8 = torch.tensor([7, 8, 9])
         meta.host_slot_mapping_scale = torch.tensor([10, 11, 12])
         meta.slot_mapping_fp8 = torch.tensor([13, 14, 15])
         meta.slot_mapping_scale = torch.tensor([16, 17, 18])
-        meta.slot_mapping_fp8_fullkv = torch.tensor([19, 20, 21])
-        meta.slot_mapping_scale_fullkv = torch.tensor([22, 23, 24])
+        meta.host_draft_indexer_k_cache_block_offsets = torch.tensor([25, 26, 27])
+        meta.draft_indexer_k_cache_block_offsets = torch.tensor([28, 29, 30])
+        meta.host_draft_slot_mapping_fp8 = torch.tensor([31, 32, 33])
+        meta.host_draft_slot_mapping_scale = torch.tensor([34, 35, 36])
+        meta.draft_slot_mapping_fp8 = torch.tensor([37, 38, 39])
+        meta.draft_slot_mapping_scale = torch.tensor([40, 41, 42])
         meta._update_indexer_k_cache_block_offsets = Mock()
-        original_manager = meta.kv_cache_manager
-        original_host_indexer_offsets = meta.host_indexer_k_cache_block_offsets
-        original_host_slot_mapping_fp8 = meta.host_slot_mapping_fp8
-        original_host_slot_mapping_scale = meta.host_slot_mapping_scale
-        original_slot_mapping_fp8_fullkv = meta.slot_mapping_fp8_fullkv
-        original_slot_mapping_scale_fullkv = meta.slot_mapping_scale_fullkv
+        del meta.slot_mapping_fp8_fullkv
+        del meta.slot_mapping_scale_fullkv
+
+        target_buffers = {
+            "host_indexer_k_cache_block_offsets": meta.host_indexer_k_cache_block_offsets,
+            "indexer_k_cache_block_offsets": meta.indexer_k_cache_block_offsets,
+            "host_slot_mapping_fp8": meta.host_slot_mapping_fp8,
+            "host_slot_mapping_scale": meta.host_slot_mapping_scale,
+            "slot_mapping_fp8": meta.slot_mapping_fp8,
+            "slot_mapping_scale": meta.slot_mapping_scale,
+        }
+        draft_buffers = {
+            "host_indexer_k_cache_block_offsets": meta.host_draft_indexer_k_cache_block_offsets,
+            "indexer_k_cache_block_offsets": meta.draft_indexer_k_cache_block_offsets,
+            "host_slot_mapping_fp8": meta.host_draft_slot_mapping_fp8,
+            "host_slot_mapping_scale": meta.host_draft_slot_mapping_scale,
+            "slot_mapping_fp8": meta.draft_slot_mapping_fp8,
+            "slot_mapping_scale": meta.draft_slot_mapping_scale,
+        }
 
         def is_attn_metadata(obj, cls):
             if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata):
@@ -3653,47 +3640,79 @@ def is_attn_metadata(obj, cls):
                 "tensorrt_llm._torch.speculative.interface.isinstance",
                 side_effect=is_attn_metadata,
             ),
+            patch(
+                "tensorrt_llm._torch.speculative.interface.torch.cuda.is_current_stream_capturing",
+                return_value=False,
+            ),
             patch.object(Indexer, "recompute_slot_mappings") as recompute,
-            patch.object(
-                Indexer,
-                "recompute_context_kv_gather_mappings",
-            ) as recompute_context_gather,
         ):
             with SpecWorkerBase.draft_kv_cache_context(Mock(), meta, mgr):
-                assert meta.kv_cache_manager is mgr
-                assert meta.host_indexer_k_cache_block_offsets is not original_host_indexer_offsets
-                assert meta.host_slot_mapping_fp8 is not original_host_slot_mapping_fp8
-                assert meta.host_slot_mapping_scale is not original_host_slot_mapping_scale
-                meta._update_indexer_k_cache_block_offsets.assert_called_once_with()
-                recompute.assert_called_once_with(meta)
-                recompute_context_gather.assert_called_once_with(meta)
-                meta.host_indexer_k_cache_block_offsets.fill_(99)
-                meta.slot_mapping_fp8_fullkv = torch.tensor([100, 200, 300])
-                meta.slot_mapping_scale_fullkv = torch.tensor([400, 500, 600])
-
-        assert meta.kv_cache_manager is original_manager
-        assert meta.host_indexer_k_cache_block_offsets is original_host_indexer_offsets
-        assert meta.host_slot_mapping_fp8 is original_host_slot_mapping_fp8
-        assert meta.host_slot_mapping_scale is original_host_slot_mapping_scale
-        assert meta.slot_mapping_fp8_fullkv is original_slot_mapping_fp8_fullkv
-        assert meta.slot_mapping_scale_fullkv is original_slot_mapping_scale_fullkv
-
-
-def test_recompute_context_kv_gather_mappings_skips_generation():
-    """Generation graphs must not allocate unused full-KV mapping tensors."""
+                for name, buffer in draft_buffers.items():
+                    assert getattr(meta, name) is buffer
+                assert meta.slot_mapping_fp8_fullkv is meta.draft_slot_mapping_fp8
+                assert meta.slot_mapping_scale_fullkv is meta.draft_slot_mapping_scale
+
+        for name, buffer in target_buffers.items():
+            assert getattr(meta, name) is buffer
+        assert meta.slot_mapping_fp8_fullkv is meta.slot_mapping_fp8
+        assert meta.slot_mapping_scale_fullkv is meta.slot_mapping_scale
+        meta._update_indexer_k_cache_block_offsets.assert_called_once_with()
+        recompute.assert_called_once_with(meta)
+
+
+def test_recompute_context_kv_gather_mappings_does_not_alias_uncompressed_cached_mappings():
+    """Cached KV requires independent full-KV mappings even without compression."""
     fp8_mapping = torch.tensor([1, 2, 3])
     scale_mapping = torch.tensor([4, 5, 6])
+    fullkv_fp8_mapping = object()
+    fullkv_scale_mapping = object()
+    host_fullkv_fp8_mapping = Mock()
+    host_fullkv_scale_mapping = Mock()
+    host_fullkv_fp8_mapping.cuda.return_value = fullkv_fp8_mapping
+    host_fullkv_scale_mapping.cuda.return_value = fullkv_scale_mapping
+    computed_fp8_mapping = torch.tensor([10, 11, 12, 13])
+    computed_scale_mapping = torch.tensor([20, 21, 22, 23])
     metadata = SimpleNamespace(
         enable_context_mla_with_cached_kv=True,
+        skip_indexer_for_ctx_reqs=False,
+        slot_mapping_fp8=fp8_mapping,
+        slot_mapping_scale=scale_mapping,
         slot_mapping_fp8_fullkv=fp8_mapping,
         slot_mapping_scale_fullkv=scale_mapping,
+        host_indexer_k_cache_block_offsets=torch.tensor([[0], [1]], dtype=torch.int32),
+    )
+    indexer_params = SimpleNamespace(
+        num_contexts=2,
+        compress_ratio=1,
+        cached_kv_tokens=torch.tensor([0, 1], dtype=torch.int32),
+        kv_lens=torch.tensor([2, 2], dtype=torch.int32),
+        head_dim=128,
+        tokens_per_block=64,
+        quant_block_size=128,
+        data_bytes_per_token=128,
     )
-    indexer_params = SimpleNamespace(num_contexts=0)
-
-    Indexer.recompute_context_kv_gather_mappings(metadata, indexer_params)
 
-    assert metadata.slot_mapping_fp8_fullkv is fp8_mapping
-    assert metadata.slot_mapping_scale_fullkv is scale_mapping
+    with (
+        patch(
+            "tensorrt_llm._torch.attention_backend.sparse.dsa.torch.empty",
+            side_effect=[host_fullkv_fp8_mapping, host_fullkv_scale_mapping],
+        ),
+        patch(
+            "tensorrt_llm._torch.attention_backend.sparse.dsa._compute_slot_mappings",
+            return_value=(computed_fp8_mapping, computed_scale_mapping),
+        ) as compute_slot_mappings,
+    ):
+        Indexer.recompute_context_kv_gather_mappings(metadata, indexer_params)
+
+    assert metadata.slot_mapping_fp8_fullkv is fullkv_fp8_mapping
+    assert metadata.slot_mapping_scale_fullkv is fullkv_scale_mapping
+    assert metadata.slot_mapping_fp8_fullkv is not fp8_mapping
+    assert metadata.slot_mapping_scale_fullkv is not scale_mapping
+    host_fullkv_fp8_mapping.copy_.assert_called_once_with(computed_fp8_mapping)
+    host_fullkv_scale_mapping.copy_.assert_called_once_with(computed_scale_mapping)
+    host_fullkv_fp8_mapping.cuda.assert_called_once_with(non_blocking=True)
+    host_fullkv_scale_mapping.cuda.assert_called_once_with(non_blocking=True)
+    compute_slot_mappings.assert_called_once()
 
 
 @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
index 073190eacf71..97313b013321 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
@@ -12,6 +12,7 @@
 
 import gc
 import unittest
+from unittest.mock import Mock
 
 import torch
 
@@ -20,7 +21,7 @@
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig as KvCacheConfigV2
 from tensorrt_llm.mapping import Mapping
-from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig
+from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode
 
 DataType = tensorrt_llm.bindings.DataType
 CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType
@@ -188,6 +189,64 @@ def test_subclass_registers_index_key_only_on_sparse_layers(self):
             mgr.shutdown()
             del mgr
 
+    def test_pool_role_representative_uses_first_config_layer_with_role(self):
+        mgr = _IndexKeyOnSparseLayersV2(sparse_layer_ids=[2, 3], **_make_kwargs(num_layers=4))
+        try:
+            pool_id = int(mgr.impl.get_layer_group_id(2))
+            self.assertEqual(int(mgr.impl.get_layer_group_id(0)), pool_id)
+
+            layer_zero = mgr.kv_cache_manager_py_config.layers[0]
+            self.assertNotIn(Role.INDEX_KEY, [buffer.role for buffer in layer_zero.buffers])
+            self.assertEqual(
+                int(mgr._get_pool_layer_id(pool_id, Role.INDEX_KEY)),
+                2,
+                "the representative must be the first config-order layer that owns the role",
+            )
+        finally:
+            mgr.shutdown()
+            del mgr
+
+    def test_pool_mapping_ignores_reordered_layer_grouping(self):
+        mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4))
+        real_impl = mgr.impl
+        try:
+            self.assertEqual(mgr.num_pools, 1)
+            pool_id = 0
+            expected_layer_id = mgr._get_pool_layer_id(pool_id, Role.KEY)
+            config_group = tuple(
+                layer.layer_id
+                for layer in mgr.kv_cache_manager_py_config.layers
+                if int(real_impl.get_layer_group_id(layer.layer_id)) == pool_id
+            )
+            self.assertGreater(len(config_group), 1)
+            reordered_grouping = (config_group[1:] + config_group[:1],)
+            reordered_first_layer_id = reordered_grouping[pool_id][0]
+            self.assertNotEqual(expected_layer_id, reordered_first_layer_id)
+
+            expected_base_address = real_impl.get_mem_pool_base_address(
+                expected_layer_id, Role.KEY, PageIndexMode.SHARED
+            )
+
+            impl_proxy = Mock(wraps=real_impl)
+            impl_proxy.layer_grouping = reordered_grouping
+            mgr.impl = impl_proxy
+            mgr._prepare_page_table_tensor(index_mapper_capacity=1)
+
+            self.assertEqual(
+                impl_proxy.get_mem_pool_base_address.call_args_list[0].args,
+                (expected_layer_id, Role.KEY, PageIndexMode.SHARED),
+                "the pool pointer origin must use the config-order representative",
+            )
+            self.assertEqual(
+                int(mgr.kv_cache_pool_pointers[pool_id, 0]),
+                int(expected_base_address),
+            )
+            impl_proxy.get_page_index_scale.assert_called_once_with(expected_layer_id, Role.KEY)
+        finally:
+            mgr.impl = real_impl
+            mgr.shutdown()
+            del mgr
+
     def test_index_key_layer_has_non_zero_page_upper_bound(self):
         # Once a subclass registers INDEX_KEY, the V2 storage must wire a
         # paged pool for that role with a non-zero page-index upper bound,
diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
index 81db2977398d..fcaa753316c9 100644
--- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py
+++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
@@ -67,7 +67,7 @@
     _resolve_transceiver_runtime_auto,
     apply_model_defaults_to_llm_args,
 )
-from tensorrt_llm.mapping import Mapping
+from tensorrt_llm.mapping import CpType, Mapping
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     AttentionLayerConfig,
     BatchDesc,
@@ -841,11 +841,38 @@ def test_v2_hybrid_incompatibility_fails_without_cpp_fallback(
     creator._max_beam_width = max_beam_width
 
     with pytest.raises(NotImplementedError, match=expected):
-        creator._fallback_if_unsupported_kv_cache_manager_v2(
+        creator._validate_or_fallback_kv_cache_manager_v2(
             MambaHybridCacheManagerV2, model_config, KvCacheConfig()
         )
 
 
+def test_dsa_star_rejects_v2_and_keeps_explicit_v1_selection():
+    model_config = SimpleNamespace(
+        pretrained_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]),
+        sparse_attention_config=SimpleNamespace(algorithm="dsa"),
+    )
+    creator = object.__new__(KvCacheCreator)
+    creator._kv_connector_manager = None
+    creator._max_beam_width = 1
+    creator._mapping = SimpleNamespace(cp_config={"cp_type": CpType.STAR})
+
+    with pytest.raises(NotImplementedError, match="STAR context parallelism"):
+        creator._validate_or_fallback_kv_cache_manager_v2(
+            KVCacheManagerV2,
+            model_config,
+            KvCacheConfig(use_kv_cache_manager_v2=True),
+        )
+
+    assert (
+        creator._validate_or_fallback_kv_cache_manager_v2(
+            KVCacheManager,
+            model_config,
+            KvCacheConfig(use_kv_cache_manager_v2=False),
+        )
+        is KVCacheManager
+    )
+
+
 def _make_mgr(
     max_batch_size=4, max_draft_len=2, enable_attention_dp=False, use_replay_state_update=False
 ):
diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py
index b678377b2e6d..43eb0173654d 100644
--- a/tests/unittest/_torch/speculative/test_eagle3.py
+++ b/tests/unittest/_torch/speculative/test_eagle3.py
@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import json
 import os
 import sys
@@ -326,6 +341,38 @@ def _make_mock_kv_cache_manager(num_seqs: int) -> MagicMock:
     return mock_kv_cache_manager
 
 
+@skip_num_gpus_less_than(1)
+def test_separate_draft_flash_mla_buffers_use_manager_geometry():
+    """Target and draft FlashMLA block IDs need stable, disjoint storage."""
+    num_seqs = 3
+    target_kv_cache_manager = _make_mock_kv_cache_manager(num_seqs)
+    draft_kv_cache_manager = _make_mock_kv_cache_manager(num_seqs)
+    draft_kv_cache_manager.max_blocks_per_seq = 7
+    attn_metadata = TrtllmAttentionMetadata(
+        max_num_requests=num_seqs,
+        max_num_tokens=num_seqs * 3,
+        kv_cache_manager=target_kv_cache_manager,
+        draft_kv_cache_manager=draft_kv_cache_manager,
+        enable_flash_mla=True,
+    )
+
+    for metadata in (
+            attn_metadata,
+            attn_metadata.create_cuda_graph_metadata(num_seqs,
+                                                     max_draft_tokens=2),
+    ):
+        assert metadata.block_ids_per_seq.shape == (num_seqs, 16)
+        assert metadata.kv_block_ids_per_seq.shape == (num_seqs, 16)
+        assert metadata.draft_block_ids_per_seq.shape == (num_seqs, 7)
+        assert metadata.draft_kv_block_ids_per_seq.shape == (num_seqs, 7)
+        assert len({
+            metadata.block_ids_per_seq.data_ptr(),
+            metadata.kv_block_ids_per_seq.data_ptr(),
+            metadata.draft_block_ids_per_seq.data_ptr(),
+            metadata.draft_kv_block_ids_per_seq.data_ptr(),
+        }) == 4
+
+
 @pytest.mark.parametrize("spec_signal", [
     None, "num_extra_kv_tokens", "is_spec_decoding_enabled",
     "has_speculative_draft_tokens", "draft_kv_cache_manager"
diff --git a/tests/unittest/llmapi/test_config_database.py b/tests/unittest/llmapi/test_config_database.py
index b443c231e15e..ac88d5c58a69 100644
--- a/tests/unittest/llmapi/test_config_database.py
+++ b/tests/unittest/llmapi/test_config_database.py
@@ -28,6 +28,7 @@
 from tensorrt_llm.commands.serve import _apply_fastapi_middlewares
 from tensorrt_llm.commands.serve import main as serve_main
 from tensorrt_llm.llmapi import llm_args as llm_args_module
+from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args
 
 from .yaml_validation_harness import (
     assert_no_default_valued_leaves,
@@ -108,13 +109,10 @@ def _get_default_values_for_config(config_path: Path) -> dict:
     if not model_cls or not hasattr(model_cls, "get_model_defaults"):
         return global_default
 
-    base_args = llm_args_module.TorchLlmArgs(model=model, skip_tokenizer_init=True).model_dump(
-        mode="json"
-    )
-
+    base_args = llm_args_module.TorchLlmArgs(model=model, skip_tokenizer_init=True)
     model_defaults = model_cls.get_model_defaults(base_args)
-    base_args.update(model_defaults)
-    return base_args
+    apply_model_defaults_to_llm_args(base_args, model_defaults)
+    return base_args.model_dump(mode="json")
 
 
 @pytest.fixture(autouse=True)
@@ -128,6 +126,23 @@ def get_config_id(config_path: Path) -> str:
     return str(config_path.relative_to(CONFIG_ROOT))
 
 
+def test_model_nested_defaults_preserve_global_nested_defaults(monkeypatch):
+    config_key = "test-only/deepseek-v3.yaml"
+    monkeypatch.setitem(
+        _CONFIG_PATH_TO_ENTRY,
+        config_key,
+        {
+            "model": "dummy/model",
+            "arch": "DeepseekV3ForCausalLM",
+        },
+    )
+
+    defaults = _get_default_values_for_config(REPO_ROOT / config_key)
+
+    assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True
+    assert defaults["kv_cache_config"]["enable_block_reuse"] is True
+
+
 def _assert_kv_cache_block_reuse_policy(config_dict: dict) -> None:
     assert config_dict.get("kv_cache_config", {}).get("enable_block_reuse") is not False, (
         "KV cache block reuse should not be disabled"
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index 2fa8772fd073..4b1970d3b95a 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -4052,3 +4052,40 @@ def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None:
         cfg = self._pretrained_config(["DeepseekV3ForCausalLM"], "deepseek_v3")
         _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg)
         assert args.cache_transceiver_config.transceiver_runtime == "PYTHON"
+
+
+@pytest.mark.parametrize(
+    "architecture,user_setting,expected",
+    [
+        ("DeepseekV32ForCausalLM", "auto", True),
+        ("DeepseekV32ForCausalLM", False, False),
+        ("DeepseekV3ForCausalLM", "auto", True),
+        ("GlmMoeDsaForCausalLM", "auto", True),
+        ("MistralLarge3ForCausalLM", "auto", True),
+    ],
+)
+def test_deepseek_v3_family_kv_cache_manager_v2_default(architecture,
+                                                        user_setting, expected):
+    """Models backed by the shared V3 implementation select V2 by default."""
+    from types import SimpleNamespace
+    from unittest.mock import MagicMock
+
+    from transformers import PretrainedConfig
+
+    from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader
+
+    checkpoint_loader = MagicMock()
+    checkpoint_loader.load_config.return_value = SimpleNamespace(
+        pretrained_config=PretrainedConfig(architectures=[architecture]),
+        mm_encoder_only=False,
+        spec_config=None,
+    )
+    llm_args = TorchLlmArgs(
+        model="/tmp/dummy_model",
+        kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=user_setting),
+    )
+
+    ModelLoader.load_config_and_apply_defaults("/tmp/dummy_model", llm_args,
+                                               checkpoint_loader)
+
+    assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected

From a6737495cb6424024811d915f6909c2bac9450eb Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Fri, 7 Aug 2026 00:42:41 -0700
Subject: [PATCH 06/12] [None][fix] address DSA KV cache review feedback

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 docs/source/features/kvcache.md               | 22 +++++++++++
 .../_torch/models/modeling_deepseekv3.py      |  8 +++-
 .../_torch/pyexecutor/kv_cache_manager_v2.py  | 13 +------
 .../attention/sparse/dsa/test_dsa_indexer.py  | 18 +--------
 .../test_kv_cache_v2_extra_buffers.py         | 19 +---------
 .../_torch/speculative/test_eagle3.py         | 32 ----------------
 tests/unittest/llmapi/test_llm_args.py        | 37 -------------------
 7 files changed, 33 insertions(+), 116 deletions(-)
 mode change 100755 => 100644 tensorrt_llm/_torch/models/modeling_deepseekv3.py

diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md
index 576927c59475..ad9b95b554ff 100644
--- a/docs/source/features/kvcache.md
+++ b/docs/source/features/kvcache.md
@@ -1,3 +1,8 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
 # KV Cache System
 
 The KV cache stores previously computed key-value pairs for reuse during generation in order to avoid redundant calculations. The TensorRT LLM KV cache system also supports reuse across requests and uses a suite of tools like offloading and prioritized eviction to increase reuse. It supports variable attention window sizes and Multi-Head Attention (MHA) optimization techniques such as MQA and GQA.
@@ -46,6 +51,23 @@ Many of the features in the KV cache system are optional or have user defined pr
 
 See [this example](../examples/kvcacheconfig.md) for an example of how to use KvCacheConfig to control KV cache behavior.
 
+### KV Cache Manager Version
+
+`KvCacheConfig.use_kv_cache_manager_v2` selects the V2 manager; `"auto"` uses
+the model-specific default. When a sparse-attention model selects V2,
+TensorRT LLM raises `NotImplementedError` for V2-incompatible features instead
+of silently falling back to the V1 manager.
+
+`DeepseekV3ForCausalLM`, `DeepseekV32ForCausalLM`,
+`GlmMoeDsaForCausalLM`, and `MistralLarge3ForCausalLM` select V2 by default
+when this option is `"auto"`. Set the option to `false` to select V1
+explicitly.
+
+For DSA, KV cache manager V2 does not support STAR context parallelism. Use
+another context-parallel mode, or set
+`kv_cache_config.use_kv_cache_manager_v2: false` to use the DSA V1 cache
+manager with STAR.
+
 ### Datatype
 
 Perhaps the most important property is ```dtype``` which specifies what data type is held in KV cache. The default 'auto' specifies that data type should be inferred from model config.
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
old mode 100755
new mode 100644
index c1fb856b0021..bd13a17b04d1
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -31,7 +31,10 @@
 import copy
 import math
 import os
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
+
+if TYPE_CHECKING:
+    from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
 
 import torch
 import triton
@@ -1908,7 +1911,8 @@ class DeepseekV3ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV3Model,
                                                         PretrainedConfig]):
 
     @classmethod
-    def get_model_defaults(cls, llm_args) -> dict:
+    def get_model_defaults(
+            cls, llm_args: "TorchLlmArgs") -> dict[str, dict[str, bool]]:
         """Use KV cache manager V2 for models backed by this implementation."""
         return {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
 
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index e55c8dddc5a7..a3db0d3eabdd 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -1200,15 +1200,6 @@ def _get_block_scale_role(self, role_a: DataRole) -> Optional[DataRole]:
             return None
         return Role.KEY_BLOCK_SCALE
 
-    def _get_pool_layer_id(self, pool_id: int, role: DataRole) -> LayerId:
-        """Return a stable config-order representative for a pool role."""
-        try:
-            return self._pool_layer_ids_by_role[(pool_id, role)]
-        except KeyError as error:
-            raise ValueError(
-                f"Logical layer group {pool_id} has no layer with role {role}"
-            ) from error
-
     def _build_pool_mapping_tensors(self):
         """Build the (kv_cache_pool_pointers, kv_cache_pool_mapping) tensors.
 
@@ -1246,7 +1237,7 @@ def _build_pool_mapping_tensors(self):
         else:
             for pool_id in range(self.num_pools):
                 role_a, _ = self._get_pool_roles(pool_id)
-                layer_id = self._get_pool_layer_id(pool_id, role_a)
+                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
                 )
@@ -1360,7 +1351,7 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
         )
         for pool_id in range(self.num_pools):
             role_a, role_b = self._get_pool_roles(pool_id)
-            layer_id = self._get_pool_layer_id(pool_id, role_a)
+            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(
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index c3e35107db48..caf1e64f185b 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -53,7 +53,6 @@
     split_prefill_chunks,
     transform_local_topk_and_prepare_pool_view,
 )
-from tensorrt_llm._torch.attention_backend.sparse.registry import get_sparse_attn_kv_cache_manager
 from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
 from tensorrt_llm._torch.speculative.interface import (
@@ -300,19 +299,6 @@ def create_dsa_cache_manager(
     return cache_manager, sparse_attn_config
 
 
-def test_dsa_cache_manager_selection_honors_v2_config():
-    sparse_attn_config = DeepSeekSparseAttentionConfig()
-
-    assert (
-        get_sparse_attn_kv_cache_manager(sparse_attn_config, use_kv_cache_manager_v2=False)
-        is DSACacheManager
-    )
-    assert (
-        get_sparse_attn_kv_cache_manager(sparse_attn_config, use_kv_cache_manager_v2=True)
-        is DSACacheManagerV2
-    )
-
-
 @pytest.mark.parametrize(
     "cache_manager_cls",
     [
@@ -3694,11 +3680,11 @@ def test_recompute_context_kv_gather_mappings_does_not_alias_uncompressed_cached
 
     with (
         patch(
-            "tensorrt_llm._torch.attention_backend.sparse.dsa.torch.empty",
+            "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.torch.empty",
             side_effect=[host_fullkv_fp8_mapping, host_fullkv_scale_mapping],
         ),
         patch(
-            "tensorrt_llm._torch.attention_backend.sparse.dsa._compute_slot_mappings",
+            "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer._compute_slot_mappings",
             return_value=(computed_fp8_mapping, computed_scale_mapping),
         ) as compute_slot_mappings,
     ):
diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
index 97313b013321..fd460e8b8d9b 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
@@ -189,35 +189,18 @@ def test_subclass_registers_index_key_only_on_sparse_layers(self):
             mgr.shutdown()
             del mgr
 
-    def test_pool_role_representative_uses_first_config_layer_with_role(self):
-        mgr = _IndexKeyOnSparseLayersV2(sparse_layer_ids=[2, 3], **_make_kwargs(num_layers=4))
-        try:
-            pool_id = int(mgr.impl.get_layer_group_id(2))
-            self.assertEqual(int(mgr.impl.get_layer_group_id(0)), pool_id)
-
-            layer_zero = mgr.kv_cache_manager_py_config.layers[0]
-            self.assertNotIn(Role.INDEX_KEY, [buffer.role for buffer in layer_zero.buffers])
-            self.assertEqual(
-                int(mgr._get_pool_layer_id(pool_id, Role.INDEX_KEY)),
-                2,
-                "the representative must be the first config-order layer that owns the role",
-            )
-        finally:
-            mgr.shutdown()
-            del mgr
-
     def test_pool_mapping_ignores_reordered_layer_grouping(self):
         mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4))
         real_impl = mgr.impl
         try:
             self.assertEqual(mgr.num_pools, 1)
             pool_id = 0
-            expected_layer_id = mgr._get_pool_layer_id(pool_id, Role.KEY)
             config_group = tuple(
                 layer.layer_id
                 for layer in mgr.kv_cache_manager_py_config.layers
                 if int(real_impl.get_layer_group_id(layer.layer_id)) == pool_id
             )
+            expected_layer_id = config_group[0]
             self.assertGreater(len(config_group), 1)
             reordered_grouping = (config_group[1:] + config_group[:1],)
             reordered_first_layer_id = reordered_grouping[pool_id][0]
diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py
index 43eb0173654d..f77dc74783bf 100644
--- a/tests/unittest/_torch/speculative/test_eagle3.py
+++ b/tests/unittest/_torch/speculative/test_eagle3.py
@@ -341,38 +341,6 @@ def _make_mock_kv_cache_manager(num_seqs: int) -> MagicMock:
     return mock_kv_cache_manager
 
 
-@skip_num_gpus_less_than(1)
-def test_separate_draft_flash_mla_buffers_use_manager_geometry():
-    """Target and draft FlashMLA block IDs need stable, disjoint storage."""
-    num_seqs = 3
-    target_kv_cache_manager = _make_mock_kv_cache_manager(num_seqs)
-    draft_kv_cache_manager = _make_mock_kv_cache_manager(num_seqs)
-    draft_kv_cache_manager.max_blocks_per_seq = 7
-    attn_metadata = TrtllmAttentionMetadata(
-        max_num_requests=num_seqs,
-        max_num_tokens=num_seqs * 3,
-        kv_cache_manager=target_kv_cache_manager,
-        draft_kv_cache_manager=draft_kv_cache_manager,
-        enable_flash_mla=True,
-    )
-
-    for metadata in (
-            attn_metadata,
-            attn_metadata.create_cuda_graph_metadata(num_seqs,
-                                                     max_draft_tokens=2),
-    ):
-        assert metadata.block_ids_per_seq.shape == (num_seqs, 16)
-        assert metadata.kv_block_ids_per_seq.shape == (num_seqs, 16)
-        assert metadata.draft_block_ids_per_seq.shape == (num_seqs, 7)
-        assert metadata.draft_kv_block_ids_per_seq.shape == (num_seqs, 7)
-        assert len({
-            metadata.block_ids_per_seq.data_ptr(),
-            metadata.kv_block_ids_per_seq.data_ptr(),
-            metadata.draft_block_ids_per_seq.data_ptr(),
-            metadata.draft_kv_block_ids_per_seq.data_ptr(),
-        }) == 4
-
-
 @pytest.mark.parametrize("spec_signal", [
     None, "num_extra_kv_tokens", "is_spec_decoding_enabled",
     "has_speculative_draft_tokens", "draft_kv_cache_manager"
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index 4b1970d3b95a..2fa8772fd073 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -4052,40 +4052,3 @@ def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None:
         cfg = self._pretrained_config(["DeepseekV3ForCausalLM"], "deepseek_v3")
         _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg)
         assert args.cache_transceiver_config.transceiver_runtime == "PYTHON"
-
-
-@pytest.mark.parametrize(
-    "architecture,user_setting,expected",
-    [
-        ("DeepseekV32ForCausalLM", "auto", True),
-        ("DeepseekV32ForCausalLM", False, False),
-        ("DeepseekV3ForCausalLM", "auto", True),
-        ("GlmMoeDsaForCausalLM", "auto", True),
-        ("MistralLarge3ForCausalLM", "auto", True),
-    ],
-)
-def test_deepseek_v3_family_kv_cache_manager_v2_default(architecture,
-                                                        user_setting, expected):
-    """Models backed by the shared V3 implementation select V2 by default."""
-    from types import SimpleNamespace
-    from unittest.mock import MagicMock
-
-    from transformers import PretrainedConfig
-
-    from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader
-
-    checkpoint_loader = MagicMock()
-    checkpoint_loader.load_config.return_value = SimpleNamespace(
-        pretrained_config=PretrainedConfig(architectures=[architecture]),
-        mm_encoder_only=False,
-        spec_config=None,
-    )
-    llm_args = TorchLlmArgs(
-        model="/tmp/dummy_model",
-        kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=user_setting),
-    )
-
-    ModelLoader.load_config_and_apply_defaults("/tmp/dummy_model", llm_args,
-                                               checkpoint_loader)
-
-    assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected

From 72b73d62f949705b355b068da815ef044d9c4440 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Fri, 7 Aug 2026 00:48:55 -0700
Subject: [PATCH 07/12] [None][docs] defer KV cache documentation update

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 docs/source/features/kvcache.md | 22 ----------------------
 1 file changed, 22 deletions(-)

diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md
index ad9b95b554ff..576927c59475 100644
--- a/docs/source/features/kvcache.md
+++ b/docs/source/features/kvcache.md
@@ -1,8 +1,3 @@
-<!--
-SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-SPDX-License-Identifier: Apache-2.0
--->
-
 # KV Cache System
 
 The KV cache stores previously computed key-value pairs for reuse during generation in order to avoid redundant calculations. The TensorRT LLM KV cache system also supports reuse across requests and uses a suite of tools like offloading and prioritized eviction to increase reuse. It supports variable attention window sizes and Multi-Head Attention (MHA) optimization techniques such as MQA and GQA.
@@ -51,23 +46,6 @@ Many of the features in the KV cache system are optional or have user defined pr
 
 See [this example](../examples/kvcacheconfig.md) for an example of how to use KvCacheConfig to control KV cache behavior.
 
-### KV Cache Manager Version
-
-`KvCacheConfig.use_kv_cache_manager_v2` selects the V2 manager; `"auto"` uses
-the model-specific default. When a sparse-attention model selects V2,
-TensorRT LLM raises `NotImplementedError` for V2-incompatible features instead
-of silently falling back to the V1 manager.
-
-`DeepseekV3ForCausalLM`, `DeepseekV32ForCausalLM`,
-`GlmMoeDsaForCausalLM`, and `MistralLarge3ForCausalLM` select V2 by default
-when this option is `"auto"`. Set the option to `false` to select V1
-explicitly.
-
-For DSA, KV cache manager V2 does not support STAR context parallelism. Use
-another context-parallel mode, or set
-`kv_cache_config.use_kv_cache_manager_v2: false` to use the DSA V1 cache
-manager with STAR.
-
 ### Datatype
 
 Perhaps the most important property is ```dtype``` which specifies what data type is held in KV cache. The default 'auto' specifies that data type should be inferred from model config.

From 40cdbb28eb9da1a305191c3b46f6c788b91ed1a4 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sat, 8 Aug 2026 00:45:04 -0700
Subject: [PATCH 08/12] [TensorRT-LLM] Fix KV cache V2 CI regressions

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../sparse/dsa/cache_manager.py               | 12 ++----
 tensorrt_llm/_torch/pyexecutor/_util.py       | 12 ++----
 .../_torch/pyexecutor/kv_cache_manager_v2.py  | 16 +++++---
 .../executor/test_kv_cache_estimation.py      |  2 +-
 .../test_kv_cache_v2_extra_buffers.py         | 24 ++++++++++--
 .../executor/test_mamba_cache_manager.py      | 38 +++++++++++++++++++
 6 files changed, 76 insertions(+), 28 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
index 620e5452656b..e496e8bc23df 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
@@ -421,8 +421,7 @@ def __init__(
             for local_layer_idx in range(self.num_local_layers)
         ]
         self._primary_pool_page_index_params = [
-            (int(converter.scale), int(converter.layer_offset))
-            for converter in primary_converters
+            (int(converter.scale), int(converter.layer_offset)) for converter in primary_converters
         ]
         full_indexer_local_layers = [
             local_layer_idx
@@ -463,9 +462,7 @@ def get_primary_pool_page_index_params(self, local_layer_idx: int) -> Tuple[int,
         """Return the formal V2 page scale and layer offset for sparse MLA."""
         return self._primary_pool_page_index_params[local_layer_idx]
 
-    def _extra_buffers_per_layer(
-        self, *, tokens_per_block: int
-    ) -> dict[int, List[BufferConfig]]:
+    def _extra_buffers_per_layer(self, *, tokens_per_block: int) -> dict[int, List[BufferConfig]]:
         return {
             local_layer_idx: [
                 BufferConfig(
@@ -574,10 +571,7 @@ def get_unique_primary_pool(self) -> torch.Tensor:
         element_per_container = 2 if self.dtype == DataType.NVFP4 else 1
         dtype = torch.int8 if self.dtype == DataType.NVFP4 else self.dtype
         elements_per_layer = (
-            self.tokens_per_block
-            * first_num_heads
-            * first_head_dim
-            // element_per_container
+            self.tokens_per_block * first_num_heads * first_head_dim // element_per_container
         )
         shape = [
             self.blocks_in_primary_pool,
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index a24a4bc73387..d19e0d705b29 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -129,6 +129,7 @@ def get_kv_cache_manager_cls(
     config = model_config.pretrained_config
     sparse_attn_config = model_config.sparse_attention_config
     sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None)
+    use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True
     if is_hybrid_linear(config):
         # Degenerate case: model is flagged as hybrid but the config has zero
         # mamba layers. Fall through to the standard non-hybrid routes.
@@ -137,9 +138,7 @@ def get_kv_cache_manager_cls(
                         "KV cache manager without mamba caching")
             if sparse_attn_config is not None:
                 return get_sparse_attn_kv_cache_manager(
-                    sparse_attn_config,
-                    use_kv_cache_manager_v2=kv_cache_config.
-                    use_kv_cache_manager_v2)
+                    sparse_attn_config, use_kv_cache_manager_v2=use_v2)
             return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config)
 
         if (sparse_attn_config is not None
@@ -152,8 +151,6 @@ def get_kv_cache_manager_cls(
         has_additional_snapshots = bool(
             state_config.additional_snapshot_offsets_from_start
             or state_config.additional_snapshot_offsets_from_end)
-        use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True
-
         if has_additional_snapshots and not use_v2:
             raise ValueError("Mamba additional snapshot offsets require "
                              "use_kv_cache_manager_v2=True; V1 supports only "
@@ -265,9 +262,8 @@ def get_kv_cache_manager_cls(
                 "yet model retained recurrent-state snapshots.")
         return MambaHybridCacheManagerV2
     elif sparse_attn_config is not None:
-        return get_sparse_attn_kv_cache_manager(
-            sparse_attn_config,
-            use_kv_cache_manager_v2=kv_cache_config.use_kv_cache_manager_v2)
+        return get_sparse_attn_kv_cache_manager(sparse_attn_config,
+                                                use_kv_cache_manager_v2=use_v2)
     else:
         return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config)
 
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index a3db0d3eabdd..2852ac44fd57 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -1081,14 +1081,18 @@ def append_to_kv_heads_per_layer(
             )
 
         # Both backends build layer_grouping on demand, and the layer order
-        # within a group is not part of its contract. Cache a stable
-        # config-order representative for each role in each logical group.
+        # within a group is not part of its contract. Cache a stable physical-
+        # layout representative for each role from the public pool descriptors.
         self.num_pools = len(self.impl.layer_grouping)
         self._pool_layer_ids_by_role: Dict[Tuple[int, DataRole], LayerId] = {}
-        for layer in config.layers:
-            pool_id = int(self.impl.get_layer_group_id(layer.layer_id))
-            for buffer in layer.buffers:
-                self._pool_layer_ids_by_role.setdefault((pool_id, buffer.role), layer.layer_id)
+        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
diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py
index 5c75c9186bd7..7ce5d626d944 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py
@@ -754,7 +754,7 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None:
         patch.object(creator, "_enable_kv_cache_stats", return_value=False),
         patch.object(
             creator,
-            "_fallback_if_unsupported_kv_cache_manager_v2",
+            "_validate_or_fallback_kv_cache_manager_v2",
             return_value=KVCacheManagerV2,
         ),
         patch(
diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
index fd460e8b8d9b..dd80817a3d9e 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
@@ -190,7 +190,7 @@ def test_subclass_registers_index_key_only_on_sparse_layers(self):
             del mgr
 
     def test_pool_mapping_ignores_reordered_layer_grouping(self):
-        mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4))
+        mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4, head_dim=[64, 192, 64, 192]))
         real_impl = mgr.impl
         try:
             self.assertEqual(mgr.num_pools, 1)
@@ -200,9 +200,25 @@ def test_pool_mapping_ignores_reordered_layer_grouping(self):
                 for layer in mgr.kv_cache_manager_py_config.layers
                 if int(real_impl.get_layer_group_id(layer.layer_id)) == pool_id
             )
-            expected_layer_id = config_group[0]
+            storage_variant = next(
+                variant
+                for pool_group in real_impl.pool_group_descs
+                for variant in pool_group.slot_desc.variants
+                if int(variant.layer_group_id) == pool_id
+            )
+            expected_layer_id = next(
+                int(buffer_id.layer_id)
+                for coalesced in storage_variant.coalesced_buffers
+                for buffer_id in coalesced.buffer_ids
+                if buffer_id.role == Role.KEY
+            )
             self.assertGreater(len(config_group), 1)
-            reordered_grouping = (config_group[1:] + config_group[:1],)
+            self.assertNotEqual(expected_layer_id, int(config_group[0]))
+            reordered_group = tuple(
+                layer_id for layer_id in config_group if int(layer_id) != expected_layer_id
+            )
+            reordered_group += (expected_layer_id,)
+            reordered_grouping = (reordered_group,)
             reordered_first_layer_id = reordered_grouping[pool_id][0]
             self.assertNotEqual(expected_layer_id, reordered_first_layer_id)
 
@@ -218,7 +234,7 @@ def test_pool_mapping_ignores_reordered_layer_grouping(self):
             self.assertEqual(
                 impl_proxy.get_mem_pool_base_address.call_args_list[0].args,
                 (expected_layer_id, Role.KEY, PageIndexMode.SHARED),
-                "the pool pointer origin must use the config-order representative",
+                "the pool pointer origin must use the physical-layout representative",
             )
             self.assertEqual(
                 int(mgr.kv_cache_pool_pointers[pool_id, 0]),
diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
index fcaa753316c9..804e4a947162 100644
--- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py
+++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
@@ -873,6 +873,44 @@ def test_dsa_star_rejects_v2_and_keeps_explicit_v1_selection():
     )
 
 
+@pytest.mark.parametrize("zero_layer_hybrid", [False, True])
+@pytest.mark.parametrize(
+    ("manager_setting", "expected_use_v2"),
+    [(False, False), ("auto", False), (True, True)],
+)
+def test_sparse_cache_manager_factory_only_enables_explicit_v2(
+    monkeypatch, zero_layer_hybrid, manager_setting, expected_use_v2
+):
+    sparse_attn_config = SimpleNamespace(algorithm="dsa")
+    model_config = SimpleNamespace(
+        pretrained_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]),
+        sparse_attention_config=sparse_attn_config,
+        get_num_mamba_layers=lambda: 0,
+    )
+    expected_manager_cls = object()
+    sparse_manager_factory = MagicMock(return_value=expected_manager_cls)
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.pyexecutor._util.is_hybrid_linear",
+        lambda config: zero_layer_hybrid,
+    )
+    monkeypatch.setattr(
+        "tensorrt_llm._torch.pyexecutor._util.get_sparse_attn_kv_cache_manager",
+        sparse_manager_factory,
+    )
+
+    assert (
+        get_kv_cache_manager_cls(
+            model_config,
+            KvCacheConfig(use_kv_cache_manager_v2=manager_setting),
+        )
+        is expected_manager_cls
+    )
+    sparse_manager_factory.assert_called_once_with(
+        sparse_attn_config,
+        use_kv_cache_manager_v2=expected_use_v2,
+    )
+
+
 def _make_mgr(
     max_batch_size=4, max_draft_len=2, enable_attention_dp=False, use_replay_state_update=False
 ):

From 679b10e05fa730341605b20e74563b7e01936884 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sat, 8 Aug 2026 06:00:49 -0700
Subject: [PATCH 09/12] [TensorRT-LLM] Fix DSA MTP draft cache replay

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../attention_backend/sparse/dsa/metadata.py  | 108 +++++++++++-----
 tensorrt_llm/_torch/speculative/eagle3.py     |  11 ++
 tensorrt_llm/_torch/speculative/interface.py  |  22 +++-
 tensorrt_llm/llmapi/llm_utils.py              |  53 ++++++++
 .../attention/sparse/dsa/test_dsa_indexer.py  |  82 ++++++++++--
 .../_torch/speculative/test_eagle3.py         |  95 ++++++++++++++
 tests/unittest/_torch/test_model_config.py    | 121 ++++++++++++++++++
 7 files changed, 444 insertions(+), 48 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
index 10513814e39d..d697b1f8583c 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
@@ -574,8 +574,13 @@ def create_buffers_for_indexer(self, capture_graph=False):
         # the draft-replay context swaps these in by rebinding, so CUDA graph
         # capture bakes distinct addresses for the target and draft segments
         # and both sides can be refreshed eagerly outside the graph.
-        if self.draft_kv_cache_manager is not None and hasattr(
-            self.draft_kv_cache_manager, "index_head_dim"
+        if self.draft_kv_cache_manager is not None and all(
+            hasattr(self.draft_kv_cache_manager, attr)
+            for attr in (
+                "index_head_dim",
+                "get_pool_block_indices",
+                "indexer_k_cache_page_scale",
+            )
         ):
             self.draft_indexer_k_cache_block_offsets = self.get_empty(
                 self.cuda_graph_buffers,
@@ -637,6 +642,24 @@ def create_buffers_for_indexer(self, capture_graph=False):
             dtype=torch.int32,
             capture_graph=capture_graph,
         )
+        if self.draft_kv_cache_manager is not None and all(
+            hasattr(self.draft_kv_cache_manager, attr)
+            for attr in (
+                "index_head_dim",
+                "get_pool_block_indices",
+                "indexer_k_cache_page_scale",
+            )
+        ):
+            self.draft_block_table = self.get_empty(
+                self.cuda_graph_buffers,
+                [
+                    self.max_num_sequences,
+                    self.draft_kv_cache_manager.max_blocks_per_seq,
+                ],
+                cache_name="draft_block_table",
+                dtype=torch.int32,
+                capture_graph=capture_graph,
+            )
         self.scheduler_metadata_buffer = self.get_empty(
             self.cuda_graph_buffers,
             (self.num_sms + 1, 2),
@@ -787,6 +810,29 @@ def create_expanded_buffers(self, capture_graph=False):
             device="cpu",
             pin_memory=prefer_pinned(),
         )
+        if self.draft_kv_cache_manager is not None and all(
+            hasattr(self.draft_kv_cache_manager, attr)
+            for attr in (
+                "index_head_dim",
+                "get_pool_block_indices",
+                "indexer_k_cache_page_scale",
+            )
+        ):
+            self.draft_block_table_expanded = self.get_empty(
+                self.cuda_graph_buffers,
+                [
+                    self.max_num_sequences * (1 + self.max_draft_tokens),
+                    self.draft_kv_cache_manager.max_blocks_per_seq,
+                ],
+                cache_name="draft_block_table_expanded",
+                dtype=torch.int32,
+                capture_graph=capture_graph,
+            )
+            self.host_draft_block_table_expanded = torch.zeros_like(
+                self.draft_block_table_expanded,
+                device="cpu",
+                pin_memory=prefer_pinned(),
+            )
         self.scheduler_metadata_buffer_expanded = self.get_empty(
             self.cuda_graph_buffers,
             (self.num_sms + 1, 2),
@@ -1005,24 +1051,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
                 self.kv_lens_expanded_host[:num_tokens], non_blocking=True
             )
 
-            # Expand indexer_k_cache_block_offsets (only generation)
-            # host_indexer_k_cache_block_offsets already contains INDEX_KEY
-            # page indices from prepare_for_indexer_k_cache().
-            if self.kv_cache_manager is not None and self.num_generations > 0:
-                max_len = self.host_indexer_k_cache_block_offsets.shape[1]
-                gen_block_tensor = self.host_indexer_k_cache_block_offsets[
-                    self.num_contexts : self.num_seqs, :max_len
-                ]
-                expanded_blocks = gen_block_tensor.repeat_interleave(
-                    1 + self.max_draft_tokens, dim=0
-                )
-                self.host_block_table_expanded[:num_tokens, :max_len].copy_(
-                    expanded_blocks, non_blocking=True
-                )
-                self.block_table_expanded[:num_tokens].copy_(
-                    self.host_block_table_expanded[:num_tokens], non_blocking=True
-                )
-                self.block_table_expanded.clamp_(min=0)
+            self._refresh_expanded_block_table(1 + self.max_draft_tokens)
 
         self.expand_for_dsl = (
             use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1
@@ -1048,22 +1077,37 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor):
                 self.kv_lens_expanded_cuda[:num_tokens].copy_(
                     self.kv_lens_expanded_host[:num_tokens], non_blocking=True
                 )
-                max_len = self.host_indexer_k_cache_block_offsets.shape[1]
-                gen_block_tensor = self.host_indexer_k_cache_block_offsets[
-                    self.num_contexts : self.num_seqs, :max_len
-                ]
-                expanded_blocks = gen_block_tensor.repeat_interleave(expand_factor, dim=0)
-                self.host_block_table_expanded[:num_tokens, :max_len].copy_(
-                    expanded_blocks, non_blocking=True
-                )
-                self.block_table_expanded[:num_tokens].copy_(
-                    self.host_block_table_expanded[:num_tokens], non_blocking=True
-                )
-                self.block_table_expanded.clamp_(min=0)
+                self._refresh_expanded_block_table(expand_factor)
         else:
             self.dsl_expand_factor = 1
             self.dsl_atom = 1 + self.max_draft_tokens
 
+    def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None):
+        """Refresh the active cache's expanded INDEX_KEY page table."""
+        if self.kv_cache_manager is None or self.num_generations == 0:
+            return
+        if repeat_factor is None:
+            if self.use_expanded_buffers_for_mtp:
+                repeat_factor = 1 + self.max_draft_tokens
+            elif self.expand_for_dsl and self.dsl_expand_factor > 1:
+                repeat_factor = self.dsl_expand_factor
+            else:
+                return
+
+        num_tokens = self.num_generations * repeat_factor
+        max_len = self.host_indexer_k_cache_block_offsets.shape[1]
+        gen_block_tensor = self.host_indexer_k_cache_block_offsets[
+            self.num_contexts : self.num_seqs, :max_len
+        ]
+        expanded_blocks = gen_block_tensor.repeat_interleave(repeat_factor, dim=0)
+        self.host_block_table_expanded[:num_tokens, :max_len].copy_(
+            expanded_blocks, non_blocking=True
+        )
+        self.block_table_expanded[:num_tokens].copy_(
+            self.host_block_table_expanded[:num_tokens], non_blocking=True
+        )
+        self.block_table_expanded.clamp_(min=0)
+
     def prepare_for_indexer_k_cache(self):
         if self.kv_cache_manager is None:
             return
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index c201c76592ed..6117ef0649fe 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -957,6 +957,17 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 attn_metadata, draft_kv_cache_manager) as draft_attn_metadata:
             attn_metadata = draft_attn_metadata
             inputs["attn_metadata"] = draft_attn_metadata
+            if (reuse_mtp_topk and draft_kv_cache_manager is not None
+                    and hasattr(draft_kv_cache_manager, "index_head_dim") and
+                    hasattr(draft_kv_cache_manager, "get_pool_block_indices")
+                    and hasattr(draft_kv_cache_manager,
+                                "indexer_k_cache_page_scale")):
+                # Overlap scheduling corrects kv_lens_cuda from the runtime
+                # accepted-token counts inside the captured graph. The target
+                # forward refreshes target slot mappings during that correction;
+                # refresh once more after rebinding so the first draft forward
+                # writes the separate DSA indexer cache at the same positions.
+                attn_metadata.on_update_kv_lens()
             for i in range(runtime_draft_len):
                 if reuse_mtp_topk:
                     attn_metadata.set_skip_topk(i > 0)
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index a061e3b0821e..49305fc24252 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -172,8 +172,14 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
 
     from ..attention_backend.sparse.dsa import (DSAtrtllmAttentionMetadata,
                                                 Indexer)
+
+    # DeepSeek-V4 metadata inherits DSA metadata, but its cache manager uses a
+    # different dual-pool layout. Only enter the DSA replay path when the
+    # manager exposes DSA's indexer-cache page-mapping contract.
     if (isinstance(attn_metadata, DSAtrtllmAttentionMetadata)
-            and hasattr(draft_kv_cache_manager, 'index_head_dim')):
+            and hasattr(draft_kv_cache_manager, 'index_head_dim')
+            and hasattr(draft_kv_cache_manager, 'get_pool_block_indices')
+            and hasattr(draft_kv_cache_manager, 'indexer_k_cache_page_scale')):
         m = attn_metadata
         saved['saved_dsa_state'] = {
             'host_indexer_k_cache_block_offsets':
@@ -183,6 +189,9 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
             'host_slot_mapping_scale': m.host_slot_mapping_scale,
             'slot_mapping_fp8': m.slot_mapping_fp8,
             'slot_mapping_scale': m.slot_mapping_scale,
+            'block_table': m.block_table,
+            'block_table_expanded': m.block_table_expanded,
+            'host_block_table_expanded': m.host_block_table_expanded,
         }
         # The cached-KV feature owns these references even when an optimized
         # path aliases them to slot_mapping_*. With the feature disabled, the
@@ -206,6 +215,10 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
         m.slot_mapping_fp8 = m.draft_slot_mapping_fp8
         m.host_slot_mapping_scale = m.host_draft_slot_mapping_scale
         m.slot_mapping_scale = m.draft_slot_mapping_scale
+        m.block_table = m.draft_block_table
+        m.block_table_expanded = m.draft_block_table_expanded
+        m.host_block_table_expanded = m.host_draft_block_table_expanded
+        m._invalidate_pool_view_cache()
         # Recording a capture executes no kernels, so the draft mappings only
         # need refreshing when the transfers actually run: eager forwards
         # (warmup) and the pre-replay call from model_engine. The per-step
@@ -213,7 +226,8 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
         # device from the rebound block-offset buffer.
         # kv_cache_manager was already swapped to the draft manager above.
         if not torch.cuda.is_current_stream_capturing():
-            m._update_indexer_k_cache_block_offsets()
+            m.prepare_for_indexer_k_cache()
+            m._refresh_expanded_block_table()
             Indexer.recompute_slot_mappings(m)
         Indexer.recompute_context_kv_gather_mappings(m)
     return saved
@@ -248,6 +262,10 @@ def restore_attn_metadata_after_draft_replay(attn_metadata, saved_state):
         m.host_slot_mapping_scale = saved_dsa['host_slot_mapping_scale']
         m.slot_mapping_fp8 = saved_dsa['slot_mapping_fp8']
         m.slot_mapping_scale = saved_dsa['slot_mapping_scale']
+        m.block_table = saved_dsa['block_table']
+        m.block_table_expanded = saved_dsa['block_table_expanded']
+        m.host_block_table_expanded = saved_dsa['host_block_table_expanded']
+        m._invalidate_pool_view_cache()
         if 'slot_mapping_fp8_fullkv' in saved_dsa:
             m.slot_mapping_fp8_fullkv = saved_dsa['slot_mapping_fp8_fullkv']
             m.slot_mapping_scale_fullkv = saved_dsa['slot_mapping_scale_fullkv']
diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py
index a8bd94bdb797..4ff146188a18 100644
--- a/tensorrt_llm/llmapi/llm_utils.py
+++ b/tensorrt_llm/llmapi/llm_utils.py
@@ -10,6 +10,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import transformers
+from pydantic import BaseModel
 
 from .._utils import global_mpi_rank, local_mpi_rank, mpi_rank
 # yapf: disable
@@ -529,6 +530,52 @@ def _deep_merge(base: Dict[str, Any], *overlays: Dict[str,
     return result
 
 
+def _capture_pydantic_fields_set(value: Any) -> Any:
+    """Capture explicit-field state for nested Pydantic configurations."""
+    if isinstance(value, BaseModel):
+        return (
+            "model",
+            set(value.model_fields_set),
+            {
+                field_name:
+                _capture_pydantic_fields_set(getattr(value, field_name))
+                for field_name in type(value).model_fields
+            },
+        )
+    if isinstance(value, (list, tuple)):
+        return ("sequence",
+                [_capture_pydantic_fields_set(item) for item in value])
+    if isinstance(value, dict):
+        return (
+            "mapping",
+            {
+                key: _capture_pydantic_fields_set(item)
+                for key, item in value.items()
+            },
+        )
+    return None
+
+
+def _restore_pydantic_fields_set(value: Any, state: Any) -> None:
+    """Restore explicit-field state after validating reconstructed configs."""
+    if state is None:
+        return
+
+    kind = state[0]
+    if kind == "model" and isinstance(value, BaseModel):
+        object.__setattr__(value, "__pydantic_fields_set__", set(state[1]))
+        for field_name, child_state in state[2].items():
+            _restore_pydantic_fields_set(getattr(value, field_name),
+                                         child_state)
+    elif kind == "sequence" and isinstance(value, (list, tuple)):
+        for item, child_state in zip(value, state[1]):
+            _restore_pydantic_fields_set(item, child_state)
+    elif kind == "mapping" and isinstance(value, dict):
+        for key, child_state in state[1].items():
+            if key in value:
+                _restore_pydantic_fields_set(value[key], child_state)
+
+
 def apply_model_defaults_to_llm_args(
         llm_args: 'TorchLlmArgs',
         model_defaults_dict: Dict[str, Any]) -> Dict[str, Any]:
@@ -549,6 +596,7 @@ def apply_model_defaults_to_llm_args(
             "config the user did not turn on. Declare a runtime preference "
             "via get_preferred_transceiver_runtime() instead.")
 
+    explicit_fields_state = _capture_pydantic_fields_set(llm_args)
     user_overrides = llm_args.model_dump(exclude_unset=True)
     base_state = llm_args.model_dump()
     merged_state = _deep_merge(base_state, model_defaults_dict, user_overrides)
@@ -557,6 +605,11 @@ def apply_model_defaults_to_llm_args(
 
     for field_name in llm_args.model_fields:
         setattr(llm_args, field_name, getattr(new_args, field_name))
+    # Reconstructing from a complete model_dump marks every nested default as
+    # explicit. Preserve the original user-intent metadata after validation so
+    # later checkpoint resolution can still distinguish omitted values from
+    # overrides (for example DeepSeek index_topk).
+    _restore_pydantic_fields_set(llm_args, explicit_fields_state)
 
     def _compute_applied(defaults: Dict[str, Any],
                          overrides: Dict[str, Any]) -> Dict[str, Any]:
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index caf1e64f185b..db329e1b48db 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -451,7 +451,9 @@ def __init__(self, head_dim):
     mla_params = MLAParams(sparse_params.index_head_dim)
 
     # Mock RotaryEmbedding since we're only testing cache management, not rope functionality
-    with patch("tensorrt_llm._torch.attention_backend.sparse.dsa.RotaryEmbedding") as mock_rope:
+    with patch(
+        "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.RotaryEmbedding"
+    ) as mock_rope:
         # Create a mock instance with a simple forward method
         mock_rope_instance = Mock()
         mock_rope_instance.forward = Mock(side_effect=lambda pos_ids, tensors: tensors)
@@ -3539,9 +3541,7 @@ def _make_mock_metadata():
     @staticmethod
     def _make_mock_draft_manager():
         """Create a mock draft KV cache manager with host block offsets."""
-        mgr = Mock(name="draft_kv_cache_manager")
-        mgr.host_kv_cache_block_offsets = torch.tensor([100, 200, 300])
-        return mgr
+        return SimpleNamespace(host_kv_cache_block_offsets=torch.tensor([100, 200, 300]))
 
     def test_prepare_swaps_and_restore_recovers(self):
         """Test that prepare swaps KV manager and restore recovers original state."""
@@ -3576,11 +3576,14 @@ def test_prepare_swaps_and_restore_recovers(self):
         torch.testing.assert_close(meta.kv_cache_block_offsets, original_offsets)
         torch.testing.assert_close(meta.host_kv_cache_block_offsets, original_host_offsets)
 
-    def test_draft_context_recomputes_and_restores_dsa_metadata(self):
+    @pytest.mark.parametrize("capturing", [False, True])
+    def test_draft_context_recomputes_and_restores_dsa_metadata(self, capturing):
         """Draft context rebinds preallocated DSA buffers and restores lazy aliases."""
         meta = self._make_mock_metadata()
         mgr = self._make_mock_draft_manager()
         mgr.index_head_dim = 128
+        mgr.get_pool_block_indices = Mock()
+        mgr.indexer_k_cache_page_scale = 1
         meta.enable_flash_mla = False
         meta.enable_context_mla_with_cached_kv = False
         meta.host_indexer_k_cache_block_offsets = torch.tensor([1, 2, 3])
@@ -3589,12 +3592,20 @@ def test_draft_context_recomputes_and_restores_dsa_metadata(self):
         meta.host_slot_mapping_scale = torch.tensor([10, 11, 12])
         meta.slot_mapping_fp8 = torch.tensor([13, 14, 15])
         meta.slot_mapping_scale = torch.tensor([16, 17, 18])
-        meta.host_draft_indexer_k_cache_block_offsets = torch.tensor([25, 26, 27])
-        meta.draft_indexer_k_cache_block_offsets = torch.tensor([28, 29, 30])
-        meta.host_draft_slot_mapping_fp8 = torch.tensor([31, 32, 33])
-        meta.host_draft_slot_mapping_scale = torch.tensor([34, 35, 36])
-        meta.draft_slot_mapping_fp8 = torch.tensor([37, 38, 39])
-        meta.draft_slot_mapping_scale = torch.tensor([40, 41, 42])
+        meta.block_table = torch.tensor([19, 20, 21])
+        meta.block_table_expanded = torch.tensor([22, 23, 24])
+        meta.host_block_table_expanded = torch.tensor([25, 26, 27])
+        meta.host_draft_indexer_k_cache_block_offsets = torch.tensor([28, 29, 30])
+        meta.draft_indexer_k_cache_block_offsets = torch.tensor([31, 32, 33])
+        meta.host_draft_slot_mapping_fp8 = torch.tensor([34, 35, 36])
+        meta.host_draft_slot_mapping_scale = torch.tensor([37, 38, 39])
+        meta.draft_slot_mapping_fp8 = torch.tensor([40, 41, 42])
+        meta.draft_slot_mapping_scale = torch.tensor([43, 44, 45])
+        meta.draft_block_table = torch.tensor([46, 47, 48])
+        meta.draft_block_table_expanded = torch.tensor([49, 50, 51])
+        meta.host_draft_block_table_expanded = torch.tensor([52, 53, 54])
+        meta.prepare_for_indexer_k_cache = Mock()
+        meta._refresh_expanded_block_table = Mock()
         meta._update_indexer_k_cache_block_offsets = Mock()
         del meta.slot_mapping_fp8_fullkv
         del meta.slot_mapping_scale_fullkv
@@ -3606,6 +3617,9 @@ def test_draft_context_recomputes_and_restores_dsa_metadata(self):
             "host_slot_mapping_scale": meta.host_slot_mapping_scale,
             "slot_mapping_fp8": meta.slot_mapping_fp8,
             "slot_mapping_scale": meta.slot_mapping_scale,
+            "block_table": meta.block_table,
+            "block_table_expanded": meta.block_table_expanded,
+            "host_block_table_expanded": meta.host_block_table_expanded,
         }
         draft_buffers = {
             "host_indexer_k_cache_block_offsets": meta.host_draft_indexer_k_cache_block_offsets,
@@ -3614,6 +3628,9 @@ def test_draft_context_recomputes_and_restores_dsa_metadata(self):
             "host_slot_mapping_scale": meta.host_draft_slot_mapping_scale,
             "slot_mapping_fp8": meta.draft_slot_mapping_fp8,
             "slot_mapping_scale": meta.draft_slot_mapping_scale,
+            "block_table": meta.draft_block_table,
+            "block_table_expanded": meta.draft_block_table_expanded,
+            "host_block_table_expanded": meta.host_draft_block_table_expanded,
         }
 
         def is_attn_metadata(obj, cls):
@@ -3628,7 +3645,7 @@ def is_attn_metadata(obj, cls):
             ),
             patch(
                 "tensorrt_llm._torch.speculative.interface.torch.cuda.is_current_stream_capturing",
-                return_value=False,
+                return_value=capturing,
             ),
             patch.object(Indexer, "recompute_slot_mappings") as recompute,
         ):
@@ -3642,8 +3659,45 @@ def is_attn_metadata(obj, cls):
             assert getattr(meta, name) is buffer
         assert meta.slot_mapping_fp8_fullkv is meta.slot_mapping_fp8
         assert meta.slot_mapping_scale_fullkv is meta.slot_mapping_scale
-        meta._update_indexer_k_cache_block_offsets.assert_called_once_with()
-        recompute.assert_called_once_with(meta)
+        if capturing:
+            meta.prepare_for_indexer_k_cache.assert_not_called()
+            meta._refresh_expanded_block_table.assert_not_called()
+            recompute.assert_not_called()
+        else:
+            meta.prepare_for_indexer_k_cache.assert_called_once_with()
+            meta._refresh_expanded_block_table.assert_called_once_with()
+            recompute.assert_called_once_with(meta)
+        meta._update_indexer_k_cache_block_offsets.assert_not_called()
+
+    def test_non_dsa_manager_skips_inherited_dsa_replay_state(self):
+        """A DSA metadata subclass must not impose DSA layout on its manager."""
+        meta = self._make_mock_metadata()
+        mgr = self._make_mock_draft_manager()
+        target_manager = meta.kv_cache_manager
+        target_offsets = meta.kv_cache_block_offsets
+        target_host_offsets = meta.host_kv_cache_block_offsets
+        # DeepSeek-V4 cache managers expose this field but not the DSA page
+        # mapping contract used by _update_indexer_k_cache_block_offsets().
+        mgr.index_head_dim = 128
+
+        def is_attn_metadata(obj, cls):
+            if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata):
+                return obj is meta
+            return builtins.isinstance(obj, cls)
+
+        with patch(
+            "tensorrt_llm._torch.speculative.interface.isinstance",
+            side_effect=is_attn_metadata,
+        ):
+            saved = prepare_attn_metadata_for_draft_replay(meta, mgr)
+
+        assert saved is not None
+        assert "saved_dsa_state" not in saved
+
+        restore_attn_metadata_after_draft_replay(meta, saved)
+        assert meta.kv_cache_manager is target_manager
+        assert meta.kv_cache_block_offsets is target_offsets
+        assert meta.host_kv_cache_block_offsets is target_host_offsets
 
 
 def test_recompute_context_kv_gather_mappings_does_not_alias_uncompressed_cached_mappings():
diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py
index f77dc74783bf..52e50f9f1b15 100644
--- a/tests/unittest/_torch/speculative/test_eagle3.py
+++ b/tests/unittest/_torch/speculative/test_eagle3.py
@@ -18,6 +18,7 @@
 import sys
 import tempfile
 import unittest
+from contextlib import contextmanager
 from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import MagicMock
@@ -47,6 +48,100 @@
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 
 
+def test_mtp_eagle_refreshes_runtime_draft_slot_mappings_before_forward(
+) -> None:
+    """The first MTP-Eagle draft forward must use runtime-corrected mappings."""
+    events = []
+    target_manager = object()
+    target_block_table = object()
+    draft_block_table = object()
+    draft_manager = SimpleNamespace(
+        index_head_dim=128,
+        get_pool_block_indices=lambda *args, **kwargs: None,
+        indexer_k_cache_page_scale=1,
+    )
+
+    class _AttentionMetadata:
+
+        def __init__(self):
+            self.seq_lens_cuda = torch.tensor([1], dtype=torch.int32)
+            self.num_ctx_tokens = 0
+            self.kv_cache_manager = target_manager
+            self.block_table = target_block_table
+
+        def set_in_mtp_draft_loop(self, active):
+            events.append(("set_in_mtp_draft_loop", active))
+
+        def set_mtp_num_accepted(self, value):
+            events.append(("set_mtp_num_accepted", value))
+
+        def set_skip_topk(self, value):
+            events.append(("set_skip_topk", value))
+
+        def on_update_kv_lens(self):
+            assert self.kv_cache_manager is draft_manager
+            assert self.block_table is draft_block_table
+            events.append(("on_update_kv_lens", draft_manager))
+
+    attn_metadata = _AttentionMetadata()
+
+    @contextmanager
+    def draft_context(metadata, manager):
+        assert metadata is attn_metadata
+        assert manager is draft_manager
+        assert metadata.kv_cache_manager is target_manager
+        assert metadata.block_table is target_block_table
+        metadata.kv_cache_manager = draft_manager
+        metadata.block_table = draft_block_table
+        events.append(("enter_draft_context", manager))
+        try:
+            yield metadata
+        finally:
+            metadata.kv_cache_manager = target_manager
+            metadata.block_table = target_block_table
+
+    class _StopAfterFirstDraftForward(Exception):
+        pass
+
+    def run_draft_forward(*args, **kwargs):
+        assert attn_metadata.kv_cache_manager is draft_manager
+        assert attn_metadata.block_table is draft_block_table
+        events.append(("run_draft_forward", draft_manager))
+        raise _StopAfterFirstDraftForward
+
+    worker = SimpleNamespace(
+        is_mtp_eagle=True,
+        draft_kv_cache_context=draft_context,
+        _run_draft_forward=run_draft_forward,
+    )
+    spec_metadata = SimpleNamespace(runtime_draft_len=1)
+    inputs = {"position_ids": torch.tensor([0], dtype=torch.int32)}
+
+    from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker
+
+    with pytest.raises(_StopAfterFirstDraftForward):
+        Eagle3OneModelWorker._forward_linear_draft_loop(
+            worker,
+            inputs,
+            attn_metadata,
+            spec_metadata,
+            draft_model=object(),
+            draft_kv_cache_manager=draft_manager,
+            num_contexts=0,
+            batch_size=1,
+            num_accepted_tokens=torch.tensor([1], dtype=torch.int32),
+            original_all_rank_num_tokens=None,
+        )
+
+    event_names = [name for name, _ in events]
+    assert event_names.index("enter_draft_context") < event_names.index(
+        "on_update_kv_lens")
+    assert event_names.index("on_update_kv_lens") < event_names.index(
+        "run_draft_forward")
+    assert attn_metadata.kv_cache_manager is target_manager
+    assert attn_metadata.block_table is target_block_table
+
+
 def test_dynamic_tree_metadata_forces_target_mask_prepare_each_step() -> None:
     metadata = TrtllmAttentionMetadata(
         seq_lens=None,
diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py
index 6e3daea5e695..a9e71de4843a 100644
--- a/tests/unittest/_torch/test_model_config.py
+++ b/tests/unittest/_torch/test_model_config.py
@@ -1,15 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import json
 import struct
 import types
 
 import pytest
 import torch
+import transformers
 
 from tensorrt_llm._torch.model_config import _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT, ModelConfig
 from tensorrt_llm._torch.pyexecutor.model_loader import (
     validate_and_set_kv_cache_quant,
     validate_encoder_decoder_kv_cache_config,
 )
+from tensorrt_llm.llmapi.llm_args import (
+    DeepSeekSparseAttentionConfig,
+    DeepSeekV4SparseAttentionConfig,
+    TorchLlmArgs,
+)
+from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
 
@@ -230,6 +240,117 @@ def test_deepseek_v4_missing_compress_ratios_raises(tmp_path, monkeypatch):
         ModelConfig.from_pretrained(str(tmp_path))
 
 
+@pytest.mark.parametrize(
+    "user_index_topk,expected_index_topk",
+    [(None, 2048), (1024, 1024)],
+)
+def test_deepseek_v32_model_defaults_preserve_index_topk_override(
+    tmp_path, monkeypatch, user_index_topk, expected_index_topk
+):
+    """Applying unrelated model defaults must preserve nested user intent."""
+    from tensorrt_llm._torch import model_config as model_config_module
+
+    pretrained_config = transformers.PretrainedConfig(
+        architectures=["DeepseekV32ForCausalLM"],
+        index_n_heads=64,
+        index_head_dim=128,
+        index_topk=2048,
+        indexer_rope_interleave=False,
+        dtype="bfloat16",
+        intermediate_size=1024,
+    )
+    monkeypatch.setattr(
+        model_config_module,
+        "load_pretrained_config",
+        lambda *args, **kwargs: pretrained_config,
+    )
+
+    sparse_kwargs = {}
+    if user_index_topk is not None:
+        sparse_kwargs["index_topk"] = user_index_topk
+    llm_args = TorchLlmArgs(
+        model=str(tmp_path),
+        sparse_attention_config=DeepSeekSparseAttentionConfig(**sparse_kwargs),
+    )
+    apply_model_defaults_to_llm_args(
+        llm_args, {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+    )
+
+    model_config = ModelConfig.from_pretrained(
+        str(tmp_path), sparse_attention_config=llm_args.sparse_attention_config
+    )
+    sparse_config = model_config.sparse_attention_config
+    assert sparse_config.index_topk == expected_index_topk
+    assert sparse_config.needs_separate_short_long_cuda_graphs()
+    assert sparse_config.seq_len_threshold == expected_index_topk
+
+
+def test_model_defaults_preserve_explicit_nested_none(tmp_path):
+    """An explicit None must remain distinct from an omitted nested field."""
+    llm_args = TorchLlmArgs(
+        model=str(tmp_path),
+        sparse_attention_config=DeepSeekSparseAttentionConfig(index_topk=None),
+    )
+    assert "index_topk" in llm_args.sparse_attention_config.model_fields_set
+
+    apply_model_defaults_to_llm_args(
+        llm_args, {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+    )
+
+    assert llm_args.sparse_attention_config.index_topk is None
+    assert "index_topk" in llm_args.sparse_attention_config.model_fields_set
+
+
+@pytest.mark.parametrize(
+    "sparse_kwargs,expected_index_topk",
+    [
+        pytest.param({}, 1024, id="omitted"),
+        pytest.param({"index_topk": 512}, 512, id="explicit"),
+    ],
+)
+def test_deepseek_v4_model_defaults_preserve_index_topk_override(
+    tmp_path, monkeypatch, sparse_kwargs, expected_index_topk
+):
+    """V4's default top-k must not hide a checkpoint or explicit override."""
+    from tensorrt_llm._torch import model_config as model_config_module
+
+    pretrained_config = transformers.PretrainedConfig(
+        architectures=["DeepseekV4ForCausalLM"],
+        index_n_heads=64,
+        index_head_dim=128,
+        index_topk=1024,
+        indexer_rope_interleave=False,
+        compress_ratios=[1, 4, 128, 4],
+        window_size=128,
+        num_hidden_layers=4,
+        dtype="bfloat16",
+        intermediate_size=1024,
+    )
+    monkeypatch.setattr(
+        model_config_module,
+        "load_pretrained_config",
+        lambda *args, **kwargs: pretrained_config,
+    )
+
+    llm_args = TorchLlmArgs(
+        model=str(tmp_path),
+        sparse_attention_config=DeepSeekV4SparseAttentionConfig(
+            index_head_dim=128, **sparse_kwargs
+        ),
+    )
+    apply_model_defaults_to_llm_args(
+        llm_args, {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+    )
+    assert ("index_topk" in llm_args.sparse_attention_config.model_fields_set) is bool(
+        sparse_kwargs
+    )
+
+    model_config = ModelConfig.from_pretrained(
+        str(tmp_path), sparse_attention_config=llm_args.sparse_attention_config
+    )
+    assert model_config.sparse_attention_config.index_topk == expected_index_topk
+
+
 def test_model_config_sets_is_encoder_decoder_from_pretrained_config():
     model_config = ModelConfig(
         pretrained_config=make_pretrained_config(

From b309e17f99237af1be6af6526ea6769c944fd90e Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sun, 9 Aug 2026 01:17:29 -0700
Subject: [PATCH 10/12] refactor: identify native DSA cache managers

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../attention_backend/sparse/dsa/__init__.py  |  3 +-
 .../sparse/dsa/cache_manager.py               |  5 +++
 .../attention_backend/sparse/dsa/metadata.py  | 32 +++--------------
 tensorrt_llm/_torch/speculative/eagle3.py     | 20 +++++------
 tensorrt_llm/_torch/speculative/interface.py  | 10 +++---
 .../attention/sparse/dsa/test_dsa_indexer.py  | 34 ++++++++++++++++---
 .../_torch/speculative/test_eagle3.py         | 14 ++++----
 7 files changed, 62 insertions(+), 56 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
index 685c2e35bfda..80908e08b0ac 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/__init__.py
@@ -4,7 +4,7 @@
 """DeepSeek Sparse Attention backend package."""
 
 from .backend import DSATrtllmAttention
-from .cache_manager import DSACacheManager, DSACacheManagerV2
+from .cache_manager import DSACacheManager, DSACacheManagerV2, is_dsa_cache_manager
 from .indexer import (
     _DG_SCHEDULE_BLOCK_KV,
     HAS_FAST_HADAMARD,
@@ -45,6 +45,7 @@
     "_select_indexer_compress_ratio",
     "build_req_idx_per_token",
     "compute_cu_seqlen_kv_bounds_with_cache",
+    "is_dsa_cache_manager",
     "rotate_activation",
     "split_prefill_chunks",
     "transform_local_topk_and_prepare_pool_view",
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
index e496e8bc23df..ed19c6ea869e 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/cache_manager.py
@@ -615,3 +615,8 @@ def shutdown(self) -> None:
         self.indexer_k_cache_pool_per_layer = []
         self._unique_primary_pool = None
         super().shutdown()
+
+
+def is_dsa_cache_manager(cache_manager: object) -> bool:
+    """Return whether a manager uses native DSA indexer page mapping."""
+    return isinstance(cache_manager, (DSACacheManager, DSACacheManagerV2))
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
index d697b1f8583c..a41631cda0e1 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa/metadata.py
@@ -17,6 +17,7 @@
 from tensorrt_llm._utils import get_sm_version, prefer_pinned
 from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata
 
+from .cache_manager import is_dsa_cache_manager
 from .indexer import (
     _DG_SCHEDULE_BLOCK_KV,
     Indexer,
@@ -137,10 +138,8 @@ def __init__(self, *args, **kwargs):
 
     def __post_init__(self):
         """Allocate indexer K-cache buffers and heuristic TopK metadata."""
-        from .cache_manager import DSACacheManager, DSACacheManagerV2
-
         super().__post_init__()
-        if not isinstance(self.kv_cache_manager, (DSACacheManager, DSACacheManagerV2)):
+        if not is_dsa_cache_manager(self.kv_cache_manager):
             has_deepseek_v4_cache_interface = all(
                 hasattr(self.kv_cache_manager, attr)
                 for attr in ("compressed_block_sizes", "get_cache_indices")
@@ -574,14 +573,7 @@ def create_buffers_for_indexer(self, capture_graph=False):
         # the draft-replay context swaps these in by rebinding, so CUDA graph
         # capture bakes distinct addresses for the target and draft segments
         # and both sides can be refreshed eagerly outside the graph.
-        if self.draft_kv_cache_manager is not None and all(
-            hasattr(self.draft_kv_cache_manager, attr)
-            for attr in (
-                "index_head_dim",
-                "get_pool_block_indices",
-                "indexer_k_cache_page_scale",
-            )
-        ):
+        if is_dsa_cache_manager(self.draft_kv_cache_manager):
             self.draft_indexer_k_cache_block_offsets = self.get_empty(
                 self.cuda_graph_buffers,
                 [self.max_num_sequences, self.draft_kv_cache_manager.max_blocks_per_seq],
@@ -642,14 +634,7 @@ def create_buffers_for_indexer(self, capture_graph=False):
             dtype=torch.int32,
             capture_graph=capture_graph,
         )
-        if self.draft_kv_cache_manager is not None and all(
-            hasattr(self.draft_kv_cache_manager, attr)
-            for attr in (
-                "index_head_dim",
-                "get_pool_block_indices",
-                "indexer_k_cache_page_scale",
-            )
-        ):
+        if is_dsa_cache_manager(self.draft_kv_cache_manager):
             self.draft_block_table = self.get_empty(
                 self.cuda_graph_buffers,
                 [
@@ -810,14 +795,7 @@ def create_expanded_buffers(self, capture_graph=False):
             device="cpu",
             pin_memory=prefer_pinned(),
         )
-        if self.draft_kv_cache_manager is not None and all(
-            hasattr(self.draft_kv_cache_manager, attr)
-            for attr in (
-                "index_head_dim",
-                "get_pool_block_indices",
-                "indexer_k_cache_page_scale",
-            )
-        ):
+        if is_dsa_cache_manager(self.draft_kv_cache_manager):
             self.draft_block_table_expanded = self.get_empty(
                 self.cuda_graph_buffers,
                 [
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 6117ef0649fe..34d8f150a866 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -939,6 +939,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                                    num_accepted_tokens,
                                    original_all_rank_num_tokens):
         """Linear draft loop, unified for Eagle3 and MTP Eagle."""
+        from ..attention_backend.sparse.dsa import (DSAtrtllmAttentionMetadata,
+                                                    is_dsa_cache_manager)
+
         runtime_draft_len = spec_metadata.runtime_draft_len
         num_gens = batch_size - num_contexts
         next_draft_tokens = []
@@ -946,9 +949,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
             attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
         position_ids = inputs["position_ids"]
 
-        reuse_mtp_topk = (self.is_mtp_eagle
-                          and hasattr(attn_metadata, "set_skip_topk"))
-        if reuse_mtp_topk:
+        uses_dsa_mtp_metadata = self.is_mtp_eagle and isinstance(
+            attn_metadata, DSAtrtllmAttentionMetadata)
+        if uses_dsa_mtp_metadata:
             attn_metadata.set_in_mtp_draft_loop(True)
             # Accepted counts let the indexer stash each gen's last-accepted row.
             attn_metadata.set_mtp_num_accepted(num_accepted_tokens)
@@ -957,11 +960,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 attn_metadata, draft_kv_cache_manager) as draft_attn_metadata:
             attn_metadata = draft_attn_metadata
             inputs["attn_metadata"] = draft_attn_metadata
-            if (reuse_mtp_topk and draft_kv_cache_manager is not None
-                    and hasattr(draft_kv_cache_manager, "index_head_dim") and
-                    hasattr(draft_kv_cache_manager, "get_pool_block_indices")
-                    and hasattr(draft_kv_cache_manager,
-                                "indexer_k_cache_page_scale")):
+            if uses_dsa_mtp_metadata and is_dsa_cache_manager(
+                    draft_kv_cache_manager):
                 # Overlap scheduling corrects kv_lens_cuda from the runtime
                 # accepted-token counts inside the captured graph. The target
                 # forward refreshes target slot mappings during that correction;
@@ -969,7 +969,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 # writes the separate DSA indexer cache at the same positions.
                 attn_metadata.on_update_kv_lens()
             for i in range(runtime_draft_len):
-                if reuse_mtp_topk:
+                if uses_dsa_mtp_metadata:
                     attn_metadata.set_skip_topk(i > 0)
                 # Run draft model (mode-specific via helper). The helper
                 # passes ``all_rank_num_tokens`` as a kwarg so the draft model
@@ -1167,7 +1167,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                 }
         next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
 
-        if reuse_mtp_topk:
+        if uses_dsa_mtp_metadata:
             attn_metadata.set_skip_topk(False)
             attn_metadata.set_in_mtp_draft_loop(False)
             attn_metadata.set_mtp_num_accepted(None)
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 49305fc24252..8e51c1379a7f 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -171,15 +171,13 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
         attn_metadata.prepare_flash_mla()
 
     from ..attention_backend.sparse.dsa import (DSAtrtllmAttentionMetadata,
-                                                Indexer)
+                                                Indexer, is_dsa_cache_manager)
 
     # DeepSeek-V4 metadata inherits DSA metadata, but its cache manager uses a
-    # different dual-pool layout. Only enter the DSA replay path when the
-    # manager exposes DSA's indexer-cache page-mapping contract.
+    # different dual-pool layout. Only native DSA cache managers use the DSA
+    # draft-replay buffers below.
     if (isinstance(attn_metadata, DSAtrtllmAttentionMetadata)
-            and hasattr(draft_kv_cache_manager, 'index_head_dim')
-            and hasattr(draft_kv_cache_manager, 'get_pool_block_indices')
-            and hasattr(draft_kv_cache_manager, 'indexer_k_cache_page_scale')):
+            and is_dsa_cache_manager(draft_kv_cache_manager)):
         m = attn_metadata
         saved['saved_dsa_state'] = {
             'host_indexer_k_cache_block_offsets':
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index db329e1b48db..e06ceb1767c0 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -40,6 +40,9 @@
     PositionalEmbeddingParams,
     RopeParams,
 )
+from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import (
+    DeepseekV4CacheManager,
+)
 from tensorrt_llm._torch.attention_backend.sparse.dsa import (
     DSABackendForwardArgs,
     DSACacheManager,
@@ -50,6 +53,7 @@
     _effective_compress_ratio_divisor,
     _select_indexer_compress_ratio,
     compute_cu_seqlen_kv_bounds_with_cache,
+    is_dsa_cache_manager,
     split_prefill_chunks,
     transform_local_topk_and_prepare_pool_view,
 )
@@ -83,6 +87,28 @@ def has_deep_gemm():
         return False
 
 
+@pytest.mark.parametrize("cache_manager_cls", [DSACacheManager, DSACacheManagerV2])
+def test_is_dsa_cache_manager_accepts_native_managers(cache_manager_cls):
+    assert is_dsa_cache_manager(object.__new__(cache_manager_cls))
+
+
+@pytest.mark.parametrize(
+    "cache_manager",
+    [
+        None,
+        object.__new__(DeepseekV4CacheManager),
+        SimpleNamespace(
+            index_head_dim=128,
+            get_pool_block_indices=Mock(),
+            indexer_k_cache_page_scale=1,
+        ),
+    ],
+    ids=["none", "deepseek-v4", "structural-lookalike"],
+)
+def test_is_dsa_cache_manager_rejects_non_native_managers(cache_manager):
+    assert not is_dsa_cache_manager(cache_manager)
+
+
 def test_metadata_cache_geometry_comes_from_sparse_metadata_params():
     sparse_config = DeepSeekV4SparseAttentionConfig(
         compress_ratios=[1, 4, 128],
@@ -3580,10 +3606,8 @@ def test_prepare_swaps_and_restore_recovers(self):
     def test_draft_context_recomputes_and_restores_dsa_metadata(self, capturing):
         """Draft context rebinds preallocated DSA buffers and restores lazy aliases."""
         meta = self._make_mock_metadata()
-        mgr = self._make_mock_draft_manager()
-        mgr.index_head_dim = 128
-        mgr.get_pool_block_indices = Mock()
-        mgr.indexer_k_cache_page_scale = 1
+        mgr = object.__new__(DSACacheManagerV2)
+        mgr.host_kv_cache_block_offsets = torch.tensor([100, 200, 300])
         meta.enable_flash_mla = False
         meta.enable_context_mla_with_cached_kv = False
         meta.host_indexer_k_cache_block_offsets = torch.tensor([1, 2, 3])
@@ -3679,6 +3703,8 @@ def test_non_dsa_manager_skips_inherited_dsa_replay_state(self):
         # DeepSeek-V4 cache managers expose this field but not the DSA page
         # mapping contract used by _update_indexer_k_cache_block_offsets().
         mgr.index_head_dim = 128
+        mgr.get_pool_block_indices = Mock()
+        mgr.indexer_k_cache_page_scale = 1
 
         def is_attn_metadata(obj, cls):
             if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata):
diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py
index 52e50f9f1b15..3c0ca12ef94a 100644
--- a/tests/unittest/_torch/speculative/test_eagle3.py
+++ b/tests/unittest/_torch/speculative/test_eagle3.py
@@ -31,6 +31,8 @@
                         skip_pre_blackwell)
 
 from tensorrt_llm import LLM, SamplingParams
+from tensorrt_llm._torch.attention_backend.sparse.dsa import (
+    DSACacheManagerV2, DSAtrtllmAttentionMetadata)
 from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata
 from tensorrt_llm._torch.metadata import KVCacheParams
 from tensorrt_llm._torch.pyexecutor._util import \
@@ -55,17 +57,13 @@ def test_mtp_eagle_refreshes_runtime_draft_slot_mappings_before_forward(
     target_manager = object()
     target_block_table = object()
     draft_block_table = object()
-    draft_manager = SimpleNamespace(
-        index_head_dim=128,
-        get_pool_block_indices=lambda *args, **kwargs: None,
-        indexer_k_cache_page_scale=1,
-    )
+    draft_manager = object.__new__(DSACacheManagerV2)
 
-    class _AttentionMetadata:
+    class _AttentionMetadata(DSAtrtllmAttentionMetadata):
 
         def __init__(self):
-            self.seq_lens_cuda = torch.tensor([1], dtype=torch.int32)
-            self.num_ctx_tokens = 0
+            self._seq_lens_cuda = torch.tensor([1], dtype=torch.int32)
+            self._num_ctx_tokens = 0
             self.kv_cache_manager = target_manager
             self.block_table = target_block_table
 

From 3353119cce1e6e448101d434974f6de276709f4c Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sun, 9 Aug 2026 03:46:26 -0700
Subject: [PATCH 11/12] [None][refactor] resolve KV cache defaults via
 preferences

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../_torch/models/modeling_deepseekv3.py      |  14 +-
 .../_torch/models/modeling_deepseekv4.py      |  32 ++-
 .../_torch/models/modeling_kimi_k25.py        |   9 +-
 .../_torch/models/modeling_minimaxm2.py       |  13 +-
 .../_torch/models/modeling_nemotron_h.py      |   8 +-
 .../_torch/models/modeling_qwen3_5.py         |  11 +-
 .../_torch/models/modeling_qwen3_next.py      |  16 +-
 tensorrt_llm/_torch/models/modeling_utils.py  |  50 +++-
 .../_torch/pyexecutor/model_loader.py         |  38 +--
 tensorrt_llm/llmapi/llm_utils.py              | 164 +++++++------
 .../cache_transceiver_precheck/README.md      |   2 +-
 .../precheck_config.py                        |   4 +-
 .../run_precheck.py                           |  17 +-
 .../executor/test_mamba_cache_manager.py      |   9 +-
 .../modeling/test_modeling_deepseekv4.py      |  11 +-
 .../_torch/modeling/test_modeling_kimi_k25.py |   7 +
 tests/unittest/_torch/test_model_config.py    | 118 ---------
 tests/unittest/llmapi/test_config_database.py |  28 ++-
 tests/unittest/llmapi/test_llm_args.py        | 229 ++++++++++++++----
 .../test_cache_transceiver_precheck_config.py |   2 +-
 .../test_cache_transceiver_precheck_run.py    |   8 +-
 21 files changed, 450 insertions(+), 340 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
index bd13a17b04d1..9755f0fc2806 100644
--- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
@@ -31,10 +31,7 @@
 import copy
 import math
 import os
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
-
-if TYPE_CHECKING:
-    from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 import torch
 import triton
@@ -1911,10 +1908,11 @@ class DeepseekV3ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV3Model,
                                                         PretrainedConfig]):
 
     @classmethod
-    def get_model_defaults(
-            cls, llm_args: "TorchLlmArgs") -> dict[str, dict[str, bool]]:
-        """Use KV cache manager V2 for models backed by this implementation."""
-        return {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+    def get_preferred_kv_cache_manager_version(cls,
+                                               pretrained_config: Any = None
+                                               ) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for this model implementation."""
+        return "V2"
 
     @classmethod
     def get_preferred_transceiver_runtime(
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py
index a4260e15dbb0..79e059a0a857 100644
--- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py
@@ -31,10 +31,7 @@
 import copy
 import math
 import os
-from typing import TYPE_CHECKING, Dict, List, Optional
-
-if TYPE_CHECKING:
-    from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
+from typing import Dict, List, Literal, Optional
 
 import torch
 import triton
@@ -2533,14 +2530,25 @@ def forward(
 @register_auto_model("DeepseekV4ForCausalLM")
 class DeepseekV4ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV4Model, PretrainedConfig]):
     @classmethod
-    def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
-        return {
-            "kv_cache_config": {
-                "tokens_per_block": 128,
-                "use_kv_cache_manager_v2": True,
-                "enable_swa_scratch_reuse": True,
-            }
-        }
+    def get_preferred_kv_cache_tokens_per_block(
+        cls, pretrained_config: object | None = None
+    ) -> int:
+        """Prefer 128-token KV-cache blocks for DeepSeek-V4."""
+        return 128
+
+    @classmethod
+    def get_preferred_kv_cache_swa_scratch_reuse(
+        cls, pretrained_config: object | None = None
+    ) -> bool:
+        """Prefer SWA scratch reuse for DeepSeek-V4."""
+        return True
+
+    @classmethod
+    def get_preferred_kv_cache_manager_version(
+        cls, pretrained_config: object | None = None
+    ) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for DeepSeek-V4."""
+        return "V2"
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig]):
         model_config = _normalize_deepseek_v4_nvfp4_mixed_precision_config(model_config)
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py
index 8eed050eafe2..723f242cb613 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py
@@ -1518,9 +1518,12 @@ class KimiK25ForConditionalGeneration(PreTrainedModel):
     _LANG_PREFIX = "language_model."
 
     @classmethod
-    def get_model_defaults(cls, llm_args: Any) -> dict:
-        """Use the C++ KV cache manager V2 by default."""
-        return {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+    def get_preferred_kv_cache_manager_version(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for Kimi K2.5."""
+        return "V2"
 
     @classmethod
     def get_preferred_transceiver_runtime(
diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm2.py b/tensorrt_llm/_torch/models/modeling_minimaxm2.py
index e35f8bf1b73b..2ca84ee1f2a0 100644
--- a/tensorrt_llm/_torch/models/modeling_minimaxm2.py
+++ b/tensorrt_llm/_torch/models/modeling_minimaxm2.py
@@ -13,15 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import Dict, List, Literal, Optional
 
 import torch
 from torch import nn
 from transformers import PretrainedConfig
 
-if TYPE_CHECKING:
-    from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
-
 from tensorrt_llm._ipc_utils import can_access_peer
 from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType
 from tensorrt_llm.mapping import Mapping
@@ -403,9 +400,11 @@ def forward(
 @register_auto_model("MiniMaxM2ForCausalLM")
 class MiniMaxM2ForCausalLM(DecoderModelForCausalLM[MiniMaxM2Model, PretrainedConfig]):
     @classmethod
-    def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
-        """Use KV cache manager V2 by default."""
-        return {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+    def get_preferred_kv_cache_manager_version(
+        cls, pretrained_config: object | None = None
+    ) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for MiniMax M2."""
+        return "V2"
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig]):
         super().__init__(
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
index 4484ee8315b4..bd5915343b24 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
@@ -991,10 +991,16 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
         return {
             "kv_cache_config": {
                 "enable_block_reuse": False,
-                "use_kv_cache_manager_v2": True,
             }
         }
 
+    @classmethod
+    def get_preferred_kv_cache_manager_version(cls,
+                                               pretrained_config: object
+                                               | None = None) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for the hybrid state layout."""
+        return "V2"
+
     @classmethod
     def get_preferred_transceiver_runtime(cls,
                                           pretrained_config: object
diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py
index 8acf89ab927b..b71fe8495928 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py
@@ -701,10 +701,17 @@ def get_model_defaults(cls, llm_args):
         # model class (this VLM wrapper), not on the inner decoder. Both
         # inner LMs (`Qwen3_5MoeForCausalLM` / `Qwen3_5ForCausalLM`) inherit
         # `Qwen3NextForCausalLM`'s defaults unchanged, so delegate to it to
-        # propagate the hybrid manager selection and keep block reuse disabled
-        # until a recurrent-state snapshot policy is configured.
+        # keep block reuse disabled until a recurrent-state snapshot policy is
+        # configured.
         return Qwen3NextForCausalLM.get_model_defaults(llm_args)
 
+    @classmethod
+    def get_preferred_kv_cache_manager_version(
+        cls, pretrained_config: object | None = None
+    ) -> Literal["V2"]:
+        """Match the hybrid text decoder's KV cache manager preference."""
+        return "V2"
+
     @classmethod
     def get_preferred_transceiver_runtime(
         cls, pretrained_config: object | None = None
diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py
index 268429d038f3..aa2626bf1da7 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py
@@ -1053,13 +1053,15 @@ def __init__(
 
     @classmethod
     def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict:
-        """Use V2 by default; explicit V1 selections remain supported."""
-        return {
-            "kv_cache_config": {
-                "enable_block_reuse": False,
-                "use_kv_cache_manager_v2": True,
-            }
-        }
+        """Disable block reuse until a snapshot policy is configured."""
+        return {"kv_cache_config": {"enable_block_reuse": False}}
+
+    @classmethod
+    def get_preferred_kv_cache_manager_version(cls,
+                                               pretrained_config: object
+                                               | None = None) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for the hybrid state layout."""
+        return "V2"
 
     @classmethod
     def get_preferred_transceiver_runtime(cls,
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index 1b4b5f84e52e..fd6feb66d520 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -658,13 +658,55 @@ def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict:
         The returned dict is deep-merged with the user's llm_args, with
         user-set values taking priority over these defaults.
 
-        Note: ``cache_transceiver_config`` is rejected here (enforced at
-        load time) — the deep-merge could materialize or silently enable a
-        transceiver config the user did not turn on. Use
-        :meth:`get_preferred_transceiver_runtime` instead.
+        Note: ``cache_transceiver_config`` and
+        ``kv_cache_config.use_kv_cache_manager_v2`` are rejected here
+        (enforced at load time). Use
+        :meth:`get_preferred_transceiver_runtime` and
+        :meth:`get_preferred_kv_cache_manager_version` instead.
         """
         return {}
 
+    @classmethod
+    def get_preferred_kv_cache_manager_version(
+            cls,
+            pretrained_config: Any = None) -> Optional[Literal["V1", "V2"]]:
+        """Return the model's preferred KV cache manager version.
+
+        The preference is adopted only when the user leaves
+        ``kv_cache_config.use_kv_cache_manager_v2`` at ``"auto"``. Return
+        ``None`` to defer to the global default (V1).
+
+        Args:
+            pretrained_config: The loaded Hugging Face config. Shared model
+                implementations can inspect it to select a preference for the
+                original checkpoint architecture.
+        """
+        return None
+
+    @classmethod
+    def get_preferred_kv_cache_tokens_per_block(cls,
+                                                pretrained_config: Any = None
+                                                ) -> Optional[int]:
+        """Return the model's preferred KV-cache block size.
+
+        The preference is adopted only when the user does not explicitly set
+        ``kv_cache_config.tokens_per_block``. Return ``None`` to keep the
+        global default.
+        """
+        return None
+
+    @classmethod
+    def get_preferred_kv_cache_swa_scratch_reuse(cls,
+                                                 pretrained_config: Any = None
+                                                 ) -> Optional[bool]:
+        """Return the model's preferred SWA scratch-reuse setting.
+
+        The preference is adopted only when the user does not explicitly set
+        ``kv_cache_config.enable_swa_scratch_reuse``. Return ``None`` to keep
+        the global default.
+        """
+        return None
+
     @classmethod
     def get_preferred_transceiver_runtime(
             cls,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 7a852bf9bfd1..d5a216136148 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -23,7 +23,8 @@
                                           ExecutorMemoryType,
                                           ModelExpressConfig,
                                           SparseAttentionConfig, TorchLlmArgs)
-from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
+from tensorrt_llm.llmapi.llm_utils import (_apply_kv_cache_config_preferences,
+                                           _resolve_kv_cache_manager_v2_auto,
                                            _resolve_transceiver_runtime_auto,
                                            apply_model_defaults_to_llm_args)
 from tensorrt_llm.logger import logger
@@ -413,6 +414,7 @@ def load_config_and_apply_defaults(
             # No config to resolve a model class from; still resolve the
             # "auto" sentinel so it never leaks past config loading.
             _resolve_transceiver_runtime_auto(llm_args)
+            _resolve_kv_cache_manager_v2_auto(llm_args)
             return llm_args
 
         config_kwargs = {
@@ -428,9 +430,25 @@ def load_config_and_apply_defaults(
         config = checkpoint_loader.load_config(checkpoint_dir, **config_kwargs)
 
         model_cls = AutoModelForCausalLM._resolve_class(config)
-        use_kv_cache_manager_v2 = (
+        original_kv_cache_manager_setting = (
             llm_args.kv_cache_config.use_kv_cache_manager_v2)
 
+        # Preferences follow the checkpoint's original architecture:
+        # _resolve_class may rewrite it to an execution class (e.g.
+        # MTPDraftModelForCausalLM), which must not drop the target model's
+        # preferences.
+        preference_cls = model_cls
+        architectures = getattr(config.pretrained_config, 'architectures', None)
+        if architectures:
+            preference_cls = get_registered_model_class(
+                architectures[0]) or model_cls
+
+        applied_kv_preferences = _apply_kv_cache_config_preferences(
+            llm_args, preference_cls, config.pretrained_config)
+        if applied_kv_preferences:
+            logger.info("Applied KV-cache preferences for %s: %s",
+                        preference_cls.__name__, applied_kv_preferences)
+
         # model_cls is None when the architecture is unknown/unsupported.
         model_defaults = {}
         if model_cls and hasattr(model_cls, 'get_model_defaults'):
@@ -452,23 +470,13 @@ def load_config_and_apply_defaults(
             update_spec_config_from_model_config(llm_args.speculative_config,
                                                  config.pretrained_config)
 
-        # The transceiver preference follows the checkpoint's original
-        # architecture: _resolve_class may rewrite it to an execution class
-        # (e.g. MTPDraftModelForCausalLM), which must not drop the target
-        # model's preference.
-        preference_cls = model_cls
-        architectures = getattr(config.pretrained_config, 'architectures', None)
-        if architectures:
-            preference_cls = get_registered_model_class(
-                architectures[0]) or model_cls
-
         # Resolve "auto" sentinel values after model defaults are applied.
         _resolve_transceiver_runtime_auto(llm_args, preference_cls,
                                           config.pretrained_config)
-        _resolve_kv_cache_manager_v2_auto(
-            llm_args, model_defaults, original_setting=use_kv_cache_manager_v2)
+        _resolve_kv_cache_manager_v2_auto(llm_args, preference_cls,
+                                          config.pretrained_config)
         _validate_and_adjust_mamba_snapshot_config(config, llm_args)
-        if use_kv_cache_manager_v2 == "auto":
+        if original_kv_cache_manager_setting == "auto":
             logger.info(
                 "Resolved use_kv_cache_manager_v2='auto' to %s for %s",
                 llm_args.kv_cache_config.use_kv_cache_manager_v2,
diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py
index 4ff146188a18..0de67d373542 100644
--- a/tensorrt_llm/llmapi/llm_utils.py
+++ b/tensorrt_llm/llmapi/llm_utils.py
@@ -10,7 +10,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import transformers
-from pydantic import BaseModel
 
 from .._utils import global_mpi_rank, local_mpi_rank, mpi_rank
 # yapf: disable
@@ -530,52 +529,6 @@ def _deep_merge(base: Dict[str, Any], *overlays: Dict[str,
     return result
 
 
-def _capture_pydantic_fields_set(value: Any) -> Any:
-    """Capture explicit-field state for nested Pydantic configurations."""
-    if isinstance(value, BaseModel):
-        return (
-            "model",
-            set(value.model_fields_set),
-            {
-                field_name:
-                _capture_pydantic_fields_set(getattr(value, field_name))
-                for field_name in type(value).model_fields
-            },
-        )
-    if isinstance(value, (list, tuple)):
-        return ("sequence",
-                [_capture_pydantic_fields_set(item) for item in value])
-    if isinstance(value, dict):
-        return (
-            "mapping",
-            {
-                key: _capture_pydantic_fields_set(item)
-                for key, item in value.items()
-            },
-        )
-    return None
-
-
-def _restore_pydantic_fields_set(value: Any, state: Any) -> None:
-    """Restore explicit-field state after validating reconstructed configs."""
-    if state is None:
-        return
-
-    kind = state[0]
-    if kind == "model" and isinstance(value, BaseModel):
-        object.__setattr__(value, "__pydantic_fields_set__", set(state[1]))
-        for field_name, child_state in state[2].items():
-            _restore_pydantic_fields_set(getattr(value, field_name),
-                                         child_state)
-    elif kind == "sequence" and isinstance(value, (list, tuple)):
-        for item, child_state in zip(value, state[1]):
-            _restore_pydantic_fields_set(item, child_state)
-    elif kind == "mapping" and isinstance(value, dict):
-        for key, child_state in state[1].items():
-            if key in value:
-                _restore_pydantic_fields_set(value[key], child_state)
-
-
 def apply_model_defaults_to_llm_args(
         llm_args: 'TorchLlmArgs',
         model_defaults_dict: Dict[str, Any]) -> Dict[str, Any]:
@@ -596,20 +549,28 @@ def apply_model_defaults_to_llm_args(
             "config the user did not turn on. Declare a runtime preference "
             "via get_preferred_transceiver_runtime() instead.")
 
-    explicit_fields_state = _capture_pydantic_fields_set(llm_args)
+    kv_cache_defaults = model_defaults_dict.get("kv_cache_config")
+    if (kv_cache_defaults is not None
+            and not isinstance(kv_cache_defaults, dict)):
+        raise ValueError(
+            "Model default 'kv_cache_config' must be a dict. Runtime "
+            "preferences must use the corresponding get_preferred_* hook.")
+    if (kv_cache_defaults is not None
+            and "use_kv_cache_manager_v2" in kv_cache_defaults):
+        raise ValueError(
+            "Model defaults must not contain "
+            "'kv_cache_config.use_kv_cache_manager_v2'. Declare a runtime "
+            "preference via get_preferred_kv_cache_manager_version() "
+            "instead.")
+
     user_overrides = llm_args.model_dump(exclude_unset=True)
     base_state = llm_args.model_dump()
     merged_state = _deep_merge(base_state, model_defaults_dict, user_overrides)
 
     new_args = llm_args.__class__(**merged_state)
 
-    for field_name in llm_args.model_fields:
+    for field_name in type(llm_args).model_fields:
         setattr(llm_args, field_name, getattr(new_args, field_name))
-    # Reconstructing from a complete model_dump marks every nested default as
-    # explicit. Preserve the original user-intent metadata after validation so
-    # later checkpoint resolution can still distinguish omitted values from
-    # overrides (for example DeepSeek index_topk).
-    _restore_pydantic_fields_set(llm_args, explicit_fields_state)
 
     def _compute_applied(defaults: Dict[str, Any],
                          overrides: Dict[str, Any]) -> Dict[str, Any]:
@@ -650,11 +611,60 @@ def _two_model_spec_dec_decoding_type(
     return getattr(spec_config, "decoding_type", "speculative decoding")
 
 
-def _resolve_kv_cache_manager_v2_auto(
+def _apply_kv_cache_config_preferences(
         llm_args: 'TorchLlmArgs',
-        model_defaults_dict: Dict[str, Any],
-        original_setting: Optional[Union[bool, str]] = None) -> bool:
-    """Resolve the KV cache manager auto setting after model defaults are applied.
+        model_cls: Optional[type] = None,
+        pretrained_config: Any = None) -> Dict[str, Any]:
+    """Apply scalar KV-cache preferences that the user left unspecified."""
+    if model_cls is None:
+        return {}
+
+    kv_cache_config = llm_args.kv_cache_config
+    user_overrides = kv_cache_config.model_dump(exclude_unset=True)
+    applied = {}
+
+    get_tokens_per_block = getattr(model_cls,
+                                   'get_preferred_kv_cache_tokens_per_block',
+                                   None)
+    if ("tokens_per_block" not in user_overrides
+            and get_tokens_per_block is not None):
+        preferred_tokens_per_block = get_tokens_per_block(pretrained_config)
+        if preferred_tokens_per_block is not None:
+            if (type(preferred_tokens_per_block) is not int
+                    or preferred_tokens_per_block <= 0):
+                raise ValueError(
+                    f"{model_cls.__name__}."
+                    "get_preferred_kv_cache_tokens_per_block() must return a "
+                    f"positive int or None, got {preferred_tokens_per_block!r}."
+                )
+            kv_cache_config.tokens_per_block = preferred_tokens_per_block
+            applied["tokens_per_block"] = preferred_tokens_per_block
+
+    get_swa_scratch_reuse = getattr(model_cls,
+                                    'get_preferred_kv_cache_swa_scratch_reuse',
+                                    None)
+    if ("enable_swa_scratch_reuse" not in user_overrides
+            and get_swa_scratch_reuse is not None):
+        preferred_swa_scratch_reuse = get_swa_scratch_reuse(pretrained_config)
+        if preferred_swa_scratch_reuse is not None:
+            if type(preferred_swa_scratch_reuse) is not bool:
+                raise ValueError(
+                    f"{model_cls.__name__}."
+                    "get_preferred_kv_cache_swa_scratch_reuse() must return "
+                    f"bool or None, got {preferred_swa_scratch_reuse!r}.")
+            kv_cache_config.enable_swa_scratch_reuse = (
+                preferred_swa_scratch_reuse)
+            applied["enable_swa_scratch_reuse"] = preferred_swa_scratch_reuse
+
+    if "tokens_per_block" in applied:
+        llm_args.validate_helix_tokens_per_block()
+    return applied
+
+
+def _resolve_kv_cache_manager_v2_auto(llm_args: 'TorchLlmArgs',
+                                      model_cls: Optional[type] = None,
+                                      pretrained_config: Any = None) -> bool:
+    """Resolve the KV cache manager auto setting from the model preference.
 
     A model default of V2 is demoted to V1 for routes V2 cannot serve; an
     explicit user value otherwise wins. The compatibility arms are:
@@ -674,8 +684,7 @@ def _resolve_kv_cache_manager_v2_auto(
     -- sparse attention picks its class from the algorithm alone -- keep a V2
     manager after the demotion, so the arm does not protect them.
     """
-    setting = (llm_args.kv_cache_config.use_kv_cache_manager_v2
-               if original_setting is None else original_setting)
+    setting = llm_args.kv_cache_config.use_kv_cache_manager_v2
     if setting != "auto":
         if setting is True:
             decoding_type = _two_model_spec_dec_decoding_type(llm_args)
@@ -690,40 +699,43 @@ def _resolve_kv_cache_manager_v2_auto(
                     "one-model variant of this decoding mode.")
         return setting
 
-    kv_cache_defaults = model_defaults_dict.get("kv_cache_config", {})
-    model_default = (kv_cache_defaults.get("use_kv_cache_manager_v2", False)
-                     if isinstance(kv_cache_defaults, dict) else False)
-    if model_default == "auto":
-        model_default = False
-    if not isinstance(model_default, bool):
+    preferred_version = None
+    if model_cls is not None:
+        get_preferred = getattr(model_cls,
+                                'get_preferred_kv_cache_manager_version', None)
+        if get_preferred is not None:
+            preferred_version = get_preferred(pretrained_config)
+    if preferred_version not in (None, "V1", "V2"):
         raise ValueError(
-            "Model default kv_cache_config.use_kv_cache_manager_v2 must be "
-            f"True, False, or 'auto', got {model_default!r}.")
+            f"{model_cls.__name__}.get_preferred_kv_cache_manager_version() "
+            f"must return 'V1', 'V2', or None, got {preferred_version!r}.")
+
+    use_v2 = preferred_version == "V2"
 
     transceiver_config = llm_args.cache_transceiver_config
-    if (model_default and transceiver_config is not None
+    if (use_v2 and transceiver_config is not None
             and transceiver_config.backend is not None):
         effective_backend, _ = transceiver_config._resolve_default_backend()
         runtime = transceiver_config.transceiver_runtime
         if effective_backend != "NIXL" or runtime != "PYTHON":
             logger.info(
-                "KV cache manager V2 is the model default, but disaggregated "
+                "KV cache manager V2 is the model preference, but disaggregated "
                 "serving uses transceiver_runtime=%r with backend=%r; "
                 "falling back to V1.", runtime, effective_backend)
-            model_default = False
+            use_v2 = False
 
-    if model_default:
+    if use_v2:
         decoding_type = _two_model_spec_dec_decoding_type(llm_args)
         if decoding_type is not None:
             logger.info(
-                "KV cache manager V2 is the model default, but %s runs the "
+                "KV cache manager V2 is the model preference, but %s runs the "
                 "draft model in a separate engine and V2 sizes both KV cache "
                 "managers from the full max_gpu_total_bytes budget; falling "
                 "back to V1.", decoding_type)
-            model_default = False
+            use_v2 = False
 
-    llm_args.kv_cache_config.use_kv_cache_manager_v2 = model_default
-    return model_default
+    llm_args.kv_cache_config.use_kv_cache_manager_v2 = use_v2
+    return use_v2
 
 
 def _resolve_transceiver_runtime_auto(llm_args: 'TorchLlmArgs',
diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md
index ddf2514b7a19..0722e0bf1618 100644
--- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md
+++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md
@@ -16,7 +16,7 @@ starts.
 | Same UCX env vars (incl. the `unset UCX_TLS` cases) | `jenkins/scripts/perf/submit.py` builds the precheck commands from the **same** `ucx_tls_cmd` + `$CTX/GEN_WORKER_ENV_VARS` strings as the worker steps; `slurm_precheck_run.sh` sources the same `slurm_env_setup.sh` (the `UCX_TLS=tcp` fixup) as `slurm_run.sh`. |
 | Same instance count / parallelism | One precheck `srun` per ctx/gen server with the same `-N/--ntasks/--ntasks-per-node/--mpi=pmix` and the same node slices (`-w`) as the real server steps (`slurm_launch_draft.sh`). TP/PP/CP/attention-DP come from the same `worker_config`. |
 | Same transceiver config | `CacheTransceiverConfig(**yaml["worker_config"][role]["cache_transceiver_config"])` — the yaml block is passed through verbatim (backend, `max_tokens_in_buffer`, timeouts, ...). |
-| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves against the model class's `get_model_defaults()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. |
+| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves via `get_preferred_kv_cache_manager_version()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. |
 
 Asymmetric layouts (ctx dep4 → gen dep16, ctx pp8 → gen tp32, ...) are
 supported: data is seeded per (request, **global** layer) and constant along
diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
index 8219a2581d9c..53b79919a9cf 100644
--- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
+++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
@@ -307,8 +307,8 @@ def resolve_plan(cfg, benchmark_mode="e2e"):
         plan[f"{role}_kv_dtype"] = str(kv_cfg.get("dtype", "auto"))
         # Tri-state, matching KvCacheConfig's pydantic default: explicit
         # True/False from the yaml wins; absent means "auto", which the
-        # driver resolves against the model class's get_model_defaults() at
-        # runtime — exactly like serving (_resolve_kv_cache_manager_v2_auto).
+        # driver resolves against the model class's manager preference at
+        # runtime, exactly like serving (_resolve_kv_cache_manager_v2_auto).
         plan[f"{role}_use_kv_cache_manager_v2"] = kv_cfg.get("use_kv_cache_manager_v2", "auto")
     plan["fingerprint"] = plan_fingerprint(plan)
     return plan
diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
index 146c40173cc3..ec5ceb32ff58 100644
--- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
+++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
@@ -321,7 +321,7 @@ def resolve_model_prefs(model_dir, side, cache_cfg):
     """Mirror serving's model-preference resolution (PR #15823 semantics).
 
     - use_kv_cache_manager_v2 == "auto" (yaml absent): adopt the model
-      class's get_model_defaults() value, default False
+      class's get_preferred_kv_cache_manager_version() value, default V1
       (llm_utils._resolve_kv_cache_manager_v2_auto).
     - cache_cfg.transceiver_runtime == "auto": adopt
       model_cls.get_preferred_transceiver_runtime(), NIXL-gated, via the
@@ -349,25 +349,16 @@ def resolve_model_prefs(model_dir, side, cache_cfg):
 
     setting = side["use_kv_cache_manager_v2"]
     if setting == "auto":
-        defaults = {}
-        if model_cls is not None:
-            try:
-                defaults = model_cls.get_model_defaults(None) or {}
-            except Exception as e:  # noqa: BLE001 - model hooks may need llm_args
-                print(
-                    f"[precheck] WARNING: get_model_defaults failed ({e!r}); assuming V1",
-                    flush=True,
-                )
         try:
             # The REAL serving resolver, via the same shim pattern as the
             # runtime resolution below -- one owner for the 'auto' semantics.
             # cache_transceiver_config feeds the resolver's disagg gating
-            # (a V2 model default requires the NIXL Python transceiver).
+            # (a V2 model preference requires the NIXL Python transceiver).
             shim = types.SimpleNamespace(
                 kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto"),
                 cache_transceiver_config=cache_cfg,
             )
-            use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, defaults))
+            use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, model_cls, hf_view))
         except Exception as e:  # noqa: BLE001 - fall back like a missing model
             print(
                 f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); assuming V1", flush=True
@@ -820,7 +811,7 @@ def __init__(self, args, plan, side, comm):
         self.kvm = None
         self.xcvr = None
         self.runtime = "CPP"
-        # Resolved in setup(): "auto" needs the model class (get_model_defaults).
+        # Resolved in setup(): "auto" needs the model preference hook.
         self.use_v2 = False
         self.mapping = None
         self.llm_request_state = None
diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
index 804e4a947162..78e262fae150 100644
--- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py
+++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
@@ -515,7 +515,7 @@ def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots(
         kv_cache_config=kv_cache_config,
     )
 
-    assert _resolve_kv_cache_manager_v2_auto(llm_args, {}) is False
+    assert _resolve_kv_cache_manager_v2_auto(llm_args) is False
     assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
     with pytest.raises(ValueError, match="use_kv_cache_manager_v2=True"):
         get_kv_cache_manager_cls(
@@ -696,7 +696,7 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
         model_defaults = model_cls.get_model_defaults(llm_args)
         apply_model_defaults_to_llm_args(llm_args, model_defaults)
         _resolve_transceiver_runtime_auto(llm_args, model_cls)
-        _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults, original_setting="auto")
+        _resolve_kv_cache_manager_v2_auto(llm_args, model_cls)
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
         assert llm_args.kv_cache_config.enable_block_reuse is False
         assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON"
@@ -736,8 +736,7 @@ def test_qwen3_gdn_replay_defaults_to_v2_cache_manager(
     apply_model_defaults_to_llm_args(llm_args, model_defaults)
     _resolve_kv_cache_manager_v2_auto(
         llm_args,
-        model_defaults,
-        original_setting=manager_setting,
+        Qwen3NextForCausalLM,
     )
 
     assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected_v2
@@ -763,7 +762,7 @@ def test_kimi_defaults_to_mixed_manager(monkeypatch: pytest.MonkeyPatch) -> None
     llm_args = TorchLlmArgs(model="/tmp/dummy_model")
     model_defaults = KimiLinearForCausalLM.get_model_defaults(llm_args)
     apply_model_defaults_to_llm_args(llm_args, model_defaults)
-    resolved = _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults, original_setting="auto")
+    resolved = _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM)
 
     assert resolved is False
     assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
index 1aba431efa05..8cdefcfba19a 100644
--- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
@@ -156,13 +156,10 @@ class LlmArgs:
 
     defaults = DeepseekV4ForCausalLM.get_model_defaults(LlmArgs())
 
-    assert defaults == {
-        "kv_cache_config": {
-            "tokens_per_block": 128,
-            "use_kv_cache_manager_v2": True,
-            "enable_swa_scratch_reuse": True,
-        }
-    }
+    assert defaults == {}
+    assert DeepseekV4ForCausalLM.get_preferred_kv_cache_tokens_per_block() == 128
+    assert DeepseekV4ForCausalLM.get_preferred_kv_cache_swa_scratch_reuse() is True
+    assert DeepseekV4ForCausalLM.get_preferred_kv_cache_manager_version() == "V2"
 
 
 def test_deepseek_v4_weight_remap_for_mxfp4_routed_experts():
diff --git a/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py b/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py
index 232893dfca3e..bd1c64562311 100644
--- a/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py
+++ b/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py
@@ -611,6 +611,13 @@ def test_prefers_python_transceiver(self):
             KimiK25ForConditionalGeneration.get_preferred_transceiver_runtime(), "PYTHON"
         )
 
+    def test_prefers_kv_cache_manager_v2(self):
+        """Kimi-K2.5 prefers KV cache manager V2."""
+        self.assertEqual(
+            KimiK25ForConditionalGeneration.get_preferred_kv_cache_manager_version(),
+            "V2",
+        )
+
 
 # ---------------------------------------------------------------------------
 # E2E Smoke Test — requires GPU + model checkpoint
diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py
index a9e71de4843a..a6d6fd755976 100644
--- a/tests/unittest/_torch/test_model_config.py
+++ b/tests/unittest/_torch/test_model_config.py
@@ -7,19 +7,12 @@
 
 import pytest
 import torch
-import transformers
 
 from tensorrt_llm._torch.model_config import _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT, ModelConfig
 from tensorrt_llm._torch.pyexecutor.model_loader import (
     validate_and_set_kv_cache_quant,
     validate_encoder_decoder_kv_cache_config,
 )
-from tensorrt_llm.llmapi.llm_args import (
-    DeepSeekSparseAttentionConfig,
-    DeepSeekV4SparseAttentionConfig,
-    TorchLlmArgs,
-)
-from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
 
@@ -240,117 +233,6 @@ def test_deepseek_v4_missing_compress_ratios_raises(tmp_path, monkeypatch):
         ModelConfig.from_pretrained(str(tmp_path))
 
 
-@pytest.mark.parametrize(
-    "user_index_topk,expected_index_topk",
-    [(None, 2048), (1024, 1024)],
-)
-def test_deepseek_v32_model_defaults_preserve_index_topk_override(
-    tmp_path, monkeypatch, user_index_topk, expected_index_topk
-):
-    """Applying unrelated model defaults must preserve nested user intent."""
-    from tensorrt_llm._torch import model_config as model_config_module
-
-    pretrained_config = transformers.PretrainedConfig(
-        architectures=["DeepseekV32ForCausalLM"],
-        index_n_heads=64,
-        index_head_dim=128,
-        index_topk=2048,
-        indexer_rope_interleave=False,
-        dtype="bfloat16",
-        intermediate_size=1024,
-    )
-    monkeypatch.setattr(
-        model_config_module,
-        "load_pretrained_config",
-        lambda *args, **kwargs: pretrained_config,
-    )
-
-    sparse_kwargs = {}
-    if user_index_topk is not None:
-        sparse_kwargs["index_topk"] = user_index_topk
-    llm_args = TorchLlmArgs(
-        model=str(tmp_path),
-        sparse_attention_config=DeepSeekSparseAttentionConfig(**sparse_kwargs),
-    )
-    apply_model_defaults_to_llm_args(
-        llm_args, {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
-    )
-
-    model_config = ModelConfig.from_pretrained(
-        str(tmp_path), sparse_attention_config=llm_args.sparse_attention_config
-    )
-    sparse_config = model_config.sparse_attention_config
-    assert sparse_config.index_topk == expected_index_topk
-    assert sparse_config.needs_separate_short_long_cuda_graphs()
-    assert sparse_config.seq_len_threshold == expected_index_topk
-
-
-def test_model_defaults_preserve_explicit_nested_none(tmp_path):
-    """An explicit None must remain distinct from an omitted nested field."""
-    llm_args = TorchLlmArgs(
-        model=str(tmp_path),
-        sparse_attention_config=DeepSeekSparseAttentionConfig(index_topk=None),
-    )
-    assert "index_topk" in llm_args.sparse_attention_config.model_fields_set
-
-    apply_model_defaults_to_llm_args(
-        llm_args, {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
-    )
-
-    assert llm_args.sparse_attention_config.index_topk is None
-    assert "index_topk" in llm_args.sparse_attention_config.model_fields_set
-
-
-@pytest.mark.parametrize(
-    "sparse_kwargs,expected_index_topk",
-    [
-        pytest.param({}, 1024, id="omitted"),
-        pytest.param({"index_topk": 512}, 512, id="explicit"),
-    ],
-)
-def test_deepseek_v4_model_defaults_preserve_index_topk_override(
-    tmp_path, monkeypatch, sparse_kwargs, expected_index_topk
-):
-    """V4's default top-k must not hide a checkpoint or explicit override."""
-    from tensorrt_llm._torch import model_config as model_config_module
-
-    pretrained_config = transformers.PretrainedConfig(
-        architectures=["DeepseekV4ForCausalLM"],
-        index_n_heads=64,
-        index_head_dim=128,
-        index_topk=1024,
-        indexer_rope_interleave=False,
-        compress_ratios=[1, 4, 128, 4],
-        window_size=128,
-        num_hidden_layers=4,
-        dtype="bfloat16",
-        intermediate_size=1024,
-    )
-    monkeypatch.setattr(
-        model_config_module,
-        "load_pretrained_config",
-        lambda *args, **kwargs: pretrained_config,
-    )
-
-    llm_args = TorchLlmArgs(
-        model=str(tmp_path),
-        sparse_attention_config=DeepSeekV4SparseAttentionConfig(
-            index_head_dim=128, **sparse_kwargs
-        ),
-    )
-    apply_model_defaults_to_llm_args(
-        llm_args, {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
-    )
-    assert ("index_topk" in llm_args.sparse_attention_config.model_fields_set) is bool(
-        sparse_kwargs
-    )
-
-    model_config = ModelConfig.from_pretrained(
-        str(tmp_path), sparse_attention_config=llm_args.sparse_attention_config
-    )
-    assert model_config.sparse_attention_config.index_topk == expected_index_topk
-
-
 def test_model_config_sets_is_encoder_decoder_from_pretrained_config():
     model_config = ModelConfig(
         pretrained_config=make_pretrained_config(
diff --git a/tests/unittest/llmapi/test_config_database.py b/tests/unittest/llmapi/test_config_database.py
index ac88d5c58a69..b1f836346adf 100644
--- a/tests/unittest/llmapi/test_config_database.py
+++ b/tests/unittest/llmapi/test_config_database.py
@@ -15,6 +15,7 @@
 """L0 tests for validating curated and database YAML configs against TorchLlmArgs."""
 
 from pathlib import Path
+from types import SimpleNamespace
 from unittest import mock
 
 import pytest
@@ -28,7 +29,11 @@
 from tensorrt_llm.commands.serve import _apply_fastapi_middlewares
 from tensorrt_llm.commands.serve import main as serve_main
 from tensorrt_llm.llmapi import llm_args as llm_args_module
-from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args
+from tensorrt_llm.llmapi.llm_utils import (
+    _apply_kv_cache_config_preferences,
+    _resolve_kv_cache_manager_v2_auto,
+    apply_model_defaults_to_llm_args,
+)
 
 from .yaml_validation_harness import (
     assert_no_default_valued_leaves,
@@ -106,12 +111,17 @@ def _get_default_values_for_config(config_path: Path) -> dict:
     model = entry["model"]
     arch = entry["arch"]
     model_cls = get_registered_model_class(arch)
-    if not model_cls or not hasattr(model_cls, "get_model_defaults"):
+    if not model_cls:
         return global_default
 
     base_args = llm_args_module.TorchLlmArgs(model=model, skip_tokenizer_init=True)
-    model_defaults = model_cls.get_model_defaults(base_args)
-    apply_model_defaults_to_llm_args(base_args, model_defaults)
+    pretrained_config = SimpleNamespace(architectures=[arch])
+    _apply_kv_cache_config_preferences(base_args, model_cls, pretrained_config)
+    get_model_defaults = getattr(model_cls, "get_model_defaults", None)
+    model_defaults = get_model_defaults(base_args) if get_model_defaults else {}
+    if model_defaults:
+        apply_model_defaults_to_llm_args(base_args, model_defaults)
+    _resolve_kv_cache_manager_v2_auto(base_args, model_cls, pretrained_config)
     return base_args.model_dump(mode="json")
 
 
@@ -126,14 +136,18 @@ def get_config_id(config_path: Path) -> str:
     return str(config_path.relative_to(CONFIG_ROOT))
 
 
-def test_model_nested_defaults_preserve_global_nested_defaults(monkeypatch):
-    config_key = "test-only/deepseek-v3.yaml"
+@pytest.mark.parametrize(
+    "arch",
+    ["DeepseekV3ForCausalLM", "KimiK25ForConditionalGeneration"],
+)
+def test_model_preferences_preserve_global_nested_defaults(monkeypatch, arch):
+    config_key = f"test-only/{arch}.yaml"
     monkeypatch.setitem(
         _CONFIG_PATH_TO_ENTRY,
         config_key,
         {
             "model": "dummy/model",
-            "arch": "DeepseekV3ForCausalLM",
+            "arch": arch,
         },
     )
 
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index 2fa8772fd073..79fe5737df75 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -58,7 +58,8 @@
                                           update_llm_args_with_extra_dict,
                                           update_llm_args_with_extra_options)
 # fmt: on
-from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
+from tensorrt_llm.llmapi.llm_utils import (_apply_kv_cache_config_preferences,
+                                           _resolve_kv_cache_manager_v2_auto,
                                            _resolve_transceiver_runtime_auto,
                                            apply_model_defaults_to_llm_args)
 from tensorrt_llm.llmapi.mm_encoder import MultimodalEncoder
@@ -513,34 +514,49 @@ def test_post_processor_hook_rejected_with_skip_tokenizer_init():
 
 
 @pytest.mark.cpu_only
-class TestModelDefaults:
-    """Test suite for model-specific default overrides functionality."""
+class TestKvCacheManagerV2AutoResolution:
+    """Test model preferences for the KV cache manager version."""
 
-    def test_compute_applied_llm_defaults_simple_field(self):
-        model_defaults = {"enable_chunked_prefill": False}
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
-        applied = apply_model_defaults_to_llm_args(llm_args, model_defaults)
-        assert applied == model_defaults
+    class _PreferV2:
+
+        @classmethod
+        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+            return "V2"
+
+    class _PreferV1:
+
+        @classmethod
+        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+            return "V1"
+
+    class _InvalidPreference:
+
+        @classmethod
+        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+            return "V3"
+
+    class _ArchitectureSensitive:
+
+        @classmethod
+        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+            architectures = getattr(pretrained_config, "architectures", [])
+            return "V2" if architectures == ["PreferV2ForCausalLM"] else "V1"
 
     @pytest.mark.parametrize("explicit_auto", [False, True])
-    def test_kv_cache_manager_v2_auto_uses_model_default(self, explicit_auto):
+    def test_auto_uses_model_preference(self, explicit_auto):
         kv_cache_config = (KvCacheConfig(use_kv_cache_manager_v2="auto")
                            if explicit_auto else KvCacheConfig())
         llm_args = TorchLlmArgs(model="/tmp/dummy_model",
                                 kv_cache_config=kv_cache_config)
-        model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
 
-        apply_model_defaults_to_llm_args(llm_args, model_defaults)
-        _resolve_kv_cache_manager_v2_auto(llm_args,
-                                          model_defaults,
-                                          original_setting="auto")
+        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
 
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
 
-    def test_kv_cache_manager_v2_auto_falls_back_to_false(self):
+    def test_auto_without_preference_falls_back_to_v1(self):
         llm_args = TorchLlmArgs(model="/tmp/dummy_model")
 
-        _resolve_kv_cache_manager_v2_auto(llm_args, {})
+        _resolve_kv_cache_manager_v2_auto(llm_args)
 
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
 
@@ -552,54 +568,158 @@ def test_kv_cache_manager_v2_auto_falls_back_to_false(self):
             ("MPI", None),
         ],
     )
-    def test_kv_cache_manager_v2_auto_falls_back_for_incompatible_disagg(
-            self, backend, runtime):
+    def test_auto_v2_falls_back_for_incompatible_disagg(self, backend, runtime):
         llm_args = TorchLlmArgs(
             model="/tmp/dummy_model",
             cache_transceiver_config=CacheTransceiverConfig(
                 backend=backend, transceiver_runtime=runtime),
         )
-        model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
 
-        apply_model_defaults_to_llm_args(llm_args, model_defaults)
-        _resolve_kv_cache_manager_v2_auto(llm_args,
-                                          model_defaults,
-                                          original_setting="auto")
+        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
 
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
 
-    def test_kv_cache_manager_v2_auto_keeps_python_nixl_model_default(self):
+    def test_auto_v2_keeps_python_nixl_preference(self):
         llm_args = TorchLlmArgs(
             model="/tmp/dummy_model",
             cache_transceiver_config=CacheTransceiverConfig(
                 backend="NIXL", transceiver_runtime="PYTHON"),
         )
-        model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
 
-        apply_model_defaults_to_llm_args(llm_args, model_defaults)
-        _resolve_kv_cache_manager_v2_auto(llm_args,
-                                          model_defaults,
-                                          original_setting="auto")
+        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
 
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
 
-    @pytest.mark.parametrize("user_setting", [False, True])
-    def test_kv_cache_manager_v2_explicit_value_overrides_model_default(
-            self, user_setting):
+    @pytest.mark.parametrize(("user_setting", "model_cls"), [
+        pytest.param(False, _PreferV2, id="explicit-v1"),
+        pytest.param(True, _PreferV1, id="explicit-v2"),
+    ])
+    def test_explicit_value_overrides_model_preference(self, user_setting,
+                                                       model_cls):
         llm_args = TorchLlmArgs(
             model="/tmp/dummy_model",
             kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=user_setting))
-        model_defaults = {
-            "kv_cache_config": {
-                "use_kv_cache_manager_v2": not user_setting
-            }
-        }
 
-        apply_model_defaults_to_llm_args(llm_args, model_defaults)
-        _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults)
+        _resolve_kv_cache_manager_v2_auto(llm_args, model_cls)
 
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is user_setting
 
+    def test_pretrained_config_is_forwarded_to_preference(self):
+        from types import SimpleNamespace
+
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
+        pretrained_config = SimpleNamespace(
+            architectures=["PreferV2ForCausalLM"])
+
+        _resolve_kv_cache_manager_v2_auto(llm_args, self._ArchitectureSensitive,
+                                          pretrained_config)
+
+        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
+
+    def test_invalid_model_preference_raises(self):
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
+
+        with pytest.raises(ValueError, match="must return 'V1', 'V2', or None"):
+            _resolve_kv_cache_manager_v2_auto(llm_args, self._InvalidPreference)
+
+    @pytest.mark.parametrize("architecture", [
+        "DeepseekV3ForCausalLM",
+        "DeepseekV32ForCausalLM",
+        "GlmMoeDsaForCausalLM",
+        "MistralLarge3ForCausalLM",
+        "DeepseekV4ForCausalLM",
+        "KimiK25ForConditionalGeneration",
+        "MiniMaxM2ForCausalLM",
+        "NemotronHForCausalLM",
+        "NemotronHPuzzleForCausalLM",
+        "Qwen3NextForCausalLM",
+        "Qwen3_5MoeForCausalLM",
+        "Qwen3_5ForCausalLM",
+        "Qwen3_5MoeForConditionalGeneration",
+        "Qwen3_5ForConditionalGeneration",
+    ])
+    def test_registered_models_prefer_v2(self, architecture):
+        from tensorrt_llm._torch.models.modeling_utils import \
+            MODEL_CLASS_MAPPING
+
+        model_cls = MODEL_CLASS_MAPPING[architecture]
+        assert model_cls.get_preferred_kv_cache_manager_version() == "V2"
+
+
+@pytest.mark.cpu_only
+class TestKvCacheConfigPreferences:
+    """Test scalar KV-cache preferences outside generic model defaults."""
+
+    def test_deepseek_v4_preferences_preserve_unrelated_nested_config(self):
+        from tensorrt_llm._torch.models.modeling_deepseekv4 import \
+            DeepseekV4ForCausalLM
+
+        sparse_config = DeepSeekV4SparseAttentionConfig(index_head_dim=128)
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model",
+                                sparse_attention_config=sparse_config)
+
+        applied = _apply_kv_cache_config_preferences(llm_args,
+                                                     DeepseekV4ForCausalLM)
+
+        assert applied == {
+            "tokens_per_block": 128,
+            "enable_swa_scratch_reuse": True,
+        }
+        assert llm_args.kv_cache_config.tokens_per_block == 128
+        assert llm_args.kv_cache_config.enable_swa_scratch_reuse is True
+        assert llm_args.sparse_attention_config is sparse_config
+        assert sparse_config.model_dump(exclude_unset=True) == {
+            "index_head_dim": 128
+        }
+
+    def test_explicit_values_override_deepseek_v4_preferences(self):
+        from tensorrt_llm._torch.models.modeling_deepseekv4 import \
+            DeepseekV4ForCausalLM
+
+        llm_args = TorchLlmArgs(
+            model="/tmp/dummy_model",
+            kv_cache_config=KvCacheConfig(tokens_per_block=32,
+                                          enable_swa_scratch_reuse=False),
+        )
+
+        applied = _apply_kv_cache_config_preferences(llm_args,
+                                                     DeepseekV4ForCausalLM)
+
+        assert applied == {}
+        assert llm_args.kv_cache_config.tokens_per_block == 32
+        assert llm_args.kv_cache_config.enable_swa_scratch_reuse is False
+
+
+@pytest.mark.cpu_only
+class TestModelDefaults:
+    """Test suite for model-specific default overrides functionality."""
+
+    def test_compute_applied_llm_defaults_simple_field(self):
+        model_defaults = {"enable_chunked_prefill": False}
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
+        applied = apply_model_defaults_to_llm_args(llm_args, model_defaults)
+        assert applied == model_defaults
+
+    def test_kv_cache_manager_version_preference_is_rejected(self):
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
+
+        with pytest.raises(ValueError,
+                           match="get_preferred_kv_cache_manager_version"):
+            apply_model_defaults_to_llm_args(
+                llm_args,
+                {"kv_cache_config": {
+                    "use_kv_cache_manager_v2": True
+                }})
+
+    def test_kv_cache_config_object_cannot_bypass_preference_guard(self):
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
+
+        with pytest.raises(ValueError, match="get_preferred_"):
+            apply_model_defaults_to_llm_args(llm_args, {
+                "kv_cache_config":
+                KvCacheConfig(use_kv_cache_manager_v2=True)
+            })
+
     @pytest.mark.parametrize(
         "defaults_dict,should_raise,error_contains",
         [
@@ -3572,12 +3692,16 @@ def test_rejects_non_bool(self):
 
 
 class _PreferPythonTransceiverModel:
-    """Fake model class opting into the Python transceiver."""
+    """Fake model class preferring Python transceiver and KV manager V2."""
 
     @classmethod
     def get_model_defaults(cls, llm_args):
         return {}
 
+    @classmethod
+    def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+        return "V2"
+
     @classmethod
     def get_preferred_transceiver_runtime(cls, pretrained_config=None):
         return "PYTHON"
@@ -3694,6 +3818,11 @@ def test_hybrid_snapshot_policy_controls_block_reuse(
             assert "no Mamba state snapshot policy" in warnings[0]
 
     def test_hybrid_fixed_snapshot_rejects_auto_resolved_v1(self, monkeypatch):
+        from tensorrt_llm._torch.models.modeling_utils import \
+            MODEL_CLASS_MAPPING
+
+        monkeypatch.setitem(MODEL_CLASS_MAPPING, "Qwen3NextForCausalLM",
+                            _NoModelDefaults)
         args = TorchLlmArgs(
             model="/tmp/dummy_model",
             kv_cache_config=KvCacheConfig(
@@ -3793,13 +3922,9 @@ def test_default_backend_env_override_falls_back_to_v1_cpp(
             monkeypatch.delenv(env_var, raising=False)
         monkeypatch.setenv(backend_env, "1")
         args = self._disagg_args(backend="DEFAULT")
-        model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
-        apply_model_defaults_to_llm_args(args, model_defaults)
 
         _resolve_transceiver_runtime_auto(args, _PreferPythonTransceiverModel)
-        _resolve_kv_cache_manager_v2_auto(args,
-                                          model_defaults,
-                                          original_setting="auto")
+        _resolve_kv_cache_manager_v2_auto(args, _PreferPythonTransceiverModel)
 
         assert args.cache_transceiver_config.transceiver_runtime is None
         assert args.kv_cache_config.use_kv_cache_manager_v2 is False
@@ -3880,6 +4005,7 @@ def test_model_loader_resolves_auto_without_checkpoint_loader(self):
         ModelLoader.load_config_and_apply_defaults("/tmp/dummy_model", args,
                                                    None)
         assert args.cache_transceiver_config.transceiver_runtime is None
+        assert args.kv_cache_config.use_kv_cache_manager_v2 is False
 
     @staticmethod
     def _fake_checkpoint_loader(architectures):
@@ -3904,6 +4030,7 @@ def test_model_loader_full_chain_adopts_model_preference(self, monkeypatch):
         model_loader_mod.ModelLoader.load_config_and_apply_defaults(
             "/tmp/dummy_model", args, fake_loader)
         assert args.cache_transceiver_config.transceiver_runtime == "PYTHON"
+        assert args.kv_cache_config.use_kv_cache_manager_v2 is True
 
     @pytest.mark.parametrize("arch,expected", [
         ("ModelAForCausalLM", "PYTHON"),
@@ -3963,6 +4090,7 @@ def get_preferred_transceiver_runtime(cls, pretrained_config=None):
         model_loader_mod.ModelLoader.load_config_and_apply_defaults(
             "/tmp/dummy_model", args, fake_loader)
         assert args.cache_transceiver_config.transceiver_runtime == "PYTHON"
+        assert args.kv_cache_config.use_kv_cache_manager_v2 is True
 
     def test_model_loader_full_chain_aggregated_stays_none(self, monkeypatch):
         """Full chain in aggregated mode: no transceiver config appears."""
@@ -3978,6 +4106,7 @@ def test_model_loader_full_chain_aggregated_stays_none(self, monkeypatch):
         model_loader_mod.ModelLoader.load_config_and_apply_defaults(
             "/tmp/dummy_model", args, fake_loader)
         assert args.cache_transceiver_config is None
+        assert args.kv_cache_config.use_kv_cache_manager_v2 is True
 
     def test_resolve_default_backend_env_priority(self, monkeypatch):
         env_vars = ("TRTLLM_USE_NIXL_KVCACHE", "TRTLLM_USE_UCX_KVCACHE",
@@ -4004,8 +4133,8 @@ def test_resolve_default_backend_env_priority(self, monkeypatch):
             backend="UCX")._resolve_default_backend() == ("UCX", None)
 
 
-class TestDeepseekTransceiverPreference:
-    """Per-architecture preferred KV-cache transceiver runtime.
+class TestDeepseekRuntimePreferences:
+    """DeepSeek KV-cache manager and transceiver preferences.
 
     DeepseekV3ForCausalLM and DeepseekV32ForCausalLM prefer the Python KV-cache
     transceiver, while GlmMoeDsaForCausalLM (GLM 5.2) requires the C++ transceiver
@@ -4041,6 +4170,12 @@ def test_prefers_python_without_config(self) -> None:
         assert DeepseekV3ForCausalLM.get_preferred_transceiver_runtime(
         ) == "PYTHON"
 
+    def test_prefers_kv_cache_manager_v2(self) -> None:
+        from tensorrt_llm._torch.models.modeling_deepseekv3 import \
+            DeepseekV3ForCausalLM
+        assert DeepseekV3ForCausalLM.get_preferred_kv_cache_manager_version(
+        ) == "V2"
+
     def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None:
         """DeepseekV3ForCausalLM on NIXL adopts the Python transceiver from 'auto'."""
         from tensorrt_llm._torch.models.modeling_deepseekv3 import \
diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py
index 99dc35706f03..252493d1dba2 100644
--- a/tests/unittest/others/test_cache_transceiver_precheck_config.py
+++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py
@@ -280,7 +280,7 @@ def test_addr_file_owner_only(self, tmp_path):
 
 def test_use_kv_cache_manager_v2_flags():
     # Absent -> "auto" (the driver resolves it against the model's
-    # get_model_defaults at runtime, like serving).
+    # get_preferred_kv_cache_manager_version at runtime, like serving).
     plan = pcfg.resolve_plan(_disagg_yaml())
     assert plan["ctx_use_kv_cache_manager_v2"] == "auto"
     assert plan["gen_use_kv_cache_manager_v2"] == "auto"
diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py
index ed060d1e9053..6d336998b06d 100644
--- a/tests/unittest/others/test_cache_transceiver_precheck_run.py
+++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py
@@ -389,11 +389,11 @@ def test_kv_cache_manager_ctor_kwargs(self, api, manager_attr):
     def test_serving_resolvers(self, api):
         import inspect
 
-        # The driver calls resolve_kv_cache_manager_v2_auto(shim, defaults):
-        # the first two params are fixed, anything added later must default.
+        # The driver calls the serving resolver with a shim, model class, and
+        # pretrained config. Only the shim is required.
         v2 = inspect.signature(api.resolve_kv_cache_manager_v2_auto).parameters
-        assert list(v2)[:2] == ["llm_args", "model_defaults_dict"]
-        assert all(p.default is not inspect.Parameter.empty for p in list(v2.values())[2:])
+        assert list(v2)[:3] == ["llm_args", "model_cls", "pretrained_config"]
+        assert all(p.default is not inspect.Parameter.empty for p in list(v2.values())[1:])
         rt = inspect.signature(api.resolve_transceiver_runtime_auto).parameters
         assert list(rt)[:1] == ["llm_args"] and len(rt) >= 3
 

From f3146d7ae71aee295c5c4503aad4744b4a528c2f Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sun, 9 Aug 2026 09:11:33 -0700
Subject: [PATCH 12/12] [None][refactor] finish KV cache preference migration

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../_torch/models/modeling_deepseekv3.py      |   0
 .../_torch/models/modeling_deepseekv4.py      |  24 +-
 .../_torch/models/modeling_gpt_oss.py         |  28 +-
 .../_torch/models/modeling_nemotron_h.py      |   4 +-
 tensorrt_llm/_torch/models/modeling_utils.py  |  35 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   2 +-
 .../_torch/pyexecutor/model_loader.py         |   9 +-
 tensorrt_llm/llmapi/llm_args.py               |   4 +-
 tensorrt_llm/llmapi/llm_utils.py              |  70 +---
 .../defs/accuracy/test_llm_api_pytorch.py     |   5 +-
 .../run_precheck.py                           |   3 +-
 .../attention/sparse/dsa/test_dsa_indexer.py  | 336 +++++-------------
 .../test_kv_cache_v2_extra_buffers.py         |  84 ++---
 .../executor/test_mamba_cache_manager.py      |  95 +----
 .../modeling/test_modeling_deepseekv4.py      |  16 +-
 .../_torch/modeling/test_modeling_gpt_oss.py  |  23 +-
 .../_torch/modeling/test_modeling_kimi_k25.py |   7 -
 .../modeling/test_modeling_qwen3_5_vl_moe.py  |   6 +-
 .../_torch/speculative/test_eagle3.py         |  60 +---
 tests/unittest/_torch/test_model_config.py    |   3 -
 tests/unittest/llmapi/test_config_database.py |  45 +--
 tests/unittest/llmapi/test_llm_args.py        | 312 ++++++----------
 .../test_cache_transceiver_precheck_run.py    |  20 ++
 23 files changed, 338 insertions(+), 853 deletions(-)
 mode change 100644 => 100755 tensorrt_llm/_torch/models/modeling_deepseekv3.py

diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py
old mode 100644
new mode 100755
diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py
index 79e059a0a857..7bbcd0ce9186 100644
--- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py
+++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py
@@ -31,7 +31,10 @@
 import copy
 import math
 import os
-from typing import Dict, List, Literal, Optional
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional
+
+if TYPE_CHECKING:
+    from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
 
 import torch
 import triton
@@ -2530,18 +2533,13 @@ def forward(
 @register_auto_model("DeepseekV4ForCausalLM")
 class DeepseekV4ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV4Model, PretrainedConfig]):
     @classmethod
-    def get_preferred_kv_cache_tokens_per_block(
-        cls, pretrained_config: object | None = None
-    ) -> int:
-        """Prefer 128-token KV-cache blocks for DeepSeek-V4."""
-        return 128
-
-    @classmethod
-    def get_preferred_kv_cache_swa_scratch_reuse(
-        cls, pretrained_config: object | None = None
-    ) -> bool:
-        """Prefer SWA scratch reuse for DeepSeek-V4."""
-        return True
+    def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
+        return {
+            "kv_cache_config": {
+                "tokens_per_block": 128,
+                "enable_swa_scratch_reuse": True,
+            }
+        }
 
     @classmethod
     def get_preferred_kv_cache_manager_version(
diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py
index 6ba492929cd5..d333379c78a9 100644
--- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py
+++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py
@@ -1,4 +1,7 @@
-from typing import TYPE_CHECKING, Any, Dict, Literal, Optional
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Dict, Literal, Optional
 
 import torch
 from torch import nn
@@ -32,9 +35,6 @@
 from .modeling_speculative import SpecDecOneEngineForCausalLM
 from .modeling_utils import DecoderModel, filter_weights, register_auto_model
 
-if TYPE_CHECKING:
-    from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
-
 # Use TinyGEMM when the number of tokens is not larger than this threshold
 MIN_LATENCY_TINYGEMM_NUM_TOKENS = 128
 
@@ -556,8 +556,10 @@ def forward(
 class GptOssForCausalLM(SpecDecOneEngineForCausalLM[Transformer, GptOssConfig]):
 
     @classmethod
-    def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
-        """Select KV cache manager V2 by default.
+    def get_preferred_kv_cache_manager_version(cls,
+                                               pretrained_config: Any = None
+                                               ) -> Literal["V2"]:
+        """Prefer KV cache manager V2 for the model's VSWA layout.
 
         GPT-OSS applies a sliding window to every other layer
         (see ``AttentionBlock.__init__``), so the KV cache is VSWA: two
@@ -566,15 +568,13 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
         sliding-window and full-attention pools independently instead of
         statically dividing memory between them.
 
-        Users keep full control: an explicit
-        ``kv_cache_config.use_kv_cache_manager_v2`` otherwise wins over this
-        default. Two-model speculative decoding is the exception, since V2
-        sizes both the target and draft KV cache managers from the full
-        budget: ``auto`` demotes this default to V1 there and an explicit
-        ``True`` is rejected, both in
-        ``llm_utils._resolve_kv_cache_manager_v2_auto``.
+        The preference is adopted only when the user leaves
+        ``kv_cache_config.use_kv_cache_manager_v2`` at ``"auto"``. Two-model
+        speculative decoding demotes it to V1 because V2 sizes both the target
+        and draft KV cache managers from the full budget; an explicit ``True``
+        is rejected by ``llm_utils._resolve_kv_cache_manager_v2_auto``.
         """
-        return {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+        return "V2"
 
     @classmethod
     def get_preferred_transceiver_runtime(
diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
index bd5915343b24..7bc4e125c968 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
@@ -985,8 +985,8 @@ def load_weights(self,
     def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict:
         """Model-specific defaults for NemotronH.
 
-        Uses KV cache manager V2 for the hybrid state layout. Block reuse
-        remains opt-in because it also requires a Mamba snapshot policy.
+        Block reuse remains opt-in because it also requires a Mamba snapshot
+        policy.
         """
         return {
             "kv_cache_config": {
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index fd6feb66d520..bb81db4cf27a 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -658,11 +658,10 @@ def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict:
         The returned dict is deep-merged with the user's llm_args, with
         user-set values taking priority over these defaults.
 
-        Note: ``cache_transceiver_config`` and
-        ``kv_cache_config.use_kv_cache_manager_v2`` are rejected here
-        (enforced at load time). Use
-        :meth:`get_preferred_transceiver_runtime` and
-        :meth:`get_preferred_kv_cache_manager_version` instead.
+        Note: ``cache_transceiver_config`` is rejected here (enforced at
+        load time) — the deep-merge could materialize or silently enable a
+        transceiver config the user did not turn on. Use
+        :meth:`get_preferred_transceiver_runtime` instead.
         """
         return {}
 
@@ -674,7 +673,7 @@ def get_preferred_kv_cache_manager_version(
 
         The preference is adopted only when the user leaves
         ``kv_cache_config.use_kv_cache_manager_v2`` at ``"auto"``. Return
-        ``None`` to defer to the global default (V1).
+        ``None`` to use the built-in V1 fallback.
 
         Args:
             pretrained_config: The loaded Hugging Face config. Shared model
@@ -683,30 +682,6 @@ def get_preferred_kv_cache_manager_version(
         """
         return None
 
-    @classmethod
-    def get_preferred_kv_cache_tokens_per_block(cls,
-                                                pretrained_config: Any = None
-                                                ) -> Optional[int]:
-        """Return the model's preferred KV-cache block size.
-
-        The preference is adopted only when the user does not explicitly set
-        ``kv_cache_config.tokens_per_block``. Return ``None`` to keep the
-        global default.
-        """
-        return None
-
-    @classmethod
-    def get_preferred_kv_cache_swa_scratch_reuse(cls,
-                                                 pretrained_config: Any = None
-                                                 ) -> Optional[bool]:
-        """Return the model's preferred SWA scratch-reuse setting.
-
-        The preference is adopted only when the user does not explicitly set
-        ``kv_cache_config.enable_swa_scratch_reuse``. Return ``None`` to keep
-        the global default.
-        """
-        return None
-
     @classmethod
     def get_preferred_transceiver_runtime(
             cls,
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index d19e0d705b29..ea73b8e37479 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -651,7 +651,7 @@ def _validate_or_fallback_kv_cache_manager_v2(
                         f"{incompat_str}; CppMambaHybridCacheManager does not "
                         "provide a compatible fallback. Use max_beam_width=1 "
                         "and disable the KV connector.")
-                # Plain V2 (explicitly enabled or selected by a model default):
+                # Plain V2 (explicitly enabled or selected by a model preference):
                 # V2 was a preference, not a structural requirement, so we can
                 # safely fall back to V1.
                 logger.warning(
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index d5a216136148..30e775a15d35 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -23,8 +23,7 @@
                                           ExecutorMemoryType,
                                           ModelExpressConfig,
                                           SparseAttentionConfig, TorchLlmArgs)
-from tensorrt_llm.llmapi.llm_utils import (_apply_kv_cache_config_preferences,
-                                           _resolve_kv_cache_manager_v2_auto,
+from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
                                            _resolve_transceiver_runtime_auto,
                                            apply_model_defaults_to_llm_args)
 from tensorrt_llm.logger import logger
@@ -443,12 +442,6 @@ def load_config_and_apply_defaults(
             preference_cls = get_registered_model_class(
                 architectures[0]) or model_cls
 
-        applied_kv_preferences = _apply_kv_cache_config_preferences(
-            llm_args, preference_cls, config.pretrained_config)
-        if applied_kv_preferences:
-            logger.info("Applied KV-cache preferences for %s: %s",
-                        preference_cls.__name__, applied_kv_preferences)
-
         # model_cls is None when the architecture is unknown/unsupported.
         model_defaults = {}
         if model_cls and hasattr(model_cls, 'get_model_defaults'):
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 1afb2cc25994..b7527d0ffd95 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -3913,8 +3913,8 @@ class KvCacheConfig(StrictBaseModel, PybindMirror):
         status="prototype",
         description=
         "Whether to use the KV cache manager v2 (experimental). 'auto' uses "
-        "the model-specific default and falls back to False when the model "
-        "does not specify one.")
+        "the model-specific preference and falls back to False when the model "
+        "does not declare one.")
 
     # This is a pure python field, not a pybind field. It is only for the Pytorch backend.
     enable_swa_scratch_reuse: bool = Field(
diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py
index 0de67d373542..a21241b6bcf0 100644
--- a/tensorrt_llm/llmapi/llm_utils.py
+++ b/tensorrt_llm/llmapi/llm_utils.py
@@ -549,20 +549,6 @@ def apply_model_defaults_to_llm_args(
             "config the user did not turn on. Declare a runtime preference "
             "via get_preferred_transceiver_runtime() instead.")
 
-    kv_cache_defaults = model_defaults_dict.get("kv_cache_config")
-    if (kv_cache_defaults is not None
-            and not isinstance(kv_cache_defaults, dict)):
-        raise ValueError(
-            "Model default 'kv_cache_config' must be a dict. Runtime "
-            "preferences must use the corresponding get_preferred_* hook.")
-    if (kv_cache_defaults is not None
-            and "use_kv_cache_manager_v2" in kv_cache_defaults):
-        raise ValueError(
-            "Model defaults must not contain "
-            "'kv_cache_config.use_kv_cache_manager_v2'. Declare a runtime "
-            "preference via get_preferred_kv_cache_manager_version() "
-            "instead.")
-
     user_overrides = llm_args.model_dump(exclude_unset=True)
     base_state = llm_args.model_dump()
     merged_state = _deep_merge(base_state, model_defaults_dict, user_overrides)
@@ -611,62 +597,12 @@ def _two_model_spec_dec_decoding_type(
     return getattr(spec_config, "decoding_type", "speculative decoding")
 
 
-def _apply_kv_cache_config_preferences(
-        llm_args: 'TorchLlmArgs',
-        model_cls: Optional[type] = None,
-        pretrained_config: Any = None) -> Dict[str, Any]:
-    """Apply scalar KV-cache preferences that the user left unspecified."""
-    if model_cls is None:
-        return {}
-
-    kv_cache_config = llm_args.kv_cache_config
-    user_overrides = kv_cache_config.model_dump(exclude_unset=True)
-    applied = {}
-
-    get_tokens_per_block = getattr(model_cls,
-                                   'get_preferred_kv_cache_tokens_per_block',
-                                   None)
-    if ("tokens_per_block" not in user_overrides
-            and get_tokens_per_block is not None):
-        preferred_tokens_per_block = get_tokens_per_block(pretrained_config)
-        if preferred_tokens_per_block is not None:
-            if (type(preferred_tokens_per_block) is not int
-                    or preferred_tokens_per_block <= 0):
-                raise ValueError(
-                    f"{model_cls.__name__}."
-                    "get_preferred_kv_cache_tokens_per_block() must return a "
-                    f"positive int or None, got {preferred_tokens_per_block!r}."
-                )
-            kv_cache_config.tokens_per_block = preferred_tokens_per_block
-            applied["tokens_per_block"] = preferred_tokens_per_block
-
-    get_swa_scratch_reuse = getattr(model_cls,
-                                    'get_preferred_kv_cache_swa_scratch_reuse',
-                                    None)
-    if ("enable_swa_scratch_reuse" not in user_overrides
-            and get_swa_scratch_reuse is not None):
-        preferred_swa_scratch_reuse = get_swa_scratch_reuse(pretrained_config)
-        if preferred_swa_scratch_reuse is not None:
-            if type(preferred_swa_scratch_reuse) is not bool:
-                raise ValueError(
-                    f"{model_cls.__name__}."
-                    "get_preferred_kv_cache_swa_scratch_reuse() must return "
-                    f"bool or None, got {preferred_swa_scratch_reuse!r}.")
-            kv_cache_config.enable_swa_scratch_reuse = (
-                preferred_swa_scratch_reuse)
-            applied["enable_swa_scratch_reuse"] = preferred_swa_scratch_reuse
-
-    if "tokens_per_block" in applied:
-        llm_args.validate_helix_tokens_per_block()
-    return applied
-
-
 def _resolve_kv_cache_manager_v2_auto(llm_args: 'TorchLlmArgs',
                                       model_cls: Optional[type] = None,
                                       pretrained_config: Any = None) -> bool:
     """Resolve the KV cache manager auto setting from the model preference.
 
-    A model default of V2 is demoted to V1 for routes V2 cannot serve; an
+    A model preference for V2 is demoted to V1 for routes V2 cannot serve; an
     explicit user value otherwise wins. The compatibility arms are:
 
     - Disaggregated serving: hybrid Mamba V2 requires the Python transceiver
@@ -676,8 +612,8 @@ def _resolve_kv_cache_manager_v2_auto(llm_args: 'TorchLlmArgs',
       with its own KV cache manager. ``build_managers`` hands that manager the
       target's ``kv_cache_config`` unsplit, and V2 capacity is governed solely
       by ``max_gpu_total_bytes``, so both managers size their pools from the
-      full budget. The model default falls back to V1, and an explicit ``True``
-      is rejected rather than deferred to that allocation.
+      full budget. The model preference falls back to V1, and an explicit
+      ``True`` is rejected rather than deferred to that allocation.
 
     The fallback only reaches models whose manager class is selected by
     ``use_kv_cache_manager_v2``. Models routed to a V2 manager unconditionally
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index d7aa17fda956..c2835eca30bf 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -7959,8 +7959,7 @@ def test_nvfp4(self, tp_size, ep_size):
         # precision.
         model_name = "zai-org/GLM-5.2"
         model_path = f"{llm_models_root()}/GLM-5.2-NVFP4"
-        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7,
-                                        use_kv_cache_manager_v2=True)
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7)
 
         pytorch_config = dict(
             disable_overlap_scheduler=False,
@@ -7976,9 +7975,9 @@ def test_nvfp4(self, tp_size, ep_size):
                  pipeline_parallel_size=1,
                  moe_expert_parallel_size=ep_size,
                  kv_cache_config=kv_cache_config,
-                 max_batch_size=128,
                  max_seq_len=8192,
                  **pytorch_config) as llm:
+            assert llm.args.kv_cache_config.use_kv_cache_manager_v2 is True
             assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4
             task = GSM8K(model_name)
             task.evaluate(llm)
diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
index ec5ceb32ff58..e8dcea152fc9 100644
--- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
+++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
@@ -321,7 +321,7 @@ def resolve_model_prefs(model_dir, side, cache_cfg):
     """Mirror serving's model-preference resolution (PR #15823 semantics).
 
     - use_kv_cache_manager_v2 == "auto" (yaml absent): adopt the model
-      class's get_preferred_kv_cache_manager_version() value, default V1
+      class's get_preferred_kv_cache_manager_version() value, falling back to V1
       (llm_utils._resolve_kv_cache_manager_v2_auto).
     - cache_cfg.transceiver_runtime == "auto": adopt
       model_cls.get_preferred_transceiver_runtime(), NIXL-gated, via the
@@ -357,6 +357,7 @@ def resolve_model_prefs(model_dir, side, cache_cfg):
             shim = types.SimpleNamespace(
                 kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto"),
                 cache_transceiver_config=cache_cfg,
+                speculative_config=None,
             )
             use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, model_cls, hf_view))
         except Exception as e:  # noqa: BLE001 - fall back like a missing model
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index e06ceb1767c0..6ae2544df242 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -40,9 +40,6 @@
     PositionalEmbeddingParams,
     RopeParams,
 )
-from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import (
-    DeepseekV4CacheManager,
-)
 from tensorrt_llm._torch.attention_backend.sparse.dsa import (
     DSABackendForwardArgs,
     DSACacheManager,
@@ -53,14 +50,13 @@
     _effective_compress_ratio_divisor,
     _select_indexer_compress_ratio,
     compute_cu_seqlen_kv_bounds_with_cache,
-    is_dsa_cache_manager,
     split_prefill_chunks,
     transform_local_topk_and_prepare_pool_view,
 )
-from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata
+from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata
+from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role
 from tensorrt_llm._torch.speculative.interface import (
-    SpecWorkerBase,
     prepare_attn_metadata_for_draft_replay,
     restore_attn_metadata_after_draft_replay,
 )
@@ -87,28 +83,6 @@ def has_deep_gemm():
         return False
 
 
-@pytest.mark.parametrize("cache_manager_cls", [DSACacheManager, DSACacheManagerV2])
-def test_is_dsa_cache_manager_accepts_native_managers(cache_manager_cls):
-    assert is_dsa_cache_manager(object.__new__(cache_manager_cls))
-
-
-@pytest.mark.parametrize(
-    "cache_manager",
-    [
-        None,
-        object.__new__(DeepseekV4CacheManager),
-        SimpleNamespace(
-            index_head_dim=128,
-            get_pool_block_indices=Mock(),
-            indexer_k_cache_page_scale=1,
-        ),
-    ],
-    ids=["none", "deepseek-v4", "structural-lookalike"],
-)
-def test_is_dsa_cache_manager_rejects_non_native_managers(cache_manager):
-    assert not is_dsa_cache_manager(cache_manager)
-
-
 def test_metadata_cache_geometry_comes_from_sparse_metadata_params():
     sparse_config = DeepSeekV4SparseAttentionConfig(
         compress_ratios=[1, 4, 128],
@@ -325,76 +299,37 @@ def create_dsa_cache_manager(
     return cache_manager, sparse_attn_config
 
 
-@pytest.mark.parametrize(
-    "cache_manager_cls",
-    [
-        pytest.param(DSACacheManager, id="v1"),
-        pytest.param(DSACacheManagerV2, id="v2"),
-    ],
-)
-def test_transform_local_topk_uses_page_mapping_for_pp_local_layer(cache_manager_cls):
-    """Flatten sparse indices with the active manager's PP-local page mapping."""
-    num_local_layers = 3
-    global_layer_idx = 5
-    local_layer_idx = 1
-    tokens_per_block = 4
-    page_index_scale = 3
-    layer_offset = 1
-
-    if cache_manager_cls is DSACacheManager:
-        cache_manager = SimpleNamespace(num_local_layers=num_local_layers)
-    else:
-        cache_manager = SimpleNamespace(_primary_pool_page_index_params=[(3, 0), (3, 1), (3, 2)])
-    cache_manager.layer_offsets = {global_layer_idx: local_layer_idx}
-    get_page_index_params = MethodType(
-        cache_manager_cls.get_primary_pool_page_index_params,
+def test_transform_local_topk_uses_v2_page_mapping():
+    """Use the V2 pool scale and offset when flattening DSA indices."""
+    cache_manager = SimpleNamespace(
+        _primary_pool_page_index_params=[(3, 1)],
+        get_primary_pool_page_index_params=None,
+    )
+    cache_manager.get_primary_pool_page_index_params = MethodType(
+        DSACacheManagerV2.get_primary_pool_page_index_params,
         cache_manager,
     )
-    cache_manager.get_primary_pool_page_index_params = Mock(wraps=get_page_index_params)
-
     req_idx = torch.tensor([0], dtype=torch.int32)
     block_table = torch.tensor([[3, 7]], dtype=torch.int32)
     topk_indices = torch.tensor([[6]], dtype=torch.int32)
-    pool_view = torch.empty(0)
     metadata = SimpleNamespace(
         kv_cache_manager=cache_manager,
         _ensure_pool_view_cached=Mock(),
         _cached_block_table_ctx=block_table,
         _cached_req_idx_ctx=req_idx,
-        _cached_tokens_per_block=tokens_per_block,
-        _cached_pool_view=pool_view,
+        _cached_tokens_per_block=4,
+        _cached_pool_view=torch.empty(0),
     )
-    attention = SimpleNamespace(layer_idx=global_layer_idx, local_layer_idx=None)
-    pp_local_layer_idx = TrtllmAttention.get_local_layer_idx(attention, metadata)
-    expected_flat_index = (
-        7 * (page_index_scale * tokens_per_block) + (layer_offset * tokens_per_block) + 2
-    )
-    expected_indices = torch.tensor([[expected_flat_index]], dtype=torch.int32)
+    expected_indices = torch.tensor([[90]], dtype=torch.int32)
 
     with patch(
         "torch.ops.trtllm.convert_req_index_to_global",
         return_value=expected_indices,
-    ) as convert_req_index_to_global:
-        global_indices, returned_pool_view = transform_local_topk_and_prepare_pool_view(
-            topk_indices,
-            metadata,
-            layer_idx=pp_local_layer_idx,
-        )
+    ) as convert:
+        actual, _ = transform_local_topk_and_prepare_pool_view(topk_indices, metadata, layer_idx=0)
 
-    assert pp_local_layer_idx == local_layer_idx
-    cache_manager.get_primary_pool_page_index_params.assert_called_once_with(local_layer_idx)
-    convert_req_index_to_global.assert_called_once_with(
-        req_idx,
-        block_table,
-        topk_indices,
-        tokens_per_block,
-        1,
-        page_index_scale * tokens_per_block,
-        layer_offset,
-    )
-    assert torch.equal(global_indices, expected_indices)
-    assert returned_pool_view is pool_view
-    metadata._ensure_pool_view_cached.assert_called_once_with()
+    torch.testing.assert_close(actual, expected_indices)
+    convert.assert_called_once_with(req_idx, block_table, topk_indices, 4, 1, 12, 1)
 
 
 def test_dsa_cache_manager_v2_respects_shared_indexer_layer_mask():
@@ -406,6 +341,19 @@ def test_dsa_cache_manager_v2_respects_shared_indexer_layer_mask():
         num_hidden_layers=num_layers,
         index_topk_pattern=["F", "S", "F"],
     )
+    sparse_config = DeepSeekSparseAttentionConfig(index_head_dim=head_dim)
+    model_config = SimpleNamespace(
+        pretrained_config=SimpleNamespace(),
+        sparse_attention_config=sparse_config,
+    )
+    assert (
+        get_kv_cache_manager_cls(
+            model_config,
+            KvCacheConfig(use_kv_cache_manager_v2=True),
+        )
+        is DSACacheManagerV2
+    )
+
     cache_manager, _ = create_dsa_cache_manager(
         batch_size=2,
         head_dim=head_dim,
@@ -419,7 +367,6 @@ def test_dsa_cache_manager_v2_respects_shared_indexer_layer_mask():
     try:
         assert cache_manager.indexer_k_cache_local_layer_mask == [True, False, True]
         assert cache_manager.indexer_k_cache_page_scale == 2
-        assert cache_manager.indexer_k_cache_pool_per_layer[1] is None
         assert cache_manager.get_indexer_k_cache_buffers(0).shape == (
             cache_manager.blocks_in_primary_pool * cache_manager.indexer_k_cache_page_scale,
             tokens_per_block,
@@ -428,20 +375,19 @@ def test_dsa_cache_manager_v2_respects_shared_indexer_layer_mask():
         )
 
         indexer_bytes_per_token = head_dim + 4
-        expected_cache_bytes = 0
-        for local_layer_idx, owns_indexer in enumerate([True, False, True]):
+        has_indexer_buffers = []
+        for local_layer_idx in range(num_layers):
             layer_config = cache_manager.kv_cache_manager_py_config.layers[local_layer_idx]
-            has_indexer_buffer = any(
-                buffer.role == Role.INDEX_KEY for buffer in layer_config.buffers
+            has_indexer_buffers.append(
+                any(buffer.role == Role.INDEX_KEY for buffer in layer_config.buffers)
             )
-            assert has_indexer_buffer is owns_indexer
-
-            key_bytes = cache_manager.get_layer_bytes_per_token(local_layer_idx, Role.KEY)
-            all_bytes = cache_manager.get_layer_bytes_per_token(local_layer_idx, Role.ALL)
-            expected_all_bytes = key_bytes + indexer_bytes_per_token if owns_indexer else key_bytes
-            assert all_bytes == expected_all_bytes
-            expected_cache_bytes += expected_all_bytes
+        assert has_indexer_buffers == [True, False, True]
 
+        key_bytes = sum(
+            cache_manager.get_layer_bytes_per_token(local_layer_idx, Role.KEY)
+            for local_layer_idx in range(num_layers)
+        )
+        expected_cache_bytes = key_bytes + 2 * indexer_bytes_per_token
         assert cache_manager.get_cache_bytes_per_token() == expected_cache_bytes
         with pytest.raises(AssertionError, match="shared-indexer layer"):
             cache_manager.get_indexer_k_cache_buffers(1)
@@ -3445,6 +3391,9 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip):
     )
 
     Indexer.prepare(metadata)
+    if not enable_indexer_skip:
+        assert metadata.slot_mapping_fp8_fullkv is not metadata.slot_mapping_fp8
+        assert metadata.slot_mapping_scale_fullkv is not metadata.slot_mapping_scale
     indexer._update_k_cache(k_fp8, k_scale, metadata)
 
     # Test custom kernel
@@ -3571,23 +3520,25 @@ def _make_mock_draft_manager():
 
     def test_prepare_swaps_and_restore_recovers(self):
         """Test that prepare swaps KV manager and restore recovers original state."""
-        from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata
-
         meta = self._make_mock_metadata()
         mgr = self._make_mock_draft_manager()
+        # DeepSeek-V4 metadata inherits DSA metadata, while its manager exposes
+        # index_head_dim without implementing DSA's full page-mapping contract.
+        mgr.index_head_dim = 128
+        mgr.get_pool_block_indices = Mock()
+        mgr.indexer_k_cache_page_scale = 1
         original_kv_mgr = meta.kv_cache_manager
         original_offsets = meta.kv_cache_block_offsets.clone()
         original_host_offsets = meta.host_kv_cache_block_offsets.clone()
 
+        def is_attn_metadata(obj, cls):
+            if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata):
+                return obj is meta
+            return builtins.isinstance(obj, cls)
+
         with patch(
             "tensorrt_llm._torch.speculative.interface.isinstance",
-            side_effect=lambda obj, cls: (
-                obj is meta
-                if cls is TrtllmAttentionMetadata
-                else False
-                if cls.__name__ == "DSAtrtllmAttentionMetadata"
-                else builtins.isinstance(obj, cls)
-            ),
+            side_effect=is_attn_metadata,
         ):
             saved = prepare_attn_metadata_for_draft_replay(meta, mgr)
 
@@ -3602,61 +3553,34 @@ def test_prepare_swaps_and_restore_recovers(self):
         torch.testing.assert_close(meta.kv_cache_block_offsets, original_offsets)
         torch.testing.assert_close(meta.host_kv_cache_block_offsets, original_host_offsets)
 
-    @pytest.mark.parametrize("capturing", [False, True])
-    def test_draft_context_recomputes_and_restores_dsa_metadata(self, capturing):
-        """Draft context rebinds preallocated DSA buffers and restores lazy aliases."""
+    def test_native_dsa_replay_swaps_and_restores_buffers(self):
+        """Switch native DSA metadata to draft buffers and restore it."""
         meta = self._make_mock_metadata()
         mgr = object.__new__(DSACacheManagerV2)
         mgr.host_kv_cache_block_offsets = torch.tensor([100, 200, 300])
-        meta.enable_flash_mla = False
         meta.enable_context_mla_with_cached_kv = False
-        meta.host_indexer_k_cache_block_offsets = torch.tensor([1, 2, 3])
-        meta.indexer_k_cache_block_offsets = torch.tensor([4, 5, 6])
-        meta.host_slot_mapping_fp8 = torch.tensor([7, 8, 9])
-        meta.host_slot_mapping_scale = torch.tensor([10, 11, 12])
-        meta.slot_mapping_fp8 = torch.tensor([13, 14, 15])
-        meta.slot_mapping_scale = torch.tensor([16, 17, 18])
-        meta.block_table = torch.tensor([19, 20, 21])
-        meta.block_table_expanded = torch.tensor([22, 23, 24])
-        meta.host_block_table_expanded = torch.tensor([25, 26, 27])
-        meta.host_draft_indexer_k_cache_block_offsets = torch.tensor([28, 29, 30])
-        meta.draft_indexer_k_cache_block_offsets = torch.tensor([31, 32, 33])
-        meta.host_draft_slot_mapping_fp8 = torch.tensor([34, 35, 36])
-        meta.host_draft_slot_mapping_scale = torch.tensor([37, 38, 39])
-        meta.draft_slot_mapping_fp8 = torch.tensor([40, 41, 42])
-        meta.draft_slot_mapping_scale = torch.tensor([43, 44, 45])
-        meta.draft_block_table = torch.tensor([46, 47, 48])
-        meta.draft_block_table_expanded = torch.tensor([49, 50, 51])
-        meta.host_draft_block_table_expanded = torch.tensor([52, 53, 54])
-        meta.prepare_for_indexer_k_cache = Mock()
-        meta._refresh_expanded_block_table = Mock()
-        meta._update_indexer_k_cache_block_offsets = Mock()
+        meta._invalidate_pool_view_cache = Mock()
+        buffer_pairs = {
+            "host_indexer_k_cache_block_offsets": "host_draft_indexer_k_cache_block_offsets",
+            "indexer_k_cache_block_offsets": "draft_indexer_k_cache_block_offsets",
+            "host_slot_mapping_fp8": "host_draft_slot_mapping_fp8",
+            "host_slot_mapping_scale": "host_draft_slot_mapping_scale",
+            "slot_mapping_fp8": "draft_slot_mapping_fp8",
+            "slot_mapping_scale": "draft_slot_mapping_scale",
+            "block_table": "draft_block_table",
+            "block_table_expanded": "draft_block_table_expanded",
+            "host_block_table_expanded": "host_draft_block_table_expanded",
+        }
+        target_buffers = {}
+        draft_buffers = {}
+        for target_name, draft_name in buffer_pairs.items():
+            target_buffers[target_name] = object()
+            draft_buffers[target_name] = object()
+            setattr(meta, target_name, target_buffers[target_name])
+            setattr(meta, draft_name, draft_buffers[target_name])
         del meta.slot_mapping_fp8_fullkv
         del meta.slot_mapping_scale_fullkv
 
-        target_buffers = {
-            "host_indexer_k_cache_block_offsets": meta.host_indexer_k_cache_block_offsets,
-            "indexer_k_cache_block_offsets": meta.indexer_k_cache_block_offsets,
-            "host_slot_mapping_fp8": meta.host_slot_mapping_fp8,
-            "host_slot_mapping_scale": meta.host_slot_mapping_scale,
-            "slot_mapping_fp8": meta.slot_mapping_fp8,
-            "slot_mapping_scale": meta.slot_mapping_scale,
-            "block_table": meta.block_table,
-            "block_table_expanded": meta.block_table_expanded,
-            "host_block_table_expanded": meta.host_block_table_expanded,
-        }
-        draft_buffers = {
-            "host_indexer_k_cache_block_offsets": meta.host_draft_indexer_k_cache_block_offsets,
-            "indexer_k_cache_block_offsets": meta.draft_indexer_k_cache_block_offsets,
-            "host_slot_mapping_fp8": meta.host_draft_slot_mapping_fp8,
-            "host_slot_mapping_scale": meta.host_draft_slot_mapping_scale,
-            "slot_mapping_fp8": meta.draft_slot_mapping_fp8,
-            "slot_mapping_scale": meta.draft_slot_mapping_scale,
-            "block_table": meta.draft_block_table,
-            "block_table_expanded": meta.draft_block_table_expanded,
-            "host_block_table_expanded": meta.host_draft_block_table_expanded,
-        }
-
         def is_attn_metadata(obj, cls):
             if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata):
                 return obj is meta
@@ -3669,116 +3593,16 @@ def is_attn_metadata(obj, cls):
             ),
             patch(
                 "tensorrt_llm._torch.speculative.interface.torch.cuda.is_current_stream_capturing",
-                return_value=capturing,
+                return_value=True,
             ),
-            patch.object(Indexer, "recompute_slot_mappings") as recompute,
-        ):
-            with SpecWorkerBase.draft_kv_cache_context(Mock(), meta, mgr):
-                for name, buffer in draft_buffers.items():
-                    assert getattr(meta, name) is buffer
-                assert meta.slot_mapping_fp8_fullkv is meta.draft_slot_mapping_fp8
-                assert meta.slot_mapping_scale_fullkv is meta.draft_slot_mapping_scale
-
-        for name, buffer in target_buffers.items():
-            assert getattr(meta, name) is buffer
-        assert meta.slot_mapping_fp8_fullkv is meta.slot_mapping_fp8
-        assert meta.slot_mapping_scale_fullkv is meta.slot_mapping_scale
-        if capturing:
-            meta.prepare_for_indexer_k_cache.assert_not_called()
-            meta._refresh_expanded_block_table.assert_not_called()
-            recompute.assert_not_called()
-        else:
-            meta.prepare_for_indexer_k_cache.assert_called_once_with()
-            meta._refresh_expanded_block_table.assert_called_once_with()
-            recompute.assert_called_once_with(meta)
-        meta._update_indexer_k_cache_block_offsets.assert_not_called()
-
-    def test_non_dsa_manager_skips_inherited_dsa_replay_state(self):
-        """A DSA metadata subclass must not impose DSA layout on its manager."""
-        meta = self._make_mock_metadata()
-        mgr = self._make_mock_draft_manager()
-        target_manager = meta.kv_cache_manager
-        target_offsets = meta.kv_cache_block_offsets
-        target_host_offsets = meta.host_kv_cache_block_offsets
-        # DeepSeek-V4 cache managers expose this field but not the DSA page
-        # mapping contract used by _update_indexer_k_cache_block_offsets().
-        mgr.index_head_dim = 128
-        mgr.get_pool_block_indices = Mock()
-        mgr.indexer_k_cache_page_scale = 1
-
-        def is_attn_metadata(obj, cls):
-            if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata):
-                return obj is meta
-            return builtins.isinstance(obj, cls)
-
-        with patch(
-            "tensorrt_llm._torch.speculative.interface.isinstance",
-            side_effect=is_attn_metadata,
         ):
             saved = prepare_attn_metadata_for_draft_replay(meta, mgr)
 
-        assert saved is not None
-        assert "saved_dsa_state" not in saved
-
+        assert "saved_dsa_state" in saved
+        assert meta.slot_mapping_fp8 is draft_buffers["slot_mapping_fp8"]
         restore_attn_metadata_after_draft_replay(meta, saved)
-        assert meta.kv_cache_manager is target_manager
-        assert meta.kv_cache_block_offsets is target_offsets
-        assert meta.host_kv_cache_block_offsets is target_host_offsets
-
-
-def test_recompute_context_kv_gather_mappings_does_not_alias_uncompressed_cached_mappings():
-    """Cached KV requires independent full-KV mappings even without compression."""
-    fp8_mapping = torch.tensor([1, 2, 3])
-    scale_mapping = torch.tensor([4, 5, 6])
-    fullkv_fp8_mapping = object()
-    fullkv_scale_mapping = object()
-    host_fullkv_fp8_mapping = Mock()
-    host_fullkv_scale_mapping = Mock()
-    host_fullkv_fp8_mapping.cuda.return_value = fullkv_fp8_mapping
-    host_fullkv_scale_mapping.cuda.return_value = fullkv_scale_mapping
-    computed_fp8_mapping = torch.tensor([10, 11, 12, 13])
-    computed_scale_mapping = torch.tensor([20, 21, 22, 23])
-    metadata = SimpleNamespace(
-        enable_context_mla_with_cached_kv=True,
-        skip_indexer_for_ctx_reqs=False,
-        slot_mapping_fp8=fp8_mapping,
-        slot_mapping_scale=scale_mapping,
-        slot_mapping_fp8_fullkv=fp8_mapping,
-        slot_mapping_scale_fullkv=scale_mapping,
-        host_indexer_k_cache_block_offsets=torch.tensor([[0], [1]], dtype=torch.int32),
-    )
-    indexer_params = SimpleNamespace(
-        num_contexts=2,
-        compress_ratio=1,
-        cached_kv_tokens=torch.tensor([0, 1], dtype=torch.int32),
-        kv_lens=torch.tensor([2, 2], dtype=torch.int32),
-        head_dim=128,
-        tokens_per_block=64,
-        quant_block_size=128,
-        data_bytes_per_token=128,
-    )
-
-    with (
-        patch(
-            "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer.torch.empty",
-            side_effect=[host_fullkv_fp8_mapping, host_fullkv_scale_mapping],
-        ),
-        patch(
-            "tensorrt_llm._torch.attention_backend.sparse.dsa.indexer._compute_slot_mappings",
-            return_value=(computed_fp8_mapping, computed_scale_mapping),
-        ) as compute_slot_mappings,
-    ):
-        Indexer.recompute_context_kv_gather_mappings(metadata, indexer_params)
-
-    assert metadata.slot_mapping_fp8_fullkv is fullkv_fp8_mapping
-    assert metadata.slot_mapping_scale_fullkv is fullkv_scale_mapping
-    assert metadata.slot_mapping_fp8_fullkv is not fp8_mapping
-    assert metadata.slot_mapping_scale_fullkv is not scale_mapping
-    host_fullkv_fp8_mapping.copy_.assert_called_once_with(computed_fp8_mapping)
-    host_fullkv_scale_mapping.copy_.assert_called_once_with(computed_scale_mapping)
-    host_fullkv_fp8_mapping.cuda.assert_called_once_with(non_blocking=True)
-    host_fullkv_scale_mapping.cuda.assert_called_once_with(non_blocking=True)
-    compute_slot_mappings.assert_called_once()
+        assert meta.slot_mapping_fp8 is target_buffers["slot_mapping_fp8"]
+        assert meta.slot_mapping_fp8_fullkv is target_buffers["slot_mapping_fp8"]
 
 
 @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
index dd80817a3d9e..069c47841287 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_v2_extra_buffers.py
@@ -157,6 +157,33 @@ def test_default_hook_keeps_nvfp4_scale_buffers(self):
             mgr.shutdown()
             del mgr
 
+    def test_page_table_uses_physical_pool_representative(self):
+        mgr = KVCacheManagerV2(**_make_kwargs(head_dim=[64, 192, 64, 192]))
+        real_impl = mgr.impl
+        try:
+            pool_id = 0
+            physical_layer = mgr._pool_layer_ids_by_role[(pool_id, Role.KEY)]
+            other_layer = next(
+                layer_id
+                for layer_id in real_impl.layer_grouping[pool_id]
+                if int(layer_id) != int(physical_layer)
+            )
+            impl_proxy = Mock(wraps=real_impl)
+            impl_proxy.layer_grouping = ((other_layer, physical_layer),)
+            mgr.impl = impl_proxy
+
+            mgr._prepare_page_table_tensor(index_mapper_capacity=1)
+
+            self.assertEqual(
+                impl_proxy.get_mem_pool_base_address.call_args_list[0].args,
+                (physical_layer, Role.KEY, PageIndexMode.SHARED),
+            )
+            impl_proxy.get_page_index_scale.assert_called_once_with(physical_layer, Role.KEY)
+        finally:
+            mgr.impl = real_impl
+            mgr.shutdown()
+            del mgr
+
     def test_subclass_registers_index_key_only_on_sparse_layers(self):
         # Sparse layer convention: layers 0-2 dense (no INDEX_KEY), 3+ sparse
         # (INDEX_KEY registered). Matches MiniMax-M3.
@@ -189,63 +216,6 @@ def test_subclass_registers_index_key_only_on_sparse_layers(self):
             mgr.shutdown()
             del mgr
 
-    def test_pool_mapping_ignores_reordered_layer_grouping(self):
-        mgr = KVCacheManagerV2(**_make_kwargs(num_layers=4, head_dim=[64, 192, 64, 192]))
-        real_impl = mgr.impl
-        try:
-            self.assertEqual(mgr.num_pools, 1)
-            pool_id = 0
-            config_group = tuple(
-                layer.layer_id
-                for layer in mgr.kv_cache_manager_py_config.layers
-                if int(real_impl.get_layer_group_id(layer.layer_id)) == pool_id
-            )
-            storage_variant = next(
-                variant
-                for pool_group in real_impl.pool_group_descs
-                for variant in pool_group.slot_desc.variants
-                if int(variant.layer_group_id) == pool_id
-            )
-            expected_layer_id = next(
-                int(buffer_id.layer_id)
-                for coalesced in storage_variant.coalesced_buffers
-                for buffer_id in coalesced.buffer_ids
-                if buffer_id.role == Role.KEY
-            )
-            self.assertGreater(len(config_group), 1)
-            self.assertNotEqual(expected_layer_id, int(config_group[0]))
-            reordered_group = tuple(
-                layer_id for layer_id in config_group if int(layer_id) != expected_layer_id
-            )
-            reordered_group += (expected_layer_id,)
-            reordered_grouping = (reordered_group,)
-            reordered_first_layer_id = reordered_grouping[pool_id][0]
-            self.assertNotEqual(expected_layer_id, reordered_first_layer_id)
-
-            expected_base_address = real_impl.get_mem_pool_base_address(
-                expected_layer_id, Role.KEY, PageIndexMode.SHARED
-            )
-
-            impl_proxy = Mock(wraps=real_impl)
-            impl_proxy.layer_grouping = reordered_grouping
-            mgr.impl = impl_proxy
-            mgr._prepare_page_table_tensor(index_mapper_capacity=1)
-
-            self.assertEqual(
-                impl_proxy.get_mem_pool_base_address.call_args_list[0].args,
-                (expected_layer_id, Role.KEY, PageIndexMode.SHARED),
-                "the pool pointer origin must use the physical-layout representative",
-            )
-            self.assertEqual(
-                int(mgr.kv_cache_pool_pointers[pool_id, 0]),
-                int(expected_base_address),
-            )
-            impl_proxy.get_page_index_scale.assert_called_once_with(expected_layer_id, Role.KEY)
-        finally:
-            mgr.impl = real_impl
-            mgr.shutdown()
-            del mgr
-
     def test_index_key_layer_has_non_zero_page_upper_bound(self):
         # Once a subclass registers INDEX_KEY, the V2 storage must wire a
         # paged pool for that role with a non-zero page-index upper bound,
diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
index 78e262fae150..3b519541fbf1 100644
--- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py
+++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py
@@ -65,9 +65,8 @@
 from tensorrt_llm.llmapi.llm_utils import (
     _resolve_kv_cache_manager_v2_auto,
     _resolve_transceiver_runtime_auto,
-    apply_model_defaults_to_llm_args,
 )
-from tensorrt_llm.mapping import CpType, Mapping
+from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     AttentionLayerConfig,
     BatchDesc,
@@ -675,7 +674,7 @@ def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2)
     )
 
 
-def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
+def test_hybrid_models_prefer_v2_and_python_transceiver(monkeypatch):
     from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM
     from tensorrt_llm._torch.models.modeling_qwen3_5 import Qwen3_5VLModel
     from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM
@@ -693,12 +692,9 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
             model="/tmp/dummy_model",
             cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"),
         )
-        model_defaults = model_cls.get_model_defaults(llm_args)
-        apply_model_defaults_to_llm_args(llm_args, model_defaults)
         _resolve_transceiver_runtime_auto(llm_args, model_cls)
         _resolve_kv_cache_manager_v2_auto(llm_args, model_cls)
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
-        assert llm_args.kv_cache_config.enable_block_reuse is False
         assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON"
 
 
@@ -712,7 +708,7 @@ def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch):
         (None, False, False),
     ],
 )
-def test_qwen3_gdn_replay_defaults_to_v2_cache_manager(
+def test_qwen3_gdn_replay_uses_v2_preference(
     monkeypatch,
     replay_env,
     manager_setting,
@@ -732,8 +728,6 @@ def test_qwen3_gdn_replay_defaults_to_v2_cache_manager(
         ),
         speculative_config=MTPDecodingConfig(max_draft_len=3),
     )
-    model_defaults = Qwen3NextForCausalLM.get_model_defaults(llm_args)
-    apply_model_defaults_to_llm_args(llm_args, model_defaults)
     _resolve_kv_cache_manager_v2_auto(
         llm_args,
         Qwen3NextForCausalLM,
@@ -750,18 +744,22 @@ def test_qwen3_gdn_replay_defaults_to_v2_cache_manager(
     )
 
 
-def test_kimi_defaults_to_mixed_manager(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Kimi K3 declares no V2 default: block reuse defaults off and the
-    Mixed manager (separate KV / recurrent-state pools) is the default
-    route, which SA speculative decoding requires."""
+def test_kimi_without_v2_preference_uses_mixed_manager(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Kimi K3 uses separate KV and recurrent-state pools for SA decoding."""
     from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM
 
     monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False)
     monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False)
 
-    llm_args = TorchLlmArgs(model="/tmp/dummy_model")
-    model_defaults = KimiLinearForCausalLM.get_model_defaults(llm_args)
-    apply_model_defaults_to_llm_args(llm_args, model_defaults)
+    llm_args = TorchLlmArgs(
+        model="/tmp/dummy_model",
+        kv_cache_config=KvCacheConfig(
+            enable_block_reuse=False,
+            tokens_per_block=64,
+        ),
+    )
     resolved = _resolve_kv_cache_manager_v2_auto(llm_args, KimiLinearForCausalLM)
 
     assert resolved is False
@@ -845,71 +843,6 @@ def test_v2_hybrid_incompatibility_fails_without_cpp_fallback(
         )
 
 
-def test_dsa_star_rejects_v2_and_keeps_explicit_v1_selection():
-    model_config = SimpleNamespace(
-        pretrained_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]),
-        sparse_attention_config=SimpleNamespace(algorithm="dsa"),
-    )
-    creator = object.__new__(KvCacheCreator)
-    creator._kv_connector_manager = None
-    creator._max_beam_width = 1
-    creator._mapping = SimpleNamespace(cp_config={"cp_type": CpType.STAR})
-
-    with pytest.raises(NotImplementedError, match="STAR context parallelism"):
-        creator._validate_or_fallback_kv_cache_manager_v2(
-            KVCacheManagerV2,
-            model_config,
-            KvCacheConfig(use_kv_cache_manager_v2=True),
-        )
-
-    assert (
-        creator._validate_or_fallback_kv_cache_manager_v2(
-            KVCacheManager,
-            model_config,
-            KvCacheConfig(use_kv_cache_manager_v2=False),
-        )
-        is KVCacheManager
-    )
-
-
-@pytest.mark.parametrize("zero_layer_hybrid", [False, True])
-@pytest.mark.parametrize(
-    ("manager_setting", "expected_use_v2"),
-    [(False, False), ("auto", False), (True, True)],
-)
-def test_sparse_cache_manager_factory_only_enables_explicit_v2(
-    monkeypatch, zero_layer_hybrid, manager_setting, expected_use_v2
-):
-    sparse_attn_config = SimpleNamespace(algorithm="dsa")
-    model_config = SimpleNamespace(
-        pretrained_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]),
-        sparse_attention_config=sparse_attn_config,
-        get_num_mamba_layers=lambda: 0,
-    )
-    expected_manager_cls = object()
-    sparse_manager_factory = MagicMock(return_value=expected_manager_cls)
-    monkeypatch.setattr(
-        "tensorrt_llm._torch.pyexecutor._util.is_hybrid_linear",
-        lambda config: zero_layer_hybrid,
-    )
-    monkeypatch.setattr(
-        "tensorrt_llm._torch.pyexecutor._util.get_sparse_attn_kv_cache_manager",
-        sparse_manager_factory,
-    )
-
-    assert (
-        get_kv_cache_manager_cls(
-            model_config,
-            KvCacheConfig(use_kv_cache_manager_v2=manager_setting),
-        )
-        is expected_manager_cls
-    )
-    sparse_manager_factory.assert_called_once_with(
-        sparse_attn_config,
-        use_kv_cache_manager_v2=expected_use_v2,
-    )
-
-
 def _make_mgr(
     max_batch_size=4, max_draft_len=2, enable_attention_dp=False, use_replay_state_update=False
 ):
diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
index 8cdefcfba19a..6e80d1513ef8 100644
--- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
@@ -150,15 +150,15 @@ def test_deepseek_v4_fused_hc_default_enabled(monkeypatch):
     assert _resolve_enable_fused_hc(config) is False
 
 
-def test_deepseek_v4_model_defaults():
-    class LlmArgs:
-        pass
+def test_deepseek_v4_kv_cache_defaults_and_v2_preference():
+    defaults = DeepseekV4ForCausalLM.get_model_defaults(None)
 
-    defaults = DeepseekV4ForCausalLM.get_model_defaults(LlmArgs())
-
-    assert defaults == {}
-    assert DeepseekV4ForCausalLM.get_preferred_kv_cache_tokens_per_block() == 128
-    assert DeepseekV4ForCausalLM.get_preferred_kv_cache_swa_scratch_reuse() is True
+    assert defaults == {
+        "kv_cache_config": {
+            "tokens_per_block": 128,
+            "enable_swa_scratch_reuse": True,
+        }
+    }
     assert DeepseekV4ForCausalLM.get_preferred_kv_cache_manager_version() == "V2"
 
 
diff --git a/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py b/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py
index 00e3812ed0fb..9d4adc08987b 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import json
 import os
 import shutil
@@ -21,8 +24,7 @@
 from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig,
                                  KvCacheConfig, MoeConfig)
 from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
-from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
-                                           apply_model_defaults_to_llm_args)
+from tensorrt_llm.llmapi.llm_utils import _resolve_kv_cache_manager_v2_auto
 from tensorrt_llm.mapping import Mapping
 
 configs = """
@@ -56,31 +58,26 @@ def test_gpt_oss_prefers_python_transceiver() -> None:
 
 
 def _resolve_gpt_oss_kv_cache_manager_v2(**llm_args_kwargs) -> bool:
-    """Run GPT-OSS model defaults through the same path model loading uses."""
+    """Resolve the GPT-OSS preference through the model-loading path."""
     llm_args = TorchLlmArgs(model="/tmp/dummy_model", **llm_args_kwargs)
-    original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2
-    model_defaults = GptOssForCausalLM.get_model_defaults(llm_args)
-    apply_model_defaults_to_llm_args(llm_args, model_defaults)
-    return _resolve_kv_cache_manager_v2_auto(llm_args,
-                                             model_defaults,
-                                             original_setting=original_setting)
+    return _resolve_kv_cache_manager_v2_auto(llm_args, GptOssForCausalLM)
 
 
-def test_gpt_oss_model_defaults_select_v2():
+def test_gpt_oss_model_preference_selects_v2():
     """GPT-OSS is VSWA, so "auto" resolves to KVCacheManagerV2."""
     assert _resolve_gpt_oss_kv_cache_manager_v2() is True
 
 
 @pytest.mark.parametrize("user_setting", [False, True])
 def test_gpt_oss_explicit_setting_wins(user_setting):
-    """An explicit user value is never overridden by the model default."""
+    """An explicit user value is never overridden by the model preference."""
     assert _resolve_gpt_oss_kv_cache_manager_v2(kv_cache_config=KvCacheConfig(
         use_kv_cache_manager_v2=user_setting)) is user_setting
 
 
 def test_gpt_oss_two_model_eagle3_falls_back_to_v1():
     """Two-model Eagle3 builds a separate draft engine with its own KV cache
-    manager, and V2 sizes both from the full budget, so the model default is
+    manager, and V2 sizes both from the full budget, so the model preference is
     demoted to V1."""
     assert _resolve_gpt_oss_kv_cache_manager_v2(
         speculative_config=Eagle3DecodingConfig(
@@ -90,7 +87,7 @@ def test_gpt_oss_two_model_eagle3_falls_back_to_v1():
 
 
 def test_gpt_oss_explicit_v2_rejects_two_model_eagle3():
-    """The demotion above only applies to the model default. An explicit
+    """The demotion above only applies to the model preference. An explicit
     request for the same unsupported combination is rejected rather than
     silently honored."""
     with pytest.raises(ValueError,
diff --git a/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py b/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py
index bd1c64562311..232893dfca3e 100644
--- a/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py
+++ b/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py
@@ -611,13 +611,6 @@ def test_prefers_python_transceiver(self):
             KimiK25ForConditionalGeneration.get_preferred_transceiver_runtime(), "PYTHON"
         )
 
-    def test_prefers_kv_cache_manager_v2(self):
-        """Kimi-K2.5 prefers KV cache manager V2."""
-        self.assertEqual(
-            KimiK25ForConditionalGeneration.get_preferred_kv_cache_manager_version(),
-            "V2",
-        )
-
 
 # ---------------------------------------------------------------------------
 # E2E Smoke Test — requires GPU + model checkpoint
diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py
index 9aa8e7e11a06..36c9d8b0c9f7 100644
--- a/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py
+++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.py
@@ -29,7 +29,10 @@
 from tensorrt_llm.inputs import ContentFormat
 from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY
 from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
-from tensorrt_llm.llmapi.llm_utils import apply_model_defaults_to_llm_args
+from tensorrt_llm.llmapi.llm_utils import (
+    _resolve_kv_cache_manager_v2_auto,
+    apply_model_defaults_to_llm_args,
+)
 from tensorrt_llm.models.modeling_utils import QuantConfig
 from tensorrt_llm.quantization import QuantAlgo
 
@@ -187,6 +190,7 @@ def test_qwen35_moe_model_defaults(
         llm_args.quant_config = QuantConfig(quant_algo=quant_algo)
         defaults = model_cls.get_model_defaults(llm_args)
         apply_model_defaults_to_llm_args(llm_args, defaults)
+        _resolve_kv_cache_manager_v2_auto(llm_args, model_cls)
 
         assert llm_args.kv_cache_config.enable_block_reuse is False
         assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py
index 3c0ca12ef94a..57571e4db563 100644
--- a/tests/unittest/_torch/speculative/test_eagle3.py
+++ b/tests/unittest/_torch/speculative/test_eagle3.py
@@ -50,61 +50,43 @@
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 
 
-def test_mtp_eagle_refreshes_runtime_draft_slot_mappings_before_forward(
-) -> None:
-    """The first MTP-Eagle draft forward must use runtime-corrected mappings."""
+def test_mtp_eagle_refreshes_dsa_metadata_before_draft_forward() -> None:
+    """Refresh DSA mappings after switching to the draft cache."""
     events = []
-    target_manager = object()
-    target_block_table = object()
-    draft_block_table = object()
     draft_manager = object.__new__(DSACacheManagerV2)
 
-    class _AttentionMetadata(DSAtrtllmAttentionMetadata):
+    class _Metadata(DSAtrtllmAttentionMetadata):
 
         def __init__(self):
             self._seq_lens_cuda = torch.tensor([1], dtype=torch.int32)
             self._num_ctx_tokens = 0
-            self.kv_cache_manager = target_manager
-            self.block_table = target_block_table
 
         def set_in_mtp_draft_loop(self, active):
-            events.append(("set_in_mtp_draft_loop", active))
+            pass
 
         def set_mtp_num_accepted(self, value):
-            events.append(("set_mtp_num_accepted", value))
+            pass
 
         def set_skip_topk(self, value):
-            events.append(("set_skip_topk", value))
+            pass
 
         def on_update_kv_lens(self):
-            assert self.kv_cache_manager is draft_manager
-            assert self.block_table is draft_block_table
-            events.append(("on_update_kv_lens", draft_manager))
+            events.append("refresh")
 
-    attn_metadata = _AttentionMetadata()
+    metadata = _Metadata()
 
     @contextmanager
-    def draft_context(metadata, manager):
-        assert metadata is attn_metadata
+    def draft_context(attn_metadata, manager):
+        assert attn_metadata is metadata
         assert manager is draft_manager
-        assert metadata.kv_cache_manager is target_manager
-        assert metadata.block_table is target_block_table
-        metadata.kv_cache_manager = draft_manager
-        metadata.block_table = draft_block_table
-        events.append(("enter_draft_context", manager))
-        try:
-            yield metadata
-        finally:
-            metadata.kv_cache_manager = target_manager
-            metadata.block_table = target_block_table
+        events.append("switch")
+        yield metadata
 
     class _StopAfterFirstDraftForward(Exception):
         pass
 
     def run_draft_forward(*args, **kwargs):
-        assert attn_metadata.kv_cache_manager is draft_manager
-        assert attn_metadata.block_table is draft_block_table
-        events.append(("run_draft_forward", draft_manager))
+        events.append("forward")
         raise _StopAfterFirstDraftForward
 
     worker = SimpleNamespace(
@@ -112,17 +94,15 @@ def run_draft_forward(*args, **kwargs):
         draft_kv_cache_context=draft_context,
         _run_draft_forward=run_draft_forward,
     )
-    spec_metadata = SimpleNamespace(runtime_draft_len=1)
-    inputs = {"position_ids": torch.tensor([0], dtype=torch.int32)}
 
     from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker
 
     with pytest.raises(_StopAfterFirstDraftForward):
         Eagle3OneModelWorker._forward_linear_draft_loop(
             worker,
-            inputs,
-            attn_metadata,
-            spec_metadata,
+            {"position_ids": torch.tensor([0], dtype=torch.int32)},
+            metadata,
+            SimpleNamespace(runtime_draft_len=1),
             draft_model=object(),
             draft_kv_cache_manager=draft_manager,
             num_contexts=0,
@@ -131,13 +111,7 @@ def run_draft_forward(*args, **kwargs):
             original_all_rank_num_tokens=None,
         )
 
-    event_names = [name for name, _ in events]
-    assert event_names.index("enter_draft_context") < event_names.index(
-        "on_update_kv_lens")
-    assert event_names.index("on_update_kv_lens") < event_names.index(
-        "run_draft_forward")
-    assert attn_metadata.kv_cache_manager is target_manager
-    assert attn_metadata.block_table is target_block_table
+    assert events == ["switch", "refresh", "forward"]
 
 
 def test_dynamic_tree_metadata_forces_target_mask_prepare_each_step() -> None:
diff --git a/tests/unittest/_torch/test_model_config.py b/tests/unittest/_torch/test_model_config.py
index a6d6fd755976..6e3daea5e695 100644
--- a/tests/unittest/_torch/test_model_config.py
+++ b/tests/unittest/_torch/test_model_config.py
@@ -1,6 +1,3 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
 import json
 import struct
 import types
diff --git a/tests/unittest/llmapi/test_config_database.py b/tests/unittest/llmapi/test_config_database.py
index b1f836346adf..b443c231e15e 100644
--- a/tests/unittest/llmapi/test_config_database.py
+++ b/tests/unittest/llmapi/test_config_database.py
@@ -15,7 +15,6 @@
 """L0 tests for validating curated and database YAML configs against TorchLlmArgs."""
 
 from pathlib import Path
-from types import SimpleNamespace
 from unittest import mock
 
 import pytest
@@ -29,11 +28,6 @@
 from tensorrt_llm.commands.serve import _apply_fastapi_middlewares
 from tensorrt_llm.commands.serve import main as serve_main
 from tensorrt_llm.llmapi import llm_args as llm_args_module
-from tensorrt_llm.llmapi.llm_utils import (
-    _apply_kv_cache_config_preferences,
-    _resolve_kv_cache_manager_v2_auto,
-    apply_model_defaults_to_llm_args,
-)
 
 from .yaml_validation_harness import (
     assert_no_default_valued_leaves,
@@ -111,18 +105,16 @@ def _get_default_values_for_config(config_path: Path) -> dict:
     model = entry["model"]
     arch = entry["arch"]
     model_cls = get_registered_model_class(arch)
-    if not model_cls:
+    if not model_cls or not hasattr(model_cls, "get_model_defaults"):
         return global_default
 
-    base_args = llm_args_module.TorchLlmArgs(model=model, skip_tokenizer_init=True)
-    pretrained_config = SimpleNamespace(architectures=[arch])
-    _apply_kv_cache_config_preferences(base_args, model_cls, pretrained_config)
-    get_model_defaults = getattr(model_cls, "get_model_defaults", None)
-    model_defaults = get_model_defaults(base_args) if get_model_defaults else {}
-    if model_defaults:
-        apply_model_defaults_to_llm_args(base_args, model_defaults)
-    _resolve_kv_cache_manager_v2_auto(base_args, model_cls, pretrained_config)
-    return base_args.model_dump(mode="json")
+    base_args = llm_args_module.TorchLlmArgs(model=model, skip_tokenizer_init=True).model_dump(
+        mode="json"
+    )
+
+    model_defaults = model_cls.get_model_defaults(base_args)
+    base_args.update(model_defaults)
+    return base_args
 
 
 @pytest.fixture(autouse=True)
@@ -136,27 +128,6 @@ def get_config_id(config_path: Path) -> str:
     return str(config_path.relative_to(CONFIG_ROOT))
 
 
-@pytest.mark.parametrize(
-    "arch",
-    ["DeepseekV3ForCausalLM", "KimiK25ForConditionalGeneration"],
-)
-def test_model_preferences_preserve_global_nested_defaults(monkeypatch, arch):
-    config_key = f"test-only/{arch}.yaml"
-    monkeypatch.setitem(
-        _CONFIG_PATH_TO_ENTRY,
-        config_key,
-        {
-            "model": "dummy/model",
-            "arch": arch,
-        },
-    )
-
-    defaults = _get_default_values_for_config(REPO_ROOT / config_key)
-
-    assert defaults["kv_cache_config"]["use_kv_cache_manager_v2"] is True
-    assert defaults["kv_cache_config"]["enable_block_reuse"] is True
-
-
 def _assert_kv_cache_block_reuse_policy(config_dict: dict) -> None:
     assert config_dict.get("kv_cache_config", {}).get("enable_block_reuse") is not False, (
         "KV cache block reuse should not be disabled"
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index 79fe5737df75..ceef1c8de786 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -58,8 +58,7 @@
                                           update_llm_args_with_extra_dict,
                                           update_llm_args_with_extra_options)
 # fmt: on
-from tensorrt_llm.llmapi.llm_utils import (_apply_kv_cache_config_preferences,
-                                           _resolve_kv_cache_manager_v2_auto,
+from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
                                            _resolve_transceiver_runtime_auto,
                                            apply_model_defaults_to_llm_args)
 from tensorrt_llm.llmapi.mm_encoder import MultimodalEncoder
@@ -513,183 +512,6 @@ def test_post_processor_hook_rejected_with_skip_tokenizer_init():
     TorchLlmArgs(model="/tmp/dummy_model", skip_tokenizer_init=True)
 
 
-@pytest.mark.cpu_only
-class TestKvCacheManagerV2AutoResolution:
-    """Test model preferences for the KV cache manager version."""
-
-    class _PreferV2:
-
-        @classmethod
-        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
-            return "V2"
-
-    class _PreferV1:
-
-        @classmethod
-        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
-            return "V1"
-
-    class _InvalidPreference:
-
-        @classmethod
-        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
-            return "V3"
-
-    class _ArchitectureSensitive:
-
-        @classmethod
-        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
-            architectures = getattr(pretrained_config, "architectures", [])
-            return "V2" if architectures == ["PreferV2ForCausalLM"] else "V1"
-
-    @pytest.mark.parametrize("explicit_auto", [False, True])
-    def test_auto_uses_model_preference(self, explicit_auto):
-        kv_cache_config = (KvCacheConfig(use_kv_cache_manager_v2="auto")
-                           if explicit_auto else KvCacheConfig())
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model",
-                                kv_cache_config=kv_cache_config)
-
-        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
-
-        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
-
-    def test_auto_without_preference_falls_back_to_v1(self):
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
-
-        _resolve_kv_cache_manager_v2_auto(llm_args)
-
-        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
-
-    @pytest.mark.parametrize(
-        ("backend", "runtime"),
-        [
-            ("NIXL", "CPP"),
-            ("UCX", None),
-            ("MPI", None),
-        ],
-    )
-    def test_auto_v2_falls_back_for_incompatible_disagg(self, backend, runtime):
-        llm_args = TorchLlmArgs(
-            model="/tmp/dummy_model",
-            cache_transceiver_config=CacheTransceiverConfig(
-                backend=backend, transceiver_runtime=runtime),
-        )
-
-        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
-
-        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
-
-    def test_auto_v2_keeps_python_nixl_preference(self):
-        llm_args = TorchLlmArgs(
-            model="/tmp/dummy_model",
-            cache_transceiver_config=CacheTransceiverConfig(
-                backend="NIXL", transceiver_runtime="PYTHON"),
-        )
-
-        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
-
-        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
-
-    @pytest.mark.parametrize(("user_setting", "model_cls"), [
-        pytest.param(False, _PreferV2, id="explicit-v1"),
-        pytest.param(True, _PreferV1, id="explicit-v2"),
-    ])
-    def test_explicit_value_overrides_model_preference(self, user_setting,
-                                                       model_cls):
-        llm_args = TorchLlmArgs(
-            model="/tmp/dummy_model",
-            kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=user_setting))
-
-        _resolve_kv_cache_manager_v2_auto(llm_args, model_cls)
-
-        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is user_setting
-
-    def test_pretrained_config_is_forwarded_to_preference(self):
-        from types import SimpleNamespace
-
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
-        pretrained_config = SimpleNamespace(
-            architectures=["PreferV2ForCausalLM"])
-
-        _resolve_kv_cache_manager_v2_auto(llm_args, self._ArchitectureSensitive,
-                                          pretrained_config)
-
-        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
-
-    def test_invalid_model_preference_raises(self):
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
-
-        with pytest.raises(ValueError, match="must return 'V1', 'V2', or None"):
-            _resolve_kv_cache_manager_v2_auto(llm_args, self._InvalidPreference)
-
-    @pytest.mark.parametrize("architecture", [
-        "DeepseekV3ForCausalLM",
-        "DeepseekV32ForCausalLM",
-        "GlmMoeDsaForCausalLM",
-        "MistralLarge3ForCausalLM",
-        "DeepseekV4ForCausalLM",
-        "KimiK25ForConditionalGeneration",
-        "MiniMaxM2ForCausalLM",
-        "NemotronHForCausalLM",
-        "NemotronHPuzzleForCausalLM",
-        "Qwen3NextForCausalLM",
-        "Qwen3_5MoeForCausalLM",
-        "Qwen3_5ForCausalLM",
-        "Qwen3_5MoeForConditionalGeneration",
-        "Qwen3_5ForConditionalGeneration",
-    ])
-    def test_registered_models_prefer_v2(self, architecture):
-        from tensorrt_llm._torch.models.modeling_utils import \
-            MODEL_CLASS_MAPPING
-
-        model_cls = MODEL_CLASS_MAPPING[architecture]
-        assert model_cls.get_preferred_kv_cache_manager_version() == "V2"
-
-
-@pytest.mark.cpu_only
-class TestKvCacheConfigPreferences:
-    """Test scalar KV-cache preferences outside generic model defaults."""
-
-    def test_deepseek_v4_preferences_preserve_unrelated_nested_config(self):
-        from tensorrt_llm._torch.models.modeling_deepseekv4 import \
-            DeepseekV4ForCausalLM
-
-        sparse_config = DeepSeekV4SparseAttentionConfig(index_head_dim=128)
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model",
-                                sparse_attention_config=sparse_config)
-
-        applied = _apply_kv_cache_config_preferences(llm_args,
-                                                     DeepseekV4ForCausalLM)
-
-        assert applied == {
-            "tokens_per_block": 128,
-            "enable_swa_scratch_reuse": True,
-        }
-        assert llm_args.kv_cache_config.tokens_per_block == 128
-        assert llm_args.kv_cache_config.enable_swa_scratch_reuse is True
-        assert llm_args.sparse_attention_config is sparse_config
-        assert sparse_config.model_dump(exclude_unset=True) == {
-            "index_head_dim": 128
-        }
-
-    def test_explicit_values_override_deepseek_v4_preferences(self):
-        from tensorrt_llm._torch.models.modeling_deepseekv4 import \
-            DeepseekV4ForCausalLM
-
-        llm_args = TorchLlmArgs(
-            model="/tmp/dummy_model",
-            kv_cache_config=KvCacheConfig(tokens_per_block=32,
-                                          enable_swa_scratch_reuse=False),
-        )
-
-        applied = _apply_kv_cache_config_preferences(llm_args,
-                                                     DeepseekV4ForCausalLM)
-
-        assert applied == {}
-        assert llm_args.kv_cache_config.tokens_per_block == 32
-        assert llm_args.kv_cache_config.enable_swa_scratch_reuse is False
-
-
 @pytest.mark.cpu_only
 class TestModelDefaults:
     """Test suite for model-specific default overrides functionality."""
@@ -700,26 +522,6 @@ def test_compute_applied_llm_defaults_simple_field(self):
         applied = apply_model_defaults_to_llm_args(llm_args, model_defaults)
         assert applied == model_defaults
 
-    def test_kv_cache_manager_version_preference_is_rejected(self):
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
-
-        with pytest.raises(ValueError,
-                           match="get_preferred_kv_cache_manager_version"):
-            apply_model_defaults_to_llm_args(
-                llm_args,
-                {"kv_cache_config": {
-                    "use_kv_cache_manager_v2": True
-                }})
-
-    def test_kv_cache_config_object_cannot_bypass_preference_guard(self):
-        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
-
-        with pytest.raises(ValueError, match="get_preferred_"):
-            apply_model_defaults_to_llm_args(llm_args, {
-                "kv_cache_config":
-                KvCacheConfig(use_kv_cache_manager_v2=True)
-            })
-
     @pytest.mark.parametrize(
         "defaults_dict,should_raise,error_contains",
         [
@@ -817,6 +619,111 @@ def get_model_defaults(cls, llm_args):
         assert "enable_block_reuse" in error_str or "max_tokens" in error_str
 
 
+@pytest.mark.cpu_only
+class TestKvCacheManagerV2AutoResolution:
+    """Test model preferences for the KV cache manager version."""
+
+    class _PreferV2:
+
+        @classmethod
+        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+            return "V2"
+
+    class _PreferV1:
+
+        @classmethod
+        def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+            return "V1"
+
+    @pytest.mark.parametrize("explicit_auto", [False, True])
+    def test_auto_uses_model_preference(self, explicit_auto):
+        kv_cache_config = (KvCacheConfig(use_kv_cache_manager_v2="auto")
+                           if explicit_auto else KvCacheConfig())
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model",
+                                kv_cache_config=kv_cache_config)
+
+        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
+
+        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
+
+    def test_auto_without_preference_falls_back_to_v1(self):
+        llm_args = TorchLlmArgs(model="/tmp/dummy_model")
+
+        _resolve_kv_cache_manager_v2_auto(llm_args)
+
+        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
+
+    @pytest.mark.parametrize(
+        ("backend", "runtime"),
+        [
+            ("NIXL", "CPP"),
+            ("UCX", None),
+            ("MPI", None),
+        ],
+    )
+    def test_auto_v2_falls_back_for_incompatible_disagg(self, backend, runtime):
+        llm_args = TorchLlmArgs(
+            model="/tmp/dummy_model",
+            cache_transceiver_config=CacheTransceiverConfig(
+                backend=backend, transceiver_runtime=runtime),
+        )
+
+        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
+
+        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False
+
+    def test_auto_v2_keeps_python_nixl_preference(self):
+        llm_args = TorchLlmArgs(
+            model="/tmp/dummy_model",
+            cache_transceiver_config=CacheTransceiverConfig(
+                backend="NIXL", transceiver_runtime="PYTHON"),
+        )
+
+        _resolve_kv_cache_manager_v2_auto(llm_args, self._PreferV2)
+
+        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True
+
+    @pytest.mark.parametrize(("user_setting", "model_cls"), [
+        pytest.param(False, _PreferV2, id="explicit-v1"),
+        pytest.param(True, _PreferV1, id="explicit-v2"),
+    ])
+    def test_explicit_value_overrides_model_preference(self, user_setting,
+                                                       model_cls):
+        llm_args = TorchLlmArgs(
+            model="/tmp/dummy_model",
+            kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=user_setting))
+
+        _resolve_kv_cache_manager_v2_auto(llm_args, model_cls)
+
+        assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is user_setting
+
+    def test_registered_models_prefer_v2(self):
+        from tensorrt_llm._torch.models.modeling_utils import \
+            get_registered_model_class
+
+        architectures = (
+            "DeepseekV3ForCausalLM",
+            "DeepseekV32ForCausalLM",
+            "GlmMoeDsaForCausalLM",
+            "GptOssForCausalLM",
+            "MistralLarge3ForCausalLM",
+            "DeepseekV4ForCausalLM",
+            "KimiK25ForConditionalGeneration",
+            "MiniMaxM2ForCausalLM",
+            "NemotronHForCausalLM",
+            "NemotronHPuzzleForCausalLM",
+            "Qwen3NextForCausalLM",
+            "Qwen3_5MoeForCausalLM",
+            "Qwen3_5ForCausalLM",
+            "Qwen3_5MoeForConditionalGeneration",
+            "Qwen3_5ForConditionalGeneration",
+        )
+        for architecture in architectures:
+            model_cls = get_registered_model_class(architecture)
+            assert model_cls is not None
+            assert model_cls.get_preferred_kv_cache_manager_version() == "V2"
+
+
 @pytest.mark.cpu_only
 def test_KvCacheConfig_declaration():
     assert KvCacheConfig().mamba_state_cache_interval is None
@@ -4106,7 +4013,6 @@ def test_model_loader_full_chain_aggregated_stays_none(self, monkeypatch):
         model_loader_mod.ModelLoader.load_config_and_apply_defaults(
             "/tmp/dummy_model", args, fake_loader)
         assert args.cache_transceiver_config is None
-        assert args.kv_cache_config.use_kv_cache_manager_v2 is True
 
     def test_resolve_default_backend_env_priority(self, monkeypatch):
         env_vars = ("TRTLLM_USE_NIXL_KVCACHE", "TRTLLM_USE_UCX_KVCACHE",
@@ -4170,12 +4076,6 @@ def test_prefers_python_without_config(self) -> None:
         assert DeepseekV3ForCausalLM.get_preferred_transceiver_runtime(
         ) == "PYTHON"
 
-    def test_prefers_kv_cache_manager_v2(self) -> None:
-        from tensorrt_llm._torch.models.modeling_deepseekv3 import \
-            DeepseekV3ForCausalLM
-        assert DeepseekV3ForCausalLM.get_preferred_kv_cache_manager_version(
-        ) == "V2"
-
     def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None:
         """DeepseekV3ForCausalLM on NIXL adopts the Python transceiver from 'auto'."""
         from tensorrt_llm._torch.models.modeling_deepseekv3 import \
diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py
index 6d336998b06d..d6b736a4d4bd 100644
--- a/tests/unittest/others/test_cache_transceiver_precheck_run.py
+++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py
@@ -397,6 +397,26 @@ def test_serving_resolvers(self, api):
         rt = inspect.signature(api.resolve_transceiver_runtime_auto).parameters
         assert list(rt)[:1] == ["llm_args"] and len(rt) >= 3
 
+    def test_model_preference_resolver_shim_supports_v2(self, api, monkeypatch):
+        class _PreferV2:
+            @classmethod
+            def get_preferred_kv_cache_manager_version(cls, pretrained_config=None):
+                return "V2"
+
+        monkeypatch.setattr(rp, "load_internal_apis", lambda: api)
+        monkeypatch.setattr(
+            rp,
+            "_lookup_model_cls",
+            lambda model_dir: (_PreferV2, types.SimpleNamespace()),
+        )
+        cache_cfg = api.CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON")
+
+        assert rp.resolve_model_prefs(
+            "/tmp/dummy_model",
+            {"use_kv_cache_manager_v2": "auto"},
+            cache_cfg,
+        )
+
     def test_enum_members(self, api):
         for enum, members in (
             (api.DataType, ("FP8", "HALF", "BF16")),