Add caller-owned workspace sizing helper

[lesj0610]

📌 Description

Add a caller-owned workspace sizing helper so integrations can allocate the FlashInfer float workspace to the actual planned requirement instead of always reserving the default workspace size.

This PR:

• Adds count-only support to AlignedAllocator so plan allocation offsets can be reused for sizing without touching workspace memory.
• Adds batch decode and batch prefill workspace_size FFI functions that return required float and int workspace bytes.
• Exposes workspace_size on generated and JIT batch decode/prefill Python modules, with None for TRTLLM-gen where the helper is not available.
• Adds CUDA tests that verify decode and FA2 prefill can plan with buffers sized exactly from workspace_size.
• Leaves plan_info layout and runtime run paths unchanged.

🔍 Related Issues

N/A

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Validation performed:

• pre-commit run --all-files
• FLASHINFER_DISABLE_VERSION_CHECK=1 MAX_JOBS=4 FLASHINFER_NVCC_THREADS=2 python -m pytest tests/attention/test_workspace_size.py -q -s -x --tb=short

Not yet verified:

• Full CUDA/JIT regression matrix across all dtype, mask, sliding-window, positional-encoding, and backend combinations.

Reviewer Notes

This PR is independent against main.

Summary by CodeRabbit

• New Features
  • Added workspace-size reporting for paged-KV batched decode and KV-cache prefill.
  • Planning modules/wrappers now expose workspace_size to let callers size buffers before running.
• Bug Fixes
  • Improved workspace sizing/allocation behavior for more accurate buffer requirements.
  • Clarified that required float workspace buffers must be 16-byte aligned.
• Tests
  • Added CUDA-only tests validating decode and prefill workspace sizing and planning with exact-sized buffers.

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

[lesj0610]

Upstream merged.