[Bugfix][Rollout] Record prefix_cache_hit_rate from vLLM usage

[aoshen02]

Problem

rollout/prefix_cache_hit_rate always logs 0 on the vLLM path (slime/sglang reports it fine). The metric pipeline (Sample.PrefixCacheInfo.add in vime/utils/types.py, _compute_prefix_cache_metrics in vime/ray/rollout.py) is byte-identical to slime — the sglang→vllm divergence broke the wiring on two ends:

1. Parser — vllm_rollout.py / vllm_streaming_rollout.py built meta from usage.prompt_tokens + usage.completion_tokens but never read the cached count, so PrefixCacheInfo.add saw cached_tokens=0 every time. The hit-rate numerator was structurally pinned to 0 even though the denominator (prompt_tokens) was populated. vLLM reports the count nested as usage.prompt_tokens_details.cached_tokens; surface it.

2. Server — vLLM only emits prompt_tokens_details when the OpenAI frontend is started with --enable-prompt-tokens-details (default off). This flag lives on FrontendArgs, not AsyncEngineArgs, so it is not reachable via the --vllm-* auto-forwarder and must be passed explicitly when assembling the vllm serve command.

vLLM source confirmation

• /inference/v1/generate → ServingTokens.serve_tokens sets usage.prompt_tokens_details.cached_tokens = final_res.num_cached_tokens only if enable_prompt_tokens_details (which is wired from args.enable_prompt_tokens_details, CLI default False).
• Prefix caching is ON by default (enable_prefix_caching=True), so the count is meaningful for grouped RL rollouts (shared prompt across n samples).

History

This is a re-application of #83 (b502fed), which was reverted by #123 (92c3014) only to "split the change out cleanly" during the slime→vime sync, then never re-landed. Adapted to current main (inlined meta build; old _vllm_meta_from_generate_choice is gone) and extended to the streaming rollout path, which postdates #83 and had the identical gap.

This is the legitimate sglang→vllm translation point; the metric/types layer stays byte-for-byte slime.

Changes (16 insertions, 3 files)

• vime/backends/vllm_utils/vllm_engine.py — add --enable-prompt-tokens-details to the launch command (always-on, mirroring slime's always-available cached_tokens; negligible cost).
• vime/rollout/vllm_rollout.py — meta["cached_tokens"] = (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0).
• vime/rollout/vllm_streaming_rollout.py — same extraction (parity).

Cosmetic metric only — does not affect training or the train_rollout_logprob_abs_diff consistency metric.

Test

• py_compile on all three files: OK.
• Standalone logic check: usage with prompt_tokens_details.cached_tokens → hit_rate=0.500 (cached 150 / prompt 300); missing key and explicit null both → cached=0, no crash.
• tests/test_sample.py exercises update_from_meta_info with cached_tokens already in meta (accumulation logic, unchanged) — not affected.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request refactors and cleans up arguments in arguments.py, adds validation checks for prefill_num_servers and vllm_config, enables prefix-cache accounting by passing --enable-prompt-tokens-details to the vLLM command, extracts cached_tokens in rollout generation, and removes the monkey-patching of dist.gather_object. Feedback is provided to replace the assert statements in validate_args with standard if statements raising ValueError to ensure robust validation even when Python is run with optimization flags.

[gemini-code-assist]

Using assert statements for validating command-line arguments or configuration is an anti-pattern in Python. Assertions can be globally disabled in production environments when Python is run with optimization flags (e.g., python -O or PYTHONOPTIMIZE=1), which would silently bypass these validation checks.

To ensure robust validation, these checks should be implemented using standard if statements that raise a ValueError.

Suggested change

	assert not (
	getattr(args, "prefill_num_servers", None) is not None and getattr(args, "rollout_external", False)
	), "prefill_num_servers cannot be set with --rollout-external-engine-addrs."
	assert not (
	getattr(args, "vllm_config", None) is not None and getattr(args, "rollout_external", False)
	), "vllm_config cannot be set with --rollout-external-engine-addrs."
	assert not (
	getattr(args, "vllm_config", None) is not None and getattr(args, "prefill_num_servers", None) is not None
	), "vllm_config and prefill_num_servers are mutually exclusive. Use server_groups in the YAML config instead."
	if getattr(args, "prefill_num_servers", None) is not None and getattr(args, "rollout_external", False):
	raise ValueError("prefill_num_servers cannot be set with --rollout-external-engine-addrs.")
	if getattr(args, "vllm_config", None) is not None and getattr(args, "rollout_external", False):
	raise ValueError("vllm_config cannot be set with --rollout-external-engine-addrs.")
	if getattr(args, "vllm_config", None) is not None and getattr(args, "prefill_num_servers", None) is not None:
	raise ValueError("vllm_config and prefill_num_servers are mutually exclusive. Use server_groups in the YAML config instead.")

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #33103784 | 📁 Comparing 18c103b against latest (fa0b6e9)

  🔍 Preview build  

26 files changed · ± 26 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/fault-tolerance.html
• advanced/megatron-config.html
• advanced/pd-disaggregation.html
• advanced/reproducibility.html
• advanced/speculative-decoding.html
• advanced/vllm-config.html
• and 16 more...