[perf] Vectorize generate_batch_meta's force_fetch readiness check

[Chase-Rong]

Problem

The force_fetch branch decided per-sample readiness with a Python loop:

for i, global_idx in enumerate(batch_global_indexes):
    if global_idx < partition.production_status.shape[0]:
        sample_status = partition.production_status[global_idx, field_indices]
        if torch.all(sample_status == 1):
            production_status[i] = 1

production_status is an int8 tensor and field_indices is a Python list, so every
iteration pays:

• a torch advanced-index that re-converts field_indices to a tensor and allocates a
  fresh result tensor (1 ATen dispatch),
• == 1, allocating a bool tensor (2nd dispatch),
• torch.all(...) reducing to a 0-dim tensor (3rd dispatch),
• an implicit .item() when if converts that 0-dim tensor to a Python bool.

That is ~3 dispatches, 2 tensor allocations, one list→tensor conversion and one
tensor→scalar conversion per sample, to compare 20 bytes of int8. Measured
~52us per sample — the actual comparison is nanoseconds, so essentially all of it is
per-op overhead. On the controller's single request-handling thread that came to
~40ms for 512 samples and ~100ms for 1024.

Separately, get_field_custom_backend_meta tested field membership against a list
once per (sample, field) pair, i.e. O(N x F^2) where F is the registered field count.

Change

One batched gather instead of the loop: row_idx[:, None] with col_idx pulls the
whole (batch, fields) submatrix in a single index, then one == 1 and one
.all(dim=1) row-reduction. The bounds check becomes a vectorised rows < shape[0]
instead of a per-iteration Python if.

The number of bytes read is unchanged. What goes away is the fixed cost paid once per
sample: dispatches drop from ~3N to ~6, list→tensor conversions from N to 1, and the
Python loop disappears entirely.

Also make the membership test use a set, making that helper O(N x F).

Measurements

Ascend NPU + openYuanrong, verl GRPO, 28 registered fields, on main at 750719e with
this PR applied alone (the other two perf branches reverted to baseline). Medians;
40+ calls per bucket per run:

	baseline	this PR
generate_batch_meta, 512 samples	40.4-44.4ms	7.1ms	~6.0x
generate_batch_meta, 1024 samples	100.0ms	13.0ms	7.7x
get_field_custom_backend_meta, 512	7.3-7.8ms	3.5ms	~2.1x
get_field_custom_backend_meta, 1024	32.8ms	6.8ms	4.8x

Split of the 512-sample saving: 29.5ms from the loop, 3.8ms from the set. At 1024 it
is 61.0ms and 26.0ms — the helper's share grows because its baseline scales
super-linearly (7.3 -> 32.8ms for 2x the samples).

The ratio barely moves with batch size (6.0x vs 7.7x) because the old cost was
O(N) x large constant and the new one is O(N) x small constant. Two independent
baseline runs in the same session measured 40.43 and 44.44ms for the 512 bucket
(+-2% around the mean), so the ratio is well outside run-to-run noise.

Distribution matters more than the mean: before, nearly every batch call took
40-100ms; after, they are consistently under 20ms.

Scope and correctness

Three call sites benefit: kv_retrieve_meta, _handle_get_partition_meta_request,
and get_metadata(mode="force_fetch") directly.

Behaviour is unchanged, including the out-of-range guard for global indexes beyond
production_status and the empty field_indices case. Verified byte-identical
production_status against the old loop for fully-ready, half-ready, none-ready,
out-of-range and empty-data_fields inputs.

[Copilot]

Pull request overview

This PR improves controller-side metadata generation performance by removing per-sample Python loops in the force_fetch readiness calculation and by reducing per-field membership-check overhead when assembling custom backend metadata.

Changes:

• Vectorizes generate_batch_meta(..., mode="force_fetch") readiness computation via a batched row/column gather and row-wise reduction.
• Optimizes get_field_custom_backend_meta field filtering by converting field_names to a set for O(1) membership checks.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[ascend-robot]

CLA Signature Guide

@Chase-Rong , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
[b87c522 [perf] Vectorize generate_batch...](b87c522)	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.