[Core] Profile persistent FlashInfer workspace before KV sizing

[lesj0610]

Purpose

A FlashInfer wrapper allocates its dedicated int workspace on its first plan(), and the shared float arena grows on demand, so on a fresh engine neither exists yet when determine_available_memory() runs. The closing measurement of memory_profiling therefore does not see them, and KV cache sizing is computed without them.

This PR adds a lifecycle that materializes that workspace inside the profiling window, keeps it alive with a lease until the closing measurement has read it, and locks the shared arena once the last warmup that could still size it has run.

AI assistance: Codex, Claude

What this changes

Three contracts.

Fail-closed support declaration. AttentionMetadataBuilder.get_persistent_workspace_memory_profiling_support() returns REQUIRED, NEUTRAL or UNSUPPORTED, and defaults to UNSUPPORTED. The new lifecycle runs only when every builder of the model declares REQUIRED or NEUTRAL and at least one declares REQUIRED. NEUTRAL means a builder has no reservation of its own but can be carried by the common lease; a builder that has not been taught the contract keeps the default and the model stays off the path. FlashInfer declares REQUIRED; GDN declares NEUTRAL so a hybrid model can still take it.

Lease lifetime. The shared arena is owned by the global workspace manager, but the int workspace is owned by its wrapper alone. Releasing the lease inside the profiling block frees it before the closing measurement total_consumed is derived from, so the lease has to outlive that block. The order is: initial profiling snapshot, persistent workspace reservation and profile run, closing profiling measurement, lease release, CUDA graph profiling, then rebind, warmup and lock as driven by the runner lifecycle.

Conservative reservation. This PR does not query per-plan workspace sizes. FlashInfer reserves its existing default arena and materializes the wrappers the enabled prefill and decode routes will use, while memory profiling is measuring allocator use, so the int workspace each one allocates falls inside the measured window. Direct trtllm-gen's module-global workspace is materialized the same way.

A model that does not opt in does not enter the new profiling-time reservation or post-warmup lock path; the existing capture-time behavior is unchanged.

Accounting for that persistent allocation can reduce the available KV capacity. The actual change has not been measured on this branch.

Test Plan

Static checks:

• py_compile on the 12 changed Python files
• ruff check
• ruff format --check
• git diff --check
• pre-commit run mypy-3.10

Focused test:

python -m pytest \
  tests/v1/worker/test_cudagraph_memory_profiling.py \
  tests/v1/worker/test_gpu_worker.py -q

Test Result

Static checks passed. The focused run is 69 passed, 0 failed, 0 skipped.

It covers the support tri-state and its fail-closed default, REQUIRED + NEUTRAL composition on a mixed FlashInfer/GDN model, and the lease keeping a builder's int workspace alive through the closing measurement and released before CUDA graph profiling. It also covers the V1 and V2 ordering around that measurement, per-ubatch arena isolation, and arena size and pointer stability after the lock.

Notes

This PR deliberately leaves mm-prefix on the unsupported profiling path. It fixes the shared lifecycle for builders that opt in; mm-prefix will need a separate integration that materializes and retains its dedicated wrapper before it can safely opt in.

The results above come from CPU-only tests, against stubs and real WorkspaceManager instances. Not exercised on this branch:

• an actual GPU profiling run, CUDA graph capture, or serving
• the real KV cache memory and token count this reservation yields, and therefore the size of the capacity change described above
• DCP, which answers UNSUPPORTED and is covered statically only
• backend selection and other suites outside tests/v1/worker/

---

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. Not needed; this is an internal FlashInfer workspace accounting change.

[chatgpt-codex-connector]

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

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 25a67f76-68c9-4dbf-a743-72df579e9264

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}