[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.

Start by reading the HybridEngineRollout.generate() path and the HybridEngine work from #8027, then define the smallest reproducible profiling utility for the rollout stages listed in the issue. Done means reporting baseline rollout timings and memory metrics, with correctness coverage for single and multiple samples, padding and masking, EOS behavior, and response_start_idx.