feat(rl): TokenSpeed control endpoint — gRPC/ZMQ engines behind the RL control plane

[key4ng]

Description

Problem

SMG's M1 RL control plane could only drive HTTP engines: control calls went to the worker's data URL, so a TokenSpeed engine served over gRPC or ZMQ answered every /v1/rl call with 422 unsupported_connection_mode, its capabilities showed all-false, and the gRPC generate path reported weight_version only from dispatch-time labels because the proto carried no engine-reported version. Running a real trainer (slime) against TokenSpeed engines behind SMG also surfaced two data-plane gaps: the HTTP passthrough forwarded SMG's internal "model":"unknown" placeholder to the engine, and the gRPC /generate required a model field and answered a single prompt as a list, neither of which slime's SGLang-native client can use.

Solution

A worker's control endpoint is modeled as a well-known label (rl.control_url), discovered from TokenSpeed's gRPC server info (or supplied on registration), with wildcard hosts resolved to the worker's host. The RL crate proxies control calls to that URL for gRPC/ZMQ workers and keeps HTTP workers controlling themselves; a non-HTTP worker without an endpoint now fails with 422 no_control_endpoint. Capabilities come from rl.* labels the engine advertises, with a static TokenSpeed row (pause_modes: wait,abort, update_from: distributed) as the fallback. The proto gains an optional weight_version on generate chunks and completions, the servicer stamps it, and the gateway prefers the engine-reported version over the dispatch label. smg serve stamps rl.control_url on the TokenSpeed workers it launches. For slime's data plane, the wildcard model is never serialized upstream, and the gRPC /generate answers a single prompt as one object and defaults an absent model to the fleet's single served model. A trainer-side NCCL refit example (examples/rl/refit_from_trainer.py) covers TokenSpeed, whose engines refuse disk and tensor refits with 501.

Companion TokenSpeed PR: lightseekorg/tokenspeed rl/control-endpoint (control-app host/auth/route semantics, distributed-only advertisement, NCCL split guard, SGLang-shaped sidecar /get_server_info).

Design: docs/superpowers/specs/2026-09-21-rl-tokenspeed-control-endpoint-design.md (local); acceptance write-up: docs/superpowers/rl/tokenspeed/acceptance-h200.md.

Changes

• crates/rl: control_url on the worker view, resolve_control_url, no_control_endpoint error, capability labels with the static TokenSpeed row, tp_size falling back to attn_tp_size; model_gateway/src/rl_adapter.rs hands out a control client for gRPC/ZMQ workers.
• crates/grpc_client proto 0.4.21: GenerateStreamChunk.weight_version = 7, GenerateComplete.weight_version = 12; grpc_servicer 0.12.1 stamps them, merges the engine's rl_advertisement() into server info, redacts secrets; TOKENSPEED_GRPC_KEYS lifts rl.* keys into labels.
• gRPC generate: effective_weight_version (engine > dispatch label > "default") in the processor and the four streaming sites; single-prompt object shape and single-model default for model-less requests.
• HTTP router / protocol: GenerateRequest.model skips serialization when it holds the wildcard placeholder.
• bindings/python: Worker.control_url, smg serve stamps rl.control_url for TokenSpeed launches (--rl-control-port).
• examples/rl/refit_from_trainer.py + layout test; mock_worker gains server_args / weight_version; new integration test rl_tokenspeed_control_endpoint_test.rs; e2e test_rl_control_plane.py parametrized over SGLang-HTTP and TokenSpeed-gRPC.
• Docs: docs/guides/rl-tokenspeed.md, crates/rl/{README,NOTES,COUPLING}.md, examples/rl/README.md.

Test Plan

On an 8x H200 node (details and every command in docs/superpowers/rl/tokenspeed/acceptance-h200.md):

• cargo +nightly fmt --all -- --check; cargo clippy --all-targets --all-features -- -D warnings on Rust 1.98.0 (CI toolchain, OpenCV container); cargo test -p smg-rl -p openai-protocol -p smg incl. rl_control_plane_test, rl_tokenspeed_control_endpoint_test, grpc_pd_fanout_test, tenant_rate_limiting_grpc_test, zmq_backend_test: 2376 passed, 0 failed.
• grpc_servicer tests 19 passed; bindings tests, e2e helper tests, ruff, mypy (one pre-existing error in simple_eval_common.py).
• e2e pytest e2e_test/router/test_rl_control_plane.py -k TokenSpeed 6/6 (TokenSpeed gRPC engines) and -k 'Sglang or Disabled' 6/6 (SGLang HTTP, lmsysorg/sglang:v0.5.20).
• Acceptance item 1 on two TokenSpeed gRPC engines: discovery with control_url and label-sourced capabilities; two trainer-driven NCCL refits through /v1/rl with the next /generate reporting weight_version 1 then 2; fan-out overhead at most 1 ms over 12 fan-outs; killing an engine mid-fan-out gives 207 naming it (upstream_unreachable), breakers closed, load 0, restart gives 200; --enable-rl off gives 404 on /v1/rl/* with no smg_rl_ series.
• Acceptance item 2, unpatched slime a3f50097 (GRPO, GSM8K, Qwen3-0.6B, 200 steps, same seed) with two ts serve engines: slime's own sgl-router vs SMG --enable-rl: reward over the last 50 steps 0.771 vs 0.778 (σ ≈ 0.08), rollout throughput 6949 vs 7741 tok/GPU/s, step time 6.42 vs 6.79 s (+5.7 %, one run per arm). A labeled two-file slime patch (skip registration, abort via engine addresses) ran slime's rollouts over SMG's gRPC data plane for 10 steps.

Before: gRPC TokenSpeed workers answered 422 on every /v1/rl route, capabilities all false, weight_version never left "default"; slime's model-less /generate through SMG got 404 model_not_found on both families.

Checklist

• cargo +nightly fmt passes
• cargo clippy --all-targets --all-features -- -D warnings passes
• (Optional) Documentation updated
• (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

🤖 Generated with Claude Code

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}

[claude]

🟡 Nit: effective_weight_version returns an owned String, so the four call sites in regular/streaming.rs (lines 1135, 1168, 1363, 1410) now allocate per streamed chunk where the old code embedded &ctx.weight_version with no allocation — this is the per-token hot path. All inputs are already &strs that outlive the json! construction, so the helper can return a borrowed value and let the one non-streaming caller (processor.rs) add .to_string():

pub(crate) fn effective_weight_version<'a>(
    reported: Option<&'a str>,
    dispatch: Option<&'a str>,
) -> &'a str {
    reported.filter(|v| !v.is_empty()).or(dispatch).unwrap_or("default")
}

[claude]

🟡 Nit: the /generate path now prefers the engine-stamped version, but the OpenAI-shaped responses built from the same ProtoGenerateComplete still report only the dispatch-time label as system_fingerprint (processor.rs:321, processor.rs:945, and the four maybe_system_fingerprint(dispatch.weight_version…) sites in harmony/). After a trainer-driven NCCL refit on TokenSpeed the dispatch label never updates, so /generate will report v2 while /v1/chat/completions keeps reporting the stale registration-time value indefinitely. If that split is intentional scope (slime only reads /generate), fine — but the same effective_weight_version(complete.weight_version(), …) one-liner would keep the two surfaces consistent, and it may be worth a code comment either way so the asymmetry reads as deliberate.

[claude]

Reviewed the full diff: RL control-endpoint model (crates/rl), gateway adapter and client-cache wiring, proto weight_version plumbing across the servicer and the four streaming sites, the single-prompt/model-less /generate shape, smg serve label stamping, and the tests/example.

No blocking issues found. The control-URL wildcard resolution, the no_control_endpoint error surface, the stale-stub guard in the servicer, and the untagged-worker filter in the model-less default are all carefully handled and well tested.

Summary of findings:

• 🔴 Important: 0
• 🟡 Nit: 2 — per-chunk String allocation in the streaming hot path via effective_weight_version; system_fingerprint on the OpenAI-shaped gRPC responses still reports only the dispatch label, so it diverges from /generate after a TokenSpeed refit
• 🟣 Pre-existing: 0

[key4ng]

Superseded by the stacked drafts #2666 (weight_version provenance) → #2667 (control endpoints) → #2668 (slime data plane, e2e, trainer refit example): same content rebased on main, without the acceptance write-up and without the Co-Authored-By trailers that failed the commit-message check.