From c24eab390b5982e3527ba8a9aa21b46ecc6296e1 Mon Sep 17 00:00:00 2001 From: Dmitry Malishev Date: Tue, 4 Aug 2026 22:00:17 +0300 Subject: [PATCH] fix: compose the walk's attention mask explicitly (old-ggml CUDA compat) The ABot walk's recompute path reached ggml's fused masked soft-max, whose CUDA kernel in the registry's 2026-07-03 ggml line (50cf5630) rejects the 2D-mask head broadcast: the walk aborted on its first block with "ggml_cuda_compute_forward: SOFT_MAX failed, invalid argument". The engine submodule's newer ggml accepts it, which is why the addon has carried a ggml overlay repin - blocking the registry from staying on its sanctioned ggml REF. WanSelfAttention::forward now routes the masked case (only the ABot walk ever passes a mask; every other Wan model uses the nullptr default) through forward_kv with no cached context - the KV-cache formulation that composes scale -> add(mask) -> soft_max explicitly from primitive ops supported by every ggml revision in play. forward_kv was built to mirror the fused path op-for-op and both formulations were already gated against the goldens. Validated against ggml 50cf5630 (the registry REF): - CPU golden replay, recompute path, 4 blocks vs the PyTorch reference: WALK PASS at cosine 0.999951 / 0.998994 / 0.997401 / 0.992991 - identical to the fused path's gated results, i.e. the two formulations are numerically equivalent - CUDA (RTX 4050): recompute and KV walks both complete with coherent frames - previously the first block aborted in SOFT_MAX on this exact ggml revision Co-Authored-By: Claude Fable 5 --- src/wan.hpp | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/wan.hpp b/src/wan.hpp index cca0d41d5..9f6256e42 100644 --- a/src/wan.hpp +++ b/src/wan.hpp @@ -1397,6 +1397,18 @@ namespace WAN { // x: [N, n_token, dim] // pe: [n_token, d_head/2, 2, 2] // return [N, n_token, dim] + if (mask != nullptr) { + // Masked self-attention (only the ABot causal walk passes a + // mask): compose the mask explicitly (scale -> add -> soft_max) + // via the KV-cache formulation with no cached context, instead + // of the fused masked soft_max - the CUDA soft_max kernel of + // the ggml revision the registry's 2026-07-03 port line ships + // (50cf5630) rejects the fused op's 2D-mask head broadcast + // ("SOFT_MAX failed, invalid argument"). forward_kv mirrors + // the fused path op-for-op; both formulations pass the + // golden-replay gates. + return forward_kv(ctx, x, pe, mask, nullptr, nullptr, nullptr, nullptr); + } int64_t N = x->ne[2]; int64_t n_token = x->ne[1];