[Ai-assisted] CLC work stealing

[drisspg]

Stacked PRs:

• ->[Ai-assisted] CLC work stealing #2218
• Nicer headdim error message #2227

---

[Draft][Ai-assisted] CLC work stealing

Not for land yet; the scheduling leaking abstraction all over the place is bad, im going to find a better way to encapsulate that.

I had to help alot on this one even though claude and codex did most of the setup.

Example work steal "trace" recreated from printf logs

Perf run

We would expect the highest gain for the most imbalanced workloads with the current scheduling for flex; We kind of see that e.g. alibi + causal are the same and dont currently have lpt schedule set. Document mask also sees a nice boost which makes sense

I dont really know why noop (fully dense fa4 with no sparsity and no scoremod) takes a hit for hdim 64 but only on NON - GQA path, the pattern looks too regular to be chance

What is needed to figure out before land

1. A more unified API and a mechanism for turning on with the env var. think we should universally enable for Flex use cases. I think fwd only is fine for now but likely will want bwd integration
2. I have been debugging the weirdest race condition for the 128x128 test. I have narrowed it down somewhat. My current working theory is that repsonse_ptr gets allocated with random smem data. We have num_tiles < num_sms, so only initial work is needed. We query clc, if we print after the consumer_wait in warp15 we see that we clc says no more work all invalid. The other consumer warps are not properly syncing and some of them end up pulling the random reponse data before its actually beeen populated. Racecheck shows me some error but not helping find the source of this race..
3. The register spills for no-op MHA case are weird. I also spent some time debugging. NCU points pretty much to a huge registor spill. However there is not really any good reason for this to be happening in this case and not others (AFAIK). I dumped the ptx and then compiled with 13.1 ptxas and it showed no spills in the sass (claude was helping here). Im not 100% convinced this is just random ptxas edge case but leaning that way. Also the ptxas patch thing didnt seem to be realy working .. something else to look at.

[tridao]

I like the direction, CLC is the right thing to do

[tridao]

Cc @tzadouri who’s thinking about scheduling and persistence