Match FlashInfer TP8 workspace initialization

runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_allreduce.py

@@ -147,9 +147,28 @@ def main() -> None:
                     f"no FlashInfer fusion limit for SM{capability}, TP{world_size}"
                 ) from error
             if input_tensor.nbytes <= fusion_limit_bytes:
+                from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+                    get_fi_ar_workspace,
+                )
+
                 fi_comm = communicator.fi_ar_comm
                 if fi_comm is None or fi_comm.disabled:
                     raise SystemExit("FlashInfer all-reduce communicator is unavailable")
+                max_token_num = fusion_limit_bytes // (
+                    args.hidden_dim * input_tensor.element_size()
+                )
+                workspace = get_fi_ar_workspace(
+                    world_size=world_size,
+                    rank=rank,
+                    max_token_num=max_token_num,
+                    hidden_dim=args.hidden_dim,
+                    dtype=input_tensor.dtype,
+                    group=get_tp_group().device_group,
+                )
+                if workspace is None:
+                    raise SystemExit(
+                        "FlashInfer rejected the serving-matched integer workspace"
+                    )
                 if not fi_comm.should_use_fi_ar(input_tensor):
                     raise SystemExit("FlashInfer rejected a fusion-eligible payload")
                 return (