[Bug]: flakey test found in #7874

[noooop]

#7874 adjusted chunked prefill the scheduling order

fp8_e4m3 model FAILED

FAILED basic_correctness/test_chunked_prefill.py::test_models_with_fp8_kv_cache[True-1-False-4-4-fp8_e4m3-nm-testing/Qwen2-1.5B-Instruct-FP8-K-V] - AssertionError: Test7:

FAILED basic_correctness/test_chunked_prefill.py::test_models_with_fp8_kv_cache[True-1-True-4-4-fp8_e4m3-nm-testing/Qwen2-1.5B-Instruct-FP8-K-V] - AssertionError: Test7:

but bf16 model PASS

tests/basic_correctness/test_chunked_prefill.py::test_models_with_fp8_kv_cache[True-1-False-4-4-auto-Qwen/Qwen2-1.5B-Instruct]
tests/basic_correctness/test_chunked_prefill.py::test_models_with_fp8_kv_cache[True-1-True-4-4-auto-Qwen/Qwen2-1.5B-Instruct]

Test 7 is stuck on the resolution of fp8_e4m3，

🐛 Describe the bug

as @jon-chuang said:

The test using top k log probs may have been bound to be flakey. Perhaps testing style like this is more reliable especially given hardware differences or drift across kernels. #8013

(However note that this is just top 1 log probs I.e. greedy decode so not sure if even this testing strategy will be reliable)

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.

[jon-chuang]

Will try setting higher NUM_LOGPROBS for fp8e4m3 as the test is simply whether sampled for one is in top K log probs of reference and vice versa