[Ai-assisted] CLC work stealing

[drisspg]

Stacked PRs:

• Add to varlen #2346
• ->[Ai-assisted] CLC work stealing #2218

---

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

[drisspg]

@jayhshah this is what im getting with the benchmark:

This is an aws node, and im somewhat dubious of how badly it thermally throttles. But I also have some ncu compares of dense and you can potentially less tail latency with the clc in the pm samples

Normal schedule:

CLC:

NCU reported roughly 3% faster in this particular case

[jayhshah]

I threw up a branch dynamic-persistent-with-semaphore with the classic (pre CLC) way of doing dynamic persistent scheduler with a semaphore, it should be useful for ablations and eventually the sm90 kernel. Benchmark shows the expected improvement for small sequence length (holding batch * seqlen constant) but regression against SingleTileLPTScheduler for large seqlen; this is more pronounced for MHA than GQA. It's also clear that the LPT swizzle needs to be tweaked for large seqlen, at least batch and head should be treated differently.

dynamic_persistent_bench.txt

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, you can upgrade your account or add credits to your account and enable them for code reviews in your settings.

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, you can upgrade your account or add credits to your account and enable them for code reviews in your settings.

[drisspg]

Alright got 2cta working, still just forward. One thing I noticed is we are using 2cta even for M<256 which doesnt seem right

[drisspg]

Okay a few follow ups; one thing I found really helpful is these log messages + script to construct the trace from logs. I think having a FA4_LOG=0,1,2,3 and different types of logging, e.g 1 is host side log should have no perf impacts, and 3 is full debug prints in kernel using cute.printf make sense to me. I can remove from this PR but curious if others agree / think its worth adding this mechanism

Being able to explicitly disable 2cta and CLC is quite convenient, currently you need to opt in .

Also I was bad with naming and use FA4_<var> instead of FLASH_ATTENTION_<var> it is much easier to dtype :)

I think there are wins on dense cases.

But regardless I think this is an a mostly landable state and worth some input

[drisspg]

Okay

back to better state -> install the 13 verison of cutedsl explicitly

[drisspg]

cleans up logs a little

[tridao]

LGTM let's merge when it's ready