togethercomputer / togethercomputer/xorl
Reorganize src/xorl/ops: quarantine vendored code, separate kernels / layers / objectives, one home for the exact-contract family
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 46
- Forks
- 1
- Avg merge
- 4h 16m
- Merged PRs (30d)
- 4
Description
Problem
src/xorl/ops is the largest package in the tree (225 files, ~72k lines) and has become a grab-bag of four different kinds of code sharing one namespace:
- Vendored third-party trees mixed with first-party code.
ops/quack/(59 files, ~28k lines) andops/linear_attention/flashqla/are vendored,ops/bi_families_v2.pyis vendored byte-identical into the serving engine and sha256-gated — together they are more than half ofops/by volume, yet nothing structural distinguishes them from code we own. This has practical costs: a tree-wideruff --fixrecently "fixed" 45 vendored quack files, and every reader burns time discovering what is editable. - The exact/serving-parity contract family is scattered flat at the root under four naming conventions.
batch_invariant_ops.py(2,188 lines),bi_families_v2.py,bi_gemm_configs.py,exact_sampling_transforms.py,rope_class_b.py,canonical_moe_leaf.py,canonical_moe_cast.py,fused_silu_and_mul.py,kernel_config_pin.py,block_fp8_native.py— thebi_/exact_/canonical_/class_bprefixes all mean the same thing ("byte-pinned serving-parity program"), andbatch_invariant_ops.pyalone mixes lm-head, norm, trunk-linear, and router contracts in one module. - Abstraction levels are mixed.
ops/contains raw kernels (glm5_kernels/,dsv4/), autograd ops, fullnn.Modulelayers (GatedDeltaNetunderops/linear_attention/layers/,Mamba2Mixerunderops/ssm/), and complete RL objective functions (policy_loss,grpo_loss,cispo_loss,importance_sampling_loss,opd_*inops/loss/). The objectives are consumed by the trainer/runner, not by models — they are not "ops". - Model families have two homes.
ops/glm5_kernels/andops/dsv4/duplicate the family split that already exists undermodels/transformers/{glm5,deepseek_v4}/.
Reference points
- Megatron-core (the canonical first-party kernel library):
fusions/(flat, one fused op per file) is separate fromtransformer/(layers) which is separate frommodels/; third-party integration lives inextensions/; parallelism-specific ops intensor_parallel/. Kernels ≠ layers ≠ models, and each is one directory. - slime / miles (thin RL orchestration): they own no kernels at all — training compute is a
backends/megatron_utilsboundary and serving isbackends/sglang_utils. XoRL is unusual in owning both the orchestration and a kernel library, which is whyops/ballooned. The transferable lesson is the explicit boundary: our serving-parity programs are effectively a vendored contract surface with SGLang and deserve the same clear edge that slime gives its backends.
Target layout
src/xorl/ops/
_vendored/ # third-party, lint-excluded centrally, never hand-edited
quack/ # moved from ops/quack (VENDORED.md pattern generalized)
flashqla/ # moved from ops/linear_attention/flashqla
exact/ # the serving-parity byte-contract programs, one naming convention
sampling_transforms.py # from ops/exact_sampling_transforms.py
rope_class_b.py, canonical_moe_leaf.py, canonical_moe_cast.py,
one_round_swiglu.py # from fused_silu_and_mul.py (exact half)
block_fp8_native.py, kernel_config_pin.py, bi_gemm_configs.py
batch_invariant/ # batch_invariant_ops.py split by concern:
lm_head.py, norms.py, trunk_linear.py, router_gemm.py
loss/ # CE/logprob KERNELS only (per_token_ce, compiled_cross_entropy,
# bi_fused_lm_head, sampling_transform_ce, fused_linear_logprob,
# vocab_parallel_*)
moe/ # + ep_kernels/ merged in (deepep sort/scatter are MoE dispatch)
linear_attention/ # pure kernels only (chunk scan, conv, gating)
ssm/ # pure kernels only
quantize/
src/xorl/objectives/ # RL losses out of ops/loss: policy_loss, grpo_loss, cispo_loss,
# importance_sampling_loss, opd_loss, opd_streaming_kl,
# reducers, loss_output, causallm_loss (the LossOutput-level API)
src/xorl/models/layers/ # layer classes out of ops: GatedDeltaNet (+ fused_norm_gate
# modules), Mamba2Mixer
src/xorl/models/transformers/glm5/kernels/ # from ops/glm5_kernels
src/xorl/models/transformers/deepseek_v4/kernels/ # from ops/dsv4
What deliberately does NOT move
ops/bi_families_v2.py— vendored byte-identical into SGLang and sha256-gated (tests/ops/test_bi_golden_gates.py, pre-commit formatting carve-out). It stays at its exact path;ops/exact/re-exports it. Moving it means re-pinning the golden gates on both sides for zero benefit.- Names of heavily-referenced modules keep their import paths alive: every move leaves a one-line re-export stub at the old path for at least one deprecation cycle. #65 resolved 84 conflicts caused by cross-cutting churn; we do not repeat that.
Phasing (sequenced against in-flight work)
Each phase is one PR, branched from main, landed before the next starts. Phases avoid files owned by open PRs until those merge.
- Phase 0 — mark, don't move (can start now): generalize the
flashqla/VENDORED.mdpattern toquack/andbi_families_v2.py(banner comment), centralize the vendored lint/format excludes in one pre-commit/ruff block, and addops/README.mddocumenting the taxonomy above. Zero import churn. - Phase 1 — vendored quarantine (after #77 lands, since it touches adjacent files): physically move
quack/andflashqla/underops/_vendored/with re-export stubs; update thepyproject.tomlvendoring note and any sync tooling paths. Purely mechanical; biggest visual win (~30k lines out of the first-party namespace). - Phase 2 — objectives out of ops/loss (after #74 lands, since #74 reworks
ops/loss/): createxorl/objectives/, move the RL objective modules, leave the CE/logprob kernel stack inops/loss/. Consumers: trainer, model_runner, PP loss factory. - Phase 3 —
ops/exact/gathering +batch_invariant_ops.pysplit (coordinate with #71 point 4, which touches the same surface): one naming convention for the serving-parity family;bi_families_v2.pystays put and is re-exported. - Phase 4 — one home per model family:
ops/glm5_kernels/andops/dsv4/intomodels/transformers/{glm5,deepseek_v4}/kernels/; layer classes (GatedDeltaNet,Mamba2Mixer, fused norm-gate modules) intomodels/layers/; kernels stay inops/. - Phase 5 — small merges:
ep_kernels/intomoe/, retire grab-bagutils.pynames, delete empty shells (ops/group_gemmis already removed by #77).
Migration mechanics
- Move + stub:
git mv, then a stub module at the old path doingfrom xorl.ops._vendored.quack import * # moved; remove after <date>so out-of-tree configs, notebooks, and the path-pinning source-inspection tests keep working during the window. - Import rewrites inside the repo are codemodded (ruff/libcst) in the same PR as the move.
- Tests move with their subjects;
tests/distributedfile-shard claims inpr-test-cpu.ymlare updated in the same PR (enforced bytests/test_cpu_test_shards.py). - Each phase's PR must be green on the full CPU matrix plus the GPU suites touching the moved files.
Acceptance criteria
- No first-party lint/format tooling ever touches vendored trees (single centralized exclude).
-
ops/contains only kernels and autograd ops: nonn.Modulelayer classes, no RL objectives. - The serving-parity contract family lives under one directory with one naming convention;
batch_invariant_ops.pyno longer exists as a 2,100-line multi-concern module. - Each model family's kernels have exactly one home.
- Every old import path still resolves (stub) for one deprecation cycle, then stubs are deleted in a final sweep PR.
Adjacent debt (observed, explicitly out of scope here)
server/runner/model_runner.pyis ~8.5k lines and the real monolith of the repo; splitting it deserves its own issue.- The precision-related sibling packages (
fp8_training/,qarl/,qlora/,lora/) could arguably group under one namespace, but they are small, self-contained, and not worth the churn now. sim/(~21k lines) is self-contained and fine where it is.
Refs: #65 (churn cost), #71 (point 4 shares files with Phase 3), #74, #77 (sequencing). Structure comparisons: Megatron-core fusions/transformer/extensions; slime/miles backend boundaries.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Begin with Phase 0 and read flashqla/VENDORED.md, ops/README.md, pyproject.toml, and the pre-commit/ruff configuration. Check tests/ops/test_bi_golden_gates.py and tests/test_cpu_test_shards.py before any later phase. Done means vendored paths are marked and centrally excluded without import churn, while the full CPU matrix remains green.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, tooling
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100