From a5f35ea5b92076367206c4f6a4fb80110767bb41 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:10:34 +0200 Subject: [PATCH] Veda: split the kernel call into index, quantize and attend timings; cheaper host prep with equal output --- modules/loaders/tests/test_veda.py | 13 +++++++++++ modules/loaders/veda_attention.py | 11 ++++++--- modules/loaders/veda_core.py | 5 +++-- modules/loaders/veda_kernel.py | 36 +++++++++++++++++++----------- 4 files changed, 47 insertions(+), 18 deletions(-) diff --git a/modules/loaders/tests/test_veda.py b/modules/loaders/tests/test_veda.py index caddd4db..9c683f97 100644 --- a/modules/loaders/tests/test_veda.py +++ b/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): diff --git a/modules/loaders/veda_attention.py b/modules/loaders/veda_attention.py index 92b8d247..17a2c718 100644 --- a/modules/loaders/veda_attention.py +++ b/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']} | " diff --git a/modules/loaders/veda_core.py b/modules/loaders/veda_core.py index da470891..802f5596 100644 --- a/modules/loaders/veda_core.py +++ b/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] diff --git a/modules/loaders/veda_kernel.py b/modules/loaders/veda_kernel.py index 0d8c4221..bf3b7511 100644 --- a/modules/loaders/veda_kernel.py +++ b/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