[WS1] Close embedding and LM-head invariance coverage

[inaniloquentee]

Summary

• close the remaining issue [WS1][kernels] Batch-invariant embedding + LM head projection #151 coverage gap by adding RL-batch embedding invariance tests for fixed probe tokens, row permutation, padding mutation, and embedding backward layout invariance
• make the native LM-head reference use row-wise fixed-K projection so logits are bitwise batch-invariant, and update the LM-head docs/tests accordingly
• wire embedding and lm_head into the [WS1] Ground-truth harness + numerical contract for batch-invariant ops #108 operator spec path, including shared helper gradient checks
• make DeepSpeed linear_logp backend lookup pass the explicit worker device instead of relying on global device_ctx

Relationship to existing PRs

• Supersedes the useful coverage intent from [WS1]feat: add embedding invariance and LM-head linear_logp routing #190 without carrying its merge conflicts.
• Completes the [FEAT][kernels] route LM-head vocab projection through deterministic GEMM #194 device-routing/review gap on top of current main without carrying its stale failing diff.

Closes #151

Validation

• PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q tests/test_issue151_embedding_lm_head_invariance.py tests/test_embedding.py tests/test_lm_head.py tests/test_linear_logp.py tests/test_op_checks.py tests/test_operator_inputs.py tests/test_deepspeed_training_worker.py rl_engine/tests/test_dispatch.py
• python -m black --target-version py310 --line-length=100 --check
• python -m isort --profile black --line-length=100 --check-only
• python -m ruff check
• python -m mypy --python-version 3.10 --ignore-missing-imports rl_engine/
• python -m pre_commit run --files

Summary by CodeRabbit

• New Features

  • Added operator-suite coverage for embedding and language-model head, including gradient validation.
  • Kernel selection now supports explicitly specifying a device to choose the appropriate backend.

• Bug Fixes

  • Improved language-model head projection behavior for stable, row-wise fixed-K semantics and strict fp32 correctness, including correct zero-row handling.
  • Ensured embedding and language-model head outputs/gradients are invariant to batch layout, row permutation, and padded-token tail mutations.

• Documentation

  • Rewrote the language-model head operator contract to clarify semantics and dtype/accuracy/reference behavior.

• Tests

  • Added invariance, dispatch/device-selection, and lm-head fixed-K correctness tests.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: eb724d64-89d5-4a71-a6f3-037630996749

📥 Commits

Reviewing files that changed from the base of the PR and between 171c4d6 and f5673a6.

📒 Files selected for processing (4)

• docs/operators/lm_head.md
• rl_engine/kernels/ops/pytorch/linear/lm_head.py
• tests/test_issue151_embedding_lm_head_invariance.py
• tests/test_lm_head.py

🚧 Files skipped from review as they are similar to previous changes (4)

• tests/test_lm_head.py
• tests/test_issue151_embedding_lm_head_invariance.py
• rl_engine/kernels/ops/pytorch/linear/lm_head.py
• docs/operators/lm_head.md

---

📝 Walkthrough

Walkthrough

The PR changes the native LM-head reference to fixed-K row-wise projection, adds explicit-device kernel dispatch, registers embedding and LM-head operator checks, and adds batch-layout, padding, permutation, gradient, and log-probability invariance tests.

Changes

Batch-invariant forward path

Layer / File(s)	Summary
Fixed-K LM-head projection and contract docs/operators/lm_head.md, rl_engine/kernels/ops/pytorch/linear/lm_head.py, tests/test_lm_head.py	LM-head projection now uses row-wise fixed-K torch.mv computation, with updated dtype, shape, purity, dispatch, accuracy, and validation semantics.
Device-aware linear-logp dispatch rl_engine/kernels/registry.py, rl_engine/executors/deepspeed_trainer.py, rl_engine/tests/test_dispatch.py, tests/test_deepspeed_training_worker.py	Kernel lookup accepts an explicit device, resolves CUDA or ROCm accordingly, and DeepSpeed passes the worker device for non-CPU linear-logp lookup.
Embedding and LM-head operator harness contracts rl_engine/kernels/gtest/operator_specs.py, rl_engine/kernels/gtest/operator_inputs.py, tests/test_operator_inputs.py, tests/test_op_checks.py	Embedding and LM-head specifications, normalized hidden dimensions, deterministic inputs, and gradient-enabled operator-suite tests are added.
Embedding and log-probability invariance validation tests/test_issue151_embedding_lm_head_invariance.py	Tests verify embedding and LM-head/log-probability results across batch layouts, row permutations, padding mutations, devices, dtypes, and embedding gradients.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• RL-Align/RL-Kernel#170 — Directly overlaps the NativeLMHeadOp implementation and validation.
• RL-Align/RL-Kernel#122 — Shares KernelRegistry and linear-logp device/backend dispatch changes.
• RL-Align/RL-Kernel#169 — Relates to the NativeEmbeddingOp contract and operator harness coverage.

Suggested reviewers: kjldefeated, flink-ddd

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 15.91% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: embedding and LM-head batch-invariance coverage.
Linked Issues check	✅ Passed	The PR covers embedding invariance, fixed-K LM-head projection, LM-head→logprob handoff, #108 harness wiring, and layout-invariant backward checks.
Out of Scope Changes check	✅ Passed	No clear out-of-scope code changes stand out; the registry, DeepSpeed, and docs updates all support the stated embedding/LM-head invariance work.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/issue-151-complete

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@rl_engine/kernels/ops/pytorch/linear/lm_head.py`:
- Around line 118-119: Update the empty-input branch in the LM-head
implementation around flat_hidden to return an empty matmul result connected to
hidden and weight instead of hidden.new_empty, preserving autograd for zero-row
inputs. Also add the equivalent zero-row branch to the reference snippets in
tests/test_lm_head.py lines 94-100 and docs/operators/lm_head.md lines 57-61 so
they avoid torch.stack([]); the implementation site is
rl_engine/kernels/ops/pytorch/linear/lm_head.py lines 118-119.

In `@tests/test_issue151_embedding_lm_head_invariance.py`:
- Around line 240-271: Update
test_lm_head_linear_logp_handoff_is_layout_invariant_for_rl_batches to reuse the
existing CUDA/bfloat16 parameterization, covering both CPU/float32 and
CUDA/bfloat16 execution. Before calling NativeLinearLogpOp, alter the hidden
rows selected by ~mask while leaving valid rows unchanged, then retain the
reference and layout-invariance assertions so padded hidden values cannot affect
results.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: daceb832-b7f6-4a13-8c93-19c721132703

📥 Commits

Reviewing files that changed from the base of the PR and between 063a261 and 171c4d6.

📒 Files selected for processing (12)

• docs/operators/lm_head.md
• rl_engine/executors/deepspeed_trainer.py
• rl_engine/kernels/gtest/operator_inputs.py
• rl_engine/kernels/gtest/operator_specs.py
• rl_engine/kernels/ops/pytorch/linear/lm_head.py
• rl_engine/kernels/registry.py
• rl_engine/tests/test_dispatch.py
• tests/test_deepspeed_training_worker.py
• tests/test_issue151_embedding_lm_head_invariance.py
• tests/test_lm_head.py
• tests/test_op_checks.py
• tests/test_operator_inputs.py