You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Recognize trainable pre-norm attention chains as graph-fusion candidates #871
Pattern\n\nExtend the synchronized-composite detector proposed in #870 to recognize trainable pre-norm attention boundaries hidden across shallow call edges:\n\nLayerNorm -> Q/K/V MatMul siblings -> packed noncausal SDPA/MHA -> output MatMul -> residual Add.\n\nThe high-value shape is repeated inside a transformer-depth loop on a synchronous accelerator backend. Forward-only analysis is incomplete: the matching VJP also crosses projection, SDPA-backward, output-projection, and LayerNorm-backward boundaries.\n\n## Independently validated GoAI case\n\nThe batch-8 ViT path used seven synchronous Metal operations around attention. A bounded shape-keyed MPSGraph now executes the forward as one submission and the explicit seven-input backward as one submission. Runtime epsilon remains a feed rather than a cache-key dimension.\n\nApple M2 Pro, F32, batch 8, sequence 65, dimension 128, 4 heads:\n\n- boundary median speedups across three fresh-process order-alternated count-7 campaigns: 2.968x, 2.986x, 3.059x;\n- depth-4 ViT training-step median speedups: 2.110x, 2.062x, 2.062x;\n- weakest of 21 aligned full-step pairs: 1.793x;\n- output, all seven gradients, full-model logits, all parameter gradients, and input immutability passed parity gates.\n\nThe control disabled only the new attention fusion, leaving the already-merged FFN graph fusion active in both arms. This isolates the incremental attention leverage.\n\n## Detector direction\n\nBuild a shallow interprocedural dispatch DAG rather than only a linear call sequence. Rank a candidate higher when:\n\n- one normalized activation fans out into three same-shape projection MatMuls;\n- the projections reconverge at a packed MHA/SDPA op, followed by output projection and residual Add;\n- the chain repeats in a depth loop;\n- the backend is synchronous or exposes commit-and-wait native calls;\n- stable shape, dtype, batch, head count, and epsilon make graph caching bounded;\n- matching backward dispatches exist for every differentiable input.\n\nReport forward and backward eager-op counts, the depth multiplier, feature gates, and a required same-binary end-to-end benchmark. Suppress or downgrade when bias, LoRA, explicit masks, causal behavior, dynamic shapes, device crossings, or missing fused backward capability change the semantics.\n\nThis is a concrete attention-specific specialization of #870 rather than evidence that dispatch-count reduction alone wins; the earlier 51-to-5 ViT preparation flattening remained neutral.
Activity