fix: specify dtype with begin_forward aka plan

[zhyncs]

Motivation

fix #2313 cc @yzh119

v0.1.6 and nightly

Modifications

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.

[zhyncs]

@james-p-xu after this PR, you can use the nightly flashinfer for testing
also cc @BBuf

[zhyncs]

ref (prefill)
https://github.com/flashinfer-ai/flashinfer/blob/1b998c317a1705ae9300ae35a530d01efc0f81c9/python/flashinfer/prefill.py#L865-L889
https://github.com/flashinfer-ai/flashinfer/blob/1b998c317a1705ae9300ae35a530d01efc0f81c9/python/flashinfer/prefill.py#L1509-L1529

[zhyncs]

nightly pr test https://github.com/sgl-project/sglang/actions/runs/12221709641

[zhyncs]

v0.1.6 https://github.com/sgl-project/sglang/actions/runs/12221450009/job/34090501035
nightly https://github.com/sgl-project/sglang/actions/runs/12221709641/job/34091112429

Based on the results of the performance test, nightly's performance is worse than v0.1.6. @yzh119