From 171c4d69961e047a782459d7e02e90eb44e767dc Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Wed, 22 Jul 2026 16:36:12 +0800 Subject: [PATCH 1/5] test(ws1): close embedding and lm-head invariance coverage Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/lm_head.md | 111 +++---- rl_engine/executors/deepspeed_trainer.py | 6 +- rl_engine/kernels/gtest/operator_inputs.py | 12 +- rl_engine/kernels/gtest/operator_specs.py | 20 ++ .../kernels/ops/pytorch/linear/lm_head.py | 18 +- rl_engine/kernels/registry.py | 54 +++- rl_engine/tests/test_dispatch.py | 22 ++ tests/test_deepspeed_training_worker.py | 25 +- ...t_issue151_embedding_lm_head_invariance.py | 271 ++++++++++++++++++ tests/test_lm_head.py | 54 ++-- tests/test_op_checks.py | 99 +++++++ tests/test_operator_inputs.py | 19 ++ 12 files changed, 590 insertions(+), 121 deletions(-) create mode 100644 tests/test_issue151_embedding_lm_head_invariance.py diff --git a/docs/operators/lm_head.md b/docs/operators/lm_head.md index a7608385..d9dab7f1 100644 --- a/docs/operators/lm_head.md +++ b/docs/operators/lm_head.md @@ -1,18 +1,20 @@ # LM Head -The lm_head operator projects hidden states back to vocabulary logits — the final -layer of the Qwen3/Llama stack. It is a **WS1 ground-truth reference** (issue #108): -a pure-PyTorch definition of the "correct answer" that downstream fused CUDA/Triton +The lm_head operator projects hidden states back to vocabulary logits, the final +layer of the Qwen3/Llama stack. It is a WS1 ground-truth reference for issue #108: +a pure-PyTorch definition of the correct answer that downstream fused CUDA/Triton kernels are validated against. -- **LM Head** (`NativeLMHeadOp`): `out = hidden @ weight.t() (+ bias)`. +- **LM Head** (`NativeLMHeadOp`): mathematically `out = hidden @ weight.t() (+ bias)`. + The native reference implements this as row-wise fixed-K GEMV projections so the + reference path is batch-invariant. For Qwen3-8B the weight is the output projection `[vocab=151936, hidden=4096]` in the -HF `nn.Linear` `[out, in]` convention, so it is transposed internally. It is -**independent** from the embedding table (`tie_word_embeddings=false`) — the two -weights are not shared — and Qwen3 has **no bias** (`bias=None`). +HF `nn.Linear` `[out, in]` convention. It is independent from the embedding table +(`tie_word_embeddings=false`), and Qwen3 has no bias (`bias=None`). ## Entry Point + ```python from rl_engine.kernels.registry import kernel_registry @@ -24,80 +26,57 @@ logits = lm_head(hidden, weight, bias=b) # optional [vocab] bias The op exposes the WS1 dual-path contract: -- `forward(...)` — projects in the input dtype, returns the input dtype (Axis-B accuracy - candidate / dtype-behavior path). -- `forward_fp32(...)` — upcasts to fp32, accumulates in fp32, returns fp32 (the - ground-truth golden path). The matmul runs with autocast disabled and CUDA TF32 - turned off, so it stays a true fp32 reference regardless of the caller's ambient - precision context (the global `allow_tf32` flag is saved and restored around it). +- `forward(...)` projects in the input dtype and returns the input dtype. +- `forward_fp32(...)` upcasts to fp32, accumulates in fp32, and returns fp32. The + fixed-K projection runs with autocast disabled and CUDA TF32 turned off, so it + stays a true fp32 reference regardless of the caller's ambient precision context. ## Backends | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeLMHeadOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | — | — | Planned: downstream fused kernels validate against this reference. | +| CUDA / ROCm / Triton | N/A | N/A | Planned: downstream fused kernels validate against this reference. | ## Tensor Contract | Argument | Shape | Dtype | Requirements | | --- | --- | --- | --- | -| `hidden` | `[B, S, hidden]` (any leading dims) | float (fp16/bf16/fp32) | Hidden states (Qwen3-8B `hidden=4096`). | -| `weight` | `[vocab, hidden]` | float | Output projection in HF `[out, in]` layout; transposed internally. Qwen3-8B `[151936, 4096]`. | -| `bias` | `[vocab]` or `None` | float | Optional; Qwen3 has none (`None`). | -| output | `hidden.shape[:-1] + (vocab,)` | `forward`: hidden dtype · `forward_fp32`: float32 | Logits. | - -Output dtype follows `hidden`. Pure function — no randomness, no in-place mutation, -device/dtype follow the inputs. - -> **Difference from the bare `matmul` op**: lm_head takes the weight in HF `[out, in]` -> layout and transposes it internally (`weight.t()`); the `matmul` op computes a bare -> `a @ b` with no transpose. Do not use them interchangeably. +| `hidden` | `[B, S, hidden]` or any leading dims | fp16/bf16/fp32 | Hidden states. | +| `weight` | `[vocab, hidden]` | fp16/bf16/fp32 | Output projection in HF `[out, in]` layout. | +| `bias` | `[vocab]` or `None` | fp16/bf16/fp32 | Optional; Qwen3 uses `None`. | +| output | `hidden.shape[:-1] + (vocab,)` | `forward`: hidden dtype; `forward_fp32`: fp32 | Logits. | -## Dispatch Behavior - -`kernel_registry.get_op("lm_head")` resolves through the `OpBackend` priority map. On -`cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_LM_HEAD`), so every device dispatches to the fp32 reference. When fused -kernels land, they are prepended to the priority list and the native op becomes the fallback. +Output dtype follows `hidden`. The op is pure: no randomness and no in-place mutation. ## Accuracy Reference semantics (`forward_fp32`): ```python -out = hidden.float() @ weight.float().t() +flat_hidden = hidden.float().reshape(-1, hidden.size(-1)) +rows = [torch.mv(weight.float(), row) for row in flat_hidden] +out = torch.stack(rows).reshape(*hidden.shape[:-1], weight.size(0)) if bias is not None: out = out + bias.float() ``` - **Ground truth**: `forward_fp32` accumulates in and returns fp32, with autocast and - CUDA TF32 disabled so it is a true fp32 reference even if the caller has TF32 or - autocast enabled. + CUDA TF32 disabled. - **Dtype path**: `forward` runs the projection in the input dtype. Because this is a - reduction over `hidden`, low-precision accumulation **drifts** from the fp32 reference - (unlike the lossless embedding gather). Unlike `forward_fp32`, this path intentionally - follows the ambient precision context (it is the dtype-behavior path): with an fp32 - input it is bitwise-equal to the ground truth **when ambient TF32/autocast is off**, but - on a TF32-enabled GPU it tracks real hardware behavior and may drift. bf16/fp16 are - always checked with a tolerance. -- **Axis-B — accuracy tolerance**: measured as max absolute error relative to the output - peak magnitude. On a SMALL load point with the real `hidden=4096` reduction length, - bf16 drifts ~0.3–0.4% of peak and fp16 ~0.05%. Elementwise `rtol` is not used: many - logits are near zero while the accumulated error tracks the reduction length, not the - output value. -- **Axis-A — batch invariance**: a row's logits are independent of the rest of the batch, - so the output is bitwise-identical regardless of batch size or padding (`torch.equal`, - `atol=0`) — **provided the reduction order is fixed**. Multi-threaded CPU GEMM splits - the `hidden` reduction across threads by the `M = batch*seq` dimension, which silently - breaks bitwise batch invariance for large `hidden`; the tests pin a single thread to fix - the order. On GPU, cuBLAS likewise splits K by `M`, so a bitwise batch-invariant GEMM is - a downstream kernel concern, not a free property of `torch.matmul`. - -## Performance Notes - -Reference operator — no fused kernel or benchmark yet. Downstream fused kernels carry their -own benchmarks and are measured against this reference for correctness. + reduction over `hidden`, low-precision accumulation drifts from the fp32 reference and + is checked with tolerance. +- **Axis-A batch invariance**: a row's logits are bitwise-identical regardless of batch + size or padding. The native reference enforces this by flattening leading dimensions + and projecting each row through the same GEMV-shaped K reduction instead of relying on + batched GEMM, whose reduction tree can change with `M = batch * seq`. + +## Dispatch Behavior + +`kernel_registry.get_op("lm_head")` resolves through the `OpBackend` priority map. On +`cuda`, `rocm`, and `cpu`, the only registered backend today is the PyTorch native op +(`PYTORCH_NATIVE_LM_HEAD`), so every device dispatches to this reference. When fused +kernels land, they should preserve the same batch-invariant contract. ## Tests @@ -105,23 +84,13 @@ own benchmarks and are measured against this reference for correctness. python -m pytest tests/test_lm_head.py -v ``` -Covers: fp32 correctness vs naive matmul (bitwise, with ambient TF32 pinned off), -`forward_fp32` precision-context safety (true fp32 under ambient autocast + restores the -global TF32 flag on CPU; numerically beats a TF32 matmul on GPU), bf16/fp16 dtype-path -accuracy (relative-to-peak tolerance, with `bias`), output shape, bias semantics, Axis-A -batch invariance (slice + padding, single-thread reduction, all dtypes), input purity, -gradient flow to `hidden`/`weight` (closed-form check), registry dispatch, and a GPU-only -smoke test at the real Qwen3-8B dims (`vocab=151936, hidden=4096`) that skips when CUDA or -GPU memory is unavailable. +Covers fp32 correctness vs the fixed-K reference, precision-context safety, bf16/fp16 +accuracy, output shape, bias semantics, Axis-A batch invariance, input purity, gradient +flow to `hidden` and `weight`, registry dispatch, and a GPU-only smoke test at the real +Qwen3-8B dimensions. ## Implementation Files - `rl_engine/kernels/ops/pytorch/linear/lm_head.py` - `rl_engine/kernels/registry.py` - `tests/test_lm_head.py` - -## Known Limitations - -- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). -- Axis-A bitwise batch invariance holds only with a fixed reduction order (single-thread on - CPU); a batch-invariant GEMM on GPU is a downstream concern. diff --git a/rl_engine/executors/deepspeed_trainer.py b/rl_engine/executors/deepspeed_trainer.py index 9eb276c8..7b59bebd 100644 --- a/rl_engine/executors/deepspeed_trainer.py +++ b/rl_engine/executors/deepspeed_trainer.py @@ -246,7 +246,9 @@ def _next_published_weight_version(self, consumed_weight_version: int) -> int: self._latest_published_weight_version = published return published - def _export_zero3_full_state_model(self) -> tuple[torch.nn.Module, Mapping[str, Any]]: + def _export_zero3_full_state_model( + self, + ) -> tuple[torch.nn.Module, Mapping[str, Any]]: model = getattr(self.engine, "module", self.model) rank = self._engine_rank() if rank != 0: @@ -439,7 +441,7 @@ def _linear_logp_op_for_device(device: torch.device | str) -> Any: resolved = torch.device(device) if resolved.type == "cpu": return NativeLinearLogpOp() - return kernel_registry.get_op("linear_logp") + return kernel_registry.get_op("linear_logp", device=resolved) def _unwrap_training_model(engine: Any, fallback_model: torch.nn.Module) -> torch.nn.Module: diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f124cafb..4a0cee96 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -55,8 +55,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", - "embedding": f"{batch}x{seq}x{vocab}x{DEFAULT_HIDDEN}", - "lm_head": f"{batch}x{seq}x{vocab}", + "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", + "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", } try: @@ -169,9 +169,10 @@ def _make_embedding_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { "token_ids": _token_ids((batch, seq), vocab, args, device), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 0), } @@ -180,9 +181,10 @@ def _make_lm_head_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { - "hidden": _floating_tensor((batch, seq, DEFAULT_HIDDEN), args, dtype, device, 0), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 1), + "hidden": _floating_tensor((batch, seq, hidden_dim), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 1), "bias": None, } diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 55a4a203..f1a4ee49 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -57,6 +57,26 @@ def _load_object(path: str) -> Any: }, grad_input_names=("hidden", "lm_head_weight"), ), + "embedding": OperatorSpec( + name="embedding", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + }, + grad_input_names=("weight",), + ), + "lm_head": OperatorSpec( + name="lm_head", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + }, + grad_input_names=("hidden", "weight"), + ), } diff --git a/rl_engine/kernels/ops/pytorch/linear/lm_head.py b/rl_engine/kernels/ops/pytorch/linear/lm_head.py index 4604fee7..daa5dfd9 100644 --- a/rl_engine/kernels/ops/pytorch/linear/lm_head.py +++ b/rl_engine/kernels/ops/pytorch/linear/lm_head.py @@ -28,7 +28,9 @@ class NativeLMHeadOp: input dtype and therefore drifts from the fp32 ``forward_fp32`` ground truth. Axis-B accuracy uses a tolerance (``torch.allclose``), not bitwise equality. Axis-A batch invariance still holds bitwise within a single dtype - (each output row reduces over ``hidden`` independently of the batch). + because each flattened output row is projected through the same GEMV-shaped + reduction over ``hidden``. The reduction path therefore depends on K and V, + not on how many other rows share the batch. """ def __init__(self) -> None: @@ -101,16 +103,24 @@ def _lm_head( """ h = hidden.to(compute_dtype) w = weight.to(compute_dtype) - # [..., hidden] @ [hidden, vocab] -> [..., vocab]; weight is [vocab, hidden] (HF [out, in]). if strict_fp32: with NativeLMHeadOp._strict_fp32_matmul(h.device.type): - out = h @ w.t() + out = NativeLMHeadOp._fixed_k_projection(h, w) else: - out = h @ w.t() + out = NativeLMHeadOp._fixed_k_projection(h, w) if bias is not None: out = out + bias.to(compute_dtype) return out.to(output_dtype) + @staticmethod + def _fixed_k_projection(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + flat_hidden = hidden.reshape(-1, hidden.size(-1)) + if flat_hidden.size(0) == 0: + return hidden.new_empty(*hidden.shape[:-1], weight.size(0)) + rows = [torch.mv(weight, row) for row in flat_hidden] + flat_out = torch.stack(rows, dim=0) + return flat_out.reshape(*hidden.shape[:-1], weight.size(0)) + @staticmethod @contextmanager def _strict_fp32_matmul(device_type: str): diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..705b3b69 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -6,6 +6,8 @@ from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +import torch + from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -163,11 +165,18 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], - "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], + "attn": [ + OpBackend.FLASH_ATTN, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_ATTN, + ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], @@ -179,7 +188,11 @@ def __init__(self): "rope": [OpBackend.PYTORCH_NATIVE_ROPE], }, "rocm": { - "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], + "logp": [ + OpBackend.ROCM_AITER, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_NATIVE, + ], "logp_deterministic": [OpBackend.PYTORCH_NATIVE], "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [ @@ -191,7 +204,10 @@ def __init__(self): "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], @@ -237,7 +253,11 @@ def _adjust_priority_from_env(self): OpBackend.ROCM_FLASH_ATTN, OpBackend.TRITON_GENERIC, ] - elif rocm_attn_backend and rocm_attn_backend not in {"native", "pytorch", "sdpa"}: + elif rocm_attn_backend and rocm_attn_backend not in { + "native", + "pytorch", + "sdpa", + }: logger.warning( "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.", rocm_attn_backend, @@ -281,16 +301,11 @@ def _adjust_priority_for_hardware(self): except Exception as e: logger.warning(f"Failed to probe device capability: {e}") - def get_op(self, op_type: str) -> Any: + def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: """Core distribution logic: Automatically select the best operator based on hardware and priority. """ - if device_ctx.is_rocm: - platform = "rocm" - elif device_ctx.device_type == "cuda": - platform = "cuda" - else: - platform = "cpu" + platform = self._platform_for_device(device) candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: @@ -314,6 +329,21 @@ def get_op(self, op_type: str) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def _platform_for_device(self, device: torch.device | str | None) -> str: + if device is None: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + resolved = torch.device(device) + if resolved.type == "cuda": + return "rocm" if torch.version.hip is not None else "cuda" + if resolved.type in self._priority_map: + return resolved.type + return "cpu" + def _load_backend(self, backend: OpBackend) -> Optional[Type]: """Dynamic loading technique: Import modules only when needed and check environment dependencies. diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 7f2839e1..3d61a068 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -42,6 +42,28 @@ def test_rocm_attention_native_sdpa_opt_out(monkeypatch): assert registry._priority_map["rocm"]["attn"][1] == OpBackend.ROCM_FLASH_ATTN +def test_registry_explicit_device_selects_device_platform(monkeypatch): + registry = KernelRegistry() + loaded = [] + + class DummyOp: + pass + + def fake_load_backend(backend): + loaded.append(backend) + return DummyOp + + monkeypatch.setattr(registry, "_load_backend", fake_load_backend) + + op = registry.get_op("logp", device="cuda:0") + + assert isinstance(op, DummyOp) + expected = ( + OpBackend.ROCM_AITER if torch.version.hip is not None else OpBackend.CUDA_FUSED_LOGP_GENERIC + ) + assert loaded[0] == expected + + def test_executor_flow(): executor = RolloutExecutor() mock_input_ids = torch.ones((1, 16), dtype=torch.long) diff --git a/tests/test_deepspeed_training_worker.py b/tests/test_deepspeed_training_worker.py index 33b7d769..d0a61e81 100644 --- a/tests/test_deepspeed_training_worker.py +++ b/tests/test_deepspeed_training_worker.py @@ -393,7 +393,9 @@ def test_extract_hidden_states_rejects_ambiguous_multi_tensor_tuple(): _extract_hidden_states((torch.randn(2, 3, 11), torch.randn(2, 3, 5))) -def test_deepspeed_training_worker_routes_linear_logp_and_zeroes_masked_targets(monkeypatch): +def test_deepspeed_training_worker_routes_linear_logp_and_zeroes_masked_targets( + monkeypatch, +): _install_fake_deepspeed(monkeypatch) from rl_engine.executors import deepspeed_trainer @@ -463,6 +465,23 @@ def test_deepspeed_training_worker_routes_linear_logp_and_zeroes_masked_targets( assert math.isfinite(result.metrics["loss"]) +def test_deepspeed_linear_logp_device_lookup_uses_explicit_worker_device(monkeypatch): + from rl_engine.executors import deepspeed_trainer + + sentinel = object() + call = {} + + def fake_get_op(op_type, *, device=None): + call["op_type"] = op_type + call["device"] = device + return sentinel + + monkeypatch.setattr(deepspeed_trainer.kernel_registry, "get_op", fake_get_op) + + assert deepspeed_trainer._linear_logp_op_for_device("cuda:0") is sentinel + assert call == {"op_type": "linear_logp", "device": torch.device("cuda:0")} + + def test_deepspeed_training_worker_rejects_ignore_index_in_model_inputs(monkeypatch): _install_fake_deepspeed(monkeypatch) from rl_engine.executors import deepspeed_trainer @@ -503,7 +522,9 @@ def test_deepspeed_training_worker_rejects_ignore_index_in_model_inputs(monkeypa worker.train(_rollout()) -def test_deepspeed_zero3_training_gathers_lm_head_parameters_during_backward(monkeypatch): +def test_deepspeed_zero3_training_gathers_lm_head_parameters_during_backward( + monkeypatch, +): _install_fake_deepspeed_with_gather(monkeypatch) from rl_engine.executors import deepspeed_trainer diff --git a/tests/test_issue151_embedding_lm_head_invariance.py b/tests/test_issue151_embedding_lm_head_invariance.py new file mode 100644 index 00000000..6f36ee78 --- /dev/null +++ b/tests/test_issue151_embedding_lm_head_invariance.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.linear.embedding import NativeEmbeddingOp +from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp +from rl_engine.testing import ( + SyntheticRLKernelBatch, + make_synthetic_rl_kernel_batch, + selected_logprobs_reference, +) + +VOCAB_SIZE = 4096 +HIDDEN_DIM = 128 +PROMPT_PROBE_POS = 1 +COMPLETION_PROBE_POS = 5 +PROMPT_PROBE_TOKEN = 1234 +COMPLETION_PROBE_TOKEN = 2345 + +BATCH_LAYOUTS = ( + dict( + num_prompts=1, + samples_per_prompt=2, + prompt_len=4, + completion_len=6, + vocab_size=VOCAB_SIZE, + valid_density=1.0, + seed=151, + ), + dict( + num_prompts=2, + samples_per_prompt=3, + prompt_len=4, + completion_len=8, + vocab_size=VOCAB_SIZE, + valid_density=0.5, + seed=152, + ), + dict( + num_prompts=3, + samples_per_prompt=4, + prompt_len=4, + completion_len=10, + vocab_size=VOCAB_SIZE, + valid_density=0.75, + seed=153, + ), +) + +CUDA_CASE = pytest.param( + "cuda", + torch.bfloat16, + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available"), +) + + +def _make_embedding_weight(*, device: str, dtype: torch.dtype, seed: int) -> torch.Tensor: + generator = torch.Generator(device=torch.device(device)) + generator.manual_seed(seed) + return torch.randn( + VOCAB_SIZE, + HIDDEN_DIM, + device=device, + dtype=torch.float32, + generator=generator, + ).to(dtype=dtype) + + +def _stamp_probe_tokens(batch: SyntheticRLKernelBatch) -> SyntheticRLKernelBatch: + completion_offset = COMPLETION_PROBE_POS - batch.prompt_len + if completion_offset < 0 or completion_offset >= batch.completion_len: + raise ValueError("completion probe position must fall inside completion tokens") + + input_ids = batch.input_ids.clone() + token_ids = batch.token_ids.clone() + completion_mask = batch.completion_mask.clone() + attention_mask = batch.attention_mask.clone() + + input_ids[:, PROMPT_PROBE_POS] = PROMPT_PROBE_TOKEN + input_ids[:, COMPLETION_PROBE_POS] = COMPLETION_PROBE_TOKEN + token_ids[:, completion_offset] = COMPLETION_PROBE_TOKEN + completion_mask[:, completion_offset] = True + attention_mask[:, COMPLETION_PROBE_POS] = True + valid_indices = completion_mask.reshape(-1).nonzero(as_tuple=False).squeeze(-1) + + return replace( + batch, + input_ids=input_ids, + token_ids=token_ids, + completion_mask=completion_mask, + attention_mask=attention_mask, + valid_indices=valid_indices, + metadata={ + **batch.metadata, + "valid_tokens": int(completion_mask.sum().item()), + "valid_density": float(completion_mask.float().mean().item()), + }, + ) + + +def _make_layout( + layout: dict[str, int | float], *, device: str, dtype: torch.dtype +) -> SyntheticRLKernelBatch: + batch = make_synthetic_rl_kernel_batch(device=device, dtype=dtype, **layout) + return _stamp_probe_tokens(batch) + + +def _permute_rows(batch: SyntheticRLKernelBatch, perm: torch.Tensor) -> SyntheticRLKernelBatch: + completion_mask = batch.completion_mask.index_select(0, perm) + return replace( + batch, + input_ids=batch.input_ids.index_select(0, perm), + attention_mask=batch.attention_mask.index_select(0, perm), + prompt_mask=batch.prompt_mask.index_select(0, perm), + completion_mask=completion_mask, + token_ids=batch.token_ids.index_select(0, perm), + rewards=batch.rewards.index_select(0, perm), + advantages=batch.advantages.index_select(0, perm), + old_logps=batch.old_logps.index_select(0, perm), + ref_logps=batch.ref_logps.index_select(0, perm), + valid_indices=completion_mask.reshape(-1).nonzero(as_tuple=False).squeeze(-1), + ) + + +def _assert_probe_vectors( + output: torch.Tensor, + *, + batch_size: int, + prompt_reference: torch.Tensor, + completion_reference: torch.Tensor, +) -> None: + assert torch.equal( + output[:, PROMPT_PROBE_POS, :], + prompt_reference.expand(batch_size, -1), + ) + assert torch.equal( + output[:, COMPLETION_PROBE_POS, :], + completion_reference.expand(batch_size, -1), + ) + + +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_embedding_lookup_is_bitwise_identical_across_rl_batch_layouts( + device: str, dtype: torch.dtype +) -> None: + op = NativeEmbeddingOp() + weight = _make_embedding_weight(device=device, dtype=dtype, seed=2026) + prompt_reference = weight[PROMPT_PROBE_TOKEN].detach() + completion_reference = weight[COMPLETION_PROBE_TOKEN].detach() + + for layout in BATCH_LAYOUTS: + batch = _make_layout(layout, device=device, dtype=dtype) + output = op.forward(batch.input_ids, weight) + _assert_probe_vectors( + output, + batch_size=batch.batch_size, + prompt_reference=prompt_reference, + completion_reference=completion_reference, + ) + + +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_embedding_lookup_is_row_order_invariant_under_rl_permutation( + device: str, dtype: torch.dtype +) -> None: + op = NativeEmbeddingOp() + weight = _make_embedding_weight(device=device, dtype=dtype, seed=2027) + batch = _make_layout(BATCH_LAYOUTS[2], device=device, dtype=dtype) + + perm = torch.arange(batch.batch_size - 1, -1, -1, device=torch.device(device)) + original = op.forward(batch.input_ids, weight) + permuted = op.forward(_permute_rows(batch, perm).input_ids, weight) + + assert torch.equal(permuted, original.index_select(0, perm)) + + +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_embedding_lookup_active_tokens_ignore_padding_tail_mutations( + device: str, dtype: torch.dtype +) -> None: + op = NativeEmbeddingOp() + weight = _make_embedding_weight(device=device, dtype=dtype, seed=2028) + batch = _make_layout(BATCH_LAYOUTS[1], device=device, dtype=dtype) + inactive = ~batch.attention_mask + assert bool(inactive.any().item()) + + mutated_input_ids = batch.input_ids.clone() + generator = torch.Generator(device=torch.device(device)) + generator.manual_seed(404) + random_tokens = torch.randint( + 0, + VOCAB_SIZE, + mutated_input_ids.shape, + device=device, + generator=generator, + dtype=torch.long, + ) + mutated_input_ids[inactive] = random_tokens[inactive] + + baseline = op.forward(batch.input_ids, weight) + candidate = op.forward(mutated_input_ids, weight) + + assert torch.equal(candidate[batch.attention_mask], baseline[batch.attention_mask]) + + +def test_embedding_backward_weight_gradient_is_layout_invariant() -> None: + torch.manual_seed(2030) + op = NativeEmbeddingOp() + base_token_ids = torch.tensor([2, 5, 1, 5, 2, 3], dtype=torch.long) + base_upstream = torch.randn(6, HIDDEN_DIM) + layouts = [ + ((2, 3), [0, 1, 2, 3, 4, 5]), + ((3, 2), [5, 2, 1, 0, 4, 3]), + ((1, 6), [3, 0, 5, 2, 1, 4]), + ] + + canonical_grad = None + for lead_shape, order in layouts: + order_t = torch.tensor(order, dtype=torch.long) + token_ids = base_token_ids.index_select(0, order_t).reshape(lead_shape) + upstream = base_upstream.index_select(0, order_t).reshape(*lead_shape, HIDDEN_DIM) + weight = _make_embedding_weight(device="cpu", dtype=torch.float32, seed=2030) + weight.requires_grad_(True) + + (op.forward(token_ids, weight) * upstream).sum().backward() + grad = weight.grad.detach().clone() + + if canonical_grad is None: + canonical_grad = grad + else: + assert torch.allclose(grad, canonical_grad, atol=1e-6) + + +def test_lm_head_linear_logp_handoff_is_layout_invariant_for_rl_batches() -> None: + torch.manual_seed(2031) + op = NativeLinearLogpOp() + weight = torch.randn(29, 7) + bias = torch.randn(29) + base_hidden = torch.randn(6, 7) + base_target = torch.tensor([3, 7, 1, 9, 4, 6], dtype=torch.long) + base_mask = torch.tensor([True, False, True, True, False, True], dtype=torch.bool) + layouts = [ + ((2, 3), [0, 1, 2, 3, 4, 5]), + ((3, 2), [5, 1, 3, 0, 4, 2]), + ((1, 6), [2, 4, 1, 5, 0, 3]), + ] + + canonical = None + for lead_shape, order in layouts: + order_t = torch.tensor(order, dtype=torch.long) + hidden = base_hidden.index_select(0, order_t).reshape(*lead_shape, -1) + target = base_target.index_select(0, order_t).reshape(lead_shape) + mask = base_mask.index_select(0, order_t).reshape(lead_shape) + safe_target = target.masked_fill(~mask, 0) + + logps = op(hidden, weight, safe_target, bias).masked_fill(~mask, 0.0) + logits = torch.nn.functional.linear(hidden.float(), weight.float(), bias.float()) + expected = selected_logprobs_reference(logits, target, mask=mask) + + assert torch.allclose(logps, expected, atol=1e-5) + recovered = logps.reshape(-1)[torch.argsort(order_t)] + if canonical is None: + canonical = recovered + else: + assert torch.allclose(recovered, canonical, atol=1e-6) diff --git a/tests/test_lm_head.py b/tests/test_lm_head.py index 15571e2b..dc36cb6d 100644 --- a/tests/test_lm_head.py +++ b/tests/test_lm_head.py @@ -12,12 +12,10 @@ bitwise -- elementwise rtol is useless here because many logits are near zero while the accumulated error tracks the reduction length, not the output value. - * Axis-A (batch invariance): still bitwise within a single dtype, but only - once the CPU reduction order is pinned. Multi-threaded CPU GEMM splits the - K (=hidden) reduction across threads differently depending on the M - (=batch*seq) dimension, which silently breaks bitwise batch invariance for - large hidden. ``_single_thread`` fixes the reduction order; this is the - local stand-in for the planned testing/determinism.py::deterministic_context. + * Axis-A (batch invariance): bitwise within a single dtype because the native + reference projects each output row with the same GEMV-shaped fixed-K + reduction. The reduction path is independent of the flattened M + (=batch*seq) dimension. """ import contextlib @@ -93,26 +91,34 @@ def _rand_weight(vocab=_VOCAB, hidden=_HIDDEN, *, seed, dtype=torch.float32): return torch.randn(vocab, hidden, generator=gen, dtype=dtype) -# Correctness of the fp32 ground truth: forward_fp32 == naive fp32 matmul, -# bitwise. forward_fp32 disables autocast/TF32, so it is a true fp32 reference -# unconditionally. The fp32 dtype path (forward) follows the ambient precision -# context, so it is bitwise-equal to the ground truth only when TF32/autocast is -# off (the default here); only bf16/fp16 introduce drift. -def test_native_lm_head_fp32_matches_naive_matmul(): - """forward_fp32 (and fp32 forward, TF32 off) is bitwise-equal to a naive fp32 matmul.""" +def _fixed_k_reference(hidden: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor: + flat_hidden = hidden.reshape(-1, hidden.size(-1)) + flat_out = torch.stack([torch.mv(weight, row) for row in flat_hidden], dim=0) + out = flat_out.reshape(*hidden.shape[:-1], weight.size(0)) + if bias is not None: + out = out + bias + return out + + +# Correctness of the fp32 ground truth: forward_fp32 == row-wise fixed-K fp32 +# projection, bitwise. This intentionally avoids batched GEMM as the golden +# source, because batched GEMM may choose a different K reduction tree when the +# flattened M (=batch*seq) dimension changes. +def test_native_lm_head_fp32_matches_fixed_k_reference(): + """forward_fp32 and fp32 forward are bitwise-equal to the fixed-K reference.""" hidden = _rand_hidden(2, 5, seed=1) weight = _rand_weight(seed=1) - # Pin TF32 off so the fp32 forward path and the naive reference below are both - # true fp32 regardless of the machine's global default, making this assertion - # deterministic. forward_fp32 disables TF32 internally and does not need this. + # Pin TF32 off so the fp32 forward path and fixed-K reference are both true + # fp32 regardless of the machine's global default. forward_fp32 disables TF32 + # internally and does not need this. prev_tf32 = torch.backends.cuda.matmul.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = False try: - naive = hidden.float() @ weight.float().t() - assert torch.equal(NativeLMHeadOp().forward_fp32(hidden, weight), naive) + ref = _fixed_k_reference(hidden.float(), weight.float()) + assert torch.equal(NativeLMHeadOp().forward_fp32(hidden, weight), ref) # fp32 forward path computes in fp32 too -> bitwise equal to ground truth. - assert torch.equal(NativeLMHeadOp().forward(hidden, weight), naive) + assert torch.equal(NativeLMHeadOp().forward(hidden, weight), ref) finally: torch.backends.cuda.matmul.allow_tf32 = prev_tf32 @@ -124,7 +130,7 @@ def test_forward_fp32_ignores_ambient_autocast_and_restores_tf32(): weight = _rand_weight(vocab=7, hidden=32, seed=11) gen = torch.Generator().manual_seed(11) bias = torch.randn(7, generator=gen) - ref = hidden.float() @ weight.float().t() + bias.float() + ref = _fixed_k_reference(hidden.float(), weight.float(), bias.float()) prev_tf32 = torch.backends.cuda.matmul.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = True @@ -330,8 +336,6 @@ def test_native_lm_head_qwen3_8b_real_shape(): out = op.forward_fp32(hidden, weight) assert out.shape == (2, 16, _QWEN3_VOCAB) assert out.dtype == torch.float32 - # Bitwise equal to the naive fp32 matmul (same call, same inputs). - assert torch.equal(out, hidden @ weight.t()) - # NB: Axis-A bitwise is asserted on CPU (single-thread reduction) above. - # On GPU it is NOT free -- cuBLAS also splits K by the M dimension, so a - # batch-invariant GEMM is a downstream kernel concern, not validated here. + # Axis-A bitwise is already asserted above; this smoke test validates the + # fixed-K path at real Qwen3-8B vocab width and hidden reduction length. + assert torch.equal(out[:1], op.forward_fp32(hidden[:1], weight)) diff --git a/tests/test_op_checks.py b/tests/test_op_checks.py index cbecddea..e076e106 100644 --- a/tests/test_op_checks.py +++ b/tests/test_op_checks.py @@ -3,9 +3,18 @@ from __future__ import annotations +import argparse + import torch from rl_engine.kernels.gtest.op_checks import CandidateSpec, OperatorCase, run_operator_suite +from rl_engine.kernels.gtest.operator_specs import ( + make_candidate, + make_operator_case, + operator_names, +) +from rl_engine.kernels.ops.pytorch.linear.embedding import NativeEmbeddingOp +from rl_engine.kernels.ops.pytorch.linear.lm_head import NativeLMHeadOp from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp @@ -34,6 +43,54 @@ def _logp_backward_case(name: str, *, seed: int = 0) -> OperatorCase: ) +def _embedding_case(name: str, *, seed: int = 0) -> OperatorCase: + generator = torch.Generator().manual_seed(seed) + token_ids = torch.tensor([[1, 7, 3], [7, 0, 5]], dtype=torch.long) + weight = torch.randn(11, 5, generator=generator) + return OperatorCase( + name=name, + op_class="elementwise", + dtype=torch.float32, + inputs={"token_ids": token_ids, "weight": weight}, + gold_fn=NativeEmbeddingOp().forward_fp32, + grad_input_names=("weight",), + ) + + +def _lm_head_case(name: str, *, seed: int = 0) -> OperatorCase: + generator = torch.Generator().manual_seed(seed) + hidden = torch.randn(2, 3, 5, generator=generator) + weight = torch.randn(13, 5, generator=generator) + return OperatorCase( + name=name, + op_class="reduction", + dtype=torch.float32, + inputs={"hidden": hidden, "weight": weight, "bias": None}, + gold_fn=NativeLMHeadOp().forward_fp32, + grad_input_names=("hidden", "weight"), + ) + + +def _spec_args(op: str) -> argparse.Namespace: + return argparse.Namespace( + op=op, + candidate="native", + arch_key=None, + batch=1, + seq=2, + vocab=17, + seed=123, + input_mode="constant", + constant_value=0.5, + token_value=3, + normalized_dim=8, + k_dim=8, + n_dim=8, + theta=1.0e6, + eps=1.0e-6, + ) + + def test_logp_native_candidate_suite_passes(): report = run_operator_suite( "logp", @@ -51,6 +108,48 @@ def test_logp_native_candidate_suite_passes(): assert all(case.passed for case in report.candidates[0].cases) +def test_embedding_native_candidate_suite_passes_issue_108_helper(): + report = run_operator_suite( + "embedding", + candidates=[ + CandidateSpec(name="native-embedding", backend="pytorch", fn=NativeEmbeddingOp()) + ], + cases=[_embedding_case("fp32", seed=10)], + check_grad=True, + ) + + assert report.passed + assert report.candidates[0].cases[0].outputs[1].message == "gradient:weight" + + +def test_lm_head_native_candidate_suite_passes_issue_108_helper(): + report = run_operator_suite( + "lm_head", + candidates=[CandidateSpec(name="native-lm-head", backend="pytorch", fn=NativeLMHeadOp())], + cases=[_lm_head_case("fp32", seed=11)], + check_grad=True, + ) + + assert report.passed + messages = [output.message for output in report.candidates[0].cases[0].outputs] + assert "gradient:hidden" in messages + assert "gradient:weight" in messages + + +def test_issue151_ops_pass_shared_issue_108_spec_path(): + assert {"embedding", "lm_head"}.issubset(operator_names()) + + for op_name in ("embedding", "lm_head"): + args = _spec_args(op_name) + report = run_operator_suite( + op_name, + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.float32, torch.device("cpu"))], + check_grad=True, + ) + assert report.passed + + def test_suite_reports_failure_for_bad_candidate(): def bad_logp(logits, token_ids): del token_ids diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index bb1a2220..db7aaed5 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -79,3 +79,22 @@ def test_constant_linear_logp_inputs_match_operator_contract(): assert torch.equal(inputs["lm_head_weight"], torch.full((17, 128), 0.51)) assert torch.equal(inputs["target_ids"], torch.full((1, 2), 3, dtype=torch.long)) assert inputs["bias"] is None + + +def test_constant_embedding_inputs_match_operator_contract(): + args = _args(input_mode="constant", constant_value=0.5, token_value=3) + inputs = make_operator_inputs("embedding", args, torch.float32, torch.device("cpu")) + + assert torch.equal(inputs["token_ids"], torch.full((1, 2), 3, dtype=torch.long)) + assert torch.equal(inputs["weight"], torch.full((17, 128), 0.5)) + assert operator_shape_name("embedding", args) == "1x2x17x128" + + +def test_constant_lm_head_inputs_match_operator_contract(): + args = _args(input_mode="constant", constant_value=0.5) + inputs = make_operator_inputs("lm_head", args, torch.float32, torch.device("cpu")) + + assert torch.equal(inputs["hidden"], torch.full((1, 2, 128), 0.5)) + assert torch.equal(inputs["weight"], torch.full((17, 128), 0.51)) + assert inputs["bias"] is None + assert operator_shape_name("lm_head", args) == "1x2x128x17" From f5673a6dbc0c6d6365157f883602cbdfe29eaf07 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Wed, 22 Jul 2026 23:06:39 +0800 Subject: [PATCH 2/5] fix(ws1): address lm-head invariance review gaps Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/lm_head.md | 8 +++- .../kernels/ops/pytorch/linear/lm_head.py | 3 +- ...t_issue151_embedding_lm_head_invariance.py | 43 +++++++++++++------ tests/test_lm_head.py | 20 ++++++++- 4 files changed, 57 insertions(+), 17 deletions(-) diff --git a/docs/operators/lm_head.md b/docs/operators/lm_head.md index d9dab7f1..79fd9047 100644 --- a/docs/operators/lm_head.md +++ b/docs/operators/lm_head.md @@ -55,8 +55,12 @@ Reference semantics (`forward_fp32`): ```python flat_hidden = hidden.float().reshape(-1, hidden.size(-1)) -rows = [torch.mv(weight.float(), row) for row in flat_hidden] -out = torch.stack(rows).reshape(*hidden.shape[:-1], weight.size(0)) +if flat_hidden.size(0) == 0: + flat_out = flat_hidden @ weight.float().t() +else: + rows = [torch.mv(weight.float(), row) for row in flat_hidden] + flat_out = torch.stack(rows) +out = flat_out.reshape(*hidden.shape[:-1], weight.size(0)) if bias is not None: out = out + bias.float() ``` diff --git a/rl_engine/kernels/ops/pytorch/linear/lm_head.py b/rl_engine/kernels/ops/pytorch/linear/lm_head.py index daa5dfd9..95bd62ea 100644 --- a/rl_engine/kernels/ops/pytorch/linear/lm_head.py +++ b/rl_engine/kernels/ops/pytorch/linear/lm_head.py @@ -116,7 +116,8 @@ def _lm_head( def _fixed_k_projection(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: flat_hidden = hidden.reshape(-1, hidden.size(-1)) if flat_hidden.size(0) == 0: - return hidden.new_empty(*hidden.shape[:-1], weight.size(0)) + flat_out = flat_hidden @ weight.t() + return flat_out.reshape(*hidden.shape[:-1], weight.size(0)) rows = [torch.mv(weight, row) for row in flat_hidden] flat_out = torch.stack(rows, dim=0) return flat_out.reshape(*hidden.shape[:-1], weight.size(0)) diff --git a/tests/test_issue151_embedding_lm_head_invariance.py b/tests/test_issue151_embedding_lm_head_invariance.py index 6f36ee78..58f558c7 100644 --- a/tests/test_issue151_embedding_lm_head_invariance.py +++ b/tests/test_issue151_embedding_lm_head_invariance.py @@ -237,14 +237,21 @@ def test_embedding_backward_weight_gradient_is_layout_invariant() -> None: assert torch.allclose(grad, canonical_grad, atol=1e-6) -def test_lm_head_linear_logp_handoff_is_layout_invariant_for_rl_batches() -> None: - torch.manual_seed(2031) +@pytest.mark.parametrize("device,dtype", [("cpu", torch.float32), CUDA_CASE]) +def test_lm_head_linear_logp_handoff_is_layout_invariant_for_rl_batches( + device: str, dtype: torch.dtype +) -> None: + target_device = torch.device(device) + generator = torch.Generator(device=target_device) + generator.manual_seed(2031) op = NativeLinearLogpOp() - weight = torch.randn(29, 7) - bias = torch.randn(29) - base_hidden = torch.randn(6, 7) - base_target = torch.tensor([3, 7, 1, 9, 4, 6], dtype=torch.long) - base_mask = torch.tensor([True, False, True, True, False, True], dtype=torch.bool) + weight = torch.randn(29, 7, device=target_device, dtype=dtype, generator=generator) + bias = torch.randn(29, device=target_device, dtype=dtype, generator=generator) + base_hidden = torch.randn(6, 7, device=target_device, dtype=dtype, generator=generator) + base_target = torch.tensor([3, 7, 1, 9, 4, 6], device=target_device, dtype=torch.long) + base_mask = torch.tensor( + [True, False, True, True, False, True], device=target_device, dtype=torch.bool + ) layouts = [ ((2, 3), [0, 1, 2, 3, 4, 5]), ((3, 2), [5, 1, 3, 0, 4, 2]), @@ -253,19 +260,29 @@ def test_lm_head_linear_logp_handoff_is_layout_invariant_for_rl_batches() -> Non canonical = None for lead_shape, order in layouts: - order_t = torch.tensor(order, dtype=torch.long) + order_t = torch.tensor(order, device=target_device, dtype=torch.long) hidden = base_hidden.index_select(0, order_t).reshape(*lead_shape, -1) target = base_target.index_select(0, order_t).reshape(lead_shape) mask = base_mask.index_select(0, order_t).reshape(lead_shape) safe_target = target.masked_fill(~mask, 0) + padded_hidden = hidden.clone() + padding_values = torch.randn( + padded_hidden.shape, + device=target_device, + dtype=dtype, + generator=generator, + ) + padded_hidden[~mask] = padding_values[~mask] - logps = op(hidden, weight, safe_target, bias).masked_fill(~mask, 0.0) - logits = torch.nn.functional.linear(hidden.float(), weight.float(), bias.float()) + logps = op(padded_hidden, weight, safe_target, bias).masked_fill(~mask, 0.0) + logits = torch.nn.functional.linear(padded_hidden.float(), weight.float(), bias.float()) expected = selected_logprobs_reference(logits, target, mask=mask) - assert torch.allclose(logps, expected, atol=1e-5) - recovered = logps.reshape(-1)[torch.argsort(order_t)] + assert torch.allclose(logps, expected, atol=5e-2 if dtype is torch.bfloat16 else 1e-5) + recovered = logps.reshape(-1)[torch.argsort(order_t)].float() if canonical is None: canonical = recovered else: - assert torch.allclose(recovered, canonical, atol=1e-6) + assert torch.allclose( + recovered, canonical, atol=5e-2 if dtype is torch.bfloat16 else 1e-6 + ) diff --git a/tests/test_lm_head.py b/tests/test_lm_head.py index dc36cb6d..51f7c481 100644 --- a/tests/test_lm_head.py +++ b/tests/test_lm_head.py @@ -93,7 +93,10 @@ def _rand_weight(vocab=_VOCAB, hidden=_HIDDEN, *, seed, dtype=torch.float32): def _fixed_k_reference(hidden: torch.Tensor, weight: torch.Tensor, bias=None) -> torch.Tensor: flat_hidden = hidden.reshape(-1, hidden.size(-1)) - flat_out = torch.stack([torch.mv(weight, row) for row in flat_hidden], dim=0) + if flat_hidden.size(0) == 0: + flat_out = flat_hidden @ weight.t() + else: + flat_out = torch.stack([torch.mv(weight, row) for row in flat_hidden], dim=0) out = flat_out.reshape(*hidden.shape[:-1], weight.size(0)) if bias is not None: out = out + bias @@ -209,6 +212,21 @@ def test_native_lm_head_output_shape(): assert out.shape == (3, 7, _VOCAB) +def test_native_lm_head_zero_rows_preserve_autograd_edges(): + op = NativeLMHeadOp() + hidden = torch.empty(0, _HIDDEN, requires_grad=True) + weight = _rand_weight(seed=31).requires_grad_(True) + + out = op.forward_fp32(hidden, weight) + out.sum().backward() + + assert out.shape == (0, _VOCAB) + assert hidden.grad is not None + assert hidden.grad.shape == hidden.shape + assert weight.grad is not None + assert torch.equal(weight.grad, torch.zeros_like(weight)) + + # Bias: None (Qwen3 default) is a plain matmul; a provided [vocab] bias is added. def test_native_lm_head_bias(): """bias=None is a plain matmul; a [vocab] bias is added elementwise.""" From 0e5cd3e3dd7e8b7790c51f5ab38fb29bb6d4fee4 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Wed, 29 Jul 2026 17:28:39 +0800 Subject: [PATCH 3/5] feat(ws1): add SM90 embedding and LM-head backends Add single-card Hopper/H200 CUDA forward paths for embedding gather and LM-head projection. The LM-head kernel assigns one CTA per output logit and reduces the full hidden dimension without Split-K, preserving batch-invariant reduction order. Signed-off-by: inaniloquentee <3051000145@qq.com> --- csrc/cuda/embedding_lm_head_sm90.cu | 294 ++++++++++++++++++ csrc/ops.cpp | 17 +- docs/operators/embedding.md | 42 +-- docs/operators/lm_head.md | 12 +- rl_engine/_C.pyi | 12 + rl_engine/kernels/gtest/operator_specs.py | 2 + rl_engine/kernels/ops/cuda/linear/__init__.py | 7 + .../kernels/ops/cuda/linear/embedding.py | 90 ++++++ rl_engine/kernels/ops/cuda/linear/lm_head.py | 138 ++++++++ rl_engine/kernels/registry.py | 14 + setup.py | 3 +- tests/test_embedding.py | 2 +- tests/test_kernel_registry.py | 37 +++ tests/test_lm_head.py | 2 +- tests/test_operator_inputs.py | 1 + tests/test_sm90_linear_wrappers.py | 130 ++++++++ 16 files changed, 777 insertions(+), 26 deletions(-) create mode 100644 csrc/cuda/embedding_lm_head_sm90.cu create mode 100644 rl_engine/kernels/ops/cuda/linear/__init__.py create mode 100644 rl_engine/kernels/ops/cuda/linear/embedding.py create mode 100644 rl_engine/kernels/ops/cuda/linear/lm_head.py create mode 100644 tests/test_sm90_linear_wrappers.py diff --git a/csrc/cuda/embedding_lm_head_sm90.cu b/csrc/cuda/embedding_lm_head_sm90.cu new file mode 100644 index 00000000..f7289eb2 --- /dev/null +++ b/csrc/cuda/embedding_lm_head_sm90.cu @@ -0,0 +1,294 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Single-card SM90 batch-invariant embedding and LM-head reference kernels. +// +// These kernels intentionally avoid Split-K. For LM-head, each output element +// owns the full hidden-dimension reduction inside one CTA, so the K traversal +// and reduction tree depend only on hidden_size, not on batch/sequence layout. + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; + +template +__device__ __forceinline__ float to_float(T value) { + return static_cast(value); +} + +template +__device__ __forceinline__ T from_float(float value) { + return static_cast(value); +} + +template +__device__ __forceinline__ float block_sum(float value) { + static_assert(BlockSize % 32 == 0, "BlockSize must be a warp multiple"); + __shared__ float shared[BlockSize / 32]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + + if (lane == 0) { + shared[warp] = value; + } + __syncthreads(); + + value = threadIdx.x < (BlockSize / 32) ? shared[lane] : 0.0f; + if (warp == 0) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + } + return value; +} + +template +__global__ void embedding_sm90_forward_kernel(const int64_t *__restrict__ token_ids, + const input_t *__restrict__ weight, + output_t *__restrict__ output, + int64_t num_tokens, int64_t hidden_size, + int64_t vocab_size) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = num_tokens * hidden_size; + if (idx >= total) { + return; + } + + const int64_t token_row = idx / hidden_size; + const int64_t hidden_col = idx - token_row * hidden_size; + const int64_t token_id = token_ids[token_row]; + if (token_id < 0 || token_id >= vocab_size) { + output[idx] = from_float(0.0f); + return; + } + output[idx] = static_cast(weight[token_id * hidden_size + hidden_col]); +} + +template +__global__ void lm_head_sm90_forward_kernel(const scalar_t *__restrict__ hidden, + const scalar_t *__restrict__ weight, + const float *__restrict__ bias, + output_t *__restrict__ output, + int64_t num_tokens, int64_t hidden_size, + int64_t vocab_size) { + const int64_t out_idx = static_cast(blockIdx.x); + const int64_t total = num_tokens * vocab_size; + if (out_idx >= total) { + return; + } + + const int64_t token_row = out_idx / vocab_size; + const int64_t vocab_col = out_idx - token_row * vocab_size; + const scalar_t *hidden_row = hidden + token_row * hidden_size; + const scalar_t *weight_row = weight + vocab_col * hidden_size; + + float acc = 0.0f; + for (int64_t k = threadIdx.x; k < hidden_size; k += blockDim.x) { + acc += to_float(hidden_row[k]) * to_float(weight_row[k]); + } + acc = block_sum(acc); + + if (threadIdx.x == 0) { + if (bias != nullptr) { + acc += bias[vocab_col]; + } + output[out_idx] = from_float(acc); + } +} + +std::vector embedding_output_sizes(torch::Tensor token_ids, int64_t hidden_size) { + std::vector sizes; + sizes.reserve(static_cast(token_ids.dim()) + 1); + for (int64_t i = 0; i < token_ids.dim(); ++i) { + sizes.push_back(token_ids.size(i)); + } + sizes.push_back(hidden_size); + return sizes; +} + +std::vector lm_head_output_sizes(torch::Tensor hidden, int64_t vocab_size) { + std::vector sizes; + sizes.reserve(static_cast(hidden.dim())); + for (int64_t i = 0; i < hidden.dim() - 1; ++i) { + sizes.push_back(hidden.size(i)); + } + sizes.push_back(vocab_size); + return sizes; +} + +bool is_supported_float_dtype(at::ScalarType dtype) { + return dtype == at::kFloat || dtype == at::kHalf || dtype == at::kBFloat16; +} + +void check_sm90_device() { + int device = 0; + C10_CUDA_CHECK(cudaGetDevice(&device)); + cudaDeviceProp prop{}; + C10_CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); + TORCH_CHECK(prop.major == 9, "SM90 embedding/lm_head kernels require Hopper, got sm_", + prop.major, prop.minor); +} + +void check_cuda_same_device(torch::Tensor a, torch::Tensor b, const char *a_name, + const char *b_name) { + TORCH_CHECK(a.is_cuda() && b.is_cuda(), a_name, " and ", b_name, + " must be CUDA tensors"); + TORCH_CHECK(a.device() == b.device(), a_name, " and ", b_name, + " must be on the same CUDA device"); +} + +torch::Tensor embedding_sm90_forward_impl(torch::Tensor token_ids, torch::Tensor weight, + bool output_fp32) { + TORCH_CHECK(token_ids.is_cuda() && weight.is_cuda(), + "token_ids and weight must be CUDA tensors"); + TORCH_CHECK(token_ids.device() == weight.device(), + "token_ids and weight must be on the same CUDA device"); + TORCH_CHECK(weight.dim() == 2, "embedding weight must be [vocab, hidden]"); + TORCH_CHECK(weight.is_contiguous(), "embedding weight must be contiguous"); + TORCH_CHECK(is_supported_float_dtype(weight.scalar_type()), + "embedding_sm90 supports fp32, fp16, and bf16 weights"); + + c10::cuda::CUDAGuard device_guard(weight.device()); + check_sm90_device(); + const int64_t vocab_size = weight.size(0); + const int64_t hidden_size = weight.size(1); + const int64_t num_tokens = token_ids.numel(); + auto ids = token_ids.reshape({num_tokens}).to(at::kLong).contiguous(); + if (num_tokens > 0) { + const int64_t min_id = ids.min().item(); + const int64_t max_id = ids.max().item(); + TORCH_CHECK(min_id >= 0 && max_id < vocab_size, + "embedding_sm90 token ids must be in [0, ", vocab_size - 1, + "], got [", min_id, ", ", max_id, "]"); + } + auto out_options = weight.options().dtype(output_fp32 ? at::kFloat : weight.scalar_type()); + auto output = torch::empty(embedding_output_sizes(token_ids, hidden_size), out_options); + if (num_tokens == 0 || hidden_size == 0) { + return output; + } + + const int64_t total = num_tokens * hidden_size; + TORCH_CHECK(total <= std::numeric_limits::max(), + "embedding_sm90 launch size exceeds CUDA grid limit"); + const int blocks = static_cast((total + kThreads - 1) / kThreads); + auto output_2d = output.reshape({num_tokens, hidden_size}); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, weight.scalar_type(), "embedding_sm90_forward", [&] { + if (output_fp32) { + embedding_sm90_forward_kernel + <<>>( + ids.data_ptr(), weight.data_ptr(), + output_2d.data_ptr(), num_tokens, hidden_size, vocab_size); + } else { + embedding_sm90_forward_kernel + <<>>( + ids.data_ptr(), weight.data_ptr(), + output_2d.data_ptr(), num_tokens, hidden_size, vocab_size); + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +torch::Tensor lm_head_sm90_forward_impl(torch::Tensor hidden, torch::Tensor weight, + torch::optional bias, + bool output_fp32) { + check_cuda_same_device(hidden, weight, "hidden", "weight"); + TORCH_CHECK(hidden.dim() >= 2, "hidden must have shape [..., hidden]"); + TORCH_CHECK(weight.dim() == 2, "lm_head weight must be [vocab, hidden]"); + TORCH_CHECK(hidden.size(-1) == weight.size(1), "hidden/weight hidden-dim mismatch"); + TORCH_CHECK(is_supported_float_dtype(hidden.scalar_type()), + "lm_head_sm90 supports fp32, fp16, and bf16 hidden states"); + TORCH_CHECK(is_supported_float_dtype(weight.scalar_type()), + "lm_head_sm90 supports fp32, fp16, and bf16 weights"); + + c10::cuda::CUDAGuard device_guard(hidden.device()); + check_sm90_device(); + const int64_t hidden_size = hidden.size(-1); + TORCH_CHECK(hidden_size > 0, "lm_head hidden dimension must be non-zero"); + const int64_t vocab_size = weight.size(0); + const int64_t num_tokens = hidden.numel() / hidden_size; + const at::ScalarType compute_dtype = output_fp32 ? at::kFloat : hidden.scalar_type(); + + auto hidden_2d = hidden.reshape({num_tokens, hidden_size}).to(compute_dtype).contiguous(); + auto weight_2d = weight.to(compute_dtype).contiguous(); + torch::Tensor bias_f; + const float *bias_ptr = nullptr; + if (bias.has_value()) { + TORCH_CHECK(bias->is_cuda(), "lm_head bias must be a CUDA tensor"); + TORCH_CHECK(bias->device() == hidden.device(), + "lm_head bias must be on the same CUDA device as hidden"); + TORCH_CHECK(bias->dim() == 1, "lm_head bias must be 1-D [vocab]"); + TORCH_CHECK(bias->numel() == vocab_size, "lm_head bias must have vocab elements"); + TORCH_CHECK(is_supported_float_dtype(bias->scalar_type()), + "lm_head_sm90 supports fp32, fp16, and bf16 bias"); + bias_f = bias->reshape({vocab_size}).to(at::kFloat).contiguous(); + bias_ptr = bias_f.data_ptr(); + } + + auto out_options = hidden.options().dtype(output_fp32 ? at::kFloat : hidden.scalar_type()); + auto output = torch::empty(lm_head_output_sizes(hidden, vocab_size), out_options); + if (num_tokens == 0 || vocab_size == 0) { + return output; + } + + const int64_t total = num_tokens * vocab_size; + TORCH_CHECK(total <= std::numeric_limits::max(), + "lm_head_sm90 launch size exceeds CUDA grid limit"); + auto output_2d = output.reshape({num_tokens, vocab_size}); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, compute_dtype, "lm_head_sm90_forward", [&] { + if (output_fp32) { + lm_head_sm90_forward_kernel + <<(total), kThreads, 0, at::cuda::getCurrentCUDAStream()>>>( + hidden_2d.data_ptr(), weight_2d.data_ptr(), + bias_ptr, output_2d.data_ptr(), num_tokens, hidden_size, + vocab_size); + } else { + lm_head_sm90_forward_kernel + <<(total), kThreads, 0, at::cuda::getCurrentCUDAStream()>>>( + hidden_2d.data_ptr(), weight_2d.data_ptr(), + bias_ptr, output_2d.data_ptr(), num_tokens, hidden_size, + vocab_size); + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +} // namespace + +torch::Tensor embedding_sm90_forward(torch::Tensor token_ids, torch::Tensor weight) { + return embedding_sm90_forward_impl(token_ids, weight, false); +} + +torch::Tensor embedding_sm90_forward_fp32(torch::Tensor token_ids, torch::Tensor weight) { + return embedding_sm90_forward_impl(token_ids, weight, true); +} + +torch::Tensor lm_head_sm90_forward(torch::Tensor hidden, torch::Tensor weight, + torch::optional bias) { + return lm_head_sm90_forward_impl(hidden, weight, bias, false); +} + +torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::Tensor weight, + torch::optional bias) { + return lm_head_sm90_forward_impl(hidden, weight, bias, true); +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 8ff8091f..dc03ab58 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -63,6 +63,14 @@ torch::Tensor linear_logp_logits_bf16_to_dlogits(torch::Tensor logits, int64_t vocab_start_index); // RoPE (rotate-half) apply for SM90; cos/sin precomputed fp32, sin_sign = +1 fwd / -1 bwd. torch::Tensor rope_apply_sm90(torch::Tensor x, torch::Tensor cos, torch::Tensor sin, double sin_sign); +torch::Tensor embedding_sm90_forward(torch::Tensor token_ids, torch::Tensor weight); +torch::Tensor embedding_sm90_forward_fp32(torch::Tensor token_ids, torch::Tensor weight); +torch::Tensor lm_head_sm90_forward(torch::Tensor hidden, + torch::Tensor weight, + torch::optional bias); +torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, + torch::Tensor weight, + torch::optional bias); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -291,9 +299,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "In-place local bf16 probs -> TP dlogits for selected log-prob backward"); m.def("linear_logp_logits_bf16_to_dlogits", &linear_logp_logits_bf16_to_dlogits, "Build bf16 dlogits from bf16 logits and fp32 lse"); - // RoPE rotate-half apply, SM90 (forward and backward share the kernel via sin_sign) m.def("rope_apply_sm90", &rope_apply_sm90, "RoPE rotate-half apply (GPT-NeoX), SM90"); + m.def("embedding_sm90_forward", &embedding_sm90_forward, + "Single-card SM90 batch-invariant embedding forward"); + m.def("embedding_sm90_forward_fp32", &embedding_sm90_forward_fp32, + "Single-card SM90 batch-invariant embedding forward with fp32 output"); + m.def("lm_head_sm90_forward", &lm_head_sm90_forward, + "Single-card SM90 batch-invariant LM-head forward"); + m.def("lm_head_sm90_forward_fp32", &lm_head_sm90_forward_fp32, + "Single-card SM90 batch-invariant LM-head forward with fp32 output"); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) diff --git a/docs/operators/embedding.md b/docs/operators/embedding.md index 58d503b7..52783d9d 100644 --- a/docs/operators/embedding.md +++ b/docs/operators/embedding.md @@ -1,14 +1,14 @@ # Token Embedding -The embedding operator maps integer token ids to their hidden-state rows — the first +The embedding operator maps integer token ids to their hidden-state rows, the first layer of the Qwen3/Llama stack. It is a **WS1 ground-truth reference** (issue #108): a pure-PyTorch definition of the "correct answer" that downstream fused CUDA/Triton kernels are validated against. -- **Embedding** (`NativeEmbeddingOp`): `out = weight[token_ids]` — a plain row gather. +- **Embedding** (`NativeEmbeddingOp`): `out = weight[token_ids]`, a plain row gather. For Qwen3-8B the table is the input embedding `[vocab=151936, hidden=4096]` and is -**independent** from the lm_head weight (`tie_word_embeddings=false`) — the two weights +**independent** from the lm_head weight (`tie_word_embeddings=false`); the two weights are not shared. ## Entry Point @@ -22,9 +22,9 @@ h = embedding(token_ids, weight) # [B, S], [vocab, hidden] -> [B, S, hidden] The op exposes the WS1 dual-path contract: -- `forward(...)` — gathers in the weight's native dtype, casts the gathered rows back to +- `forward(...)` gathers in the weight's native dtype, casts the gathered rows back to the weight dtype (Axis-B accuracy candidate / dtype-behavior path). -- `forward_fp32(...)` — native-dtype gather, then upcasts the result to fp32 (the +- `forward_fp32(...)` uses native-dtype gather, then upcasts the result to fp32 (the ground-truth golden path). ## Backends @@ -32,7 +32,8 @@ The op exposes the WS1 dual-path contract: | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeEmbeddingOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | — | — | Planned: downstream fused kernels validate against this reference. | +| CUDA SM90 (H200/Hopper) | `SM90EmbeddingOp` | `_C.embedding_sm90_forward` | Single-card batch-invariant forward backend. | +| ROCm / Triton | N/A | N/A | Falls back to the PyTorch native reference. | ## Tensor Contract @@ -40,17 +41,18 @@ The op exposes the WS1 dual-path contract: | --- | --- | --- | --- | | `token_ids` | `[B, S]` (any shape) | integer | Index dtype; cast to int64 internally. Values in `[0, vocab)`. | | `weight` | `[vocab, hidden]` | float (fp16/bf16/fp32) | Embedding table (Qwen3-8B `[151936, 4096]`). | -| output | `token_ids.shape + (hidden,)` | `forward`: weight dtype · `forward_fp32`: float32 | Gathered rows. | +| output | `token_ids.shape + (hidden,)` | `forward`: weight dtype; `forward_fp32`: float32 | Gathered rows. | Output dtype follows `weight` (the float operand); `token_ids` stay integer. Pure -function — no randomness, no in-place mutation, device/dtype follow the inputs. +function: no randomness, no in-place mutation, device/dtype follow the inputs. ## Dispatch Behavior `kernel_registry.get_op("embedding")` resolves through the `OpBackend` priority map. On -`cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_EMBEDDING`), so every device dispatches to the fp32 reference. When fused -kernels land, they are prepended to the priority list and the native op becomes the fallback. +CPU, ROCm, and CUDA devices without the SM90 extension, dispatch uses the PyTorch native op +(`PYTORCH_NATIVE_EMBEDDING`). On H200/Hopper-class builds that expose `_C.embedding_sm90_forward`, +the CUDA SM90 single-card batch-invariant backend is prepended and the native op remains +the fallback. ## Accuracy @@ -62,20 +64,21 @@ out = F.embedding(token_ids.long(), weight).to(torch.float32) - **Ground truth**: `forward_fp32` gathers in the native dtype, then upcasts to fp32. Because a gather is a lossless row copy, this is bitwise-identical to upcasting the - whole table first — but it never allocates a multi-GB fp32 copy of the full vocab + whole table first, but it never allocates a multi-GB fp32 copy of the full vocab table for a tiny lookup; only the gathered rows are upcast. - **Dtype path**: `forward` runs the same gather, then casts back to the weight dtype; it is bitwise-equal to `forward_fp32(...).to(dtype)`. -- **Lossless gather — no accuracy drift**: a row gather performs no reduction and no +- **Lossless gather, no accuracy drift**: a row gather performs no reduction and no floating-point accumulation, so the result is **bit-exact** at every dtype. There is no Axis-B tolerance to calibrate; the gathered rows equal direct indexing exactly. -- **Axis A — batch invariance**: each token's row is independent, so the output is +- **Axis A batch invariance**: each token's row is independent, so the output is bitwise-identical regardless of batch size or padding (`torch.equal`, `atol=0`). ## Performance Notes -Reference operator — no fused kernel or benchmark yet. Downstream fused kernels carry their -own benchmarks and are measured against this reference for correctness. +The SM90 backend is a simple single-card forward gather for H200/Hopper builds. It is +not a TP/vocab-parallel integration path; downstream fused kernels carry their own +benchmarks and are measured against the PyTorch reference for correctness. ## Tests @@ -92,15 +95,18 @@ test at the real Qwen3-8B dims (`vocab=151936, hidden=4096`, boundary ids `0` an ## Implementation Files - `rl_engine/kernels/ops/pytorch/linear/embedding.py` +- `rl_engine/kernels/ops/cuda/linear/embedding.py` +- `csrc/cuda/embedding_lm_head_sm90.cu` - `rl_engine/kernels/registry.py` - `tests/test_embedding.py` ## Known Limitations -- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). +- The CUDA SM90 path is single-card forward coverage only; TP/vocab-parallel and Triton + embedding backends are outside this PR. - Out-of-range token ids are not validated; callers must keep ids in `[0, vocab)`. - **GPU backward is bitwise-reproducible only under deterministic algorithms.** The - forward is a lossless gather (always reproducible), but `∂L/∂weight` is a scatter-add: + forward is a lossless gather (always reproducible), but `dL/dweight` is a scatter-add: every repeated token id (padding, common tokens) accumulates into the same `weight.grad` row. On CUDA that accumulation uses atomic adds, whose ordering is nondeterministic, so the weight gradient is not bit-exact across runs when ids collide. PyTorch documents diff --git a/docs/operators/lm_head.md b/docs/operators/lm_head.md index 79fd9047..f0acbdd9 100644 --- a/docs/operators/lm_head.md +++ b/docs/operators/lm_head.md @@ -36,7 +36,8 @@ The op exposes the WS1 dual-path contract: | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeLMHeadOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | N/A | N/A | Planned: downstream fused kernels validate against this reference. | +| CUDA SM90 (H200/Hopper) | `SM90LMHeadOp` | `_C.lm_head_sm90_forward` | Single-card batch-invariant forward backend; no Split-K. | +| ROCm / Triton | N/A | N/A | Falls back to the PyTorch native reference. | ## Tensor Contract @@ -78,9 +79,10 @@ if bias is not None: ## Dispatch Behavior `kernel_registry.get_op("lm_head")` resolves through the `OpBackend` priority map. On -`cuda`, `rocm`, and `cpu`, the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_LM_HEAD`), so every device dispatches to this reference. When fused -kernels land, they should preserve the same batch-invariant contract. +CPU, ROCm, and CUDA devices without the SM90 extension, dispatch uses the PyTorch native op +(`PYTORCH_NATIVE_LM_HEAD`). On H200/Hopper-class builds that expose `_C.lm_head_sm90_forward`, +the CUDA SM90 single-card batch-invariant backend is prepended. Its forward path assigns +one CTA to each output logit and performs the full K reduction without Split-K. ## Tests @@ -96,5 +98,7 @@ Qwen3-8B dimensions. ## Implementation Files - `rl_engine/kernels/ops/pytorch/linear/lm_head.py` +- `rl_engine/kernels/ops/cuda/linear/lm_head.py` +- `csrc/cuda/embedding_lm_head_sm90.cu` - `rl_engine/kernels/registry.py` - `tests/test_lm_head.py` diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index c28681fd..babbdc81 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -77,6 +77,18 @@ def linear_logp_logits_bf16_to_dlogits( lse: torch.Tensor, vocab_start_index: int, ) -> torch.Tensor: ... +def embedding_sm90_forward(token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: ... +def embedding_sm90_forward_fp32(token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: ... +def lm_head_sm90_forward( + hidden: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: ... +def lm_head_sm90_forward_fp32( + hidden: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: ... def fused_logp_forward_out( logits: torch.Tensor, token_ids: torch.Tensor, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index b44b125e..09304845 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -95,6 +95,7 @@ def _load_object(path: str) -> Any: gold_method="forward_fp32", candidate_paths={ "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp", }, grad_input_names=("weight",), ), @@ -105,6 +106,7 @@ def _load_object(path: str) -> Any: gold_method="forward_fp32", candidate_paths={ "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp", }, grad_input_names=("hidden", "weight"), ), diff --git a/rl_engine/kernels/ops/cuda/linear/__init__.py b/rl_engine/kernels/ops/cuda/linear/__init__.py new file mode 100644 index 00000000..aa28f08a --- /dev/null +++ b/rl_engine/kernels/ops/cuda/linear/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from rl_engine.kernels.ops.cuda.linear.embedding import SM90EmbeddingOp +from rl_engine.kernels.ops.cuda.linear.lm_head import SM90LMHeadOp + +__all__ = ["SM90EmbeddingOp", "SM90LMHeadOp"] diff --git a/rl_engine/kernels/ops/cuda/linear/embedding.py b/rl_engine/kernels/ops/cuda/linear/embedding.py new file mode 100644 index 00000000..3992ffcd --- /dev/null +++ b/rl_engine/kernels/ops/cuda/linear/embedding.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.linear.embedding import NativeEmbeddingOp +from rl_engine.utils.logger import logger + +_SUPPORTED_DTYPES = {torch.float32, torch.float16, torch.bfloat16} + + +def _is_hopper(device: torch.device) -> bool: + try: + return torch.cuda.get_device_capability(device)[0] == 9 + except Exception: + return False + + +class _SM90EmbeddingFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, token_ids: torch.Tensor, weight: torch.Tensor, output_fp32: bool): + ctx.save_for_backward(token_ids) + ctx.weight_shape = tuple(weight.shape) + ctx.weight_dtype = weight.dtype + ctx.output_fp32 = bool(output_fp32) + if output_fp32: + return _C.embedding_sm90_forward_fp32(token_ids, weight.contiguous()) + return _C.embedding_sm90_forward(token_ids, weight.contiguous()) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + (token_ids,) = ctx.saved_tensors + grad_weight = None + if ctx.needs_input_grad[1]: + ids = token_ids.reshape(-1).to(device=grad_output.device, dtype=torch.long) + hidden_size = int(ctx.weight_shape[1]) + grad_rows = grad_output.reshape(ids.numel(), hidden_size) + valid = (ids >= 0) & (ids < int(ctx.weight_shape[0])) + if not bool(valid.all().item()): + ids = ids[valid] + grad_rows = grad_rows[valid] + grad_weight = torch.zeros( + ctx.weight_shape, + device=grad_output.device, + dtype=grad_rows.dtype, + ) + grad_weight.index_add_(0, ids, grad_rows) + grad_weight = grad_weight.to(ctx.weight_dtype) + return None, grad_weight, None + + +class SM90EmbeddingOp(torch.nn.Module): + """Single-card SM90 batch-invariant CUDA embedding op.""" + + op_class = "elementwise" + is_batch_invariant = True + + def __init__(self) -> None: + super().__init__() + if not _EXT_AVAILABLE or not hasattr(_C, "embedding_sm90_forward"): + raise RuntimeError( + "embedding_sm90_forward is not compiled into the extension. " + "Rebuild on Hopper with KERNEL_ALIGN_FORCE_SM90=1." + ) + self._fallback = NativeEmbeddingOp() + logger.info("Successfully linked to precompiled _C.embedding_sm90_forward kernel.") + + def forward(self, token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if not self._can_use_sm90(token_ids, weight): + return self._fallback.forward(token_ids, weight) + return _SM90EmbeddingFunction.apply(token_ids, weight, False) + + def forward_fp32(self, token_ids: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + if not self._can_use_sm90(token_ids, weight): + return self._fallback.forward_fp32(token_ids, weight) + return _SM90EmbeddingFunction.apply(token_ids, weight, True) + + @staticmethod + def _can_use_sm90(token_ids: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + token_ids.is_cuda + and weight.is_cuda + and token_ids.device == weight.device + and _is_hopper(weight.device) + and weight.dim() == 2 + and weight.dtype in _SUPPORTED_DTYPES + ) diff --git a/rl_engine/kernels/ops/cuda/linear/lm_head.py b/rl_engine/kernels/ops/cuda/linear/lm_head.py new file mode 100644 index 00000000..f9f37cbb --- /dev/null +++ b/rl_engine/kernels/ops/cuda/linear/lm_head.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Optional + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.linear.lm_head import NativeLMHeadOp +from rl_engine.utils.logger import logger + +_SUPPORTED_DTYPES = {torch.float32, torch.float16, torch.bfloat16} + + +def _is_hopper(device: torch.device) -> bool: + try: + return torch.cuda.get_device_capability(device)[0] == 9 + except Exception: + return False + + +class _SM90LMHeadFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + hidden: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + output_fp32: bool, + ): + bias_to_save = ( + bias if bias is not None else torch.empty(0, device=hidden.device, dtype=hidden.dtype) + ) + ctx.save_for_backward(hidden, weight, bias_to_save) + ctx.has_bias = bias is not None + ctx.output_fp32 = bool(output_fp32) + weight_c = weight.contiguous() + if output_fp32: + return _C.lm_head_sm90_forward_fp32(hidden, weight_c, bias) + return _C.lm_head_sm90_forward(hidden, weight_c, bias) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + hidden, weight, bias = ctx.saved_tensors + grad_hidden = grad_weight = grad_bias = None + + hidden_2d = hidden.reshape(-1, hidden.size(-1)) + grad_2d = grad_output.reshape(-1, weight.size(0)).float() + hidden_f = hidden_2d.float() + weight_f = weight.float() + + if ctx.needs_input_grad[0]: + grad_hidden = grad_2d.matmul(weight_f).reshape_as(hidden).to(hidden.dtype) + if ctx.needs_input_grad[1]: + grad_weight = grad_2d.transpose(0, 1).matmul(hidden_f).to(weight.dtype) + if ctx.has_bias and ctx.needs_input_grad[2]: + grad_bias = grad_2d.sum(0).to(bias.dtype) + + return grad_hidden, grad_weight, grad_bias, None + + +class SM90LMHeadOp: + """Single-card SM90 batch-invariant CUDA LM-head op. + + The CUDA forward launches one CTA per output logit and performs the full K + reduction in that CTA. There is no Split-K and no cuBLAS algorithm selection + in the forward path, so a row's logits do not depend on batch layout. + """ + + op_class = "reduction" + is_batch_invariant = True + + def __init__(self) -> None: + if not _EXT_AVAILABLE or not hasattr(_C, "lm_head_sm90_forward"): + raise RuntimeError( + "lm_head_sm90_forward is not compiled into the extension. " + "Rebuild on Hopper with KERNEL_ALIGN_FORCE_SM90=1." + ) + self._fallback = NativeLMHeadOp() + logger.info("Successfully linked to precompiled _C.lm_head_sm90_forward kernel.") + + def __call__( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.forward(hidden, weight, bias=bias) + + def forward( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not self._can_use_sm90(hidden, weight, bias): + return self._fallback.forward(hidden, weight, bias=bias) + return _SM90LMHeadFunction.apply(hidden, weight, bias, False) + + def forward_fp32( + self, + hidden: torch.Tensor, + weight: torch.Tensor, + *, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not self._can_use_sm90(hidden, weight, bias): + return self._fallback.forward_fp32(hidden, weight, bias=bias) + return _SM90LMHeadFunction.apply(hidden, weight, bias, True) + + @staticmethod + def _can_use_sm90( + hidden: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + ) -> bool: + bias_ok = bias is None or ( + bias.is_cuda + and bias.device == hidden.device + and bias.dim() == 1 + and bias.dtype in _SUPPORTED_DTYPES + ) + return ( + hidden.is_cuda + and weight.is_cuda + and hidden.device == weight.device + and _is_hopper(hidden.device) + and hidden.dim() >= 2 + and weight.dim() == 2 + and hidden.size(-1) == weight.size(1) + and hidden.dtype in _SUPPORTED_DTYPES + and weight.dtype in _SUPPORTED_DTYPES + and bias_ok + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 7a5f6e79..fb2feb6f 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -103,6 +103,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp" + CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" + CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" def resolve_logp_op_type( @@ -344,6 +346,18 @@ def _adjust_priority_for_hardware(self): f"SM{cc}: fused linear-logp SM90 kernel not compiled into _C; " "using generic linear-logp backend." ) + + sm90_embedding_compiled = _EXT_AVAILABLE and hasattr(_C, "embedding_sm90_forward") + if sm90_embedding_compiled and cc_major == 9: + embedding_list = self._priority_map["cuda"]["embedding"] + if OpBackend.CUDA_SM90_EMBEDDING not in embedding_list: + embedding_list.insert(0, OpBackend.CUDA_SM90_EMBEDDING) + + sm90_lm_head_compiled = _EXT_AVAILABLE and hasattr(_C, "lm_head_sm90_forward") + if sm90_lm_head_compiled and cc_major == 9: + lm_head_list = self._priority_map["cuda"]["lm_head"] + if OpBackend.CUDA_SM90_LM_HEAD not in lm_head_list: + lm_head_list.insert(0, OpBackend.CUDA_SM90_LM_HEAD) except Exception as e: logger.warning(f"Failed to probe device capability: {e}") diff --git a/setup.py b/setup.py index a5679a85..288c23e6 100644 --- a/setup.py +++ b/setup.py @@ -152,6 +152,7 @@ def get_extensions(): "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + "csrc/cuda/embedding_lm_head_sm90.cu", # single-card batch-invariant embedding/lm-head ] enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] @@ -162,7 +163,7 @@ def get_extensions(): cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") extra_link_args.append("-lcuda") - # det_gemm SM90 (mma.sync + TMA) path — independent of the fused_logp + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" diff --git a/tests/test_embedding.py b/tests/test_embedding.py index b8069f05..151dc890 100644 --- a/tests/test_embedding.py +++ b/tests/test_embedding.py @@ -155,7 +155,7 @@ def test_embedding_gradient_flows_to_weight(): # Registry dispatch -- "embedding" resolves to NativeEmbeddingOp (matches the # PYTORCH_NATIVE_EMBEDDING entry + the per-platform priority-map additions). def test_registry_dispatches_native_embedding_op(): - assert isinstance(kernel_registry.get_op("embedding"), NativeEmbeddingOp) + assert isinstance(kernel_registry.get_op("embedding", device="cpu"), NativeEmbeddingOp) # --------------------------------------------------------------------------- # diff --git a/tests/test_kernel_registry.py b/tests/test_kernel_registry.py index 991688a8..66a09297 100644 --- a/tests/test_kernel_registry.py +++ b/tests/test_kernel_registry.py @@ -82,3 +82,40 @@ def fake_warning(message, *args): OpBackend.TRITON_GENERIC, ] assert any("Unknown RL_KERNEL_ROCM_ATTN_BACKEND=unknown" in warning for warning in warnings) + + +def test_sm90_linear_ops_prioritize_cuda_when_extension_symbols_exist(monkeypatch): + from rl_engine.kernels.ops import base as base_module + + class FakeExtension: + fused_linear_logp_sm90 = object() + embedding_sm90_forward = object() + lm_head_sm90_forward = object() + + monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda") + monkeypatch.setattr(registry_module.torch.cuda, "get_device_capability", lambda: (9, 0)) + monkeypatch.setattr(base_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(base_module, "_C", FakeExtension()) + + registry = KernelRegistry() + + assert registry._priority_map["cuda"]["embedding"][0] is OpBackend.CUDA_SM90_EMBEDDING + assert registry._priority_map["cuda"]["lm_head"][0] is OpBackend.CUDA_SM90_LM_HEAD + + +def test_sm90_linear_ops_do_not_prioritize_cuda_on_non_hopper(monkeypatch): + from rl_engine.kernels.ops import base as base_module + + class FakeExtension: + embedding_sm90_forward = object() + lm_head_sm90_forward = object() + + monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda") + monkeypatch.setattr(registry_module.torch.cuda, "get_device_capability", lambda: (8, 0)) + monkeypatch.setattr(base_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(base_module, "_C", FakeExtension()) + + registry = KernelRegistry() + + assert OpBackend.CUDA_SM90_EMBEDDING not in registry._priority_map["cuda"]["embedding"] + assert OpBackend.CUDA_SM90_LM_HEAD not in registry._priority_map["cuda"]["lm_head"] diff --git a/tests/test_lm_head.py b/tests/test_lm_head.py index 51f7c481..f0773cf0 100644 --- a/tests/test_lm_head.py +++ b/tests/test_lm_head.py @@ -313,7 +313,7 @@ def test_lm_head_gradient_flows(): # Registry dispatch -- "lm_head" resolves to NativeLMHeadOp. def test_registry_dispatches_native_lm_head_op(): """The registry resolves "lm_head" to NativeLMHeadOp.""" - assert isinstance(kernel_registry.get_op("lm_head"), NativeLMHeadOp) + assert isinstance(kernel_registry.get_op("lm_head", device="cpu"), NativeLMHeadOp) # --------------------------------------------------------------------------- # diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 0d851649..4f742734 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -35,6 +35,7 @@ def _args(**overrides): [ "rms_norm", "matmul", + "det_gemm", "attention", "logp", "linear_logp", diff --git a/tests/test_sm90_linear_wrappers.py b/tests/test_sm90_linear_wrappers.py new file mode 100644 index 00000000..fb72986e --- /dev/null +++ b/tests/test_sm90_linear_wrappers.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest +import torch + + +def _sm90_linear_available() -> bool: + if not torch.cuda.is_available(): + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + return ( + _EXT_AVAILABLE + and hasattr(_C, "embedding_sm90_forward") + and hasattr(_C, "lm_head_sm90_forward") + and torch.cuda.get_device_capability()[0] == 9 + ) + except Exception: + return False + + +requires_sm90_linear = pytest.mark.skipif( + not _sm90_linear_available(), + reason=( + "SM90 embedding/lm_head kernels require an H200/Hopper-class GPU and " + "KERNEL_ALIGN_FORCE_SM90=1." + ), +) + + +def test_sm90_embedding_wrapper_calls_extension_symbol(monkeypatch): + from rl_engine.kernels.ops.cuda.linear import embedding as embedding_module + + calls = [] + + class FakeExtension: + @staticmethod + def embedding_sm90_forward(token_ids, weight): + calls.append(("forward", token_ids, weight)) + return torch.empty((*token_ids.shape, weight.size(1)), dtype=weight.dtype) + + @staticmethod + def embedding_sm90_forward_fp32(token_ids, weight): + calls.append(("forward_fp32", token_ids, weight)) + return torch.empty((*token_ids.shape, weight.size(1)), dtype=torch.float32) + + monkeypatch.setattr(embedding_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(embedding_module, "_C", FakeExtension) + monkeypatch.setattr( + embedding_module.SM90EmbeddingOp, + "_can_use_sm90", + staticmethod(lambda token_ids, weight: True), + ) + + op = embedding_module.SM90EmbeddingOp() + token_ids = torch.tensor([[1, 2], [3, 4]]) + weight = torch.randn(8, 5) + + assert op.forward(token_ids, weight).shape == (2, 2, 5) + assert op.forward_fp32(token_ids, weight).dtype == torch.float32 + assert [name for name, *_ in calls] == ["forward", "forward_fp32"] + + +def test_sm90_lm_head_wrapper_calls_extension_symbol(monkeypatch): + from rl_engine.kernels.ops.cuda.linear import lm_head as lm_head_module + + calls = [] + + class FakeExtension: + @staticmethod + def lm_head_sm90_forward(hidden, weight, bias): + calls.append(("forward", hidden, weight, bias)) + return torch.empty((*hidden.shape[:-1], weight.size(0)), dtype=hidden.dtype) + + @staticmethod + def lm_head_sm90_forward_fp32(hidden, weight, bias): + calls.append(("forward_fp32", hidden, weight, bias)) + return torch.empty((*hidden.shape[:-1], weight.size(0)), dtype=torch.float32) + + monkeypatch.setattr(lm_head_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(lm_head_module, "_C", FakeExtension) + monkeypatch.setattr( + lm_head_module.SM90LMHeadOp, + "_can_use_sm90", + staticmethod(lambda hidden, weight, bias: True), + ) + + op = lm_head_module.SM90LMHeadOp() + hidden = torch.randn(2, 3, 7) + weight = torch.randn(11, 7) + bias = torch.randn(11) + + assert op.forward(hidden, weight, bias=bias).shape == (2, 3, 11) + assert op.forward_fp32(hidden, weight, bias=bias).dtype == torch.float32 + assert [name for name, *_ in calls] == ["forward", "forward_fp32"] + + +@requires_sm90_linear +def test_sm90_embedding_forward_matches_direct_gather_on_h200_hopper(): + from rl_engine.kernels.ops.cuda.linear.embedding import SM90EmbeddingOp + + device = torch.device("cuda") + token_ids = torch.tensor([[0, 7, 3], [5, 1, 7]], device=device) + weight = torch.randn(11, 13, device=device, dtype=torch.bfloat16) + + out = SM90EmbeddingOp().forward(token_ids, weight) + out_fp32 = SM90EmbeddingOp().forward_fp32(token_ids, weight) + + assert torch.equal(out, weight[token_ids]) + assert torch.equal(out_fp32, weight[token_ids].float()) + + +@requires_sm90_linear +def test_sm90_lm_head_forward_is_slice_batch_invariant_on_h200_hopper(): + from rl_engine.kernels.ops.cuda.linear.lm_head import SM90LMHeadOp + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(150) + hidden = torch.randn(4, 5, 64, device=device, dtype=torch.bfloat16, generator=generator) + weight = torch.randn(37, 64, device=device, dtype=torch.bfloat16, generator=generator) + bias = torch.randn(37, device=device, dtype=torch.bfloat16, generator=generator) + op = SM90LMHeadOp() + + full = op.forward(hidden, weight, bias=bias) + full_fp32 = op.forward_fp32(hidden, weight, bias=bias) + + assert torch.equal(op.forward(hidden[2:3], weight, bias=bias), full[2:3]) + assert torch.equal(op.forward_fp32(hidden[2:3], weight, bias=bias), full_fp32[2:3]) From 7e3aae64bc35c3cb1ed505b6f4a9dd6e2750f52e Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Wed, 29 Jul 2026 20:31:34 +0800 Subject: [PATCH 4/5] fix(ws1): make SM90 linear backward invariant Signed-off-by: inaniloquentee <3051000145@qq.com> --- csrc/cuda/embedding_lm_head_sm90.cu | 6 +- docs/operators/embedding.md | 24 ++-- docs/operators/lm_head.md | 12 +- .../kernels/ops/cuda/linear/embedding.py | 61 +++++++-- rl_engine/kernels/ops/cuda/linear/lm_head.py | 45 ++++++- tests/test_sm90_linear_wrappers.py | 117 +++++++++++++++++- 6 files changed, 229 insertions(+), 36 deletions(-) diff --git a/csrc/cuda/embedding_lm_head_sm90.cu b/csrc/cuda/embedding_lm_head_sm90.cu index f7289eb2..df758c3c 100644 --- a/csrc/cuda/embedding_lm_head_sm90.cu +++ b/csrc/cuda/embedding_lm_head_sm90.cu @@ -62,7 +62,7 @@ __global__ void embedding_sm90_forward_kernel(const int64_t *__restrict__ token_ const input_t *__restrict__ weight, output_t *__restrict__ output, int64_t num_tokens, int64_t hidden_size, - int64_t vocab_size) { + int64_t /*vocab_size*/) { const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; const int64_t total = num_tokens * hidden_size; if (idx >= total) { @@ -72,10 +72,6 @@ __global__ void embedding_sm90_forward_kernel(const int64_t *__restrict__ token_ const int64_t token_row = idx / hidden_size; const int64_t hidden_col = idx - token_row * hidden_size; const int64_t token_id = token_ids[token_row]; - if (token_id < 0 || token_id >= vocab_size) { - output[idx] = from_float(0.0f); - return; - } output[idx] = static_cast(weight[token_id * hidden_size + hidden_col]); } diff --git a/docs/operators/embedding.md b/docs/operators/embedding.md index 52783d9d..41818f56 100644 --- a/docs/operators/embedding.md +++ b/docs/operators/embedding.md @@ -32,7 +32,7 @@ The op exposes the WS1 dual-path contract: | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeEmbeddingOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA SM90 (H200/Hopper) | `SM90EmbeddingOp` | `_C.embedding_sm90_forward` | Single-card batch-invariant forward backend. | +| CUDA SM90 (H200/Hopper) | `SM90EmbeddingOp` | `_C.embedding_sm90_forward` | Single-card batch-invariant forward backend; deterministic duplicate-id backward in the wrapper. | | ROCm / Triton | N/A | N/A | Falls back to the PyTorch native reference. | ## Tensor Contract @@ -78,7 +78,10 @@ out = F.embedding(token_ids.long(), weight).to(torch.float32) The SM90 backend is a simple single-card forward gather for H200/Hopper builds. It is not a TP/vocab-parallel integration path; downstream fused kernels carry their own -benchmarks and are measured against the PyTorch reference for correctness. +benchmarks and are measured against the PyTorch reference for correctness. Its backward +path is intentionally conservative: token ids are sorted, duplicate ids are reduced in +a fixed order, and only unique rows are written back. That avoids CUDA atomic-add +nondeterminism for repeated token ids at the cost of throughput. ## Tests @@ -102,15 +105,8 @@ test at the real Qwen3-8B dims (`vocab=151936, hidden=4096`, boundary ids `0` an ## Known Limitations -- The CUDA SM90 path is single-card forward coverage only; TP/vocab-parallel and Triton - embedding backends are outside this PR. -- Out-of-range token ids are not validated; callers must keep ids in `[0, vocab)`. -- **GPU backward is bitwise-reproducible only under deterministic algorithms.** The - forward is a lossless gather (always reproducible), but `dL/dweight` is a scatter-add: - every repeated token id (padding, common tokens) accumulates into the same `weight.grad` - row. On CUDA that accumulation uses atomic adds, whose ordering is nondeterministic, so - the weight gradient is not bit-exact across runs when ids collide. PyTorch documents - `embedding` backward as a nondeterministic CUDA op for this reason. Since `forward_fp32` - is the backward golden source, callers that need a reproducible GPU gradient must enable - `torch.use_deterministic_algorithms(True)` (the gradient test does this). CPU backward is - always deterministic. +- The CUDA SM90 path is single-card coverage; TP/vocab-parallel and Triton embedding + backends are outside this PR. +- Token ids must be in `[0, vocab)`. The SM90 host path validates the range before the + kernel launch and raises instead of masking invalid ids. +- The deterministic SM90 backward is a reference path, not a tuned training kernel. diff --git a/docs/operators/lm_head.md b/docs/operators/lm_head.md index f0acbdd9..901d1759 100644 --- a/docs/operators/lm_head.md +++ b/docs/operators/lm_head.md @@ -36,7 +36,7 @@ The op exposes the WS1 dual-path contract: | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeLMHeadOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA SM90 (H200/Hopper) | `SM90LMHeadOp` | `_C.lm_head_sm90_forward` | Single-card batch-invariant forward backend; no Split-K. | +| CUDA SM90 (H200/Hopper) | `SM90LMHeadOp` | `_C.lm_head_sm90_forward` | Single-card batch-invariant forward backend; no Split-K; bf16 backward uses deterministic GEMM. | | ROCm / Triton | N/A | N/A | Falls back to the PyTorch native reference. | ## Tensor Contract @@ -84,6 +84,16 @@ CPU, ROCm, and CUDA devices without the SM90 extension, dispatch uses the PyTorc the CUDA SM90 single-card batch-invariant backend is prepended. Its forward path assigns one CTA to each output logit and performs the full K reduction without Split-K. +The one-CTA-per-logit design is a deliberate invariance tradeoff: CTAs re-read the +hidden row and vocab weight row independently, so large-vocab projections are expected +to be memory-bandwidth bound compared with a tiled GEMM. This path exists to preserve a +fixed hidden-dimension accumulation order for the WS1/H200 correctness gate. + +For bf16 H200 training, `SM90LMHeadOp.backward` routes `dhidden` through +`_C.det_gemm_da` and `dweight` through `_C.det_gemm_db` (`hidden.T @ dlogits`, transposed +back to the HF `[vocab, hidden]` layout). The wrapper fails fast if those deterministic +GEMM symbols are missing instead of silently falling back to cuBLAS for bf16 gradients. + ## Tests ```bash diff --git a/rl_engine/kernels/ops/cuda/linear/embedding.py b/rl_engine/kernels/ops/cuda/linear/embedding.py index 3992ffcd..2a865158 100644 --- a/rl_engine/kernels/ops/cuda/linear/embedding.py +++ b/rl_engine/kernels/ops/cuda/linear/embedding.py @@ -19,6 +19,52 @@ def _is_hopper(device: torch.device) -> bool: return False +def _deterministic_embedding_grad_weight( + ids: torch.Tensor, + grad_rows: torch.Tensor, + *, + weight_shape: tuple[int, ...], + weight_dtype: torch.dtype, +) -> torch.Tensor: + vocab_size = int(weight_shape[0]) + grad_weight = torch.zeros( + weight_shape, + device=grad_rows.device, + dtype=grad_rows.dtype, + ) + if ids.numel() == 0: + return grad_weight.to(weight_dtype) + + valid = (ids >= 0) & (ids < vocab_size) + if not bool(valid.all().item()): + ids = ids[valid] + grad_rows = grad_rows[valid] + if ids.numel() == 0: + return grad_weight.to(weight_dtype) + + order = torch.argsort(ids, stable=True) + sorted_ids = ids.index_select(0, order) + sorted_rows = grad_rows.index_select(0, order) + unique_ids, counts = torch.unique_consecutive(sorted_ids, return_counts=True) + + segment_sums = [] + start = 0 + for count in counts.tolist(): + end = start + int(count) + acc = torch.zeros( + sorted_rows.size(1), + device=sorted_rows.device, + dtype=sorted_rows.dtype, + ) + for row in range(start, end): + acc = acc + sorted_rows[row] + segment_sums.append(acc) + start = end + + grad_weight[unique_ids] = torch.stack(segment_sums, dim=0) + return grad_weight.to(weight_dtype) + + class _SM90EmbeddingFunction(torch.autograd.Function): @staticmethod def forward(ctx, token_ids: torch.Tensor, weight: torch.Tensor, output_fp32: bool): @@ -38,17 +84,12 @@ def backward(ctx, grad_output: torch.Tensor): ids = token_ids.reshape(-1).to(device=grad_output.device, dtype=torch.long) hidden_size = int(ctx.weight_shape[1]) grad_rows = grad_output.reshape(ids.numel(), hidden_size) - valid = (ids >= 0) & (ids < int(ctx.weight_shape[0])) - if not bool(valid.all().item()): - ids = ids[valid] - grad_rows = grad_rows[valid] - grad_weight = torch.zeros( - ctx.weight_shape, - device=grad_output.device, - dtype=grad_rows.dtype, + grad_weight = _deterministic_embedding_grad_weight( + ids, + grad_rows, + weight_shape=ctx.weight_shape, + weight_dtype=ctx.weight_dtype, ) - grad_weight.index_add_(0, ids, grad_rows) - grad_weight = grad_weight.to(ctx.weight_dtype) return None, grad_weight, None diff --git a/rl_engine/kernels/ops/cuda/linear/lm_head.py b/rl_engine/kernels/ops/cuda/linear/lm_head.py index f9f37cbb..de83600b 100644 --- a/rl_engine/kernels/ops/cuda/linear/lm_head.py +++ b/rl_engine/kernels/ops/cuda/linear/lm_head.py @@ -21,6 +21,16 @@ def _is_hopper(device: torch.device) -> bool: return False +def _can_use_det_gemm_backward(hidden: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + hidden.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and _EXT_AVAILABLE + and hasattr(_C, "det_gemm_da") + and hasattr(_C, "det_gemm_db") + ) + + class _SM90LMHeadFunction(torch.autograd.Function): @staticmethod def forward( @@ -50,11 +60,42 @@ def backward(ctx, grad_output: torch.Tensor): grad_2d = grad_output.reshape(-1, weight.size(0)).float() hidden_f = hidden_2d.float() weight_f = weight.float() + needs_projection_grad = ctx.needs_input_grad[0] or ctx.needs_input_grad[1] + use_det_gemm = ( + hidden_2d.size(0) > 0 + and needs_projection_grad + and _can_use_det_gemm_backward(hidden, weight) + ) + if ( + hidden_2d.size(0) > 0 + and needs_projection_grad + and hidden.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and not use_det_gemm + ): + raise RuntimeError( + "SM90LMHeadOp.backward requires _C.det_gemm_da/db for bf16 " + "batch-invariant gradients." + ) if ctx.needs_input_grad[0]: - grad_hidden = grad_2d.matmul(weight_f).reshape_as(hidden).to(hidden.dtype) + if use_det_gemm: + grad_hidden = _C.det_gemm_da( + grad_2d.to(torch.bfloat16).contiguous(), + weight.t().contiguous(), + ) + else: + grad_hidden = grad_2d.matmul(weight_f) + grad_hidden = grad_hidden.reshape_as(hidden).to(hidden.dtype) if ctx.needs_input_grad[1]: - grad_weight = grad_2d.transpose(0, 1).matmul(hidden_f).to(weight.dtype) + if use_det_gemm: + grad_weight = _C.det_gemm_db( + hidden_2d.contiguous(), + grad_2d.to(torch.bfloat16).contiguous(), + ).t() + else: + grad_weight = grad_2d.transpose(0, 1).matmul(hidden_f) + grad_weight = grad_weight.contiguous().to(weight.dtype) if ctx.has_bias and ctx.needs_input_grad[2]: grad_bias = grad_2d.sum(0).to(bias.dtype) diff --git a/tests/test_sm90_linear_wrappers.py b/tests/test_sm90_linear_wrappers.py index fb72986e..e04b70fb 100644 --- a/tests/test_sm90_linear_wrappers.py +++ b/tests/test_sm90_linear_wrappers.py @@ -39,12 +39,12 @@ class FakeExtension: @staticmethod def embedding_sm90_forward(token_ids, weight): calls.append(("forward", token_ids, weight)) - return torch.empty((*token_ids.shape, weight.size(1)), dtype=weight.dtype) + return weight[token_ids.long()] @staticmethod def embedding_sm90_forward_fp32(token_ids, weight): calls.append(("forward_fp32", token_ids, weight)) - return torch.empty((*token_ids.shape, weight.size(1)), dtype=torch.float32) + return weight[token_ids.long()].float() monkeypatch.setattr(embedding_module, "_EXT_AVAILABLE", True) monkeypatch.setattr(embedding_module, "_C", FakeExtension) @@ -63,6 +63,47 @@ def embedding_sm90_forward_fp32(token_ids, weight): assert [name for name, *_ in calls] == ["forward", "forward_fp32"] +def test_sm90_embedding_backward_accumulates_duplicate_ids_deterministically(monkeypatch): + from rl_engine.kernels.ops.cuda.linear import embedding as embedding_module + + class FakeExtension: + @staticmethod + def embedding_sm90_forward(token_ids, weight): + return weight[token_ids.long()] + + @staticmethod + def embedding_sm90_forward_fp32(token_ids, weight): + return weight[token_ids.long()].float() + + monkeypatch.setattr(embedding_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(embedding_module, "_C", FakeExtension) + monkeypatch.setattr( + embedding_module.SM90EmbeddingOp, + "_can_use_sm90", + staticmethod(lambda token_ids, weight: True), + ) + + token_ids = torch.tensor([[2, 1, 2], [3, 1, 2]]) + weight = torch.randn(5, 4, requires_grad=True) + upstream = torch.arange(token_ids.numel() * weight.size(1), dtype=torch.float32).reshape( + *token_ids.shape, weight.size(1) + ) + + prev = torch.are_deterministic_algorithms_enabled() + torch.use_deterministic_algorithms(True) + try: + (embedding_module.SM90EmbeddingOp().forward(token_ids, weight) * upstream).sum().backward() + finally: + torch.use_deterministic_algorithms(prev) + + expected = torch.zeros_like(weight) + flat_ids = token_ids.reshape(-1) + flat_upstream = upstream.reshape(-1, weight.size(1)) + for row, token_id in enumerate(flat_ids.tolist()): + expected[token_id] += flat_upstream[row] + assert torch.equal(weight.grad, expected) + + def test_sm90_lm_head_wrapper_calls_extension_symbol(monkeypatch): from rl_engine.kernels.ops.cuda.linear import lm_head as lm_head_module @@ -72,12 +113,18 @@ class FakeExtension: @staticmethod def lm_head_sm90_forward(hidden, weight, bias): calls.append(("forward", hidden, weight, bias)) - return torch.empty((*hidden.shape[:-1], weight.size(0)), dtype=hidden.dtype) + out = hidden.float().matmul(weight.float().t()) + if bias is not None: + out = out + bias.float() + return out.to(hidden.dtype) @staticmethod def lm_head_sm90_forward_fp32(hidden, weight, bias): calls.append(("forward_fp32", hidden, weight, bias)) - return torch.empty((*hidden.shape[:-1], weight.size(0)), dtype=torch.float32) + out = hidden.float().matmul(weight.float().t()) + if bias is not None: + out = out + bias.float() + return out.float() monkeypatch.setattr(lm_head_module, "_EXT_AVAILABLE", True) monkeypatch.setattr(lm_head_module, "_C", FakeExtension) @@ -97,6 +144,68 @@ def lm_head_sm90_forward_fp32(hidden, weight, bias): assert [name for name, *_ in calls] == ["forward", "forward_fp32"] +def test_sm90_lm_head_bf16_backward_routes_projection_grads_through_det_gemm(monkeypatch): + from rl_engine.kernels.ops.cuda.linear import lm_head as lm_head_module + + calls = [] + + class FakeExtension: + @staticmethod + def lm_head_sm90_forward(hidden, weight, bias): + out = hidden.float().matmul(weight.float().t()) + if bias is not None: + out = out + bias.float() + return out.to(hidden.dtype) + + @staticmethod + def lm_head_sm90_forward_fp32(hidden, weight, bias): + out = hidden.float().matmul(weight.float().t()) + if bias is not None: + out = out + bias.float() + return out.float() + + @staticmethod + def det_gemm_da(dc, b): + calls.append(("da", tuple(dc.shape), tuple(b.shape))) + return dc.float().matmul(b.float().t()).to(torch.bfloat16) + + @staticmethod + def det_gemm_db(a, dc): + calls.append(("db", tuple(a.shape), tuple(dc.shape))) + return a.float().t().matmul(dc.float()).to(torch.bfloat16) + + monkeypatch.setattr(lm_head_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(lm_head_module, "_C", FakeExtension) + monkeypatch.setattr( + lm_head_module.SM90LMHeadOp, + "_can_use_sm90", + staticmethod(lambda hidden, weight, bias: True), + ) + monkeypatch.setattr( + lm_head_module, + "_can_use_det_gemm_backward", + lambda hidden, weight: True, + ) + + hidden = torch.randn(2, 3, 5, dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn(7, 5, dtype=torch.bfloat16, requires_grad=True) + bias = torch.randn(7, dtype=torch.bfloat16, requires_grad=True) + dy = torch.randn(2, 3, 7, dtype=torch.bfloat16) + + out = lm_head_module.SM90LMHeadOp().forward(hidden, weight, bias=bias) + out.backward(dy) + + flat_hidden = hidden.detach().reshape(-1, hidden.size(-1)) + flat_dy = dy.reshape(-1, weight.size(0)) + expected_hidden = flat_dy.float().matmul(weight.detach().float()).to(torch.bfloat16) + expected_weight = flat_dy.float().t().matmul(flat_hidden.float()).to(torch.bfloat16) + + assert calls == [("da", (6, 7), (5, 7)), ("db", (6, 5), (6, 7))] + assert torch.equal(hidden.grad, expected_hidden.reshape_as(hidden)) + assert torch.equal(weight.grad, expected_weight) + assert torch.equal(bias.grad, flat_dy.float().sum(0).to(torch.bfloat16)) + + @requires_sm90_linear def test_sm90_embedding_forward_matches_direct_gather_on_h200_hopper(): from rl_engine.kernels.ops.cuda.linear.embedding import SM90EmbeddingOp From 42ae3cb845d778b191fb913ca34fecd9a034fd4d Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Wed, 29 Jul 2026 20:51:31 +0800 Subject: [PATCH 5/5] fix(ws1): address SM90 review comments Signed-off-by: inaniloquentee <3051000145@qq.com> --- csrc/cuda/embedding_lm_head_sm90.cu | 33 +++++++++++++-- .../kernels/ops/cuda/linear/embedding.py | 40 ++++++++++--------- tests/test_sm90_linear_wrappers.py | 2 +- 3 files changed, 51 insertions(+), 24 deletions(-) diff --git a/csrc/cuda/embedding_lm_head_sm90.cu b/csrc/cuda/embedding_lm_head_sm90.cu index df758c3c..ddf65f2c 100644 --- a/csrc/cuda/embedding_lm_head_sm90.cu +++ b/csrc/cuda/embedding_lm_head_sm90.cu @@ -12,13 +12,19 @@ #include #include #include +#include #include +#include #include #include namespace { constexpr int kThreads = 256; +constexpr int kMaxCachedDevices = 64; + +// Reference path: one CTA per logit is intentionally throughput-heavy for large +// vocab prefill. It exists to preserve a fixed K reduction order for WS1. template __device__ __forceinline__ float to_float(T value) { @@ -134,10 +140,29 @@ bool is_supported_float_dtype(at::ScalarType dtype) { void check_sm90_device() { int device = 0; C10_CUDA_CHECK(cudaGetDevice(&device)); - cudaDeviceProp prop{}; - C10_CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); - TORCH_CHECK(prop.major == 9, "SM90 embedding/lm_head kernels require Hopper, got sm_", - prop.major, prop.minor); + static std::array cached_major{}; + static std::array cached_minor{}; + static std::array cached_once{}; + + if (device >= 0 && device < kMaxCachedDevices) { + std::call_once(cached_once[device], [device]() { + C10_CUDA_CHECK(cudaDeviceGetAttribute(&cached_major[device], + cudaDevAttrComputeCapabilityMajor, device)); + C10_CUDA_CHECK(cudaDeviceGetAttribute(&cached_minor[device], + cudaDevAttrComputeCapabilityMinor, device)); + }); + TORCH_CHECK(cached_major[device] == 9, + "SM90 embedding/lm_head kernels require Hopper, got sm_", + cached_major[device], cached_minor[device]); + return; + } + + int major = 0; + int minor = 0; + C10_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + C10_CUDA_CHECK(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); + TORCH_CHECK(major == 9, "SM90 embedding/lm_head kernels require Hopper, got sm_", + major, minor); } void check_cuda_same_device(torch::Tensor a, torch::Tensor b, const char *a_name, diff --git a/rl_engine/kernels/ops/cuda/linear/embedding.py b/rl_engine/kernels/ops/cuda/linear/embedding.py index 2a865158..979a63de 100644 --- a/rl_engine/kernels/ops/cuda/linear/embedding.py +++ b/rl_engine/kernels/ops/cuda/linear/embedding.py @@ -36,32 +36,34 @@ def _deterministic_embedding_grad_weight( return grad_weight.to(weight_dtype) valid = (ids >= 0) & (ids < vocab_size) - if not bool(valid.all().item()): - ids = ids[valid] - grad_rows = grad_rows[valid] - if ids.numel() == 0: - return grad_weight.to(weight_dtype) + ids = ids[valid] + grad_rows = grad_rows[valid] + if ids.numel() == 0: + return grad_weight.to(weight_dtype) order = torch.argsort(ids, stable=True) sorted_ids = ids.index_select(0, order) sorted_rows = grad_rows.index_select(0, order) unique_ids, counts = torch.unique_consecutive(sorted_ids, return_counts=True) - segment_sums = [] - start = 0 - for count in counts.tolist(): - end = start + int(count) - acc = torch.zeros( - sorted_rows.size(1), - device=sorted_rows.device, - dtype=sorted_rows.dtype, - ) - for row in range(start, end): - acc = acc + sorted_rows[row] - segment_sums.append(acc) - start = end + num_unique = int(unique_ids.numel()) + max_count = int(counts.max().item()) + hidden_size = sorted_rows.size(1) + starts = torch.cumsum(counts, dim=0) - counts + slot = torch.arange(sorted_rows.size(0), device=sorted_rows.device) + slot = slot - starts.repeat_interleave(counts) + segment = torch.arange(num_unique, device=sorted_rows.device).repeat_interleave(counts) + padded = torch.zeros( + (num_unique, max_count, hidden_size), + device=sorted_rows.device, + dtype=sorted_rows.dtype, + ) + padded[segment, slot] = sorted_rows + acc = padded[:, 0] + for slot_idx in range(1, max_count): + acc = acc + padded[:, slot_idx] - grad_weight[unique_ids] = torch.stack(segment_sums, dim=0) + grad_weight[unique_ids] = acc return grad_weight.to(weight_dtype) diff --git a/tests/test_sm90_linear_wrappers.py b/tests/test_sm90_linear_wrappers.py index e04b70fb..3ee663d7 100644 --- a/tests/test_sm90_linear_wrappers.py +++ b/tests/test_sm90_linear_wrappers.py @@ -198,7 +198,7 @@ def det_gemm_db(a, dc): flat_hidden = hidden.detach().reshape(-1, hidden.size(-1)) flat_dy = dy.reshape(-1, weight.size(0)) expected_hidden = flat_dy.float().matmul(weight.detach().float()).to(torch.bfloat16) - expected_weight = flat_dy.float().t().matmul(flat_hidden.float()).to(torch.bfloat16) + expected_weight = flat_hidden.float().t().matmul(flat_dy.float()).to(torch.bfloat16).t() assert calls == [("da", (6, 7), (5, 7)), ("db", (6, 5), (6, 7))] assert torch.equal(hidden.grad, expected_hidden.reshape_as(hidden))