diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py
index 60a7fe88f..9c85e8113 100644
--- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py
+++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py
@@ -377,6 +377,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             ]
             page_indices //= self.page_size
             metadata.page_table[:, :max_seq_pages].copy_(page_indices)
+            metadata.max_seq_len_q = self.speculative_num_draft_tokens
         elif forward_mode.is_draft_extend():
             metadata = self.draft_extend_metadata[bs]
             metadata.cache_seqlens_int32.copy_(seq_lens)
@@ -614,24 +615,42 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         bmm1_scale = q_scale * k_scale * layer.scaling
         bmm2_scale = 1.0
 
-        o = flashinfer.prefill.trtllm_batch_context_with_kv_cache(
-            query=q,
-            kv_cache=kv_cache,
-            workspace_buffer=self.workspace_buffer,
-            block_tables=self.forward_metadata.page_table,
-            seq_lens=self.forward_metadata.cache_seqlens_int32,
-            max_q_len=self.forward_metadata.max_seq_len_q,
-            max_kv_len=self.max_context_len,
-            bmm1_scale=bmm1_scale,
-            bmm2_scale=bmm2_scale,
-            batch_size=forward_batch.batch_size,
-            cum_seq_lens_q=self.forward_metadata.cu_seqlens_q,
-            cum_seq_lens_kv=self.forward_metadata.cu_seqlens_k,
-            window_left=layer.sliding_window_size,
-            # TODO: add attention_sink operation or nvfp4 scale factor if needed
-            sinks=attention_sink,
-            out_dtype=self.q_data_type,  # model_runner.dtype
-        )
+        if forward_batch.forward_mode.is_target_verify():
+            o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
+                query=q,
+                kv_cache=kv_cache,
+                workspace_buffer=self.workspace_buffer,
+                block_tables=self.forward_metadata.page_table,
+                seq_lens=self.forward_metadata.cache_seqlens_int32,
+                max_seq_len=self.max_context_len,
+                bmm1_scale=bmm1_scale,
+                bmm2_scale=bmm2_scale,
+                window_left=layer.sliding_window_size,
+                # TODO: add attention_sink operation or nvfp4 scale factor if needed
+                sinks=attention_sink,
+                out_dtype=self.q_data_type,  # model_runner.dtype
+                q_len_per_req=self.forward_metadata.max_seq_len_q,
+            )
+        else:
+
+            o = flashinfer.prefill.trtllm_batch_context_with_kv_cache(
+                query=q,
+                kv_cache=kv_cache,
+                workspace_buffer=self.workspace_buffer,
+                block_tables=self.forward_metadata.page_table,
+                seq_lens=self.forward_metadata.cache_seqlens_int32,
+                max_q_len=self.forward_metadata.max_seq_len_q,
+                max_kv_len=self.max_context_len,
+                bmm1_scale=bmm1_scale,
+                bmm2_scale=bmm2_scale,
+                batch_size=forward_batch.batch_size,
+                cum_seq_lens_q=self.forward_metadata.cu_seqlens_q,
+                cum_seq_lens_kv=self.forward_metadata.cu_seqlens_k,
+                window_left=layer.sliding_window_size,
+                # TODO: add attention_sink operation or nvfp4 scale factor if needed
+                sinks=attention_sink,
+                out_dtype=self.q_data_type,  # model_runner.dtype
+            )
 
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)