diff --git a/onnxruntime/contrib_ops/cpu/transformers/sampling.cc b/onnxruntime/contrib_ops/cpu/transformers/sampling.cc
index b37dc2fdeb..9c1b432c53 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/sampling.cc
+++ b/onnxruntime/contrib_ops/cpu/transformers/sampling.cc
@@ -140,9 +140,8 @@ Status Sampling::Compute(OpKernelContext* ctx) const {
           init_greedy_state_func_ ? init_greedy_state_func_ : GenerationCpuDeviceHelper::InitGreedyState<float>,
           device_copy_func_ ? device_copy_func_ : GenerationCpuDeviceHelper::DeviceCopy<float>,
           update_gpt_feeds_func_ ? update_gpt_feeds_func_ : GenerationCpuDeviceHelper::UpdateGptFeeds<float>,
-          // TODO(hasesh): Add support for decoder masked multihead attention
-          nullptr,
-          0};
+          gpu_device_prop_,
+          gpu_device_arch_};
       ORT_RETURN_IF_ERROR(impl.Initialize());
 
       return impl.Execute(init_run_decoder_feeds_fetches_manager_, *decoder_feeds_fetches_manager_);
@@ -165,9 +164,8 @@ Status Sampling::Compute(OpKernelContext* ctx) const {
           init_greedy_state_fp16_func_,
           device_copy_func_,
           update_gpt_feeds_fp16_func_,
-          // TODO(hasesh): Add support for decoder masked multihead attention
-          nullptr,
-          0};
+          gpu_device_prop_,
+          gpu_device_arch_};
 
       ORT_RETURN_IF_ERROR(impl.Initialize());
 
diff --git a/onnxruntime/contrib_ops/cpu/transformers/sampling.h b/onnxruntime/contrib_ops/cpu/transformers/sampling.h
index bdb7f94264..87f04bf615 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/sampling.h
+++ b/onnxruntime/contrib_ops/cpu/transformers/sampling.h
@@ -41,16 +41,16 @@ class Sampling : public IControlFlowKernel {
 
   // device helpers that is same for both GPT and encoder-decoder models.
   void SetDeviceHelpers(
-      const GenerationDeviceHelper::AddToFeedsFunc& add_to_feeds_func,
       const GenerationDeviceHelper::ReorderPastStateFunc& reorder_past_state_func,
+      const GenerationDeviceHelper::AddToFeedsFunc& add_to_feeds_func,
       const GenerationDeviceHelper::TopkFunc& topk_func,
       const GenerationDeviceHelper::DeviceCopyFunc<float>& device_copy_func,
       const GenerationDeviceHelper::GreedySearchProcessLogitsFunc<float>& process_logits_func,
       const GenerationDeviceHelper::GreedySearchProcessLogitsFunc<MLFloat16>& process_logits_fp16_func,
       const GenerationDeviceHelper::InitGreedyStateFunc<float>& init_greedy_state_func,
       const GenerationDeviceHelper::InitGreedyStateFunc<MLFloat16>& init_greedy_state_fp16_func) {
-    add_to_feeds_func_ = add_to_feeds_func;
     reorder_past_state_func_ = reorder_past_state_func;
+    add_to_feeds_func_ = add_to_feeds_func;
     topk_func_ = topk_func;
     device_copy_func_ = device_copy_func;
     process_logits_func_ = process_logits_func;
@@ -66,10 +66,13 @@ class Sampling : public IControlFlowKernel {
     update_gpt_feeds_fp16_func_ = update_gpt_feeds_fp16_func;
   }
 
+  const void* gpu_device_prop_ = nullptr;
+  int gpu_device_arch_ = 0;
+
  private:
   // Device specific functions
-  GenerationDeviceHelper::AddToFeedsFunc add_to_feeds_func_;
   GenerationDeviceHelper::ReorderPastStateFunc reorder_past_state_func_;
+  GenerationDeviceHelper::AddToFeedsFunc add_to_feeds_func_;
   GenerationDeviceHelper::TopkFunc topk_func_;
   GenerationDeviceHelper::DeviceCopyFunc<float> device_copy_func_;
 
diff --git a/onnxruntime/contrib_ops/cuda/transformers/sampling.cc b/onnxruntime/contrib_ops/cuda/transformers/sampling.cc
index 3664eae249..f299279d7b 100644
--- a/onnxruntime/contrib_ops/cuda/transformers/sampling.cc
+++ b/onnxruntime/contrib_ops/cuda/transformers/sampling.cc
@@ -32,8 +32,8 @@ transformers::CudaTensorConsoleDumper g_cuda_dumper_sampling;
 
 Sampling::Sampling(const OpKernelInfo& info)
     : onnxruntime::contrib::transformers::Sampling(info) {
-  SetDeviceHelpers(GenerationCudaDeviceHelper::AddToFeeds,
-                   GenerationCudaDeviceHelper::ReorderPastState,
+  SetDeviceHelpers(GenerationCudaDeviceHelper::ReorderPastState,
+                   GenerationCudaDeviceHelper::AddToFeeds,
                    GenerationCudaDeviceHelper::TopK,
                    GenerationCudaDeviceHelper::DeviceCopy<float>,
                    GenerationCudaDeviceHelper::GreedySearchProcessLogits<float>,
@@ -45,6 +45,12 @@ Sampling::Sampling(const OpKernelInfo& info)
                        GenerationCudaDeviceHelper::UpdateGptFeeds<MLFloat16>);
 
   SetConsoleDumper(&g_cuda_dumper_sampling);
+
+  gpu_device_prop_ = &reinterpret_cast<const CUDAExecutionProvider*>(info.GetExecutionProvider())->GetDeviceProp();
+
+  gpu_device_arch_ = static_cast<const cudaDeviceProp*>(gpu_device_prop_)->major * 100 +
+                      static_cast<const cudaDeviceProp*>(gpu_device_prop_)->minor * 10;
+
 }
 
 Status Sampling::ComputeInternal(OpKernelContext* context) const {