Skip to content

[codex] fix vllm prefix caching defaults #263

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aoshen02 wants to merge 1 commit into
base: main
Choose a base branch
from
+22 −0
Draft

[codex] fix vllm prefix caching defaults #263

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions tests/utils/test_vllm_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,25 @@ def test_build_vllm_cmd_adds_sleep_mode_only_for_offload_rollout(vllm_args):
assert vllm_args.vllm_enable_sleep_mode is True


@pytest.mark.unit
def test_build_vllm_cmd_enables_prefix_cache_by_default(vllm_args):
server_args = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000)

cmd, _ = mod.build_vllm_cmd_and_env(server_args)

assert "--enable-prefix-caching" in cmd
Comment on lines +231 to +236

@gemini-code-assist gemini-code-assist Bot Jun 17, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Test Gap: Mock does not reflect production behavior

This test passes because vllm_args is a mock SimpleNamespace that lacks the vllm_enable_prefix_caching attribute. However, in production, args is parsed via vllm_parse_args(), which populates vllm_enable_prefix_caching as False by default (due to store_true action behavior).

To prevent such false positives in the future, consider updating the test fixture or adding a test case that uses parsed arguments (e.g., by calling vllm_parse_args() or simulating a parsed namespace where vllm_enable_prefix_caching is False but not explicitly provided by the user).



@pytest.mark.unit
def test_build_vllm_cmd_honors_explicit_prefix_cache_disable(vllm_args):
vllm_args.vllm_enable_prefix_caching = False
server_args = mod._compute_server_args(vllm_args, rank=0, dist_init_addr=None, host="127.0.0.1", port=8000)

cmd, _ = mod.build_vllm_cmd_and_env(server_args)

assert "--enable-prefix-caching" not in cmd


@pytest.mark.unit
def test_build_vllm_cmd_does_not_infer_sleep_mode_from_colocate(vllm_args):
vllm_args.colocate = True
Expand Down
3 changes: 3 additions & 0 deletions vime/backends/vllm_utils/vllm_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -400,6 +400,9 @@ def build_vllm_cmd_and_env(server_args: dict[str, Any]) -> tuple[list[str], dict
if getattr(args, "use_rollout_routing_replay", False):
cmd += ["--enable-return-routed-experts"]

if getattr(args, "vllm_enable_prefix_caching", True) is not False:
cmd += ["--enable-prefix-caching"]
Comment on lines +403 to +404

@gemini-code-assist gemini-code-assist Bot Jun 17, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Correctness Issue: Default-on prefix caching is broken in production

In production, args is parsed using vllm_parse_args() (defined in vime/backends/vllm_utils/arguments.py), which imports vLLM's CLI arguments. Since --enable-prefix-caching is a store_true action, its default value in argparse is False. Therefore, when the user does not pass --vllm-enable-prefix-caching on the command line, args.vllm_enable_prefix_caching is populated as False (not missing/None).

As a result:

  1. getattr(args, "vllm_enable_prefix_caching", True) returns False.
  2. False is not False evaluates to False.
  3. --enable-prefix-caching is not added to the command line, leaving prefix caching disabled by default in production.

The unit test test_build_vllm_cmd_enables_prefix_cache_by_default only passes because it uses a mock SimpleNamespace that lacks the vllm_enable_prefix_caching attribute entirely, which does not reflect real-world execution.

Suggested Fix

To fix this robustly, we should set the default value of vllm_enable_prefix_caching to True when it is not explicitly provided by the user on the CLI. This can be done in vllm_parse_args() in vime/backends/vllm_utils/arguments.py:

# In vime/backends/vllm_utils/arguments.py -> vllm_parse_args()
args, _ = parser.parse_known_args()
user_provided, raw_values = _detect_user_provided_dests(parser, sys.argv[1:])
if "vllm_enable_prefix_caching" not in user_provided:
    args.vllm_enable_prefix_caching = True

This ensures that:

  • If the user doesn't pass the flag, it defaults to True.
  • If they override it via YAML vllm_overrides, _apply_vllm_overrides will correctly set it to False.
  • In vllm_engine.py, we can then safely use a simple truthiness check.
Suggested change
if getattr(args, "vllm_enable_prefix_caching", True) is not False:
cmd += ["--enable-prefix-caching"]
if getattr(args, "vllm_enable_prefix_caching", True):
cmd += ["--enable-prefix-caching"]


# gpu_memory_utilization: no vime-forced default. In colocate, training and rollout do not
# occupy the GPU simultaneously (sleep/offload cycles), so vLLM's own default is fine. A user
# value passed via --vllm-gpu-memory-utilization is auto-forwarded by _forward_vllm_cli_args.
Expand Down