Veda: split the kernel call into index, quantize and attend timings; cheaper host prep with equal output

modules/loaders/tests/test_veda.py

@@ -474,6 +474,19 @@ def test_timing_is_off_unless_asked_and_then_adds_one_line_with_every_phase(veda
     line = next(r.getMessage() for r in caplog.records if "time in each phase" in r.getMessage())
     assert all(name in line for name in ("gather", "score", "select", "kernel", "scatter")) and "total" in line
     assert torch.equal(plain, timed)                                                   # a clock changes no number
+    caplog.clear()
+
+    def splitter(q, k, v, mask, lay, timer):
+        with timer("kernel.index"):
+            pass
+        with timer("kernel.attend"):
+            return core.reference_attend(q, k, v, mask, lay)
+    splitter.times_itself = True
+    veda_on.attend = splitter
+    with caplog.at_level(logging.INFO):
+        _last_step(tiny, _install(tiny, veda_timing=True)[0])
+    line = next(r.getMessage() for r in caplog.records if "time in each phase" in r.getMessage())
+    assert "kernel " in line and "(index " in line and ", attend " in line and line.count("total") == 1       # a kernel's own parts are shown, not added twice
 
 
 def test_two_stages_in_one_run_give_one_line_that_names_both_sizes(veda_on, folder, comfyui, tiny_h3, caplog):

modules/loaders/veda_attention.py

@@ -149,8 +149,13 @@ def summarise(state, cfg):
                      state["dense"], f" ({why})" if why else "", state["backend"] or "?", done * 100.0, "; ".join(f"{core.describe_grid(g)}, tile plan: {how}" for g, how in sizes.items()),
                      cfg["predictor_name"])
         if veda["timer"] is not None:
-            phases = veda["timer"].totals()
-            logging.info("[FunUI Veda] time in each phase this run (ms): %s | total %.0f", ", ".join(f"{k} {v:.0f}" for k, v in phases.items()), sum(phases.values()))
+            phases = veda["timer"].totals()          # a phase with a dot is a part of the one before it: not added again
+            top = [k for k in phases if "." not in k]
+            text = []
+            for k in top:
+                inner = ["%s %.0f" % (n.split(".")[1], v) for n, v in phases.items() if n.startswith(k + ".")]
+                text.append("%s %.0f%s" % (k, phases[k], " (%s)" % ", ".join(inner) if inner else ""))
+            logging.info("[FunUI Veda] time in each phase this run (ms): %s | total %.0f", ", ".join(text), sum(phases[k] for k in top))
         logging.info("[FunUI Veda] full attention on %s of %d", f"steps {sla._runs(state['dense_at'])}" if state["dense_at"] else "no step", state["n"])
         if state["failed"] is not None:
             logging.warning("[FunUI Veda] the kernel fell back to dense at least once: %s", state["failed"])
@@ -192,7 +197,7 @@ def install_veda(model, cfg, dense_fn=None, dense_label=None, attend=None):
     to = patched.model_options.get("transformer_options", {}).copy()
     to["optimized_attention_override"] = make_override(state, cfg, bundle, attend, dense_fn, dense_label)
     patched.model_options["transformer_options"] = to
-    patched.add_wrapper_with_key("diffusion_model", "funui_sla_state", sla.make_wrapper(state, cfg))
+    sla.add_wrappers(patched, state, cfg)
     cfg["asked"] = f"trained ({core.describe_budget(bundle.generated)})"
     return patched, "\n".join([
         f"SLA on | veda (EXPERIMENTAL, Triton kernel untested on a GPU) | predictor {core.describe_bundle(bundle)} | sparsity {cfg['asked']} | "

modules/loaders/veda_core.py

@@ -474,8 +474,9 @@ def attention(q, k, v, layer, spec, plan, bundle, budgets, attend, cache, stats,
                 mask = block_mask(index, keep, lay)
             with timer("gather"):
                 v_t = gather(v, lay, hs)
-            with timer("kernel"):
-                o_t = attend(q_t, k_t, v_t, mask, lay)
+            with timer("kernel"):          # a kernel that can split its own time takes the timer too; the others are called as before
+                o_t = attend(q_t, k_t, v_t, mask, lay, timer=timer) if timer is not _NO_TIMER and getattr(attend, "times_itself", False) \
+                    else attend(q_t, k_t, v_t, mask, lay)
             with timer("scatter"):
                 scatter_(out, o_t, lay, hs)
     return out[:seq_len]

modules/loaders/veda_kernel.py

@@ -16,6 +16,8 @@ What differs from upstream, as in Veda's fork:
   * a tile is two key blocks of 64 (upstream's block), expanded on the host so the loop stays flat;
   * tile-ordered [N, H, D] tensors throughout. Veda's optional TMA path (off by default there) is not ported.
 """
+import contextlib
+
 import torch
 import triton
 import triton.language as tl
@@ -81,24 +83,32 @@ def _quant(x, block, pre_scale=1.0):
     return out, scale
 
 
-def attend(q, k, v, mask, lay):
+def attend(q, k, v, mask, lay, timer=None):
     """The block-sparse attention callable veda_core.attention takes. q, k, v: [N, H', D] contiguous, bf16/fp16, tile order;
-    mask: [H', n_tiles, n_tiles] bool. Returns [N, H', D] in q's dtype."""
+    mask: [H', n_tiles, n_tiles] bool. Returns [N, H', D] in q's dtype. `timer` (veda_timing) splits the call into
+    kernel.index (the lists the kernel walks), kernel.quantize (INT8 q and k) and kernel.attend (the kernel itself)."""
+    t = timer or (lambda name: contextlib.nullcontext())
     slots, heads, dim = q.shape
-    index, count = tile_index_list(mask & lay.kv_ok, lay.widest)         # no read of the device to size the list
-    sub = TILE // KEY_BLOCK                          # kept tiles -> kept key blocks, the kernel's unit
-    blocks = ((index.to(torch.int32) * sub)[..., None] + torch.arange(sub, device=q.device, dtype=torch.int32)).flatten(-2).contiguous()
-    valid = (lay.valid_count[:, None] - torch.arange(sub, device=q.device) * KEY_BLOCK).clamp_(0, KEY_BLOCK).flatten().to(torch.int32).contiguous()
-    q8, q_scale = _quant(q, TILE, dim ** -0.5 * LOG2E)
-    k8, k_scale = _quant(k, KEY_BLOCK)
-    v16, out = v.to(torch.float16), torch.empty_like(q)
-    _attend[(slots // TILE, heads)](
-        q8, k8, v16, q_scale, k_scale, blocks, (count * sub).contiguous(), valid, out, q.stride(0), q.stride(1), v16.stride(0),
-        v16.stride(1), q_scale.stride(0), k_scale.stride(0), blocks.stride(0), blocks.stride(1), n_q_tiles=slots // TILE, D=dim,
-        BLK=KEY_BLOCK, BLOCK_M=TILE, num_warps=WARPS, num_stages=STAGES)
+    with t("kernel.index"):
+        index, count = tile_index_list(mask & lay.kv_ok, lay.widest)         # no read of the device to size the list
+        sub = TILE // KEY_BLOCK                          # kept tiles -> kept key blocks, the kernel's unit
+        blocks = ((index.to(torch.int32) * sub)[..., None] + torch.arange(sub, device=q.device, dtype=torch.int32)).flatten(-2).contiguous()
+        valid = (lay.valid_count[:, None] - torch.arange(sub, device=q.device) * KEY_BLOCK).clamp_(0, KEY_BLOCK).flatten().to(torch.int32).contiguous()
+    with t("kernel.quantize"):
+        q8, q_scale = _quant(q, TILE, dim ** -0.5 * LOG2E)
+        k8, k_scale = _quant(k, KEY_BLOCK)
+        v16, out = v.to(torch.float16), torch.empty_like(q)
+    with t("kernel.attend"):
+        _attend[(slots // TILE, heads)](
+            q8, k8, v16, q_scale, k_scale, blocks, (count * sub).contiguous(), valid, out, q.stride(0), q.stride(1), v16.stride(0),
+            v16.stride(1), q_scale.stride(0), k_scale.stride(0), blocks.stride(0), blocks.stride(1), n_q_tiles=slots // TILE, D=dim,
+            BLK=KEY_BLOCK, BLOCK_M=TILE, num_warps=WARPS, num_stages=STAGES)
     return out
 
 
+attend.times_itself = True
+
+
 def check(device, dtype=torch.bfloat16):
     """Run the kernel on a small problem (history span, target span with partial tiles, global rows, per-row budgets)
     and compare it with the fp32 reference. Returns None when it holds, else why not (a kernel that ignores the mask