togethercomputer / togethercomputer/xorl

Reorganize src/xorl/ops: quarantine vendored code, separate kernels / layers / objectives, one home for the exact-contract family

Open
#78 2 comments 0 reactions 0 assignees View on GitHub

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:

  1. Vendored third-party trees mixed with first-party code. ops/quack/ (59 files, ~28k lines) and ops/linear_attention/flashqla/ are vendored, ops/bi_families_v2.py is vendored byte-identical into the serving engine and sha256-gated — together they are more than half of ops/ by volume, yet nothing structural distinguishes them from code we own. This has practical costs: a tree-wide ruff --fix recently "fixed" 45 vendored quack files, and every reader burns time discovering what is editable.
  2. 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 — the bi_ / exact_ / canonical_ / class_b prefixes all mean the same thing ("byte-pinned serving-parity program"), and batch_invariant_ops.py alone mixes lm-head, norm, trunk-linear, and router contracts in one module.
  3. Abstraction levels are mixed. ops/ contains raw kernels (glm5_kernels/, dsv4/), autograd ops, full nn.Module layers (GatedDeltaNet under ops/linear_attention/layers/, Mamba2Mixer under ops/ssm/), and complete RL objective functions (policy_loss, grpo_loss, cispo_loss, importance_sampling_loss, opd_* in ops/loss/). The objectives are consumed by the trainer/runner, not by models — they are not "ops".
  4. Model families have two homes. ops/glm5_kernels/ and ops/dsv4/ duplicate the family split that already exists under models/transformers/{glm5,deepseek_v4}/.

Reference points

  • Megatron-core (the canonical first-party kernel library): fusions/ (flat, one fused op per file) is separate from transformer/ (layers) which is separate from models/; third-party integration lives in extensions/; parallelism-specific ops in tensor_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_utils boundary and serving is backends/sglang_utils. XoRL is unusual in owning both the orchestration and a kernel library, which is why ops/ 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.md pattern to quack/ and bi_families_v2.py (banner comment), centralize the vendored lint/format excludes in one pre-commit/ruff block, and add ops/README.md documenting the taxonomy above. Zero import churn.
  • Phase 1 — vendored quarantine (after #77 lands, since it touches adjacent files): physically move quack/ and flashqla/ under ops/_vendored/ with re-export stubs; update the pyproject.toml vendoring 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/): create xorl/objectives/, move the RL objective modules, leave the CE/logprob kernel stack in ops/loss/. Consumers: trainer, model_runner, PP loss factory.
  • Phase 3 — ops/exact/ gathering + batch_invariant_ops.py split (coordinate with #71 point 4, which touches the same surface): one naming convention for the serving-parity family; bi_families_v2.py stays put and is re-exported.
  • Phase 4 — one home per model family: ops/glm5_kernels/ and ops/dsv4/ into models/transformers/{glm5,deepseek_v4}/kernels/; layer classes (GatedDeltaNet, Mamba2Mixer, fused norm-gate modules) into models/layers/; kernels stay in ops/.
  • Phase 5 — small merges: ep_kernels/ into moe/, retire grab-bag utils.py names, delete empty shells (ops/group_gemm is already removed by #77).

Migration mechanics

  • Move + stub: git mv, then a stub module at the old path doing from 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/distributed file-shard claims in pr-test-cpu.yml are updated in the same PR (enforced by tests/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: no nn.Module layer classes, no RL objectives.
  • The serving-parity contract family lives under one directory with one naming convention; batch_invariant_ops.py no 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.py is ~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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.