Compact math derivation slide for nunchux_low_bit_attention_blackwell.html

blackwell_attention_slides/nunchux_low_bit_attention_blackwell.html

@@ -120,6 +120,16 @@
     .label { color: var(--muted); font-size: 17px; margin-top: 5px; }
     .img-wrap { border: 1px solid var(--line); border-radius: 6px; padding: 12px; background: #fff; }
     img { display: block; max-width: 100%; }
+
+    .math-note { font-size: 22px; color: var(--muted); margin-bottom: 18px; }
+    .mathbox { border-left: 8px solid var(--mit-red); background: var(--soft); padding: 18px 20px; }
+    .mathline { display: grid; grid-template-columns: 108px 1fr; gap: 14px; align-items: start; padding: 9px 0; border-bottom: 1px solid #e5e5e5; }
+    .mathline:last-child { border-bottom: 0; }
+    .mathline .tag { color: var(--mit-red); font-weight: 850; font-size: 18px; }
+    .mathline .eq { font-family: Georgia, 'Times New Roman', serif; font-size: 24px; }
+    .mathline.strong .eq { font-size: 30px; font-weight: 850; }
+    .mathline.final { background: #fff; margin-top: 10px; padding: 13px 12px; border: 2px solid var(--mit-red); }
+    .mathline.final .eq { font-size: 25px; font-weight: 850; }
     .placeholder {
       flex: 1;
       border: 2px dashed var(--mit-red);
@@ -210,47 +220,29 @@
 
 <section class="slide">
   <div class="eyebrow">Softmax approximation / 2</div>
-  <h2>Derivation: Turn log2 Probability into an E4M3 Byte</h2>
-  <div class="cols" style="grid-template-columns: 1.12fr 0.88fr;">
+  <h2>Approximation: Directly Predict the E4M3 Byte</h2>
+  <div class="cols" style="grid-template-columns: 1.05fr 0.95fr;">
     <div>
-      <h3>One line at a time</h3>
-      <div class="formula derivation"><span class="step">0</span> Variables:
-    log2_prob = log2 of the softmax probability before FP8 rounding
-              = scaled_score - row_max, so P = 2^log2_prob
-    e4m3_code = the final 8-bit FP8 E4M3 byte stored for P
-              = 8 * exponent_bits + mantissa_bits
-
-<span class="step">1</span> Softmax in base-2 form:
-    P = 2^log2_prob
-
-<span class="step">2</span> E4M3 normal value with exponent E and mantissa M:
-    P_fp8 = 2^(E - 7) * (1 + M / 8)
-
-<span class="step">3</span> Take log2:
-    log2_prob = E - 7 + log2(1 + M / 8)
-
-<span class="step">4</span> Approximate the mantissa term:
-    log2(1 + M / 8) ≈ M / 8
-
-<span class="step">5</span> Solve for the byte value:
-    8 * log2_prob ≈ 8E - 56 + M
-    e4m3_code = 8E + M ≈ 8 * log2_prob + 56
-
-<span class="step">6</span> Add bias, round, and clamp away NaN code 127:
-    e4m3_code = clamp(round(8 * log2_prob + 56 - 0.35), 0, 126)</div>
+      <h3>Idea in one sentence</h3>
+      <p class="math-note">Softmax already gives us a log-domain value; E4M3 bytes are almost linear in log-domain value. So we skip exp2 and directly write the FP8 byte.</p>
+      <div class="mathbox">
+        <div class="mathline"><span class="tag">Input</span><span class="eq">u = \log_2 P = 	ext{scaled\_score} - 	ext{row\_max}</span></div>
+        <div class="mathline"><span class="tag">E4M3</span><span class="eq">P pprox 2^{e-7}\left(1 + rac{m}{8}
+        <div class="mathline"><span class="tag">Approx.</span><span class="eq">\log_2\left(1+rac{m}{8}
+        <div class="mathline strong"><span class="tag">Result</span><span class="eq">	ext{byte} pprox 8u + 56</span></div>
+        <div class="mathline final"><span class="tag">Kernel</span><span class="eq">	ext{byte}=\operatorname{clamp}(\operatorname{round}(8u+56-0.35),0,126)</span></div>
+      </div>
+      <p class="math-note"><span class="red">126</span> avoids E4M3 NaN code 127. The <span class="red">-0.35</span> bias centers the mantissa-linearization error.</p>
     </div>
     <div>
-      <h3>Implementation</h3>
-      <pre>; Step A: affine log2_prob -> normalized code range
-fma.rn.sat.f32 y0, S0, c_eff, bias;
-fma.rn.sat.f32 y1, S1, c_eff, bias;
+      <h3>PTX shape</h3>
+      <pre>; 1. affine u -> byte range, with saturation
+fma.rn.sat.f32 y, S, c_eff, bias;
 
-; Step B: magic-add rounds to integer code
-; byte0 of the f32 bit pattern is the E4M3 code
-fma.rn.f32x2 {t0,t1}, {y0,y1}, {126,126}, {M,M};
-fma.rn.f32x2 {t2,t3}, {y2,y3}, {126,126}, {M,M};
+; 2. magic-add: rounded byte is f32 byte0
+fma.rn.f32x2 t, y, 126, MAGIC;
 
-; Step C: pack 4 codes into one u32
+; 3. pack four byte0 fields into FP8 P
 prmt.b32 p0, t0, t1, 0x0040;
 prmt.b32 p1, t2, t3, 0x0040;
 prmt.b32 P,  p0, p1, 0x5410;</pre>