diff --git a/amp.module b/amp.module
index 9999c4d..e27db3a 100644
--- a/amp.module
+++ b/amp.module
@@ -6,8 +6,9 @@
  */
 
 use Drupal\amp\EntityTypeInfo;
+use Drupal\Core\Cache\Cache;
+use Drupal\Core\Config\ConfigFactoryInterface;
 use Drupal\Core\Entity\EntityInterface;
-use Drupal\file\Entity\File;
 use Drupal\Core\Entity\Display\EntityViewDisplayInterface;
 use Drupal\Core\Entity\Entity\EntityViewDisplay;
 use Drupal\Core\Form\FormStateInterface;
@@ -15,7 +16,9 @@ use Drupal\Core\Image\ImageInterface;
 use Drupal\Core\Routing\RouteMatchInterface;
 use Drupal\Core\Url;
 use Drupal\Core\Template\Attribute;
+use Drupal\file\Entity\File;
 use Symfony\Component\HttpFoundation\RedirectResponse;
+use Symfony\Component\HttpFoundation\Request;
 
 /**
  * Implements hook_help().
@@ -34,17 +37,46 @@ function amp_help($route_name, RouteMatchInterface $route_match) {
 }
 
 /**
+ * Implements hook_entity_view_mode_alter().
+ */
+function amp_entity_view_mode_alter(&$view_mode, EntityInterface $entity, $context) {
+
+  // Get the AMP Context service.
+  $amp_context = \Drupal::service('router.amp_context');
+
+  // If on AMP route, and in full view mode, switch to AMP view mode.
+  if ($amp_context->isAmpRoute() && $view_mode == 'full') {
+    $view_mode = 'amp';
+  }
+}
+
+/**
  * Implements hook_entity_view_alter().
  */
 function amp_entity_view_alter(array &$build, EntityInterface $entity, EntityViewDisplayInterface $display) {
 
-  // Git a list of available view modes for the current entity.
+  // Get the AMP Context service.
+  $amp_context = \Drupal::service('router.amp_context');
+
+  // Get a list of available view modes for the current entity.
   $view_modes = \Drupal::entityManager()->getViewModeOptionsByBundle('node', $entity->bundle());
 
+  // Check if entity is a node with an enabled amp view mode.
   if ($entity->getEntityType()->id() == 'node' && isset($view_modes['amp'])) {
+
+    if ($build['#view_mode'] == 'amp') {
+      // Otherwise adding 'amp' query parameter in URL will have no effect.
+      $build['#cache']['contexts'] = Cache::mergeContexts($build['#cache']['contexts'], ['url.query_args:amp']);
+      // Otherwise adding 'warnfix' query parameter in URL will have no effect.
+      $build['#cache']['contexts'] = Cache::mergeContexts($build['#cache']['contexts'], ['url.query_args:warnfix']);
+      // @TODO Add cache context to vary based on the config setting for
+      // amp_library_warnings_display.
+    }
+
     if (!empty($build['#attached']['html_head_link'])) {
       foreach ($build['#attached']['html_head_link'] as $key => $config) {
         if ($build['#view_mode'] == 'amp') {
+
           if ($config[0]['rel'] != 'canonical' && $config[0]['rel'] != 'shortlink') {
             unset($build['#attached']['html_head_link'][$key]);
           }
diff --git a/amp.routing.yml b/amp.routing.yml
index 0e3fff9..f2ac285 100644
--- a/amp.routing.yml
+++ b/amp.routing.yml
@@ -6,17 +6,6 @@ amp.settings:
   requirements:
     _permission: 'administer site configuration'
 
-amp.node_amp_page:
-  path: '/node/{node}'
-  defaults:
-    _controller: '\Drupal\amp\Controller\ampNodeViewController::view'
-    _title_callback: '\Drupal\node\Controller\NodeViewController::title'
-  requirements:
-    _entity_access: 'node.view'
-    node: \d+
-  options:
-    _amp_route: FALSE
-
 amp.test_library_hello:
   path: 'admin/amp/library/test'
   defaults:
diff --git a/src/Controller/ampNodeViewController.php b/src/Controller/ampNodeViewController.php
deleted file mode 100644
index e118a36..0000000
--- a/src/Controller/ampNodeViewController.php
+++ /dev/null
@@ -1,153 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\amp\Controller\ampNodeViewController.
- */
-
-namespace Drupal\amp\Controller;
-
-use Drupal\amp\Routing\AmpContext;
-use Drupal\Core\Cache\Cache;
-use Drupal\Core\Config\ConfigFactoryInterface;
-use Drupal\Core\Entity\EntityInterface;
-use Drupal\Core\Entity\EntityManagerInterface;
-use Drupal\Core\Render\RendererInterface;
-use Drupal\node\Controller\NodeViewController;
-use Symfony\Component\DependencyInjection\ContainerInterface;
-use Symfony\Component\HttpFoundation\Request;
-/**
- * Class ampNodeViewController.
- *
- * @package Drupal\amp\Controller
- */
-class ampNodeViewController extends NodeViewController {
-
-  /**
-   * The AMP context service.
-   *
-   * @var AmpContext $ampContext
-   */
-  protected $ampContext;
-
-  /**
-   * The config factory interface.
-   *
-   * @var ConfigFactoryInterface $configFactory
-   */
-  protected $configFactory;
-
-  /**
-   * The entity manager
-   *
-   * @var \Drupal\Core\Entity\EntityManagerInterface
-   */
-  protected $entityManager;
-
-  /**
-   * The renderer service.
-   *
-   * @var \Drupal\Core\Render\RendererInterface
-   */
-  protected $renderer;
-
-  /**
-   * Creates an EntityViewController object.
-   *
-   * @param \Drupal\amp\Routing\AmpContext; $amp_context
-   *   The AMP context.
-   * @param \Drupal\Core\Config\ConfigFactoryInterface; $config_factory_interface
-   *   The config factory.
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
-   *   The entity manager.
-   * @param \Drupal\Core\Render\RendererInterface $renderer
-   *   The renderer service.
-   */
-  public function __construct(AmpContext $amp_context, ConfigFactoryInterface
-  $config_factory_interface, EntityManagerInterface $entity_manager, RendererInterface $renderer) {
-    parent::__construct($entity_manager, $renderer);
-    $this->ampContext = $amp_context;
-    $this->configFactory = $config_factory_interface;  }
-
-  /**
-   * {@inheritdoc}
-   */
-  public static function create(ContainerInterface $container) {
-    return new static(
-      $container->get('router.amp_context'),
-      $container->get('config.factory'),
-      $container->get('entity.manager'),
-      $container->get('renderer')
-    );
-  }
-
-  public function warningsOn()
-  {
-    // First check the config if library warnings are on
-    $amp_config = $this->configFactory->get('amp.settings');
-    if ($amp_config->get('amp_library_warnings_display')) {
-      return true;
-    }
-
-    // Then check the URL if library warnings are enabled
-    /** @var Request $request */
-    $request = \Drupal::request();
-    $user_wants_amp_library_warnings = $request->get('warnfix');
-    if (isset($user_wants_amp_library_warnings)) {
-      return true;
-    }
-
-    return false;
-  }
-
-  /**
-   * {@inheritdoc}
-   */
-  public function view(EntityInterface $node, $view_mode = 'full', $langcode = NULL) {
-
-    // Only use the AMP view mode for content that is AMP enabled and in full
-    // view mode.
-    if ($this->ampContext->isAmpRoute() && $view_mode == 'full') {
-      $build = parent::view($node, 'amp', $langcode);
-
-      // Otherwise adding a ?amp query parameter at the end of URL will have no effect
-      $build['#cache']['contexts'] = Cache::mergeContexts($build['#cache']['contexts'], ['url.query_args:amp']);
-      // Otherwise adding a ?warnfix query parameter at the end of URL will have no effect
-      $build['#cache']['contexts'] = Cache::mergeContexts($build['#cache']['contexts'], ['url.query_args:warnfix']);
-      if ($this->warningsOn()) {
-        $build['#cache']['keys'][] = 'amp-warnings-on';
-      }
-      else {
-        $build['#cache']['keys'][] = 'amp-warnings-off';
-      }
-    }
-    // Otherwise return the default view mode.
-    else {
-      $build = parent::view($node, $view_mode, $langcode);
-    }
-
-    foreach ($node->uriRelationships() as $rel) {
-      // Set the node path as the canonical URL to prevent duplicate content.
-      $build['#attached']['html_head_link'][] = array(
-        array(
-          'rel' => $rel,
-          'href' => $node->url($rel),
-        ),
-        TRUE,
-      );
-
-      if ($rel == 'canonical') {
-        // Set the non-aliased canonical path as a default shortlink.
-        $build['#attached']['html_head_link'][] = array(
-          array(
-            'rel' => 'shortlink',
-            'href' => $node->url($rel, array('alias' => TRUE)),
-          ),
-          TRUE,
-        );
-      }
-    }
-
-    return $build;
-  }
-}
diff --git a/src/Element/AmpAnalytics.php b/src/Element/AmpAnalytics.php
index c14beb4..1c2f236 100644
--- a/src/Element/AmpAnalytics.php
+++ b/src/Element/AmpAnalytics.php
@@ -36,6 +36,9 @@ class AmpAnalytics extends RenderElement {
         array($class, 'preRenderAnalytics'),
       ),
       '#theme' => 'amp_analytics',
+      '#cache' => [
+        'contexts' => ['url.query_args:amp']
+      ]
     );
   }
 
diff --git a/src/Element/AmpIframe.php b/src/Element/AmpIframe.php
index 8838d66..8f30d43 100644
--- a/src/Element/AmpIframe.php
+++ b/src/Element/AmpIframe.php
@@ -43,6 +43,9 @@ class AmpIframe extends ProcessedText {
         array($class, 'preRenderAmpIframe'),
       ),
       '#theme' => 'amp_iframe',
+      '#cache' => [
+        'contexts' => ['url.query_args:amp', 'url.query_args:warnfix']
+      ]
     );
   }
 
diff --git a/src/Element/AmpProcessedText.php b/src/Element/AmpProcessedText.php
index 611adb5..de95635 100644
--- a/src/Element/AmpProcessedText.php
+++ b/src/Element/AmpProcessedText.php
@@ -10,6 +10,8 @@ namespace Drupal\amp\Element;
 use Drupal\filter\Element\ProcessedText;
 use Lullabot\AMP\AMP;
 use Drupal\amp\Service\AMPService;
+use Drupal\Core\Cache\Cache;
+use Symfony\Component\HttpFoundation\Request;
 
 /**
  * Provides an amp-processed text render element.
@@ -33,10 +35,7 @@ class AmpProcessedText extends ProcessedText {
         array($class, 'preRenderAmpText'),
       ),
       '#cache' => [
-        'tags' => ['amp-warnings'],
-        // This should be bubbling up but its not
-        // Instead we have to place the #cache setting in src/Controller/ampPage
-        // 'context' => ['url.query_args:warnfix'],
+        'contexts' => ['url.query_args:amp', 'url.query_args:warnfix']
       ]
     );
   }