Motivation

Original k_buffer shape is (layer_num, size // page_size + 1, page_size, head_num, head_dim),
when enable ASCEND_USE_FIA k_buffer shape is (layer_num, (size // page_size + 1) * page_size, 1, head_num, head_dim).

The element size of self.get_key_buffer(i)[0] in def get_contiguous_buf_infos():
original is page_size * head_num * head_dim,
when enable ASCEND_USE_FIA is 1 * head_num * head_dim.

So when enable ASCEND_USE_FIA, The calculation of kv_item_lens requires the addition of * page_size.

Modifications

Fix kv_item_lens for FIA, in order to compute data size in PD transfer cases.

Accuracy Tests

python3 few_shot_gms8k.py --data-path "/PATH/TO/test.jsonl.txt" --parallel 32 --num-questions 200
Accuracy: 0.880
Invalid: 03000
Latency: 37.457

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process