[Bug]: plamo2 and phi4flash v0 are broken on main

[heheda12345]

Your current environment

The output of python collect_env.py

Your output of `python collect_env.py` here

🐛 Describe the bug

See the below comments
#22928 (comment)
#22928 (comment)

Caused by #22928
CC @tdoublep

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.