refactor: Unify and modularize decode and prefill test. (#1375)

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

## 📌 Description

Unify decode and prefill attention test for trtllm-gen by extract common
function. Support fp8/fp4 qkvo scale, o_sf, o_sf offset test.

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 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`.
- [x] 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

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

## Reviewer Notes

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

---------

Co-authored-by: Avery Yingyi Huang <[email protected]>

Author: weireweire

flashinfer/decode.py

@@ -1288,6 +1288,13 @@ def run(
 
             self._cached_module.paged_run(*run_args)
         else:
+            # trtllm-gen does not need plan info
+            if self._backend == "trtllm-gen" and self._plan_info is None:
+                plan_info: List[int] = []
+            else:
+                plan_info = self._plan_info
+            assert plan_info is not None, "plan info is not initialized"
+
             run_args = [
                 self._float_workspace_buffer,
                 self._int_workspace_buffer,

scripts/run_test_blackwell_attention_kernels.sh

@@ -7,8 +7,8 @@ pytest -s tests/test_blackwell_fmha.py
 pytest -s tests/test_deepseek_mla.py
 
 # trtllm-gen
-pytest -s tests/test_trtllm_gen_context.py
-pytest -s tests/test_trtllm_gen_decode.py
+pytest -s tests/test_trtllm_gen_attention.py
+pytest -s tests/test_trtllm_gen_mla.py
 
 # cudnn
 pytest -s tests/test_cudnn_decode.py