test: Add trtllm-gen prefill test. Fix related wrapper issue. (#1346)

<!-- .github/pull_request_template.md -->

## 📌 Description

Make the old prefill test cover the direct function call and fp8.

Also add a new prefill test, which is more comprehensive and have very
similar structure to decode test, I plan to merge prefill/decode test
after added the output_dtype API in prefill test. And we may able to
remove the old test

Also fix some wrapper issue captured.

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 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

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] 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](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

Author: weireweire

flashinfer/prefill.py

@@ -1812,6 +1812,7 @@ def run(
         q: torch.Tensor,
         paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
         *args,
+        q_scale: Optional[float] = None,
         k_scale: Optional[float] = None,
         v_scale: Optional[float] = None,
         out: Optional[torch.Tensor] = None,
@@ -1890,6 +1891,8 @@ def run(
             logits_soft_cap = 0.0
         if sm_scale is None:
             sm_scale = 1.0 / math.sqrt(q.size(-1))
+        if q_scale is not None:
+            sm_scale *= q_scale
         if k_scale is not None:
             sm_scale *= k_scale
         if rope_scale is None:
@@ -1994,8 +1997,6 @@ def run(
             ]
 
         self._cached_module.paged_run(*run_args)
-        if v_scale is not None:
-            out *= v_scale
 
         return (out, lse) if return_lse else out