diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py
index 329ca09a81..de4be883ba 100644
--- a/atom/model_ops/attentions/aiter_attention.py
+++ b/atom/model_ops/attentions/aiter_attention.py
@@ -40,7 +40,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
 
         context_lens = [seq.num_tokens for seq in seqs]
         max_seqlen_k = max(context_lens)
-        positions = context_lens
+        positions = [i - 1 for i in context_lens]
         slot_mapping = [
             seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1
             for seq in seqs
diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py
index d5ecc67d63..e6d00cff23 100644
--- a/atom/model_ops/attentions/aiter_mla.py
+++ b/atom/model_ops/attentions/aiter_mla.py
@@ -181,7 +181,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         max_q_len = 1
 
         context_lens = [seq.num_tokens for seq in seqs]
-        positions = context_lens
+        positions = [i - 1 for i in context_lens]
         slot_mapping = [
             seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1
             for seq in seqs
diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py
index 2f50165a28..7d44504297 100644
--- a/atom/model_ops/moe.py
+++ b/atom/model_ops/moe.py
@@ -1123,7 +1123,7 @@ def __init__(
                     if is_rocm_aiter_fuse_routed_scaling_factor()
                     else 1 / self.routed_scaling_factor
                 ),
-                max_num_tokens=16384,  # Default max tokens
+                max_num_tokens=atom_config.max_num_batched_tokens,
                 is_EP=self.use_ep,
             )
         if is_rocm_aiter_fusion_shared_expert_enabled():