diff --git a/cpp/kernels/fmha_v2/fmha_test.py b/cpp/kernels/fmha_v2/fmha_test.py
index 96b1f0e457e3..94bebf718645 100644
--- a/cpp/kernels/fmha_v2/fmha_test.py
+++ b/cpp/kernels/fmha_v2/fmha_test.py
@@ -178,8 +178,7 @@ def test_trtllm_context_mla_attention_fmha(dtype, s):
         check=True)
 
     # For chunked prefill, we need to enable -save-softmax (dtype: bf16, layout: separate-q-k-v).
-    # Currently fp8 kernel doesn't support saving softmax.
-    if dtype == "-bf16":
+    if dtype in ["-bf16", "-e4m3"]:
         # padding mask
         subprocess.run(
             f"bin/fmha.exe -v 0 -runs 1 -min-s 1024 -s {s} -b 8 -h 8 -d 192 -dv 128 {dtype} "
diff --git a/cpp/kernels/fmha_v2/setup.py b/cpp/kernels/fmha_v2/setup.py
index 2d8a6b416a63..24a80b8d7134 100644
--- a/cpp/kernels/fmha_v2/setup.py
+++ b/cpp/kernels/fmha_v2/setup.py
@@ -3815,124 +3815,126 @@ def enumerate_qgmma_flash_warpspec_kernels(specs,
     combinations = product([False, True], \
         [InputLayout.PACKED_QKV, InputLayout.CONTIGUOUS_Q_KV,
          InputLayout.Q_PAGED_KV, InputLayout.SEPARATE_Q_K_V],
-        [False, True])
-    for (alibi, input_layout, enable_attn_logit_softcapping) in combinations:
+        [False, True], [False, True])
+    for (alibi, input_layout, enable_attn_logit_softcapping,
+         return_softmax) in combinations:
         # alibi and bmm1_tanh_scale shouldn't be used together.
         if alibi and enable_attn_logit_softcapping:
             continue
-        # D <= 64: KV_STEP = 256
-        specs.append(
-            kernel_spec(
-                sm=sm,
-                sm_mma=90,
-                dtype=dtype,
-                seq_len=0,  # support any sequence length
-                head_size=[32, 40, 48, 64],
-                warps_m=4,  #4x1 warpgroups
-                warps_n=1,
-                version=2,
-                interleaved=False,
-                ldgsts_q=
-                False,  # for Hopper kernels, ldgsts = False signals TMA usage.
-                ldgsts_k=False,
-                ldgsts_v=False,
-                share_smem_k_v=False,
-                loop_step=64,
-                q_tile_buffers=1,  # only used by warp specialized kernels
-                has_noloop=0,
-                noloop_step=64,
-                kv_loop_step=256,
-                kv_tile_buffers=4,  # only used by warp specialized kernels
-                unroll_threshold=1,
-                has_scale_max=False,
-                flash_attention=True,
-                warp_specialization=True,
-                alibi=alibi,
-                enable_attn_logit_softcapping=enable_attn_logit_softcapping,
-                return_softmax_stats=
-                False,  # return softmax stats is not supported for fp8 now
-                scheduling_mode=scheduling_mode,
-                input_layout=input_layout,
-                sage_block_sizes=sage_block_sizes,
-                output_dtype=output_dtype))
-
-        # 64 < D <=128: KV_STEP = 128
-        specs.append(
-            kernel_spec(
-                sm=sm,
-                sm_mma=90,
-                dtype=dtype,
-                seq_len=0,  # support any sequence length
-                head_size=[80, 96, 104, 128],
-                warps_m=4,  #4x1 warpgroups
-                warps_n=1,
-                version=2,
-                interleaved=False,
-                ldgsts_q=
-                False,  # for Hopper kernels, ldgsts = False signals TMA usage.
-                ldgsts_k=False,
-                ldgsts_v=False,
-                share_smem_k_v=False,
-                loop_step=64,
-                q_tile_buffers=1,  # only used by warp specialized kernels
-                has_noloop=0,
-                noloop_step=64,
-                kv_loop_step=256,
-                kv_tile_buffers=2,  # only used by warp specialized kernels
-                unroll_threshold=1,
-                has_scale_max=False,
-                flash_attention=True,
-                warp_specialization=True,
-                alibi=alibi,
-                enable_attn_logit_softcapping=enable_attn_logit_softcapping,
-                return_softmax_stats=
-                False,  # return softmax stats is not supported for fp8 now
-                scheduling_mode=scheduling_mode,
-                input_layout=input_layout,
-                sage_block_sizes=sage_block_sizes,
-                output_dtype=output_dtype))
-
-        # 128 < D <=256: KV_STEP = 128
-        specs.append(
-            kernel_spec(
-                sm=sm,
-                sm_mma=90,
-                dtype=dtype,
-                seq_len=0,  # support any sequence length
-                head_size=[160, 192, 256],
-                warps_m=4,  #4x1 warpgroups
-                warps_n=1,
-                version=2,
-                interleaved=False,
-                ldgsts_q=
-                False,  # for Hopper kernels, ldgsts = False signals TMA usage.
-                ldgsts_k=False,
-                ldgsts_v=False,
-                share_smem_k_v=False,
-                loop_step=64,
-                q_tile_buffers=1,  # only used by warp specialized kernels
-                has_noloop=0,
-                noloop_step=64,
-                kv_loop_step=
-                128,  # use 128 kv step size to avoid register spilling
-                kv_tile_buffers=2,  # only used by warp specialized kernels
-                unroll_threshold=1,
-                has_scale_max=False,
-                flash_attention=True,
-                warp_specialization=True,
-                alibi=alibi,
-                enable_attn_logit_softcapping=enable_attn_logit_softcapping,
-                return_softmax_stats=
-                False,  # return softmax stats is not supported for fp8 now
-                scheduling_mode=scheduling_mode,
-                input_layout=input_layout,
-                sage_block_sizes=sage_block_sizes,
-                output_dtype=output_dtype))
-
-        # context MLA (192x128)
-        # we could use param 'output_dtype' of enumerate_qgmma_flash_warpspec_kernels(),
-        # but it will generate many unnecessary kernels and they are not easy to filter out.
-        for output_type in [None, 'bf16']:
+        # for normal attention, we do not need return softmax for ws fp8 kernels currently.
+        # also fp8 input and bf16 output is only needed for MLA kernel.
+        skip_combination = return_softmax or (output_dtype is not None)
+        # for context mla, we need separate qkv as input layout when returning softmax.
+        skip_mla_combination = return_softmax and input_layout != InputLayout.SEPARATE_Q_K_V
+        if not skip_combination:
+            # D <= 64: KV_STEP = 256
+            specs.append(
+                kernel_spec(
+                    sm=sm,
+                    sm_mma=90,
+                    dtype=dtype,
+                    seq_len=0,  # support any sequence length
+                    head_size=[32, 40, 48, 64],
+                    warps_m=4,  #4x1 warpgroups
+                    warps_n=1,
+                    version=2,
+                    interleaved=False,
+                    ldgsts_q=
+                    False,  # for Hopper kernels, ldgsts = False signals TMA usage.
+                    ldgsts_k=False,
+                    ldgsts_v=False,
+                    share_smem_k_v=False,
+                    loop_step=64,
+                    q_tile_buffers=1,  # only used by warp specialized kernels
+                    has_noloop=0,
+                    noloop_step=64,
+                    kv_loop_step=256,
+                    kv_tile_buffers=4,  # only used by warp specialized kernels
+                    unroll_threshold=1,
+                    has_scale_max=False,
+                    flash_attention=True,
+                    warp_specialization=True,
+                    alibi=alibi,
+                    enable_attn_logit_softcapping=enable_attn_logit_softcapping,
+                    return_softmax_stats=return_softmax,
+                    scheduling_mode=scheduling_mode,
+                    input_layout=input_layout,
+                    sage_block_sizes=sage_block_sizes,
+                    output_dtype=output_dtype))
+
+            # 64 < D <=128: KV_STEP = 128
+            specs.append(
+                kernel_spec(
+                    sm=sm,
+                    sm_mma=90,
+                    dtype=dtype,
+                    seq_len=0,  # support any sequence length
+                    head_size=[80, 96, 104, 128],
+                    warps_m=4,  #4x1 warpgroups
+                    warps_n=1,
+                    version=2,
+                    interleaved=False,
+                    ldgsts_q=
+                    False,  # for Hopper kernels, ldgsts = False signals TMA usage.
+                    ldgsts_k=False,
+                    ldgsts_v=False,
+                    share_smem_k_v=False,
+                    loop_step=64,
+                    q_tile_buffers=1,  # only used by warp specialized kernels
+                    has_noloop=0,
+                    noloop_step=64,
+                    kv_loop_step=256,
+                    kv_tile_buffers=2,  # only used by warp specialized kernels
+                    unroll_threshold=1,
+                    has_scale_max=False,
+                    flash_attention=True,
+                    warp_specialization=True,
+                    alibi=alibi,
+                    enable_attn_logit_softcapping=enable_attn_logit_softcapping,
+                    return_softmax_stats=return_softmax,
+                    scheduling_mode=scheduling_mode,
+                    input_layout=input_layout,
+                    sage_block_sizes=sage_block_sizes,
+                    output_dtype=output_dtype))
+
+            # 128 < D <=256: KV_STEP = 128
+            specs.append(
+                kernel_spec(
+                    sm=sm,
+                    sm_mma=90,
+                    dtype=dtype,
+                    seq_len=0,  # support any sequence length
+                    head_size=[160, 192, 256],
+                    warps_m=4,  #4x1 warpgroups
+                    warps_n=1,
+                    version=2,
+                    interleaved=False,
+                    ldgsts_q=
+                    False,  # for Hopper kernels, ldgsts = False signals TMA usage.
+                    ldgsts_k=False,
+                    ldgsts_v=False,
+                    share_smem_k_v=False,
+                    loop_step=64,
+                    q_tile_buffers=1,  # only used by warp specialized kernels
+                    has_noloop=0,
+                    noloop_step=64,
+                    kv_loop_step=
+                    128,  # use 128 kv step size to avoid register spilling
+                    kv_tile_buffers=2,  # only used by warp specialized kernels
+                    unroll_threshold=1,
+                    has_scale_max=False,
+                    flash_attention=True,
+                    warp_specialization=True,
+                    alibi=alibi,
+                    enable_attn_logit_softcapping=enable_attn_logit_softcapping,
+                    return_softmax_stats=return_softmax,
+                    scheduling_mode=scheduling_mode,
+                    input_layout=input_layout,
+                    sage_block_sizes=sage_block_sizes,
+                    output_dtype=output_dtype))
+
+        if not skip_mla_combination:
+            # context MLA (192x128)
             specs.append(
                 kernel_spec(
                     sm=sm,
@@ -3962,12 +3964,11 @@ def enumerate_qgmma_flash_warpspec_kernels(specs,
                     warp_specialization=True,
                     alibi=alibi,
                     enable_attn_logit_softcapping=enable_attn_logit_softcapping,
-                    return_softmax_stats=
-                    False,  # return softmax stats is not supported for fp8 now
+                    return_softmax_stats=return_softmax,
                     scheduling_mode=scheduling_mode,
                     input_layout=input_layout,
                     sage_block_sizes=sage_block_sizes,
-                    output_dtype=output_type))
+                    output_dtype=output_dtype))
 
 
 def enumerate_igmma_kernels(specs, sm=90):
@@ -6215,6 +6216,10 @@ def enumerate_kernels():
     enumerate_hgmma_flash_warpspec_kernels(specs, sm=90, dtype='fp16')
     enumerate_hgmma_flash_warpspec_kernels(specs, sm=90, dtype='bf16')
     enumerate_qgmma_flash_warpspec_kernels(specs, sm=90, dtype='e4m3')
+    enumerate_qgmma_flash_warpspec_kernels(specs,
+                                           sm=90,
+                                           dtype='e4m3',
+                                           output_dtype="bf16")
 
     # For now SageAttention only needs BF16
     # block_size_q should be divisible by 64
diff --git a/cpp/kernels/fmha_v2/src/fmha/fragment.h b/cpp/kernels/fmha_v2/src/fmha/fragment.h
index 8be16df1636d..0c0700674ee3 100644
--- a/cpp/kernels/fmha_v2/src/fmha/fragment.h
+++ b/cpp/kernels/fmha_v2/src/fmha/fragment.h
@@ -1749,6 +1749,98 @@ struct Tile_o_normalizer<Ada_qmma_e4m3_fp32_traits, Cta_tile>
 
     // Default ctor
     Tile_o_normalizer() = default;
+
+    // The fragment accumulator.
+    using Fragment_accu = Fragment_accumulator<Traits>;
+
+    // The Mma tile.
+    using Mma_tile = typename Traits::template Mma_tile<Cta_tile>;
+
+    // The number of MMAs in the M dimension.
+    enum
+    {
+        MMAS_M = Mma_tile::MMAS_M
+    };
+
+    // The number of MMAs in the N dimension.
+    enum
+    {
+        MMAS_N = Mma_tile::VALID_MMAS_N
+    };
+
+    // The number of rows per thread.
+    enum
+    {
+        ROWS_PER_THREAD = 2 * MMAS_M
+    };
+
+    // The number of registers per thread.
+    enum
+    {
+        REGS_PER_THREAD = 8
+    };
+
+    // Warps.
+    enum
+    {
+        WARPS_M = Cta_tile::WARPS_M
+    };
+
+    enum
+    {
+        WARPS_N = Cta_tile::WARPS_N
+    };
+
+    enum
+    {
+        WARPS_K = Cta_tile::WARPS_K
+    };
+
+    // softmax data bytes
+    enum
+    {
+        BYTES_PER_ELEMENT = sizeof(float)
+    };
+
+    // Update o after P * V, the only difference from the basic class is we need to dequant the sum for softmax saver.
+    inline __device__ void final_update(Fragment_accu (&acc_o)[MMAS_M][MMAS_N], float (&sum)[ROWS_PER_THREAD])
+    {
+
+        constexpr float dequant_scale = Traits::SOFTMAX_FP_DEQUANT_SCALE;
+#pragma unroll
+        for (int mi = 0; mi < MMAS_M; ++mi)
+        {
+
+            // Precompute the scaling factors for the 2 rows.
+            float beta[2];
+#pragma unroll
+            for (int ii = 0; ii < 2; ++ii)
+            {
+                // The row.
+                int jj = 2 * mi + ii;
+
+                // The diviser.
+                beta[ii] = (sum[jj] == 0.f || sum[jj] != sum[jj]) ? 1.f : 1.f / sum[jj];
+                // softmax saver need the original sum.
+                sum[jj] = sum[jj] * dequant_scale;
+            }
+
+#pragma unroll
+            for (int ni = 0; ni < MMAS_N; ++ni)
+            {
+#pragma unroll
+                for (int ii = 0; ii < REGS_PER_THREAD; ++ii)
+                {
+                    // The register for O.
+                    float acc_o_f = acc_o[mi][ni].elt(ii);
+                    // Compute the next accumulator.
+                    acc_o_f = acc_o_f * beta[(ii & 2) / 2];
+                    // Update the accumulator.
+                    acc_o[mi][ni].elt(ii) = acc_o_f;
+                }
+            }
+        }
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/cpp/kernels/fmha_v2/src/fmha/warpspec/epilogue.h b/cpp/kernels/fmha_v2/src/fmha/warpspec/epilogue.h
index 99ea1643cd0e..d76f02021a87 100644
--- a/cpp/kernels/fmha_v2/src/fmha/warpspec/epilogue.h
+++ b/cpp/kernels/fmha_v2/src/fmha/warpspec/epilogue.h
@@ -1318,6 +1318,11 @@ struct Tile_o_epilogue<Hopper_qgmma_e4m3_fp32_traits, Kernel_traits>
 #else
             float scale = global_sum_mi == 0.f ? 1.0f : 1.0f / global_sum_mi;
 #endif
+            if constexpr (Kernel_traits::RETURN_SOFTMAX_STATS)
+            {
+                // Save the dequant exp sum for softmax saver.
+                global_sum[mi] *= Traits_o::SOFTMAX_FP_DEQUANT_SCALE;
+            }
 // Assume only N has multiple MMAs (MMAS_M = 1).
 #pragma unroll
             for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
diff --git a/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp b/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp
index 7623f6a76fc5..521de8f21ce8 100644
--- a/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp
+++ b/cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp
@@ -898,11 +898,6 @@ int main(int argc, char** argv)
                 "'-save-softmax'.\n");
             exit(1);
         }
-        if (data_type == DATA_TYPE_E4M3)
-        {
-            fprintf(stderr, "Currently fp8 kernel doesn't support saving softmax.\n");
-            exit(1);
-        }
     }
     // Sanitize
     if (min_s == -1)
diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp
index 23806ec21c7c..1f7c8040ade1 100644
--- a/cpp/tensorrt_llm/common/attentionOp.cpp
+++ b/cpp/tensorrt_llm/common/attentionOp.cpp
@@ -764,8 +764,8 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
     if (mEnableContextFMHA && mFP8ContextMLA && mFmhaDispatcher->isSeparateQAndKvInput())
     {
         fp8_q_buf_size = max_num_tokens * static_cast<size_t>(total_q_dim_all_heads);
-        fp8_k_buf_size = max_num_tokens * static_cast<size_t>(total_k_dim_all_heads);
-        fp8_v_buf_size = max_num_tokens * static_cast<size_t>(total_v_dim_all_heads);
+        fp8_k_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast<size_t>(total_k_dim_all_heads);
+        fp8_v_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast<size_t>(total_v_dim_all_heads);
     }
 
     size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens;
diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h
index 6007342a1d19..e7a02f752da9 100644
--- a/cpp/tensorrt_llm/common/attentionOp.h
+++ b/cpp/tensorrt_llm/common/attentionOp.h
@@ -414,6 +414,7 @@ class AttentionOp
     bool mFP8AttenOutput = false;
     bool mFP8ContextMLA = false;
     bool mFP8GenerationMLA = false;
+    size_t mChunkPrefillBufferBatchSize = 1;
     bool mDenseContextFMHA = false;
     bool mHasFullAttentionMask = false;
     bool mIsSpecDecodingEnabled = false;
@@ -467,8 +468,8 @@ class AttentionOp
             (int8_t) mPositionEmbeddingType, mUseLognScaling, mRemovePadding, (int32_t) mMaskType,
             mBlockSparseParams.data(), mPagedKVCache, mTokensPerBlock, mKVCacheQuantMode.value(), mTpSize, mTpRank,
             mUnfuseQkvGemm, (int32_t) mType, mMaxContextLength, mQKVBiasEnabled, mCrossAttention, mMaxDistance,
-            mPosShiftEnabled, mPagedContextFMHA, mFP8ContextFMHA, mFP8AttenOutput, mDenseContextFMHA,
-            mHasFullAttentionMask, mIsSpecDecodingEnabled, mUseSpecDecoding, mIsSpecDecTree,
+            mPosShiftEnabled, mPagedContextFMHA, mFP8ContextFMHA, mChunkPrefillBufferBatchSize, mFP8AttenOutput,
+            mDenseContextFMHA, mHasFullAttentionMask, mIsSpecDecodingEnabled, mUseSpecDecoding, mIsSpecDecTree,
             mSpecDecodingIsGenerationLengthVariable, mSpecDecodingMaxGenerationLength, mIsMLAEnabled, mIsGenerationMLA,
             mUseGenFlashMLA, mMLAParams.data(), mCpSize, mCpRank, mCpGroup, mNumAttnHeads, mNumAttnKVHeads,
             mNumKVHeadsOrigin, mAttnTpSize, mAttnTpRank, mAttnCpSize, mAttnCpRank, mUlyssesMQABroadcast,
diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp
index 7c19b51f8ad7..f88ee31e18f7 100644
--- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp
+++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp
@@ -486,13 +486,13 @@ void FusedMHARunnerV2::setupLaunchParams(MHARunnerParams runnerParams)
     }
     else
     {
-        bool isHopperBF16ContextMLA = (mFixedParams.headSize == mFixedParams.headSizeV + 64) && isSm90
-            && mFixedParams.dataType == DATA_TYPE_BF16 && mFixedParams.headSizeV == 128;
+        bool isHopperContextMLA = (mFixedParams.headSize == mFixedParams.headSizeV + 64) && isSm90
+            && (mFixedParams.dataType == DATA_TYPE_BF16 || mFixedParams.dataType == DATA_TYPE_E4M3)
+            && mFixedParams.headSizeV == 128;
         mLaunchParams.supportReturnSoftmaxStats = (runnerParams.softmaxStatsPtr != nullptr
             && mLaunchParams.flash_attention && mLaunchParams.warp_specialization
-            && ((!isHopperBF16ContextMLA
-                    && mLaunchParams.attention_input_layout == AttentionInputLayout::Q_CONTIGUOUS_KV)
-                || (isHopperBF16ContextMLA
+            && ((!isHopperContextMLA && mLaunchParams.attention_input_layout == AttentionInputLayout::Q_CONTIGUOUS_KV)
+                || (isHopperContextMLA
                     && (mLaunchParams.attention_input_layout == AttentionInputLayout::SEPARATE_Q_K_V))));
     }
 }
diff --git a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu
index 215dd5a9c3c0..97fd88a50ec0 100644
--- a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu
+++ b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cu
@@ -224,8 +224,8 @@ __global__ void mergeAttnWithSoftmaxKernel(T* merged_attn, float2* merged_softma
 // k_pe_output {total_chunk_token, h=1, d_rope}
 template <typename T, typename TCache>
 __global__ void loadChunkedKVCacheForMLAKernel(T* output_kv_ptr, T* output_k_pe_ptr,
-    tensorrt_llm::kernels::KVBlockArray const kv_cache, int64_t const* cu_ctx_chunked_len, int chunked_size,
-    int chunked_idx, float const* kv_scale_quant_orig_ptr)
+    tensorrt_llm::kernels::KVBlockArray const kv_cache, int64_t const* cu_ctx_chunked_len,
+    int64_t const* chunked_ld_global_offset, float const* kv_scale_quant_orig_ptr)
 {
     static_assert(std::is_same_v<T, TCache> || std::is_same_v<TCache, __nv_fp8_e4m3>,
         "TCache must be either the same type as T or __nv_fp8_e4m3");
@@ -238,6 +238,7 @@ __global__ void loadChunkedKVCacheForMLAKernel(T* output_kv_ptr, T* output_k_pe_
     size_t const head_dim_idx = head_dim_vec_idx * KT::kElemPerLoad;
 
     int64_t const real_chunked_size = cu_ctx_chunked_len[batch_idx + 1] - cu_ctx_chunked_len[batch_idx];
+    int64_t const global_ld_offset = chunked_ld_global_offset[batch_idx];
     int64_t const global_st_offset = cu_ctx_chunked_len[batch_idx];
     if (real_chunked_size <= 0)
     {
@@ -247,45 +248,41 @@ __global__ void loadChunkedKVCacheForMLAKernel(T* output_kv_ptr, T* output_k_pe_
     for (int local_token_idx = (threadIdx.x / KT::kThreadPerHead) + blockIdx.x * KT::kTokenPerBlock;
          local_token_idx < real_chunked_size; local_token_idx += gridDim.x * KT::kTokenPerBlock)
     {
-        int token_idx_in_kv_cache = (chunked_idx * chunked_size) + local_token_idx;
-        bool const valid_token = (local_token_idx < chunked_size);
-        if (valid_token)
+        int token_idx_in_kv_cache = global_ld_offset + local_token_idx;
+
+        auto* kvSrc = reinterpret_cast<T*>(kv_cache.getKBlockPtr(batch_idx, token_idx_in_kv_cache));
+        // head_idx === 0
+        auto kvBlockIdx
+            = kv_cache.getKVLocalIdx(token_idx_in_kv_cache, 0, KT::kVecPerHead, static_cast<int>(head_dim_vec_idx));
+        auto ld_data = (reinterpret_cast<typename KT::VecT*>(kvSrc))[kvBlockIdx];
+        if (is_valid_kv)
         {
-            auto* kvSrc = reinterpret_cast<T*>(kv_cache.getKBlockPtr(batch_idx, token_idx_in_kv_cache));
-            // head_idx === 0
-            auto kvBlockIdx
-                = kv_cache.getKVLocalIdx(token_idx_in_kv_cache, 0, KT::kVecPerHead, static_cast<int>(head_dim_vec_idx));
-            auto ld_data = (reinterpret_cast<typename KT::VecT*>(kvSrc))[kvBlockIdx];
-            if (is_valid_kv)
+            // kv_output {total_chunk_token, h=1, d}
+            int const global_st_idx = global_st_offset * KT::kLoraSize + local_token_idx * KT::kLoraSize + head_dim_idx;
+            if constexpr (std::is_same_v<TCache, T>)
             {
-                // kv_output {total_chunk_token, h=1, d}
-                int const global_st_idx
-                    = global_st_offset * KT::kLoraSize + local_token_idx * KT::kLoraSize + head_dim_idx;
-                if constexpr (std::is_same_v<TCache, T>)
-                {
-                    *reinterpret_cast<typename KT::VecT*>(output_kv_ptr + global_st_idx) = ld_data;
-                }
-                else if constexpr (std::is_same_v<TCache, __nv_fp8_e4m3>)
-                {
-                    dequantCopy<T, KT::kElemPerLoad>(output_kv_ptr + global_st_idx,
-                        reinterpret_cast<__nv_fp8_e4m3 const*>(&ld_data), kv_scale_quant_orig);
-                }
+                *reinterpret_cast<typename KT::VecT*>(output_kv_ptr + global_st_idx) = ld_data;
             }
-            else
+            else if constexpr (std::is_same_v<TCache, __nv_fp8_e4m3>)
             {
-                // k_pe_output {total_chunk_token, h=1, d_rope}
-                int const global_st_idx = global_st_offset * KT::kRopeSize + local_token_idx * KT::kRopeSize
-                    + (head_dim_idx - KT::kLoraSize);
+                dequantCopy<T, KT::kElemPerLoad>(output_kv_ptr + global_st_idx,
+                    reinterpret_cast<__nv_fp8_e4m3 const*>(&ld_data), kv_scale_quant_orig);
+            }
+        }
+        else
+        {
+            // k_pe_output {total_chunk_token, h=1, d_rope}
+            int const global_st_idx
+                = global_st_offset * KT::kRopeSize + local_token_idx * KT::kRopeSize + (head_dim_idx - KT::kLoraSize);
 
-                if constexpr (std::is_same_v<TCache, T>)
-                {
-                    *reinterpret_cast<typename KT::VecT*>(output_k_pe_ptr + global_st_idx) = ld_data;
-                }
-                else if constexpr (std::is_same_v<TCache, __nv_fp8_e4m3>)
-                {
-                    dequantCopy<T, KT::kElemPerLoad>(output_k_pe_ptr + global_st_idx,
-                        reinterpret_cast<__nv_fp8_e4m3 const*>(&ld_data), kv_scale_quant_orig);
-                }
+            if constexpr (std::is_same_v<TCache, T>)
+            {
+                *reinterpret_cast<typename KT::VecT*>(output_k_pe_ptr + global_st_idx) = ld_data;
+            }
+            else if constexpr (std::is_same_v<TCache, __nv_fp8_e4m3>)
+            {
+                dequantCopy<T, KT::kElemPerLoad>(output_k_pe_ptr + global_st_idx,
+                    reinterpret_cast<__nv_fp8_e4m3 const*>(&ld_data), kv_scale_quant_orig);
             }
         }
     }
@@ -324,17 +321,17 @@ void invokeMergeAttnWithSoftmax(T* merged_attn, float* merged_softmax_stats, T c
 // load single chunk kv from kv_cache for each request
 template <typename T, typename TCache>
 void invokeMLALoadChunkedKV(T* output_kv_ptr, T* output_k_pe_ptr, KVBlockArray const& kv_cache, int const num_contexts,
-    int64_t const* cu_ctx_chunked_len, int lora_size, int rope_size, int chunked_size, int chunked_idx,
-    float const* kv_scale_quant_orig_ptr, cudaStream_t stream)
+    int64_t const* cu_ctx_chunked_len, int64_t const* chunked_ld_global_offset, int lora_size, int rope_size,
+    int max_seq_len, float const* kv_scale_quant_orig_ptr, cudaStream_t stream)
 {
     using KT = loadChunkedKVKernelTraits<TCache>;
     TLLM_CHECK_WITH_INFO(lora_size + rope_size == KT::kHeadSize, "head dim should be equal to %d", KT::kHeadSize);
     TLLM_CHECK_WITH_INFO(lora_size == KT::kLoraSize, "lora dim should be equal to %d", KT::kLoraSize);
     TLLM_CHECK_WITH_INFO(rope_size == KT::kRopeSize, "rope dim should be equal to %d", KT::kRopeSize);
     // {chunked_unit_size / token_per_block, batch_size, head_num}
-    dim3 grid(static_cast<int>(tensorrt_llm::common::divUp(chunked_size, KT::kTokenPerBlock)), num_contexts, 1);
+    dim3 grid(static_cast<int>(tensorrt_llm::common::divUp(max_seq_len, KT::kTokenPerBlock)), num_contexts, 1);
     loadChunkedKVCacheForMLAKernel<T, TCache><<<grid, KT::kBlockSize, 0, stream>>>(output_kv_ptr, output_k_pe_ptr,
-        kv_cache, cu_ctx_chunked_len, chunked_size, chunked_idx, kv_scale_quant_orig_ptr);
+        kv_cache, cu_ctx_chunked_len, chunked_ld_global_offset, kv_scale_quant_orig_ptr);
 }
 
 #define INSTANTIATE_MLA_CHUNKED_PREFILL_KERNEL(T)                                                                      \
@@ -343,11 +340,12 @@ void invokeMLALoadChunkedKV(T* output_kv_ptr, T* output_k_pe_ptr, KVBlockArray c
         int64_t const* cu_q_seq_len, int max_q_seq_len, int64_t const* merge_op, int const num_heads,                  \
         int const head_size, cudaStream_t stream);                                                                     \
     template void invokeMLALoadChunkedKV<T, T>(T * output_kv_ptr, T * output_k_pe_ptr, KVBlockArray const& kv_cache,   \
-        int const num_contexts, int64_t const* cu_ctx_chunked_len, int lora_size, int rope_size, int chunked_size,     \
-        int chunked_idx, float const* kv_scale_quant_orig_ptr, cudaStream_t stream);                                   \
+        int const num_contexts, int64_t const* cu_ctx_chunked_len, int64_t const* chunked_ld_global_offset,            \
+        int lora_size, int rope_size, int max_seq_len, float const* kv_scale_quant_orig_ptr, cudaStream_t stream);     \
     template void invokeMLALoadChunkedKV<T, __nv_fp8_e4m3>(T * output_kv_ptr, T * output_k_pe_ptr,                     \
-        KVBlockArray const& kv_cache, int const num_contexts, int64_t const* cu_ctx_chunked_len, int lora_size,        \
-        int rope_size, int chunked_size, int chunked_idx, float const* kv_scale_quant_orig_ptr, cudaStream_t stream);
+        KVBlockArray const& kv_cache, int const num_contexts, int64_t const* cu_ctx_chunked_len,                       \
+        int64_t const* chunked_ld_global_offset, int lora_size, int rope_size, int max_seq_len,                        \
+        float const* kv_scale_quant_orig_ptr, cudaStream_t stream);
 
 INSTANTIATE_MLA_CHUNKED_PREFILL_KERNEL(half);
 INSTANTIATE_MLA_CHUNKED_PREFILL_KERNEL(float);
diff --git a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh
index 3b339ac1a85b..551e6d79a56b 100644
--- a/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh
+++ b/cpp/tensorrt_llm/kernels/mlaChunkedPrefill.cuh
@@ -35,7 +35,7 @@ void invokeMergeAttnWithSoftmax(T* merged_attn, float* merged_softmax_stats, T c
 // load single chunk kv from kv_cache for each request
 template <typename T, typename TCache>
 void invokeMLALoadChunkedKV(T* output_kv_ptr, T* output_k_pe_ptr, KVBlockArray const& kv_cache, int const num_contexts,
-    int64_t const* cu_ctx_chunked_len, int lora_size, int rope_size, int chunked_size, int chunked_idx,
-    float const* kv_scale_quant_orig_ptr, cudaStream_t stream);
+    int64_t const* cu_ctx_chunked_len, int64_t const* chunked_ld_global_offset, int lora_size, int rope_size,
+    int max_seq_len, float const* kv_scale_quant_orig_ptr, cudaStream_t stream);
 } // namespace kernels
 } // namespace tensorrt_llm
diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp
index a2b6a0f27e05..0dc89684ab83 100644
--- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp
+++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp
@@ -48,7 +48,8 @@ void initBindings(nb::module_& m)
         nb::arg("q_scaling"), nb::arg("position_embedding_type"), nb::arg("rotary_embedding_dim"),
         nb::arg("rotary_embedding_base"), nb::arg("rotary_embedding_scale_type"), nb::arg("rotary_embedding_scales"),
         nb::arg("rotary_embedding_max_position_info"), nb::arg("use_paged_context_fmha"),
-        nb::arg("attention_input_type") = std::nullopt, nb::arg("is_mla_enable"), nb::arg("q_lora_rank") = std::nullopt,
+        nb::arg("attention_input_type") = std::nullopt, nb::arg("is_mla_enable"),
+        nb::arg("chunked_prefill_buffer_batch_size") = std::nullopt, nb::arg("q_lora_rank") = std::nullopt,
         nb::arg("kv_lora_rank") = std::nullopt, nb::arg("qk_nope_head_dim") = std::nullopt,
         nb::arg("qk_rope_head_dim") = std::nullopt, nb::arg("v_head_dim") = std::nullopt,
         nb::arg("mrope_rotary_cos_sin") = std::nullopt, nb::arg("mrope_position_deltas") = std::nullopt,
diff --git a/cpp/tensorrt_llm/pybind/thop/bindings.cpp b/cpp/tensorrt_llm/pybind/thop/bindings.cpp
index ea5934514e14..0ccb2cc19c23 100644
--- a/cpp/tensorrt_llm/pybind/thop/bindings.cpp
+++ b/cpp/tensorrt_llm/pybind/thop/bindings.cpp
@@ -48,7 +48,8 @@ void initBindings(pybind11::module_& m)
         py::arg("q_scaling"), py::arg("position_embedding_type"), py::arg("rotary_embedding_dim"),
         py::arg("rotary_embedding_base"), py::arg("rotary_embedding_scale_type"), py::arg("rotary_embedding_scales"),
         py::arg("rotary_embedding_max_position_info"), py::arg("use_paged_context_fmha"),
-        py::arg("attention_input_type") = std::nullopt, py::arg("is_mla_enable"), py::arg("q_lora_rank") = std::nullopt,
+        py::arg("attention_input_type") = std::nullopt, py::arg("is_mla_enable"),
+        py::arg("chunked_prefill_buffer_batch_size") = std::nullopt, py::arg("q_lora_rank") = std::nullopt,
         py::arg("kv_lora_rank") = std::nullopt, py::arg("qk_nope_head_dim") = std::nullopt,
         py::arg("qk_rope_head_dim") = std::nullopt, py::arg("v_head_dim") = std::nullopt,
         py::arg("mrope_rotary_cos_sin") = std::nullopt, py::arg("mrope_position_deltas") = std::nullopt,
diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp
index 2911cfa6e675..6a13a5ad268e 100644
--- a/cpp/tensorrt_llm/thop/attentionOp.cpp
+++ b/cpp/tensorrt_llm/thop/attentionOp.cpp
@@ -495,7 +495,8 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     double const rotary_embedding_base, int64_t const rotary_embedding_scale_type,
     std::vector<double> rotary_embedding_scales, std::vector<int64_t> rotary_embedding_max_position_info,
     bool const use_paged_context_fmha, std::optional<int64_t> attention_input_type, bool is_mla_enable,
-    std::optional<int64_t> q_lora_rank, std::optional<int64_t> kv_lora_rank, std::optional<int64_t> qk_nope_head_dim,
+    std::optional<int64_t> chunked_prefill_buffer_batch_size, std::optional<int64_t> q_lora_rank,
+    std::optional<int64_t> kv_lora_rank, std::optional<int64_t> qk_nope_head_dim,
     std::optional<int64_t> qk_rope_head_dim, std::optional<int64_t> v_head_dim,
     std::optional<torch::Tensor> mrope_rotary_cos_sin, std::optional<torch::Tensor> mrope_position_deltas,
     std::optional<int64_t> attention_chunk_size, std::optional<torch::Tensor> softmax_stats_tensor,
@@ -628,7 +629,8 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
             static_cast<int>(v_head_dim.value()), static_cast<int>(predicted_tokens_per_seq),
             static_cast<int>(layer_num)};
 
-        op->mFP8ContextMLA = tensorrt_llm::common::getSMVersion() == 120 && op->mKVCacheQuantMode.hasFp8KvCache();
+        op->mFP8ContextMLA = (tensorrt_llm::common::getSMVersion() == 120 || tensorrt_llm::common::getSMVersion() == 90)
+            && op->mKVCacheQuantMode.hasFp8KvCache();
         op->mIsGenerationMLA = head_size == op->mMLAParams.kv_lora_rank + op->mMLAParams.qk_rope_head_dim;
         op->mFP8GenerationMLA = op->mKVCacheQuantMode.hasFp8KvCache();
         // only enable flash mla on sm90 and head_size == 576 and tokens_per_block == 64
@@ -639,6 +641,10 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
         // mNumKVHeads/mHeadSize are overwritten in common/attentionOp.cpp.
         op->mNumKVHeads = 1;
         op->mHeadSize = op->mMLAParams.kv_lora_rank + op->mMLAParams.qk_rope_head_dim;
+
+        // For chunked prefill MLA, we need larger buffer size for k and v
+        op->mChunkPrefillBufferBatchSize
+            = chunked_prefill_buffer_batch_size.has_value() ? chunked_prefill_buffer_batch_size.value() : 1;
     }
 
     auto cache_key = std::make_tuple(op->data(), runner->data());
diff --git a/cpp/tensorrt_llm/thop/attentionOp.h b/cpp/tensorrt_llm/thop/attentionOp.h
index 402bd96fccd3..ec18cc46fade 100644
--- a/cpp/tensorrt_llm/thop/attentionOp.h
+++ b/cpp/tensorrt_llm/thop/attentionOp.h
@@ -54,7 +54,8 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
     double const rotary_embedding_base, int64_t const rotary_embedding_scale_type,
     std::vector<double> rotary_embedding_scales, std::vector<int64_t> rotary_embedding_max_position_info,
     bool const use_paged_context_fmha, std::optional<int64_t> attention_input_type, bool is_mla_enable,
-    std::optional<int64_t> q_lora_rank, std::optional<int64_t> kv_lora_rank, std::optional<int64_t> qk_nope_head_dim,
+    std::optional<int64_t> chunked_prefill_buffer_batch_size, std::optional<int64_t> q_lora_rank,
+    std::optional<int64_t> kv_lora_rank, std::optional<int64_t> qk_nope_head_dim,
     std::optional<int64_t> qk_rope_head_dim, std::optional<int64_t> v_head_dim,
     torch::optional<torch::Tensor> mrope_rotary_cos_sin, torch::optional<torch::Tensor> mrope_position_deltas,
     std::optional<int64_t> attention_chunk_size, std::optional<torch::Tensor> softmax_stats_tensor,
diff --git a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp
index 882c00f7a444..6dfffec54d82 100644
--- a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp
+++ b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp
@@ -50,16 +50,16 @@ void loadPagedKVCacheForMLAHelper(torch::Tensor& compressed_kv, torch::Tensor& k
 
 template <typename T, typename TCache>
 void loadChunkedKVCacheForMLAHelper(torch::Tensor& output_kv, torch::Tensor& output_k_pe, KVBlockArray& kv_cache,
-    int const num_contexts, torch::Tensor const& cu_ctx_chunked_len, int lora_size, int rope_size,
-    int const chunked_size, int const chunked_idx, float const* kv_scale_quant_orig_ptr)
+    int const num_contexts, torch::Tensor const& cu_ctx_chunked_len, torch::Tensor const& chunked_ld_global_offset,
+    int lora_size, int rope_size, int const max_seq_len, float const* kv_scale_quant_orig_ptr)
 {
     auto stream = at::cuda::getCurrentCUDAStream(output_kv.get_device());
 
     T* output_kv_ptr = static_cast<T*>(output_kv.data_ptr());
     T* output_k_pe_ptr = static_cast<T*>(output_k_pe.data_ptr());
     tensorrt_llm::kernels::invokeMLALoadChunkedKV<T, TCache>(output_kv_ptr, output_k_pe_ptr, kv_cache, num_contexts,
-        cu_ctx_chunked_len.data_ptr<int64_t>(), lora_size, rope_size, chunked_size, chunked_idx,
-        kv_scale_quant_orig_ptr, stream);
+        cu_ctx_chunked_len.data_ptr<int64_t>(), chunked_ld_global_offset.data_ptr<int64_t>(), lora_size, rope_size,
+        max_seq_len, kv_scale_quant_orig_ptr, stream);
 }
 
 template <typename T, typename TCache>
@@ -249,13 +249,13 @@ std::vector<torch::Tensor> loadPagedKVCacheForMLA(torch::ScalarType out_dtype, i
 }
 
 std::vector<torch::Tensor> loadChunkedKVCacheForMLA(torch::ScalarType out_dtype, int64_t const num_contexts,
-    int64_t const num_ctx_cached_tokens, torch::Tensor& cu_ctx_chunked_kv_lens,
-    torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers,
-    torch::Tensor const& host_kv_cache_pool_mapping, torch::optional<torch::Tensor> kv_scale_orig_quant,
-    torch::optional<torch::Tensor> kv_scale_quant_orig, int64_t const layer_idx, int64_t const lora_size,
-    int64_t const rope_size, int64_t const tokens_per_block, int64_t const chunked_size, int64_t const chunked_index,
-    int64_t const attention_window_size, int64_t const sink_token_length, int64_t const beam_width,
-    int64_t const quant_mode)
+    int64_t const num_ctx_cached_tokens, torch::Tensor const& cu_ctx_chunked_kv_lens,
+    torch::Tensor const& chunked_ld_global_offset, torch::Tensor const& kv_cache_block_offsets,
+    torch::Tensor const& host_kv_cache_pool_pointers, torch::Tensor const& host_kv_cache_pool_mapping,
+    torch::optional<torch::Tensor> kv_scale_orig_quant, torch::optional<torch::Tensor> kv_scale_quant_orig,
+    int64_t const layer_idx, int64_t const lora_size, int64_t const rope_size, int64_t const tokens_per_block,
+    int64_t const max_seq_len, int64_t const attention_window_size, int64_t const sink_token_length,
+    int64_t const beam_width, int64_t const quant_mode)
 {
     TORCH_CHECK(out_dtype == torch::kFloat16 || out_dtype == torch::kFloat32 || out_dtype == torch::kBFloat16,
         "out_dtype only support float16, float32, bfloat16");
@@ -298,12 +298,14 @@ std::vector<torch::Tensor> loadChunkedKVCacheForMLA(torch::ScalarType out_dtype,
         if (kv_cache_quant_mode.hasFp8KvCache())
         {
             loadChunkedKVCacheForMLAHelper<half, __nv_fp8_e4m3>(outputs[0], outputs[1], kv_cache_buffer, num_contexts,
-                cu_ctx_chunked_kv_lens, lora_size, rope_size, chunked_size, chunked_index, kv_scale_quant_orig_ptr);
+                cu_ctx_chunked_kv_lens, chunked_ld_global_offset, lora_size, rope_size, max_seq_len,
+                kv_scale_quant_orig_ptr);
         }
         else
         {
             loadChunkedKVCacheForMLAHelper<half, half>(outputs[0], outputs[1], kv_cache_buffer, num_contexts,
-                cu_ctx_chunked_kv_lens, lora_size, rope_size, chunked_size, chunked_index, kv_scale_quant_orig_ptr);
+                cu_ctx_chunked_kv_lens, chunked_ld_global_offset, lora_size, rope_size, max_seq_len,
+                kv_scale_quant_orig_ptr);
         }
     }
     else if (out_dtype == torch::kFloat32)
@@ -311,12 +313,14 @@ std::vector<torch::Tensor> loadChunkedKVCacheForMLA(torch::ScalarType out_dtype,
         if (kv_cache_quant_mode.hasFp8KvCache())
         {
             loadChunkedKVCacheForMLAHelper<float, __nv_fp8_e4m3>(outputs[0], outputs[1], kv_cache_buffer, num_contexts,
-                cu_ctx_chunked_kv_lens, lora_size, rope_size, chunked_size, chunked_index, kv_scale_quant_orig_ptr);
+                cu_ctx_chunked_kv_lens, chunked_ld_global_offset, lora_size, rope_size, max_seq_len,
+                kv_scale_quant_orig_ptr);
         }
         else
         {
             loadChunkedKVCacheForMLAHelper<float, float>(outputs[0], outputs[1], kv_cache_buffer, num_contexts,
-                cu_ctx_chunked_kv_lens, lora_size, rope_size, chunked_size, chunked_index, kv_scale_quant_orig_ptr);
+                cu_ctx_chunked_kv_lens, chunked_ld_global_offset, lora_size, rope_size, max_seq_len,
+                kv_scale_quant_orig_ptr);
         }
     }
     else if (out_dtype == torch::kBFloat16)
@@ -324,13 +328,13 @@ std::vector<torch::Tensor> loadChunkedKVCacheForMLA(torch::ScalarType out_dtype,
         if (kv_cache_quant_mode.hasFp8KvCache())
         {
             loadChunkedKVCacheForMLAHelper<__nv_bfloat16, __nv_fp8_e4m3>(outputs[0], outputs[1], kv_cache_buffer,
-                num_contexts, cu_ctx_chunked_kv_lens, lora_size, rope_size, chunked_size, chunked_index,
+                num_contexts, cu_ctx_chunked_kv_lens, chunked_ld_global_offset, lora_size, rope_size, max_seq_len,
                 kv_scale_quant_orig_ptr);
         }
         else
         {
             loadChunkedKVCacheForMLAHelper<__nv_bfloat16, __nv_bfloat16>(outputs[0], outputs[1], kv_cache_buffer,
-                num_contexts, cu_ctx_chunked_kv_lens, lora_size, rope_size, chunked_size, chunked_index,
+                num_contexts, cu_ctx_chunked_kv_lens, chunked_ld_global_offset, lora_size, rope_size, max_seq_len,
                 kv_scale_quant_orig_ptr);
         }
     }
@@ -503,6 +507,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         ", int num_contexts"
         ", int num_ctx_cached_tokens"
         ", Tensor cu_ctx_chunked_kv_lens"
+        ", Tensor chunked_ld_global_offset"
         ", Tensor kv_cache_block_offsets"
         ", Tensor host_kv_cache_pool_pointers"
         ", Tensor host_kv_cache_pool_mapping"
@@ -512,8 +517,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         ", int lora_size"
         ", int rope_size"
         ", int tokens_per_block"
-        ", int chunked_size"
-        ", int chunked_index"
+        ", int max_seq_len"
         ", int attention_window_size"
         ", int sink_token_length"
         ", int beam_width"
diff --git a/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu b/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu
index 8f4bd25a9a52..3e4e9a1da0a3 100644
--- a/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu
+++ b/cpp/tests/unit_tests/kernels/mlaChunkedPrefillTest.cu
@@ -23,8 +23,8 @@ namespace
 // k_pe_output {total_tokens, h=1, rope_size}
 template <typename T, typename TCache>
 void loadChunkedKVKernelRef(T* kv_output, T* k_pe_output, tensorrt_llm::kernels::KVBlockArray const& kv_cache,
-    int num_contexts, int64_t const* cu_ctx_chunked_len, int const lora_size, int const rope_size, int const chunk_size,
-    int const chunk_idx, float const* kv_scale_quant_orig_ptr)
+    int num_contexts, int64_t const* cu_ctx_chunked_len, int64_t const* chunked_ld_global_offset, int const lora_size,
+    int const rope_size, float const* kv_scale_quant_orig_ptr)
 {
     int const head_size = lora_size + rope_size;
     float const kv_scale_quant_orig = kv_scale_quant_orig_ptr ? kv_scale_quant_orig_ptr[0] : 1.0f;
@@ -33,8 +33,8 @@ void loadChunkedKVKernelRef(T* kv_output, T* k_pe_output, tensorrt_llm::kernels:
         int const chunked_len = cu_ctx_chunked_len[b + 1] - cu_ctx_chunked_len[b];
         for (int s = 0; s < chunked_len; s++)
         {
-            int const local_token_idx = chunk_idx * chunk_size + s;
-            int const ld_token_offset = (cu_ctx_chunked_len[b] + s);
+            int const local_token_idx = chunked_ld_global_offset[b] + s;
+            int const st_token_offset = (cu_ctx_chunked_len[b] + s);
 
             auto const* kv_src = reinterpret_cast<TCache const*>(kv_cache.getKBlockPtr(b, local_token_idx));
             for (int d = 0; d < head_size; d++)
@@ -52,11 +52,11 @@ void loadChunkedKVKernelRef(T* kv_output, T* k_pe_output, tensorrt_llm::kernels:
                 }
                 if (d < lora_size)
                 {
-                    kv_output[ld_token_offset * lora_size + d] = data;
+                    kv_output[st_token_offset * lora_size + d] = data;
                 }
                 else
                 {
-                    k_pe_output[ld_token_offset * rope_size + (d - lora_size)] = data;
+                    k_pe_output[st_token_offset * rope_size + (d - lora_size)] = data;
                 }
             }
         }
@@ -166,12 +166,12 @@ void selfAttentionRef(T* output, T* const Q, T* const KV, int batch_size, int nu
 // chunked_KV {total_chunk_token, 2, H, D}
 // KV {total_kv_token, 2, H, D}
 template <typename T>
-void copyRelatedChunkedKV(T* chunked_kv, T* const kv, int chunk_idx, int chunk_size, int batch_size, int num_heads,
-    int64_t* const cu_kv_seq_len, int64_t* const cu_chunked_seq_len, int head_size)
+void copyRelatedChunkedKV(T* chunked_kv, T* const kv, int64_t const* chunked_ld_global_offset, int batch_size,
+    int num_heads, int64_t* const cu_kv_seq_len, int64_t* const cu_chunked_seq_len, int head_size)
 {
     for (int b = 0; b < batch_size; b++)
     {
-        int src_global_offset = (cu_kv_seq_len[b] + chunk_idx * chunk_size) * 2 * num_heads * head_size;
+        int src_global_offset = (cu_kv_seq_len[b] + chunked_ld_global_offset[b]) * 2 * num_heads * head_size;
         int dst_global_offset = cu_chunked_seq_len[b] * 2 * num_heads * head_size;
         int copy_length = cu_chunked_seq_len[b + 1] - cu_chunked_seq_len[b];
         if (copy_length <= 0)
@@ -261,7 +261,8 @@ protected:
     tensorrt_llm::runtime::BufferManager::ITensorPtr h_kv_cache_tensor{nullptr}, h_kv_cache_tensor_ref{nullptr},
         d_kv_cache_tensor{nullptr}, h_compressed_kv_cache_tensor{nullptr}, d_compressed_kv_cache_tensor{nullptr},
         h_compressed_offset_tensor{nullptr}, d_compressed_offset_tensor{nullptr}, h_cu_kv_seq_lens{nullptr},
-        d_cu_kv_seq_lens{nullptr}, h_cu_chunk_lens{nullptr}, d_cu_chunk_lens{nullptr}, h_cu_q_seq_lens{nullptr},
+        d_cu_kv_seq_lens{nullptr}, h_cu_chunk_lens{nullptr}, d_cu_chunk_lens{nullptr},
+        h_chunked_ld_global_offset{nullptr}, d_chunked_ld_global_offset{nullptr}, h_cu_q_seq_lens{nullptr},
         d_cu_q_seq_lens{nullptr},
 
         // for kernel 1
@@ -293,6 +294,9 @@ protected:
     int mMaxBlockPerSeq{};
     bool mIsCausalMask{};
 
+    // for chunked main loop
+    std::vector<int> max_chunk_len_per_loop;
+
     std::mt19937 gen;
 
     void SetUp() override
@@ -342,6 +346,18 @@ protected:
         std::cout << std::endl;
     }
 
+    template <typename T>
+    void showHostTensor(tensorrt_llm::runtime::BufferManager::ITensorPtr& tensor, std::string const& tensor_name)
+    {
+        std::cout << "Tensor: " << tensor_name << ": \n";
+        auto* const ptr = reinterpret_cast<T*>(tensor->data());
+        for (int _ = 0; _ < tensor->getSize(); _++)
+        {
+            std::cout << static_cast<float>(ptr[_]) << " ";
+        }
+        std::cout << std::endl;
+    }
+
     int generateRandomSizeSmallerThan(int a)
     {
         if (a <= 0)
@@ -448,14 +464,10 @@ protected:
         // cu lens
         this->h_cu_kv_seq_lens = tensorrt_llm::runtime::BufferManager::pinned(
             ITensor::makeShape({this->mBatchSize + 1}), nvinfer1::DataType::kINT64);
-        this->h_cu_chunk_lens = tensorrt_llm::runtime::BufferManager::pinned(
-            ITensor::makeShape({this->mBatchSize + 1}), nvinfer1::DataType::kINT64);
         this->h_cu_q_seq_lens = tensorrt_llm::runtime::BufferManager::pinned(
             ITensor::makeShape({this->mBatchSize + 1}), nvinfer1::DataType::kINT64);
         this->d_cu_kv_seq_lens = tensorrt_llm::runtime::BufferManager::gpuSync(
             this->h_cu_kv_seq_lens->getShape(), nvinfer1::DataType::kINT64);
-        this->d_cu_chunk_lens = tensorrt_llm::runtime::BufferManager::gpuSync(
-            this->h_cu_chunk_lens->getShape(), nvinfer1::DataType::kINT64);
         this->d_cu_q_seq_lens = tensorrt_llm::runtime::BufferManager::gpuSync(
             this->h_cu_q_seq_lens->getShape(), nvinfer1::DataType::kINT64);
         {
@@ -492,10 +504,23 @@ protected:
             cudaMemcpy(this->d_cu_q_seq_lens->data(), this->h_cu_q_seq_lens->data(),
                 this->h_cu_q_seq_lens->getSizeInBytes(), cudaMemcpyHostToDevice);
 #ifdef TRTLLM_MLA_CHUNKED_PREFILL_TEST_DBG
-            this->showHostTensor<int64_t>(this->h_cu_q_seq_lens);
-            this->showHostTensor<int64_t>(this->h_cu_kv_seq_lens);
+            this->showHostTensor<int64_t>(this->h_cu_q_seq_lens, "cu_q_seq_lens");
+            this->showHostTensor<int64_t>(this->h_cu_kv_seq_lens, "cu_kv_seq_lens");
 #endif
         }
+        int const total_chunk_size = this->mChunkSize * this->mBatchSize;
+        int const total_cached_kv_len = this->mTotalKVLen - this->mTotalQLen;
+        int const chunked_loop_num = (total_cached_kv_len + total_chunk_size - 1) / total_chunk_size;
+        this->h_cu_chunk_lens = tensorrt_llm::runtime::BufferManager::pinned(
+            ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize + 1}), nvinfer1::DataType::kINT64);
+        this->h_chunked_ld_global_offset = tensorrt_llm::runtime::BufferManager::pinned(
+            ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize}), nvinfer1::DataType::kINT64);
+        this->memsetZeroHost(this->h_chunked_ld_global_offset);
+        this->d_cu_chunk_lens = tensorrt_llm::runtime::BufferManager::gpuSync(
+            this->h_cu_chunk_lens->getShape(), nvinfer1::DataType::kINT64);
+        this->d_chunked_ld_global_offset = tensorrt_llm::runtime::BufferManager::gpuSync(
+            this->h_chunked_ld_global_offset->getShape(), nvinfer1::DataType::kINT64);
+
         // kv cache
         this->mMaxBlockPerSeq = (this->mMaxSeqLen + this->mTokensPerBlock - 1) / this->mTokensPerBlock;
         int maxChunkBlockPerSeq = (this->mChunkSize + this->mTokensPerBlock - 1) / this->mTokensPerBlock;
@@ -584,7 +609,7 @@ protected:
         this->m_h_output_tensor_accum = tensorrt_llm::runtime::BufferManager::pinned(
             ITensor::makeShape({this->mTotalQLen, this->mNumHeads, this->mNopeSize}), dtype);
         this->m_h_merge_op = tensorrt_llm::runtime::BufferManager::pinned(
-            ITensor::makeShape({this->mBatchSize}), nvinfer1::DataType::kINT64);
+            ITensor::makeShape({chunked_loop_num + 1, this->mBatchSize}), nvinfer1::DataType::kINT64);
         this->m_d_q_tensor = tensorrt_llm::runtime::BufferManager::gpuSync(this->m_h_q_tensor->getShape(), dtype);
         this->m_d_kv_full_tensor
             = tensorrt_llm::runtime::BufferManager::gpuSync(this->m_h_kv_full_tensor->getShape(), dtype);
@@ -631,9 +656,132 @@ protected:
                 m_h_output_tensor_accum->getSizeInBytes(), cudaMemcpyHostToDevice, mStream->get());
             cudaStreamSynchronize(mStream->get());
         }
+        this->prepareChunkedPrefillMetaData();
         return true;
     }
 
+    int prepareChunkedPrefillMetaData()
+    {
+        using tensorrt_llm::runtime::bufferCast;
+        int const total_chunk_size = this->mChunkSize * this->mBatchSize;
+        int chunked_loop_num = (this->mTotalKVLen - this->mTotalQLen + total_chunk_size - 1) / total_chunk_size;
+
+        auto* h_merge_op = bufferCast<int64_t>(*(this->m_h_merge_op));             // {chunked_loop_num + 1, batch_size}
+        auto* h_cu_q_seq_lens_ptr = bufferCast<int64_t>(*(this->h_cu_q_seq_lens)); // {batch_size + 1}
+        auto* h_cu_kv_seq_lens_ptr = bufferCast<int64_t>(*(this->h_cu_kv_seq_lens)); // {batch_size + 1}
+        auto* h_cu_chunk_lens_ptr
+            = bufferCast<int64_t>(*(this->h_cu_chunk_lens));            // {chunked_loop_num + 1, batch_size + 1}
+        auto* h_chunked_ld_global_offset_ptr
+            = bufferCast<int64_t>(*(this->h_chunked_ld_global_offset)); // {chunked_loop_num + 1, batch_size}
+
+        this->max_chunk_len_per_loop.clear();
+        std::vector<int64_t> chunked_seq_len_vec((chunked_loop_num + 1) * (this->mBatchSize), 0);
+        // 0 -> chunked_loop_num -1
+        int remain_buffer_len = total_chunk_size;
+        int curr_loop_idx = 0;
+        int temp_max_chunk_len = 0;
+
+#define chunked_seq_len(chunked_loop_idx, b_idx) chunked_seq_len_vec[(chunked_loop_idx) * (this->mBatchSize) + (b_idx)]
+#define cu_chunked_seq_len(chunked_loop_idx, b_idx)                                                                    \
+    h_cu_chunk_lens_ptr[(chunked_loop_idx) * (this->mBatchSize + 1) + (b_idx)]
+#define chunked_ld_global_offset(chunked_loop_idx, b_idx)                                                              \
+    h_chunked_ld_global_offset_ptr[(chunked_loop_idx) * (this->mBatchSize) + (b_idx)]
+
+        for (int b = 0; b < this->mBatchSize; b++)
+        {
+            int temp_cached_kv_len = (h_cu_kv_seq_lens_ptr[b + 1] - h_cu_kv_seq_lens_ptr[b])
+                - (h_cu_q_seq_lens_ptr[b + 1] - h_cu_q_seq_lens_ptr[b]);
+            while (temp_cached_kv_len > 0)
+            {
+                auto used_buffer_len = std::min(remain_buffer_len, temp_cached_kv_len);
+                remain_buffer_len -= used_buffer_len;
+                temp_cached_kv_len -= used_buffer_len;
+                temp_max_chunk_len = std::max(temp_max_chunk_len, used_buffer_len);
+                chunked_seq_len(curr_loop_idx, b) = used_buffer_len;
+                chunked_ld_global_offset(curr_loop_idx + 1, b)
+                    = chunked_ld_global_offset(curr_loop_idx, b) + used_buffer_len;
+                if (remain_buffer_len == 0)
+                {
+                    this->max_chunk_len_per_loop.push_back(temp_max_chunk_len);
+                    temp_max_chunk_len = 0;
+                    remain_buffer_len = total_chunk_size;
+                    curr_loop_idx++;
+                }
+            }
+        }
+        if (this->max_chunk_len_per_loop.size() < chunked_loop_num)
+        {
+            this->max_chunk_len_per_loop.push_back(temp_max_chunk_len);
+        }
+        assert(this->max_chunk_len_per_loop.size() == chunked_loop_num);
+
+        // for not cached part
+        for (int b = 0; b < this->mBatchSize; b++)
+        {
+            int uncached_len = (h_cu_q_seq_lens_ptr[b + 1] - h_cu_q_seq_lens_ptr[b]);
+            chunked_seq_len(chunked_loop_num, b) = uncached_len;
+        }
+        for (int loop_idx = 0; loop_idx < chunked_loop_num + 1; loop_idx++)
+        {
+            for (int b = 0; b < this->mBatchSize; b++)
+            {
+                cu_chunked_seq_len(loop_idx, b + 1) = cu_chunked_seq_len(loop_idx, b) + chunked_seq_len(loop_idx, b);
+            }
+        }
+        // merge op
+        for (int loop_idx = 0; loop_idx < chunked_loop_num; loop_idx++)
+        {
+            for (int b = 0; b < this->mBatchSize; b++)
+            {
+                if (chunked_seq_len(loop_idx, b) != 0 && (loop_idx == 0 || chunked_seq_len(loop_idx - 1, b) == 0))
+                {
+                    h_merge_op[loop_idx * (this->mBatchSize) + b] = 2; // copy
+                }
+                else if (chunked_seq_len(loop_idx, b) != 0)
+                {
+                    h_merge_op[loop_idx * (this->mBatchSize) + b] = 1; // merge
+                }
+                else
+                {
+                    h_merge_op[loop_idx * (this->mBatchSize) + b] = 0; // skip
+                }
+            }
+        }
+        // for the last uncached part
+        for (int b = 0; b < this->mBatchSize; b++)
+        {
+            int temp_cached_kv_len = (h_cu_kv_seq_lens_ptr[b + 1] - h_cu_kv_seq_lens_ptr[b])
+                - (h_cu_q_seq_lens_ptr[b + 1] - h_cu_q_seq_lens_ptr[b]);
+            if (temp_cached_kv_len == 0)
+            {
+                h_merge_op[chunked_loop_num * (this->mBatchSize) + b] = 2; // copy
+            }
+            else
+            {
+                h_merge_op[chunked_loop_num * (this->mBatchSize) + b] = 1; // merge
+            }
+            chunked_ld_global_offset(chunked_loop_num, b) = temp_cached_kv_len;
+        }
+
+#undef chunked_seq_len
+#undef cu_chunked_seq_len
+#undef chunked_ld_global_offset
+        // copy to device
+        cudaMemcpy(this->d_cu_chunk_lens->data(), this->h_cu_chunk_lens->data(),
+            this->h_cu_chunk_lens->getSizeInBytes(), cudaMemcpyHostToDevice);
+        cudaMemcpy(this->d_chunked_ld_global_offset->data(), this->h_chunked_ld_global_offset->data(),
+            this->h_chunked_ld_global_offset->getSizeInBytes(), cudaMemcpyHostToDevice);
+        cudaMemcpy(this->m_d_merge_op->data(), this->m_h_merge_op->data(), this->m_h_merge_op->getSizeInBytes(),
+            cudaMemcpyHostToDevice);
+#ifdef TRTLLM_MLA_CHUNKED_PREFILL_TEST_DBG
+        std::cout << "chunked_loop_num: " << chunked_loop_num << '\n';
+        this->showHostTensor<int64_t>(this->m_h_merge_op, "merge_op");
+        this->showHostTensor<int64_t>(this->h_chunked_ld_global_offset, "chunked_ld_global_offset");
+        this->showHostTensor<int64_t>(this->h_cu_chunk_lens, "cu_chunk_lens");
+#endif
+        return chunked_loop_num;
+    }
+
     void PerformNormalAttention()
     {
         using tensorrt_llm::runtime::bufferCast;
@@ -661,7 +809,7 @@ protected:
         auto* h_cu_q_seq_lens_ptr = bufferCast<int64_t>(*(this->h_cu_q_seq_lens));
         auto* h_cu_kv_seq_lens_ptr = bufferCast<int64_t>(*(this->h_cu_kv_seq_lens));
         auto* h_cu_chunk_lens_ptr = bufferCast<int64_t>(*(this->h_cu_chunk_lens));
-        auto* h_merge_op = bufferCast<int64_t>(*(this->m_h_merge_op));
+        auto* h_chunked_ld_global_offset_ptr = bufferCast<int64_t>(*(this->h_chunked_ld_global_offset));
         auto* d_kv_ptr = bufferCast<DataType>(*(this->m_d_kv_full_tensor));
         auto* d_chunked_kv_ptr = bufferCast<DataType>(*(this->m_d_chunked_kv_tensor));
         auto* d_softmax_sum_ptr = bufferCast<float>(*(this->m_d_softmax_sum_tensor));
@@ -671,16 +819,14 @@ protected:
         auto* d_merge_op = bufferCast<int64_t>(*(this->m_d_merge_op));
         auto* d_cu_q_seq_lens_ptr = bufferCast<int64_t>(*(this->d_cu_q_seq_lens));
 
-        int const loop_count = (this->mMaxSeqLen + this->mChunkSize - 1) / this->mChunkSize;
+        int const total_chunk_size = this->mChunkSize * this->mBatchSize;
+        int chunked_loop_num = (this->mTotalKVLen - this->mTotalQLen + total_chunk_size - 1) / total_chunk_size;
         // do not apply mask
-        for (int _ = 0; _ < loop_count - 1; _++)
+        for (int _ = 0; _ < chunked_loop_num; _++)
         {
-            // get chunked len for each request
-            this->PrepareChunkedLen(_);
-            cudaMemcpy(d_merge_op, h_merge_op, this->m_h_merge_op->getSizeInBytes(), cudaMemcpyHostToDevice);
             // copy related kv chunk data
-            copyRelatedChunkedKV(h_chunked_kv_ptr, h_kv_ptr, _, this->mChunkSize, this->mBatchSize, this->mNumHeads,
-                h_cu_kv_seq_lens_ptr, h_cu_chunk_lens_ptr, this->mNopeSize);
+            copyRelatedChunkedKV(h_chunked_kv_ptr, h_kv_ptr, h_chunked_ld_global_offset_ptr, this->mBatchSize,
+                this->mNumHeads, h_cu_kv_seq_lens_ptr, h_cu_chunk_lens_ptr, this->mNopeSize);
             // attention
             selfAttentionRef<DataType>(h_output_ptr, h_q_ptr, h_chunked_kv_ptr, this->mBatchSize, this->mNumHeads,
                 h_cu_q_seq_lens_ptr, h_cu_chunk_lens_ptr, this->mNopeSize, true, h_softmax_sum_ptr, false);
@@ -700,19 +846,16 @@ protected:
                 cudaMemcpyDeviceToHost);
             cudaMemcpy(h_output_accum_ptr, d_output_accum_ptr, this->m_h_output_tensor->getSizeInBytes(),
                 cudaMemcpyDeviceToHost);
+            // update merge op, ld global offset, cu chunk lens ptr.
+            d_merge_op += this->mBatchSize;
+            h_cu_chunk_lens_ptr += (this->mBatchSize + 1);
+            h_chunked_ld_global_offset_ptr += this->mBatchSize;
         }
         // final round, apply causal mask.
 
         // copy the last chunked kv data
-        copyFinalChunkedKV<DataType>(h_chunked_kv_ptr, h_kv_ptr, this->mChunkSize, this->mBatchSize, this->mNumHeads,
-            h_cu_kv_seq_lens_ptr, h_cu_chunk_lens_ptr, this->mNopeSize, h_merge_op);
-        cudaMemcpy(d_merge_op, h_merge_op, this->m_h_merge_op->getSizeInBytes(), cudaMemcpyHostToDevice);
-#ifdef TRTLLM_MLA_CHUNKED_PREFILL_TEST_DBG
-        std::cout << "merge op: ";
-        this->showHostTensor<int64_t>(this->m_h_merge_op);
-        std::cout << "cu chunk lens: ";
-        this->showHostTensor<int64_t>(this->h_cu_chunk_lens);
-#endif
+        copyRelatedChunkedKV(h_chunked_kv_ptr, h_kv_ptr, h_chunked_ld_global_offset_ptr, this->mBatchSize,
+            this->mNumHeads, h_cu_kv_seq_lens_ptr, h_cu_chunk_lens_ptr, this->mNopeSize);
         // attention
         selfAttentionRef<DataType>(h_output_ptr, h_q_ptr, h_chunked_kv_ptr, this->mBatchSize, this->mNumHeads,
             h_cu_q_seq_lens_ptr, h_cu_chunk_lens_ptr, this->mNopeSize, true, h_softmax_sum_ptr, this->mIsCausalMask);
@@ -733,55 +876,6 @@ protected:
         sync_check_cuda_error(mStream->get());
     }
 
-    void PrepareChunkedLen(int chunk_idx)
-    {
-        using tensorrt_llm::runtime::bufferCast;
-        auto* h_merge_op = bufferCast<int64_t>(*(this->m_h_merge_op));
-        auto* h_cu_q_seq_lens_ptr = bufferCast<int64_t>(*(this->h_cu_q_seq_lens));
-        auto* h_cu_kv_seq_lens_ptr = bufferCast<int64_t>(*(this->h_cu_kv_seq_lens));
-        auto* h_cu_chunk_lens_ptr = bufferCast<int64_t>(*(this->h_cu_chunk_lens));
-
-        h_cu_chunk_lens_ptr[0] = 0;
-        for (int b = 0; b < this->mBatchSize; b++)
-        {
-            int curr_kv_len = h_cu_kv_seq_lens_ptr[b + 1] - h_cu_kv_seq_lens_ptr[b];
-            int used_kv_len = chunk_idx * this->mChunkSize;
-            int curr_chunk_len = std::min(this->mChunkSize, curr_kv_len - used_kv_len);
-            if (curr_chunk_len != this->mChunkSize)
-            {
-                // last chunk, we should skip it.
-                curr_chunk_len = 0;
-            }
-            else
-            {
-                if (used_kv_len + curr_chunk_len == curr_kv_len)
-                {
-                    // last chunk, we should skip it.
-                    curr_chunk_len = 0;
-                }
-            }
-            h_cu_chunk_lens_ptr[b + 1] = h_cu_chunk_lens_ptr[b] + curr_chunk_len;
-            if (chunk_idx == 0 && curr_chunk_len > 0)
-            {
-                h_merge_op[b] = 2; // only copy result
-            }
-            else if (curr_chunk_len > 0)
-            {
-                h_merge_op[b] = 1; // merge result
-            }
-            else
-            {
-                h_merge_op[b] = 0; // skip
-            }
-        }
-#ifdef TRTLLM_MLA_CHUNKED_PREFILL_TEST_DBG
-        std::cout << "merge op: ";
-        this->showHostTensor<int64_t>(this->m_h_merge_op);
-        std::cout << "cu chunk lens: ";
-        this->showHostTensor<int64_t>(this->h_cu_chunk_lens);
-#endif
-    }
-
     void PerformLoadChunkedKVRef(int chunk_idx)
     {
         using tensorrt_llm::runtime::bufferCast;
@@ -790,7 +884,9 @@ protected:
         auto* k_pe_output_ptr = bufferCast<DataType>(*(this->h_k_pe_output_ref));
         auto* compressed_kv_cache_ptr = bufferCast<TCache>(*(this->h_compressed_kv_cache_tensor));
         auto* offset_ptr = bufferCast<int32_t>(*(this->h_compressed_offset_tensor));
-        auto* h_cu_chunk_lens_ptr = bufferCast<int64_t>(*(this->h_cu_chunk_lens));
+        auto* h_cu_chunk_lens_ptr = bufferCast<int64_t>(*(this->h_cu_chunk_lens)) + chunk_idx * (this->mBatchSize + 1);
+        auto* h_chunked_ld_global_offset_ptr
+            = bufferCast<int64_t>(*(this->h_chunked_ld_global_offset)) + chunk_idx * this->mBatchSize;
         float* kv_scale_quant_orig_ptr = nullptr;
         if constexpr (std::is_same_v<TCache, __nv_fp8_e4m3>)
         {
@@ -799,10 +895,9 @@ protected:
         tensorrt_llm::kernels::KVBlockArray kv_cache(this->mBatchSize, this->mMaxBlockPerSeq, this->mTokensPerBlock,
             sizeof(TCache) * 1 * (this->mLoraSize + this->mRopeSize), 0, 0, 0, 0, compressed_kv_cache_ptr, nullptr,
             reinterpret_cast<tensorrt_llm::kernels::KVBlockArrayForContextFMHA::DataType*>(offset_ptr));
-        this->PrepareChunkedLen(chunk_idx);
 
         loadChunkedKVKernelRef<DataType, TCache>(compressed_kv_output_ptr, k_pe_output_ptr, kv_cache, this->mBatchSize,
-            h_cu_chunk_lens_ptr, this->mLoraSize, this->mRopeSize, this->mChunkSize, chunk_idx,
+            h_cu_chunk_lens_ptr, h_chunked_ld_global_offset_ptr, this->mLoraSize, this->mRopeSize,
             kv_scale_quant_orig_ptr);
     }
 
@@ -814,7 +909,9 @@ protected:
         auto* k_pe_output_ptr = bufferCast<DataType>(*(this->d_k_pe_output));
         auto* compressed_kv_cache_ptr = bufferCast<TCache>(*(this->d_compressed_kv_cache_tensor));
         auto* offset_ptr = bufferCast<int32_t>(*(this->d_compressed_offset_tensor));
-        auto* d_cu_chunk_lens_ptr = bufferCast<int64_t>(*(this->d_cu_chunk_lens));
+        auto* d_cu_chunk_lens_ptr = bufferCast<int64_t>(*(this->d_cu_chunk_lens)) + chunk_idx * (this->mBatchSize + 1);
+        auto* d_chunked_ld_global_offset_ptr
+            = bufferCast<int64_t>(*(this->d_chunked_ld_global_offset)) + chunk_idx * this->mBatchSize;
         float* kv_scale_quant_orig_ptr = nullptr;
         if constexpr (std::is_same_v<TCache, __nv_fp8_e4m3>)
         {
@@ -823,13 +920,12 @@ protected:
         tensorrt_llm::kernels::KVBlockArray kv_cache(this->mBatchSize, this->mMaxBlockPerSeq, this->mTokensPerBlock,
             sizeof(TCache) * 1 * (this->mLoraSize + this->mRopeSize), 0, 0, 0, 0, compressed_kv_cache_ptr, nullptr,
             reinterpret_cast<tensorrt_llm::kernels::KVBlockArrayForContextFMHA::DataType*>(offset_ptr));
-        this->PrepareChunkedLen(chunk_idx);
         // copy cu chunk lens to device
         cudaMemcpy(this->d_cu_chunk_lens->data(), this->h_cu_chunk_lens->data(),
             this->h_cu_chunk_lens->getSizeInBytes(), cudaMemcpyHostToDevice);
         tensorrt_llm::kernels::invokeMLALoadChunkedKV<DataType, TCache>(compressed_kv_output_ptr, k_pe_output_ptr,
-            kv_cache, this->mBatchSize, d_cu_chunk_lens_ptr, this->mLoraSize, this->mRopeSize, this->mChunkSize,
-            chunk_idx, kv_scale_quant_orig_ptr, mStream->get());
+            kv_cache, this->mBatchSize, d_cu_chunk_lens_ptr, d_chunked_ld_global_offset_ptr, this->mLoraSize,
+            this->mRopeSize, this->max_chunk_len_per_loop[chunk_idx], kv_scale_quant_orig_ptr, mStream->get());
         cudaStreamSynchronize(this->mStream->get());
         // copy result back to host
         cudaMemcpy(this->h_compressed_kv_output->data(), compressed_kv_output_ptr,
@@ -933,8 +1029,9 @@ TYPED_TEST(MlaChunkedPrefillTest, MlaChunkedLoad)
     sync_check_cuda_error(this->mStream->get());
     bool allEqual{true};
 
-    int const loop_count = (this->mMaxSeqLen + this->mChunkSize - 1) / this->mChunkSize;
-    for (int _ = 0; _ < loop_count - 1; _++)
+    int const total_chunk_size = this->mChunkSize * this->mBatchSize;
+    int const chunked_loop_num = (this->mTotalKVLen - this->mTotalQLen + total_chunk_size - 1) / total_chunk_size;
+    for (int _ = 0; _ < chunked_loop_num - 1; _++)
     {
         this->PerformLoadChunkedKVRef(_);
         sync_check_cuda_error(this->mStream->get());
diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py
index eb77fa0055d3..e960dcf6b513 100644
--- a/tensorrt_llm/_torch/attention_backend/interface.py
+++ b/tensorrt_llm/_torch/attention_backend/interface.py
@@ -25,6 +25,7 @@ class AttentionRuntimeFeatures:
     cache_reuse: bool = False
     has_speculative_draft_tokens: bool = False
     chunk_size: int = 0  # this is the chunk size for MLA chunked prefill, it will split kv cache into chunks to save global memory.
+    chunked_prefill_buffer_batch_size: int = 4  # real chunk size for MLA chunked prefill is chunked_prefill_buffer_batch_size * chunk_size.
 
 
 # The type of requests in qkv passed to attention
@@ -634,3 +635,4 @@ class MLAParams:
     qk_nope_head_dim: int = 0
     v_head_dim: int = 0
     predicted_tokens_per_seq: int = 1
+    chunked_prefill_buffer_batch_size: int = 1
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index cdca67a7b953..a317149b1b6b 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -67,6 +67,7 @@ class TrtllmAttentionWrapper:
     qk_rope_head_dim: Optional[int]
     qk_nope_head_dim: Optional[int]
     v_head_dim: Optional[int]
+    chunked_prefill_buffer_batch_size: Optional[int]
     attention_chunk_size: Optional[int]
     softmax_stats_tensor: Optional[torch.Tensor]
     use_spec_decoding: bool
@@ -187,6 +188,7 @@ def plan(
         spec_decoding_packed_mask: Optional[torch.Tensor] = None,
         spec_decoding_generation_lengths: Optional[torch.Tensor] = None,
         attention_sinks: Optional[torch.Tensor] = None,
+        chunked_prefill_buffer_batch_size: int = 1,
         **kwargs,
     ):
         """
@@ -224,6 +226,7 @@ def plan(
             mrope_config (dict): The dictionary containing the mRope configuration.
             softmax_stats_tensor (torch.Tensor): The tensor to store the softmax statistics (max/sum)
             attention_sinks (torch.Tensor): The attention sinks (additional value in the denominator of the softmax) with shape of (num_heads_q) on GPU.
+            chunked_prefill_buffer_batch_size (int): used for malloc buffer for k and v in fp8 context mla. the max input kv length is not max_num_tokens in this case. It is chunked_prefill_buffer_batch_size * max_num_tokens.
         """
         self.layer_idx = layer_idx
         self.tokens_per_block = tokens_per_block
@@ -271,6 +274,7 @@ def plan(
         self.spec_decoding_position_offsets = spec_decoding_position_offsets
         self.spec_decoding_packed_mask = spec_decoding_packed_mask
         self.spec_decoding_generation_lengths = spec_decoding_generation_lengths
+        self.chunked_prefill_buffer_batch_size = chunked_prefill_buffer_batch_size
         self.kwargs.update(kwargs)
 
     def create_output(self, q: torch.Tensor, out_dtype: torch.dtype):
@@ -470,6 +474,7 @@ def run(
             self.use_paged_context_fmha,
             self.attention_input_type,
             self.is_mla_enable,
+            self.chunked_prefill_buffer_batch_size,
             self.q_lora_rank,
             self.kv_lora_rank,
             self.qk_nope_head_dim,
@@ -858,8 +863,11 @@ def prepare_flash_mla(self) -> None:
     def pre_process_for_chunked_prefill(
         self,
         chunked_seq_len: torch.Tensor,
+        chunked_global_offset: torch.
+        Tensor,  # [chunked_loop_num + 1, num_contexts]
         cu_chunked_seq_len: torch.Tensor,
         merge_op_tensor: torch.Tensor,
+        max_chunk_len_per_loop: list[int],
         chunked_loop_num: int,
     ) -> None:
         """
@@ -868,24 +876,46 @@ def pre_process_for_chunked_prefill(
         """
         num_contexts = self.num_contexts
         chunk_size = self.runtime_features.chunk_size
-        cached_kv_lens = torch.tensor(
-            self.kv_cache_params.num_cached_tokens_per_seq,
-            dtype=torch.int,
-            device='cpu',
-        )
+        chunk_batch_size = self.runtime_features.chunked_prefill_buffer_batch_size
+        total_chunk_size = chunk_size * chunk_batch_size
+        remain_buffer_len = total_chunk_size
+        current_batch_idx = 0
+        max_chunk_len_per_loop.clear()
+        max_chunk_len = 0
+        # cal chunked_seq_len
+        for batch_idx in range(num_contexts):
+            cached_kv_len = self.kv_cache_params.num_cached_tokens_per_seq[
+                batch_idx]
+            while cached_kv_len > 0:
+                used_buffer_len = min(remain_buffer_len, cached_kv_len)
+                chunked_seq_len[current_batch_idx, batch_idx] = used_buffer_len
+                max_chunk_len = max(max_chunk_len, used_buffer_len)
+                remain_buffer_len -= used_buffer_len
+                cached_kv_len -= used_buffer_len
+                chunked_global_offset[
+                    current_batch_idx + 1, batch_idx] = chunked_global_offset[
+                        current_batch_idx,
+                        batch_idx] + chunked_seq_len[current_batch_idx,
+                                                     batch_idx]
+                if remain_buffer_len == 0:
+                    current_batch_idx += 1
+                    remain_buffer_len = total_chunk_size
+                    max_chunk_len_per_loop.append(max_chunk_len)
+                    max_chunk_len = 0
+        if len(max_chunk_len_per_loop) < chunked_loop_num:
+            max_chunk_len_per_loop.append(max_chunk_len)
+        assert len(
+            max_chunk_len_per_loop
+        ) == chunked_loop_num, f"max_chunk_len_per_loop size {len(max_chunk_len_per_loop)} != chunked_loop_num {chunked_loop_num}"
         for loop_idx in range(chunked_loop_num):
             cu_chunked_seq_len[loop_idx, 0] = 0
-            used_chunk_seq_len = loop_idx * chunk_size
-            chunked_seq_len[loop_idx, :num_contexts] = torch.clamp(
-                cached_kv_lens[:num_contexts] - used_chunk_seq_len,
-                min=0,
-                max=chunk_size)
             torch.cumsum(chunked_seq_len[loop_idx, :num_contexts],
                          dim=0,
                          dtype=torch.int64,
                          out=cu_chunked_seq_len[loop_idx, 1:num_contexts + 1])
             for s in range(num_contexts):
-                if loop_idx == 0 and chunked_seq_len[loop_idx, s] > 0:
+                if chunked_seq_len[loop_idx, s] > 0 and (
+                        loop_idx == 0 or chunked_seq_len[loop_idx - 1, s] == 0):
                     merge_op_tensor[loop_idx, s] = 2  # copy only
                 elif chunked_seq_len[loop_idx, s] > 0:
                     merge_op_tensor[loop_idx, s] = 1  # merge
@@ -894,7 +924,7 @@ def pre_process_for_chunked_prefill(
 
         # set merge op for last attn
         for s in range(num_contexts):
-            if cached_kv_lens[s] == 0:
+            if self.kv_cache_params.num_cached_tokens_per_seq[s] == 0:
                 merge_op_tensor[chunked_loop_num, s] = 2  # copy only
             else:
                 merge_op_tensor[chunked_loop_num, s] = 1  # merge
@@ -916,14 +946,16 @@ def prepare_context_mla_with_cached_kv(self,
             # currently we assume that the chunk size is the same as the max_num_tokens
             if self.runtime_features.chunked_prefill:
                 chunk_size = self.runtime_features.chunk_size
-                self.chunked_loop_num = (self.max_ctx_cached_token_len +
-                                         chunk_size - 1) // chunk_size
-                self.chunked_seq_len = torch.empty(
+                chunk_batch_size = self.runtime_features.chunked_prefill_buffer_batch_size
+                total_chunk_size = chunk_size * chunk_batch_size
+                self.chunked_loop_num = math.ceil(self.num_ctx_cached_tokens /
+                                                  total_chunk_size)
+                self.chunked_seq_len = torch.zeros(
                     (self.chunked_loop_num, self.num_seqs),
                     dtype=torch.int,
                     device='cuda',
                 )
-                self.host_chunked_seq_len = torch.empty_like(
+                self.host_chunked_seq_len = torch.zeros_like(
                     self.chunked_seq_len,
                     device='cpu',
                     pin_memory=True,
@@ -938,6 +970,17 @@ def prepare_context_mla_with_cached_kv(self,
                     device='cpu',
                     pin_memory=True,
                 )
+                self.chunked_global_offset = torch.zeros(
+                    (self.chunked_loop_num + 1, self.num_contexts),
+                    dtype=torch.int64,
+                    device='cuda',
+                )
+                self.host_chunked_global_offset = torch.zeros_like(
+                    self.chunked_global_offset,
+                    device='cpu',
+                    pin_memory=True,
+                )
+                self.max_chunk_len_per_loop = []
                 # For last chunk we use the uncached kv
                 self.merge_op_tensor = torch.empty(
                     (self.chunked_loop_num + 1, self.num_contexts),
@@ -952,8 +995,10 @@ def prepare_context_mla_with_cached_kv(self,
 
                 self.pre_process_for_chunked_prefill(
                     chunked_seq_len=self.host_chunked_seq_len,
+                    chunked_global_offset=self.host_chunked_global_offset,
                     cu_chunked_seq_len=self.host_cu_chunked_seq_len,
                     merge_op_tensor=self.host_merge_op_tensor,
+                    max_chunk_len_per_loop=self.max_chunk_len_per_loop,
                     chunked_loop_num=self.chunked_loop_num)
                 self.chunked_seq_len.copy_(self.host_chunked_seq_len,
                                            non_blocking=True)
@@ -961,6 +1006,8 @@ def prepare_context_mla_with_cached_kv(self,
                                               non_blocking=True)
                 self.merge_op_tensor.copy_(self.host_merge_op_tensor,
                                            non_blocking=True)
+                self.chunked_global_offset.copy_(
+                    self.host_chunked_global_offset, non_blocking=True)
         else:
             self.num_ctx_cached_tokens = 0
             self.max_ctx_cached_token_len = 0
@@ -1171,6 +1218,7 @@ def forward(
         output: Optional[torch.Tensor] = None,
         output_sf: Optional[torch.Tensor] = None,
         attention_sinks: Optional[torch.Tensor] = None,
+        chunked_prefill_buffer_batch_size: int = 1,
         **kwargs,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]:
         assert isinstance(
@@ -1220,7 +1268,8 @@ def forward(
             host_kv_cache_pool_pointers=metadata.host_kv_cache_pool_pointers,
             host_kv_cache_pool_mapping=metadata.host_kv_cache_pool_mapping,
             block_ids_per_seq=metadata.block_ids_per_seq,
-            workspace=None,
+            workspace=metadata.
+            workspace,  # re-enable it, if pass None to it, fp8 mla will encounter invalid cuda free issue.
             cache_indirection=metadata.cache_indirection,
             kv_scale_orig_quant=self.kv_scale_orig_quant,
             kv_scale_quant_orig=self.kv_scale_quant_orig,
@@ -1244,6 +1293,7 @@ def forward(
             spec_decoding_generation_lengths=metadata.
             spec_decoding_generation_lengths,
             attention_sinks=attention_sinks,
+            chunked_prefill_buffer_batch_size=chunked_prefill_buffer_batch_size,
         )
         out_dtype = None
         if out_scale is not None:
@@ -1350,9 +1400,10 @@ def load_paged_kv_cache_for_mla(
     def load_chunked_kv_cache_for_mla(
         self,
         metadata: TrtllmAttentionMetadata,
-        chunked_idx: int,
         num_ctx_cached_tokens: int,
         cu_chunked_seq_len: torch.Tensor,
+        chunked_global_offset: torch.Tensor,
+        chunked_max_seq_len: int,
         out_dtype: torch.dtype,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         assert out_dtype in [torch.float16, torch.bfloat16, torch.float32]
@@ -1376,6 +1427,7 @@ def load_chunked_kv_cache_for_mla(
             metadata.num_contexts,
             num_ctx_cached_tokens,
             cu_chunked_seq_len,
+            chunked_global_offset,
             metadata.kv_cache_block_offsets,
             metadata.kv_cache_manager.kv_cache_pool_pointers,
             metadata.kv_cache_manager.kv_cache_pool_mapping,
@@ -1385,8 +1437,7 @@ def load_chunked_kv_cache_for_mla(
             self.mla_params.kv_lora_rank,
             self.mla_params.qk_rope_head_dim,
             metadata.kv_cache_manager.tokens_per_block,
-            metadata.runtime_features.chunk_size,
-            chunked_idx,
+            chunked_max_seq_len,
             metadata.kv_cache_manager.max_seq_len,
             sink_token_length,
             beam_width,
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index b38eee88b9c8..c3b7e7ab4b01 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -1189,11 +1189,15 @@ def forward_context_with_chunked_prefill(
             temp_cu_chunked_seq_len = attn_metadata.cu_chunked_seq_len[loop_idx]
             total_ctx_chunked_tokens = attn_metadata.host_cu_chunked_seq_len[
                 loop_idx, attn_metadata.num_contexts]
+            chunked_global_offset = attn_metadata.chunked_global_offset[
+                loop_idx]
+            chunked_max_seq_len = attn_metadata.max_chunk_len_per_loop[loop_idx]
             chunked_compressed_kv, chunked_k_pe = trtllm_attention.load_chunked_kv_cache_for_mla(
                 metadata=attn_metadata,
-                chunked_idx=loop_idx,
                 num_ctx_cached_tokens=total_ctx_chunked_tokens,
                 cu_chunked_seq_len=temp_cu_chunked_seq_len,
+                chunked_global_offset=chunked_global_offset,
+                chunked_max_seq_len=chunked_max_seq_len,
                 out_dtype=q.dtype)
 
             # up proj to uncompressed kv
@@ -1242,6 +1246,8 @@ def forward_context_with_chunked_prefill(
                 out_scale=out_scale,
                 attention_mask=PredefinedAttentionMask.FULL,
                 softmax_stats_tensor=self.temp_softmax_stats_tensor,
+                chunked_prefill_buffer_batch_size=attn_metadata.
+                runtime_features.chunked_prefill_buffer_batch_size,
                 output=temp_attn_output,
             )
             # merge attn result
@@ -1292,6 +1298,8 @@ def forward_context_with_chunked_prefill(
             latent_cache=None,
             out_scale=out_scale,
             softmax_stats_tensor=self.temp_softmax_stats_tensor,
+            chunked_prefill_buffer_batch_size=attn_metadata.runtime_features.
+            chunked_prefill_buffer_batch_size,
             output=temp_attn_output,
         )
         temp_merge_op = attn_metadata.merge_op_tensor[chunked_loop_num]
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index cb045f537915..4901a2fca2f8 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -366,10 +366,10 @@ def drafting_loop_wrapper(model):
             )
             executor_config.kv_cache_config.enable_block_reuse = False
         if executor_config.enable_chunked_context and sm_version not in [
-                90, 100
+                90, 100, 120
         ]:
             logger.warning(
-                "Chunked Prefill for MLA can only be enabled on SM90/SM100, "
+                "Chunked Prefill for MLA can only be enabled on SM90/SM100/SM120, "
                 f"disable enable_chunked_context for SM{sm_version}")
             executor_config.enable_chunked_context = False
             model_engine.attn_runtime_features.chunked_prefill = False
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 999e9cf15bc4..67ff8eedd2aa 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -1718,13 +1718,11 @@ def test_no_kv_cache_reuse(self, quant_dtype, mtp_nextn, fp8kv,
         (True, True),
     ])
     @parametrize_with_ids("kv_cache_reuse", [True, False])
-    @parametrize_with_ids(
-        "quant_dtype",
-        [
-            pytest.param("none", marks=skip_pre_hopper),
-            # pytest.param("fp8", marks=skip_pre_hopper),
-            pytest.param("nvfp4", marks=skip_pre_blackwell)
-        ])
+    @parametrize_with_ids("quant_dtype", [
+        pytest.param("none", marks=skip_pre_hopper),
+        pytest.param("fp8", marks=skip_pre_hopper),
+        pytest.param("nvfp4", marks=skip_pre_blackwell)
+    ])
     # currently, chunked prefill is not supported for fp8 and nvfp4
     def test_chunked_prefill(self, quant_dtype, kv_cache_reuse, fp8kv,
                              overlap_scheduler):
diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml
index 9a885b7873be..a1ffbe3e3899 100644
--- a/tests/integration/test_lists/test-db/l0_h100.yml
+++ b/tests/integration/test_lists/test-db/l0_h100.yml
@@ -56,6 +56,7 @@ l0_h100:
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_no_kv_cache_reuse[quant_dtype=fp8-mtp_nextn=2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True]
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=none-kv_cache_reuse=True-fp8kv=False-overlap_scheduler=True]
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=none-kv_cache_reuse=False-fp8kv=False-overlap_scheduler=True]
+  - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_chunked_prefill[quant_dtype=fp8-kv_cache_reuse=True-fp8kv=True-overlap_scheduler=True]
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency]
   - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=False]
   - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True]