diff --git a/packages/runtime-core/__tests__/memoize.spec.ts b/packages/runtime-core/__tests__/memoize.spec.ts
new file mode 100644
index 000000000..f2e053982
--- /dev/null
+++ b/packages/runtime-core/__tests__/memoize.spec.ts
@@ -0,0 +1,54 @@
+import { h, Component, memoize, nextTick } from '../src'
+import { renderIntsance, serialize } from '@vue/runtime-test'
+
+describe('memoize', () => {
+  it('should work', async () => {
+    class App extends Component {
+      count = 1
+      render() {
+        return h('div', [
+          this.count,
+          this.count % 2
+            ? memoize(() => h('div', `A` + this.count), this, 0)
+            : null,
+          memoize(() => h('div', `B` + this.count), this, 1)
+        ])
+      }
+    }
+
+    const app = renderIntsance(App)
+    expect(serialize(app.$el)).toBe(`<div>1<div>A1</div><div>B1</div></div>`)
+
+    app.count++
+    await nextTick()
+    expect(serialize(app.$el)).toBe(`<div>2<div>B1</div></div>`)
+
+    app.count++
+    await nextTick()
+    // test remounting a memoized tree
+    expect(serialize(app.$el)).toBe(`<div>3<div>A1</div><div>B1</div></div>`)
+  })
+
+  it('should invalidate based on keys', async () => {
+    class App extends Component {
+      foo = 1
+      bar = 1
+      render() {
+        return memoize(() => h('div', this.foo + this.bar), this, 0, [this.bar])
+      }
+    }
+
+    const app = renderIntsance(App)
+    expect(serialize(app.$el)).toBe(`<div>2</div>`)
+
+    app.foo++
+    await nextTick()
+    // should not update
+    expect(serialize(app.$el)).toBe(`<div>2</div>`)
+
+    app.bar++
+    await nextTick()
+    // should update now
+    expect(serialize(app.$el)).toBe(`<div>4</div>`)
+  })
+})
diff --git a/packages/runtime-core/src/index.ts b/packages/runtime-core/src/index.ts
index cfb457e35..3737cd50f 100644
--- a/packages/runtime-core/src/index.ts
+++ b/packages/runtime-core/src/index.ts
@@ -18,6 +18,7 @@ export { createAsyncComponent } from './optional/asyncComponent'
 export { KeepAlive } from './optional/keepAlive'
 export { mixins } from './optional/mixins'
 export { EventEmitter } from './optional/eventEmitter'
+export { memoize } from './optional/memoize'
 export { withHooks, useState, useEffect } from './optional/hooks'
 
 // flags & types
diff --git a/packages/runtime-core/src/optional/memoize.ts b/packages/runtime-core/src/optional/memoize.ts
index 0bfca68bb..1659ff153 100644
--- a/packages/runtime-core/src/optional/memoize.ts
+++ b/packages/runtime-core/src/optional/memoize.ts
@@ -17,13 +17,13 @@ import { warn } from '../warning'
 
 const memoizeMap = new WeakMap()
 
-export function memoize(
-  getter: () => any,
+export function memoize<T>(
+  getter: () => T,
   instance: Component,
   id: number,
   keys?: any[]
-): any {
-  if (__DEV__ && !Array.isArray(keys)) {
+): T {
+  if (__DEV__ && arguments.length > 3 && !Array.isArray(keys)) {
     warn(
       `keys passed to v-memo or memoize must be an array. Got ${String(keys)}`
     )