diff --git a/compiler/include/zamalang/Support/KeySet.h b/compiler/include/zamalang/Support/KeySet.h
index 88ce5364c..a31e4f017 100644
--- a/compiler/include/zamalang/Support/KeySet.h
+++ b/compiler/include/zamalang/Support/KeySet.h
@@ -6,6 +6,7 @@
 
 extern "C" {
 #include "concrete-ffi.h"
+#include "zamalang/Runtime/context.h"
 }
 
 #include "zamalang/Support/ClientParameters.h"
@@ -40,6 +41,12 @@ public:
   CircuitGate inputGate(size_t pos) { return std::get<0>(inputs[pos]); }
   CircuitGate outputGate(size_t pos) { return std::get<0>(outputs[pos]); }
 
+  void generateRuntimeContext() {
+    auto ksk = std::get<1>(this->keyswitchKeys["ksk_v0"]);
+    auto bsk = std::get<1>(this->bootstrapKeys["bsk_v0"]);
+    setGlobalRuntimeContext(createRuntimeContext(ksk, bsk));
+  }
+
 protected:
   llvm::Error generateSecretKey(LweSecretKeyID id, LweSecretKeyParam param,
                                 SecretRandomGenerator *generator);
diff --git a/compiler/lib/Support/CompilerTools.cpp b/compiler/lib/Support/CompilerTools.cpp
index 5a8626217..2bc9a03d9 100644
--- a/compiler/lib/Support/CompilerTools.cpp
+++ b/compiler/lib/Support/CompilerTools.cpp
@@ -234,6 +234,9 @@ JITLambda::Argument::Argument(KeySet &keySet) : keySet(keySet) {
   for (auto i = inputs.size(); i < rawArg.size(); i++) {
     rawArg[i] = &outputs[i - inputs.size()];
   }
+
+  // Setup runtime context with appropriate keys
+  keySet.generateRuntimeContext();
 }
 
 JITLambda::Argument::~Argument() {