feat: masked layout fp4 gemm using cute-dsl

[yzh119]

📌 Description

Implement fp4 gemm (w/ masked layout) requested in sgl-project/sglang#7994
Adapted from cutlass's dense_blockscaled_gemm_persistent example, with DeepGEMM style tile-scheduler

🔍 Related Issues

sgl-project/sglang#7994

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

cc @fzyzcjy
Co-authored-by: Avery Huang [email protected]

[gemini-code-assist]

Note

Gemini is unable to generate a summary for this pull request due to the file types involved not being currently supported.

[yzh119]

There are still a lot of work to be done:

1. AOT compile.
2. API design.

Left them for future PRs, let's unblock users and test functionality first.

[fengxie]

I think it's safe to delete this docstring?

[fengxie]

It's a non-blocking comments.

This assumes static shape if it's passed by members. Just double check it's what we are expecting here? To support dynamic shape, m/k/l must be passed via run_cute_ptr's argument list as Int32 type I believe.

[yzh119]

I think currently it's okay to assume we have static shapes, the number of groups should depend on the TP/EP size and N/K are fixed, we can compile one for each cudagraph configuration. For M we can just set a maximum possible value and the kernel execution time will only depend on the value of mask_m tensor, not M.

cc @kaixih for confirmation.

[fengxie]

cool. I think it's also one of the advantage of using jit here. You can also selectively choose static shape which usually end-up with better SASS.