Skip to content

Recognize trainable pre-norm attention chains as graph-fusion candidates #871

Open

Description

@jxsl13

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions