Sync from GitHub 2026-06-12T13:28:48Z

index.html

@@ -32,6 +32,7 @@
   .row { display:flex; gap:12px; align-items:center; flex-wrap:wrap; margin:10px 0; }
   label { color:var(--dim); font-size:13px; }
   input[type=number] { width:72px; background:#1b1e25; color:var(--fg); border:1px solid #2a2e37; border-radius:6px; padding:4px 6px; font:inherit; }
+  select { background:#1b1e25; color:var(--fg); border:1px solid #2a2e37; border-radius:6px; padding:4px 6px; font:inherit; cursor:pointer; }
   button { background:var(--acc); color:#10131a; border:0; border-radius:8px; padding:8px 18px; font:inherit; font-weight:600; cursor:pointer; }
   button:disabled { opacity:.4; cursor:default; }
   #canvas { background:#1b1e25; border:1px solid #2a2e37; border-radius:8px; padding:14px; margin-top:14px; min-height:120px; white-space:pre-wrap; word-break:break-word; }
@@ -57,6 +58,12 @@
   <p><span id="status" class="status-badge loading"><span class="spinner"></span><span id="statusText">Loading…</span></span></p>
   <textarea id="prompt">Lily runs 12 km/h for 4 hours. How far in 8 hours?</textarea>
   <div class="row">
+    <label>model <select id="modelSel">
+      <option value="mdlm">MDLM · fp16</option>
+      <option value="mdlm-q4">MDLM · q4</option>
+      <option value="bd3lm">BD3LM · fp16 · block diffusion</option>
+      <option value="bd3lm-q4">BD3LM · q4 · block diffusion</option>
+    </select></label>
     <label>new tokens <input id="maxNew" type="number" value="128" step="32" min="16"></label>
     <label>steps <input id="steps" type="number" value="128" step="16" min="4"></label>
     <label>block <input id="blockSize" type="number" value="32" step="8" min="8"></label>
@@ -83,8 +90,6 @@ import { DiffusionLM } from './kohra.js';
 
 // fused fp16 (G1d): RMSNorm fused to SimplifiedLayerNormalization so WebGPU's native
 // fp16 kernels don't overflow the decomposed Pow(x,2) -> all-zero. ~1.4GB, ~9.8 tok/s.
-const HF = 'https://huggingface.co/naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX/resolve/main';
-const TOKENIZER_ID = 'naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX';
 const $ = (id) => document.getElementById(id);
 // state: 'loading' (spinner) | 'ready' | 'error'
 const status = (t, state = 'loading') => {
@@ -92,11 +97,37 @@ const status = (t, state = 'loading') => {
   $('status').className = 'status-badge ' + state;
 };
 
+// Model registry — each entry is a self-contained HF repo (graph + tokenizer + bundled
+// kohra.js). bd3lm (block diffusion) drives the block-causal sampler; q4 variants are a
+// smaller download. The dropdown picks one (reloads with ?arch=<key>); ?model= (+ optional
+// ?arch=, ?tokenizer=, ?ort=, ?opt=) overrides for local/dev. MDLM-fp16 stays the default.
+const REPO = {
+  mdlm:  'https://huggingface.co/naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX/resolve/main',
+  bd3lm: 'https://huggingface.co/naklitechie/Qwen3-0.6B-diffusion-bd3lm-ONNX/resolve/main',
+};
+const TOKENIZER = {
+  mdlm:  'naklitechie/Qwen3-0.6B-diffusion-mdlm-ONNX',
+  bd3lm: 'naklitechie/Qwen3-0.6B-diffusion-bd3lm-ONNX',
+};
+// q4 (RTN MatMulNBits) needs the ORT-web dev build with the WebGPU q4 fix + full graph opt.
+const ORT_Q4 = '1.26.0-dev.20260416-b7804b056c';
+const REGISTRY = {
+  'mdlm':     { label: 'MDLM · fp16',                  repo: 'mdlm',  file: 'model_fp16_fused.onnx',    blockCausal: false, q4: false },
+  'mdlm-q4':  { label: 'MDLM · q4 (smaller)',          repo: 'mdlm',  file: 'model_q4f16_rtn_sym.onnx', blockCausal: false, q4: true  },
+  'bd3lm':    { label: 'BD3LM · fp16 · block diff',    repo: 'bd3lm', file: 'model_fp16_fused.onnx',    blockCausal: true,  q4: false },
+  'bd3lm-q4': { label: 'BD3LM · q4 · block diff',      repo: 'bd3lm', file: 'model_q4f16_rtn_sym.onnx', blockCausal: true,  q4: true  },
+};
+
 const params = new URLSearchParams(location.search);
-const modelUrl = params.get('model') ?? `${HF}/onnx/model_fp16_fused.onnx`;
-// optional overrides for q4/ORT experiments; defaults keep the shipped demo unchanged
-const ortVersion = params.get('ort') || undefined;
-const graphOptimizationLevel = params.get('opt') || 'disabled';
+const modelOverride = params.get('model');
+const reqArch = params.get('arch') || (modelOverride?.includes('bd3lm') ? 'bd3lm' : 'mdlm');
+const arch = REGISTRY[reqArch] ? reqArch : 'mdlm';        // unknown key → safe default
+const preset = REGISTRY[arch];
+const modelUrl = modelOverride ?? `${REPO[preset.repo]}/onnx/${preset.file}`;
+const tokenizerId = params.get('tokenizer') ?? TOKENIZER[preset.repo];
+const blockCausal = preset.blockCausal;
+const ortVersion = params.get('ort') || (preset.q4 ? ORT_Q4 : undefined);
+const graphOptimizationLevel = params.get('opt') || (preset.q4 ? 'all' : 'disabled');
 
 let lm;
 
@@ -123,6 +154,7 @@ async function generate() {
       blockSize: +$('blockSize').value,
       temperature: +$('temp').value,
       threshold: +$('threshold').value || null,   // Fast-dLLM confidence-threshold decoding
+      blockCausal,                                 // bd3lm: static block-causal attention mask
       stripThink: true,
       onStep: (s) => {
         renderCanvas(s);
@@ -144,10 +176,10 @@ async function generate() {
 async function init() {
   try {
     if (!navigator.gpu) { status('This demo needs WebGPU — try Chrome/Edge 121+.', 'error'); return; }
-    status('Loading model… first run downloads ~1.4GB + compiles WebGPU shaders (cached after).');
-    lm = await DiffusionLM.from_pretrained({ model: modelUrl, tokenizer: TOKENIZER_ID, ortVersion, graphOptimizationLevel });
+    status(`Loading model… first run downloads ${preset.q4 ? '~0.7GB' : '~1.5GB'} + compiles WebGPU shaders (cached after).`);
+    lm = await DiffusionLM.from_pretrained({ model: modelUrl, tokenizer: tokenizerId, ortVersion, graphOptimizationLevel });
     window.kohra = lm; // handy for the console / preview probing
-    status('Ready — fused fp16 Qwen3-0.6B-MDLM on WebGPU. Press Generate.', 'ready');
+    status(`Ready — ${preset.q4 ? 'q4f16' : 'fused fp16'} Qwen3-0.6B-${blockCausal ? 'BD3LM (block diffusion)' : 'MDLM'} on WebGPU. Press Generate.`, 'ready');
     $('go').disabled = false;
   } catch (e) {
     console.error(e);
@@ -158,6 +190,10 @@ async function init() {
 addEventListener('error', e => status('Uncaught: ' + (e.error?.message ?? e.message), 'error'));
 addEventListener('unhandledrejection', e => status('Rejection: ' + (e.reason?.message ?? e.reason), 'error'));
 $('go').addEventListener('click', generate);
+// Model picker: reflect the active arch, and reload with the chosen preset on change
+// (drop any ?model override so the registry URL is used).
+$('modelSel').value = arch;
+$('modelSel').addEventListener('change', () => { location.search = '?arch=' + $('modelSel').value; });
 init();
 </script>
 </body>

kohra.js

@@ -35,6 +35,9 @@ const GEN_DEFAULTS = {
   threshold: null,       // Fast-dLLM: if set (e.g. 0.9), unmask ALL block positions with
                          // confidence >= threshold per step (min 1) instead of a fixed
                          // count → fewer forwards when the model is confident. null = use `steps`.
+  blockCausal: false,    // bd3lm (block diffusion): pass a static [1,1,T,T] additive
+                         // block-causal attention_mask each forward + denoise on the
+                         // physical pos//blockSize grid. false = MDLM (single-input, no mask).
   chatTemplate: true,    // wrap prompt in Qwen3 ChatML + generation prompt
   stripThink: false,     // drop a leading empty <think>…</think> block from the result
   onStep: null,          // ({ x, P, fresh, block, numBlocks, forward, elapsedMs }) => void
@@ -145,9 +148,43 @@ export class DiffusionLM {
     promptIds.forEach((t, i) => { x[i] = BigInt(t); });
     for (let i = P; i < T; i++) x[i] = MASK;
 
-    const numBlocks = Math.ceil(cfg.maxNewTokens / cfg.blockSize);
+    // Block ranges to denoise, left-to-right. MDLM: blocks start at the prompt end
+    // (P + b·blockSize). bd3lm: blocks live on the absolute pos//blockSize grid, so we
+    // denoise from the physical block containing P — its prompt prefix is fixed, its
+    // masked tail is the first generated tokens (bidirectional within that block).
+    const blocks = [];
+    if (cfg.blockCausal) {
+      for (let blk = Math.floor(P / cfg.blockSize); blk * cfg.blockSize < T; blk++) {
+        blocks.push([blk * cfg.blockSize, Math.min((blk + 1) * cfg.blockSize, T)]);
+      }
+    } else {
+      const numBlocks = Math.ceil(cfg.maxNewTokens / cfg.blockSize);
+      for (let b = 0; b < numBlocks; b++) {
+        const start = P + b * cfg.blockSize;
+        blocks.push([start, Math.min(start + cfg.blockSize, T)]);
+      }
+    }
+    const numBlocks = blocks.length;
     const stepsPerBlock = Math.ceil(cfg.steps / numBlocks);
 
+    // bd3lm's static block-causal additive mask: 0 where block(key) ≤ block(query),
+    // else -1e9. Built once and fed every forward (the graph's 2nd input, fp32). The
+    // mask never changes during sampling; block-causality is what makes the cache-free
+    // full-canvas forward give correct current-block logits despite still-masked future
+    // blocks (a future block can't influence an earlier one).
+    let attnTensor = null;
+    if (cfg.blockCausal) {
+      const m = new Float32Array(T * T);
+      for (let q = 0; q < T; q++) {
+        const bq = Math.floor(q / cfg.blockSize);
+        const row = q * T;
+        for (let k = 0; k < T; k++) {
+          m[row + k] = Math.floor(k / cfg.blockSize) <= bq ? 0 : -1e9;
+        }
+      }
+      attnTensor = new ort.Tensor('float32', m, [1, 1, T, T]);
+    }
+
     let forward = 0;
     const t0 = performance.now();
 
@@ -156,8 +193,7 @@ export class DiffusionLM {
     };
 
     for (let b = 0; b < numBlocks; b++) {
-      const start = P + b * cfg.blockSize;
-      const end = Math.min(start + cfg.blockSize, T);
+      const [start, end] = blocks[b];
       const nMasked = inBlockMasked(start, end);
       if (nMasked === 0) continue;
 
@@ -168,9 +204,9 @@ export class DiffusionLM {
 
       while (inBlockMasked(start, end) > 0) {
         if (schedule && step >= schedule.length) break;
-        const out = await this.session.run({
-          input_ids: new ort.Tensor('int64', x.slice(), [1, T]),
-        });
+        const feeds = { input_ids: new ort.Tensor('int64', x.slice(), [1, T]) };
+        if (attnTensor) feeds.attention_mask = attnTensor;
+        const out = await this.session.run(feeds);
         forward++;
         const logits = out.logits.data;       // Float32Array, [T * V] (ORT upcasts fp16)
         const V = out.logits.dims[2];