[REQUEST] OPSD Profile and improve HybridEngine rollout performance

Is your feature request related to a problem? Please describe.

For OPSD/OPD training, each iteration includes student rollout generation, teacher forward and logit caching, and student forward/distillation loss computation.

The HybridEngine-backed rollout implementation introduced in #8027 provides a useful foundation for on-policy distillation. However, the performance characteristics of the rollout path are not yet well understood. Rollout generation may become a significant part of the total training time, especially with:

• ZeRO-3
• Multiple samples per prompt
• Long prompts or responses
• HybridEngine inference-cache management
• Training/inference mode transitions
• LoRA fuse/unfuse operations

A small and reproducible profiling baseline would help identify the actual bottlenecks before making broader changes to HybridEngine or the OPSD trainer.

This proposal is also related to the OPSD work listed in the Q3 roadmap (#8104).

Describe the solution you'd like

I would like to add a lightweight profiling or benchmark utility for the HybridEngine-backed OPSD rollout path.

The initial benchmark could measure the following components where practical:

• Prompt expansion
• Prefill
• Decode
• HybridEngine parameter gathering
• Inference-cache acquisition and release
• LoRA fuse/unfuse, when enabled
• Attention-mask construction
• Rollout output post-processing
• End-to-end rollout latency

The benchmark could report:

• End-to-end rollout latency
• Rollout tokens per second
• Prefill and decode latency
• Post-processing latency
• Cache-management latency
• Peak memory usage
• Optionally, total OPSD step time

The first version should remain intentionally small and easy to reproduce. It could use a small causal language model and a limited smoke-test matrix, for example:

Parameter	Example values
Batch size	1, 4
Samples per prompt	1, 4
Prompt length	128, 512
Response length	32, 128
Decoding	Greedy and sampling
ZeRO stage	Stage 0 and Stage 3, if practical

The initial contribution would focus on profiling and benchmarking rather than optimizing multiple components at once. After reviewing the measurements with the maintainers, one measurable bottleneck could be selected for a targeted optimization.

One possible candidate is the post-processing path in HybridEngineRollout.generate(). When n_samples_per_prompt > 1, prompts are expanded with repeat_interleave(), while the prompt portion of the output attention mask is currently populated using a Python loop. If profiling shows that this is a meaningful cost, it may be replaceable with a batched tensor operation.

Any such optimization should preserve:

• Single-sample behavior
• Multiple samples per prompt
• Left-padded prompts
• Variable prompt attention masks
• EOS and padding behavior
• response_start_idx
• ZeRO-3-backed HybridEngine execution

Correctness tests should cover:

1. Single-sample rollout
2. Multiple samples per prompt
3. Left-padded prompts
4. Partially masked prompts
5. Prompt and attention-mask alignment
6. EOS and padding behavior
7. Equivalence between the reference and optimized post-processing paths

Describe alternatives you've considered

Rather than modifying HybridEngine broadly without measurements, I would prefer a profiling-first approach.

Potential follow-up areas include:

• Reusing the HybridEngine inference workspace across compatible rollout calls
• Reducing unnecessary cache release and re-acquisition operations
• Avoiding unnecessary gc.collect() or empty_cache() calls on hot paths, where safe
• Reducing repeated LoRA fuse/unfuse work
• Improving sampled-decode performance
• Measuring the complete OPSD training step instead of the rollout path alone

These are possible directions rather than commitments for the initial contribution. I would be happy to narrow the scope based on maintainer feedback.

Additional context

This proposal is intended as a profiling and benchmark follow-up to #8027. A possible incremental implementation plan is:

1. Add a small profiling or benchmark utility.
2. Add correctness coverage for relevant rollout edge cases.
3. Share initial measurements and identify the highest-priority bottleneck.
4. Implement one targeted optimization, if appropriate.
5. Add before/after measurements and document the results.

I would appreciate guidance on the following:

1. Would this profiling-first follow-up to #8027 be useful?
2. Is there a preferred model and hardware configuration for the initial benchmark?
3. Should the first contribution focus on the rollout path or the complete OPSD training step?
4. Are there known HybridEngine bottlenecks that should be prioritized, such as ZeRO-3 parameter gathering or inference-cache management?
5. Would maintainers prefer standalone benchmark timing or optional instrumentation inside the rollout implementation?

Thank you for considering this proposal. I would be happy to start with the smallest useful benchmark and adjust the scope based on maintainer feedback.