[feat] Init true on policy with qwen_dense

[maocheng23]

Summary

Initial framework introducing true-on-policy for Qwen3-dense across SGLang, Megatron, and Miles. This is one of three tightly-coupled PRs that must land together — they share a single contract identifier qwen3_dense_true_on_policy_v1 defined by a vendored schema in each repo.

Companion PRs (must land in lockstep):

• Megatron-LM: [feat] Init true on policy with qwen_dense radixark/Megatron-LM#29
• Miles: [feat] Init true on policy with qwen_dense radixark/miles#1052

Target

Bit-identical (exact-zero) logprob parity between the SGLang rollout engine and the Megatron trainer for every scored response token at TP=1, TP>1, PP>1, and Ulysses CP for Qwen3-4B (dense).

Design

Three-layer contract architecture:

1. Miles (launcher): picks contract, validates parallel layout against typed model profile, renders SGLang/Megatron/env args from a typed kernel policy, drives PPO logprob comparison.
2. SGLang (rollout numerical truth, this PR): produces logprobs using deterministic kernels — TP-invariant matmul + fixed binary tree all-reduce + fp32 log-softmax.
3. Megatron (parity target): reproduces SGLang's numerics in a differentiable training forward via SGLangSpecProvider layer classes.

The contract object owns its own runtime policy — adding a new architecture (e.g. Qwen3-MoE) is one new contract object + one new model profile entry, not edits across three repos.

In this PR (SGLang)

• New python/sglang/srt/true_on_policy/ package:
  • schema.py — vendored shared identity, byte-identical with Megatron and Miles copies
  • contracts.py — SGLangTrueOnPolicyContract with policy_for(server_args) factory; resolver is registry lookup, not branch
  • config.py — policy-aware helpers (is_true_on_policy_enabled, should_force_bfloat16_*, should_use_tp_invariant_*)
• New python/sglang/srt/tp_invariant_ops/ package: deterministic K-block fixed-tree matmul + tree all-reduce Triton kernels
• Forward-path consumers read typed runtime policy (no raw flag branching):
  • layers/{linear,layernorm,sampler,activation,communicator,logits_processor}.py
  • layers/rotary_embedding/{base,mrope}.py
  • layers/attention/vision.py
  • model_executor/{model_runner,cuda_graph_runner,forward_batch_info}.py
• Model integration:
  • models/qwen3.py — dense, parity-validated path
  • models/{qwen3_moe,qwen2,qwen2_moe,sdar,sdar_moe,step3p5}.py — defensive helper-API alignment (these models are not yet declared true-on-policy ready; alignment keeps helper API consistent for future contract additions)
• New --true-on-policy-contract CLI flag; --rl-on-policy-target demoted to legacy fallback (zero in-tree readers)
• patch_prefill_only_deterministic_inference_for_cuda_graph — env/server_args mutation gutted

Validation

• ✅ CPU unit tests green (test_on_policy_wiring.py, test_dense_deterministic_math.py)
• 🔴 GPU exact-zero E2E gate not yet run at TP=1, TP>1, PP>1, CP. This is the next task on the stack before adopting Qwen3-MoE / Qwen3-Next.

Out of scope

• Qwen3-MoE / Qwen3-Next contracts — additive after this stack lands
• Alignment harness (layer-dump comparator, first-divergence comparator) — separate follow-up PR
• Full retirement of legacy rl_on_policy_target field — zero remaining in-tree readers; cosmetic cleanup

Test plan

• CPU unit tests pass in CI
• GPU exact-zero E2E gate at TP=1
• GPU exact-zero E2E gate at TP=2
• grad_norm parity check (validates Miles _ReplicatedLossAllGatherLastDim fix)

🤖 Generated with Claude Code