Skip to content

Example of manual decomposition leading to large speedup #946

Description

@oflatt

@yihozhang mentioned being interested in such examples, and @tucker-luminal found one. Forwarding the message here:

The egglog in question is set of four rules that are roughly looking for different variations on combinations the same set of operations (Gather, Add, Mul, Cast). https://github.com/luminal-ai/luminal/blob/497b1aad9b7af7db7efeb93b46845d9f8f8919cd/crates/luminal_cuda_lite/src/kernel/hlir.rs#L2255

When they are four separate rules and I try our qwen3_moe pytorch pipeline it can take upwards of fifteen minutes in luminal

However when I write them as a four stage process with markers luminal-ai/luminal@main...tucker/lum-677-investigate-why-qwen3-moe-models-take-so-long-to-compile

One another interesting thing to note is that this is a case where these have 0 matches, so while the graph we are trying to compile does expose this problem, it’s not as though there are genuinely lots of matches that need to be found.

It completes in under a minute. Austin thought that yihong might find it interesting to look at. I have also included a script that will run the two egglog programs with the version of egglog that we are using

run_repro.sh

qwen3_moe_tiny_staged_rules_FAST.txt
qwen3_moe_tiny_flat_rules_SLOW.txt

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions