[Bugfix] Fix persistent_masked_m_silu_mul_quant tests

[varun-sundar-rabindranath]

Purpose

1. The persistent_masked_m_silu_mul_quant tests fail on main on H100/B200. This is because on H100/B200 the kernel performs ue8m0 scale ceiling, but the reference implementation does not do that. To fix this, pass an argument to the kernel to specify if it should do ue8m0 scale ceiling and only test the non ue8m0 case.
2. The kernel produces incorrect results for test input (8, 128, 128 * 33, fp8_dtype). This is because when hidden size is >= 4096 we launch a kernel with 8 warps and the kernel requires that the NUM_GROUPS (hidden_size / 128) divides evenly between all warps. For this case, the PR just falls back to the 1 warp case.

Test Plan

• Run test_silu_mul_fp8_quant_deep_gemm.py locally

Test Result

• Test passes

[gemini-code-assist]

Code Review

This pull request introduces two important bug fixes for the persistent_masked_m_silu_mul_quant kernel. The first fix correctly handles cases where the number of groups is not divisible by the number of warps by falling back to a 1-warp configuration. The second fix adds a parameter to control ue8m0 scale ceiling to align the kernel's behavior with the reference implementation in tests. The changes are logical and well-targeted. I've found one critical issue in an updated validation check that needs to be addressed.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".