Shared memory out of resource

[TechxGenus]

Thanks for sharing this great resource.
I'm trying to run some benchmarks with test_mask from examples/flex_attn.ipynb on one RTX 4090. When I set B=1,H=16,S=2048,D=128, it triggers an error:

triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 131074, Hardware limit: 101376. Reducing block sizes or `num_stages` may help.

[drisspg]

Thanks for opening this PR up! Admittedly we dont have much CI testing for PyTorch on 4090. Would you mind trying to create a minimal repro on posting it on PyTorch. Feel free to tag me in the issue

[TechxGenus]

Thanks. I create it here: pytorch/pytorch#133254