diff --git a/mindformers/core/context/build_context.py b/mindformers/core/context/build_context.py
index 16fdcc2b64e6a164fc7e70e565629d51ba183a06..359a6a17cd884e8ed67b60124045041ef7515af8 100644
--- a/mindformers/core/context/build_context.py
+++ b/mindformers/core/context/build_context.py
@@ -240,23 +240,13 @@ class MFContextOperator(MFContextConfig):
             custom_matmul_shuffle = 'on'
             lccl_deterministic = '0'
 
-        ms_alloc_conf = os.environ.get('MS_ALLOC_CONF', 'enable_vmm:False')
-        cpu_affinity = os.environ.get('CPU_AFFINITY', 'True')
-        ms_internal_disable_custom_kernel_list = ""
-        mode = self.config.get_value('context.mode', 'GRAPH_MODE')
-        if MODE.get(mode) == MODE.get('GRAPH_MODE'):
-            ms_internal_disable_custom_kernel_list = os.environ.get('MS_INTERNAL_DISABLE_CUSTOM_KERNEL_LIST',
-                                                                    'PagedAttention')
         env = {
             'HCCL_DETERMINISTIC': hccl_deterministic,
             'ASCEND_LAUNCH_BLOCKING': ascend_launch_blocking,
             'TE_PARALLEL_COMPILER': te_parallel_compiler,
             'CUSTOM_MATMUL_SHUFFLE': custom_matmul_shuffle,
             'LCCL_DETERMINISTIC': lccl_deterministic,
-            'MS_ENABLE_GRACEFUL_EXIT': '1' if self.use_graceful_exit else '0',
-            'MS_ALLOC_CONF': ms_alloc_conf,
-            'CPU_AFFINITY': cpu_affinity,
-            'MS_INTERNAL_DISABLE_CUSTOM_KERNEL_LIST': ms_internal_disable_custom_kernel_list
+            'MS_ENABLE_GRACEFUL_EXIT': '1' if self.use_graceful_exit else '0'
         }
 
         run_mode = (
@@ -267,6 +257,16 @@ class MFContextOperator(MFContextConfig):
                 and RunMode(run_mode) in [RunMode.PREDICT, RunMode.EVAL]
                 and use_past
         ):
+            ms_alloc_conf = os.environ.get('MS_ALLOC_CONF', 'enable_vmm:False')
+            cpu_affinity = os.environ.get('CPU_AFFINITY', 'True')
+            ms_internal_disable_custom_kernel_list = ""
+            mode = self.config.get_value('context.mode', 'GRAPH_MODE')
+            if MODE.get(mode) == MODE.get('GRAPH_MODE'):
+                ms_internal_disable_custom_kernel_list = os.environ.get('MS_INTERNAL_DISABLE_CUSTOM_KERNEL_LIST',
+                                                                        'PagedAttention')
+            env['MS_ALLOC_CONF'] = ms_alloc_conf
+            env['CPU_AFFINITY'] = cpu_affinity
+            env['MS_INTERNAL_DISABLE_CUSTOM_KERNEL_LIST'] = ms_internal_disable_custom_kernel_list
             env['RUN_MODE'] = run_mode
         if (
                 self.enable_mindio_ttp_save_ckpt and