diff --git a/_includes/extensions/code-highlight.html b/_includes/extensions/code-highlight.html
index f2750d8940c..c5782cc96ab 100644
--- a/_includes/extensions/code-highlight.html
+++ b/_includes/extensions/code-highlight.html
@@ -19,67 +19,73 @@
 {%- if badge_enabled -%}
 <script>
 document.addEventListener('DOMContentLoaded', function(event) {
-  var els = document.querySelectorAll('pre code')
-
-  function addLangData(block) {
-    var outer = block.parentElement.parentElement.parentElement;
-    var lang = block.getAttribute('data-lang');
-    for (var i = 0; i < outer.classList.length; i++) {
-      var cls = outer.classList[i];
+  function getLangFromClass(elem) {
+    for (cls of elem.classList) {
       if (cls.startsWith('language-')) {
-        lang = cls;
-        break;
+        return cls.substr(9);
       }
     }
-    if (!lang) {
-      cls = block.getAttribute('class');
-      lang = cls ? cls : '';
-    }
-    if (lang.startsWith('language-')) {
-      lang = lang.substr(9);
-    }
-    block.parentNode.setAttribute('data-lang', lang);
+    return null;
   }
 
-  function addBadge(block) {
-    var enabled = ('{{ badge_enabled }}' || 'true').toLowerCase();
-    if (enabled == 'true') {
-      var pre = block.parentElement;
-      pre.classList.add('badge');
-    }
-  }
+  function addBadge(preElem, lang) {
+    let badgeWrapper = document.createElement('div');
+    badgeWrapper.classList.add('badge-wrapper')
 
-  function handle(block) {
-    addLangData(block);
-    addBadge(block)
-  }
+    let badge = document.createElement('div');
+    badge.classList.add('badge');
+    badge.append(lang);
 
-  for (var i = 0; i < els.length; i++) {
-    var el = els[i];
-    handle(el);
+    preElem.replaceWith(badgeWrapper);
+    badgeWrapper.append(badge, preElem);
   }
+
+  // Kramdown code blocks (Rouge highlighter):
+  //   div.highlighter_rouge.language-XXX > div.highlight > pre.highlight > code
+  document.querySelectorAll('.highlighter-rouge').forEach((e) => {
+    let pre = e.querySelector('pre');
+    if (pre !== null) {
+      let lang = getLangFromClass(e);
+      if (lang !== null) {
+        addBadge(pre, lang);
+      }
+    }
+  });
+
+  // Liquid code blocks (any highlighter):
+  //   figure.highlight > pre > code.language-XXX[data-lang="XXX"]
+  // Kramdown code blocks (no highlighter):
+  //   pre > code.language-XXX
+  document.querySelectorAll('pre > code').forEach((e) => {
+    let lang = e.hasAttribute('data-lang')
+      ? e.getAttribute('data-lang')   // Liquid
+      : getLangFromClass(e);          // Kramdown
+
+    if (lang !== null) {
+      addBadge(e.parentElement, lang);
+    }
+  });
 });
 </script>
 
 <style>
-  /* code language badge */
-  pre.badge::before {
-    content: attr(data-lang);
-    color: {{badge_color}};
-    background-color: {{badge_background_color}};
-    padding: 0 .5em;
-    border-radius: 0 2px;
-    text-transform: {{badge_text_transform}};
-    text-align: center;
-    min-width: 32px;
-    display: inline-block;
+  .badge-wrapper {
+    position: relative;
+  }
+
+  .badge-wrapper .badge {
     position: absolute;
+    top: 0;
     right: 0;
-  }
 
-  /* fix wrong badge display for firefox browser */
-  code > table pre::before {
-    display: none;
+    min-width: 32px;
+    padding: 0 .5em;
+    border-bottom-left-radius: 2px;
+    text-align: center;
+
+    color: {{badge_color}};
+    background-color: {{badge_background_color}};
+    text-transform: {{badge_text_transform}};
   }
 </style>
 {%- endif -%}
diff --git a/_sass/yat/_base.scss b/_sass/yat/_base.scss
index be1f842e5e6..e12dc36c171 100644
--- a/_sass/yat/_base.scss
+++ b/_sass/yat/_base.scss
@@ -157,15 +157,8 @@ code {
 
 pre {
   overflow-x: auto;
-  position: relative;
   background-color: #f0f0f0;
-
-  > code {
-    display: inline-block;
-    padding: 20px!important;
-    background-color: transparent;
-    border: 0;
-  }
+  padding: 20px;
 
   table, pre {
     margin-bottom: 0;