[Bugfix] Fix block size in block_table with PCP

[Livinfly]

I discovered a potential issue regarding the missing logic of the block_table in #28718.
In the PCP rank, it should also store only one copy, so the parameter passed from MultiGroupBlockTable to BlockTable should likely be multiplied by cp_world_size. The original code seems to be missing the PCP_world_size factor.
Could you please verify this, or am I mistaken?

@pisceskkk

Purpose

fix block_table

Test Plan

NA

Test Result

NA

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in the block_table by factoring in pcp_world_size for block size calculations, which is essential for correct behavior with Prefill Context Parallelism. The logic change is sound. I have included one review comment suggesting a refactoring to address code duplication for fetching distributed world sizes. This will improve maintainability and prevent potential inconsistencies in the future.

[pisceskkk]

Yes, we did overlook the logic modifications for MultiGroupBlockTable. Thank you very much for fixing this.
CC @LucasWilkinson

[pisceskkk]

LGTM, thanks for fix!