diff --git a/flask_cache/jinja2ext.py b/flask_cache/jinja2ext.py
index 73be6ac..515a47e 100644
--- a/flask_cache/jinja2ext.py
+++ b/flask_cache/jinja2ext.py
@@ -30,7 +30,12 @@
 
 from jinja2 import nodes
 from jinja2.ext import Extension
-from flask.ext.cache import make_template_fragment_key
+from ._compat import PY2
+
+if PY2:
+    from flask.ext.cache import make_template_fragment_key
+else:
+    from flask_cache import make_template_fragment_key
 
 JINJA_CACHE_ATTR_NAME = '_template_fragment_cache'