[refactor] enable page size > 1 for fi backend

[zzh-stable]

[refactor] enable page size > 1 for fi backend

Motivation:

While reviewing the code, I noticed that the FlashInferBackend(fi.py) currently hardcodes the page_size constraint to 1 during attention computation. This triggers _flatten_cache during inference, which forcefully flattens the cache into a single-token granularity. Considering that the FlashInfer kernel actually fully supports configurable PagedAttention block sizes, I made a simple refactor to allow it to directly accept and support the system-configured page_size. This should make the overall data dimension reduction and transfer between the frontend and backend more natural.

Main Changes:

• Remove Hardcoded Constraints: Removed the assertion in FIMetadata that only supports page_size=1, and dropped the _flatten_cache dimension flattening logic for kv_cache in the forward stage.
• Refine Page Index Calculation: When concatenating physical pages for requests (prepare_metadata), implemented stride-based sampling using page_size to translate them into physical page numbers (Block IDs). Synchronously aligned the unit length of placeholder memory allocated during CUDAGraph capture, and added the calculation for the actual valid padding length of the last page (last_page_len_cpu).

Testing Status:

• Successfully ran core-related unit tests locally and completed LLM inference conversations using the FlashInfer backend without any errors on a single-node environment.
• In principle, this only involves underlying index alignment and does not affect existing scheduling logic like TP distributed parallelism. If there are any oversights or missed edge cases, please feel free to point them out. Thanks for reviewing!

[DarkSharpness]

Hi @zzh-stable . Thanks for your contribution. We don't use page-size > 1 for flashinfer, because it seems that there's clear performance gain. Also, in SGLang, we also treat the page-size = 1 for flashinfer backend. Could you do some local micro benchmarks and share the results? We should not optimize too much before profiling.

BTW, actually the page-size of attention backend should eventually be decoupled from the allocation page-size (attn-page-size should be no larger than alloc-page-size).