MCPcopy Create free account

hub / github.com/RL-Align/RL-Kernel / functions

Functions1,290 in github.com/RL-Align/RL-Kernel

⨍ Functions 1,290◇ Types & classes 193↳ Endpoints 3
↓ 137 callersMethodget
(self, ref)
tests/test_ray_actor_manager.py:88
↓ 82 callersMethoddetach
(self)
tests/test_weight_sync_bridge.py:65
↓ 61 callersMethodcontiguous
(self)
tests/test_weight_sync_bridge.py:68
↓ 49 callersMethodbackward
(ctx, grad_output)
tests/test_op_checks.py:165
↓ 43 callersFunction_model
()
tests/test_weight_sync_bridge.py:33
↓ 41 callersMethodnumel
(self)
tests/test_weight_sync_bridge.py:71
↓ 30 callersMethodpublish
( self, model: torch.nn.Module, *, weight_version: int, metadata: Opti
rl_engine/executors/bridge.py:1527
↓ 29 callersMethodget_op
Core distribution logic: Automatically select the best operator based on hardware and priority.
rl_engine/kernels/registry.py:224
↓ 26 callersFunction_rand
(shape, *, seed, dtype=torch.float32)
tests/test_rms_norm.py:17
↓ 26 callersMethodforward_fp32
Same as apply but forces float32 output for numerical stability.
rl_engine/kernels/ops/pytorch/loss/logp.py:56
↓ 26 callersFunctionselected_logprobs_reference
Reference selected-token logprobs for RL kernel validation.
rl_engine/testing/reference_ops.py:15
↓ 25 callersMethodimport_update
(self, manifest: WeightUpdateManifest)
rl_engine/executors/bridge.py:1573
↓ 23 callersFunction_rand
(shape, *, seed, dtype=torch.float32)
tests/test_swiglu.py:15
↓ 23 callersMethodforward_fp32
fp32 gold standard: internal computation and output are fp32.
rl_engine/kernels/ops/pytorch/rotary_embedding/rope.py:36
↓ 22 callersFunction_floating_tensor
( shape: tuple[int, ...], args: argparse.Namespace, dtype: torch.dtype, device: torch.device,
rl_engine/kernels/gtest/operator_inputs.py:197
↓ 21 callersFunction_kv
K/V of KV-head count, length s.
tests/test_kv_cache_attention.py:91
↓ 19 callersMethodstate_dict
( self, *, destination: _TDestination, prefix: str = "", keep_vars: bo
rl_engine/executors/deepspeed_trainer.py:650
↓ 18 callersMethodforward_fp32
Ground truth: upcast to fp32, accumulate in fp32, force fp32 output. The whole score->softmax->value path is wrapped to disable autocast and
rl_engine/kernels/ops/pytorch/attention/standard_attn.py:84
↓ 18 callersMethodforward_fp32
Ground-truth: accumulate in fp32 and force fp32 output.
rl_engine/kernels/ops/pytorch/norm/rms_norm.py:40
↓ 18 callersMethodrelease
(self, update_id: str)
rl_engine/executors/bridge.py:1869
↓ 17 callersMethod_check
(self, code: int, operation: str)
rl_engine/executors/bridge.py:1169
↓ 16 callersFunction_make_inputs
Deterministic RoPE inputs.
tests/test_rope.py:27
↓ 16 callersFunctionmake_synthetic_rl_kernel_batch
Create deterministic RL-shaped tensors for kernel tests.
rl_engine/testing/rl_batch.py:87
↓ 15 callersMethodimport_update
(self, manifest: WeightUpdateManifest)
rl_engine/executors/bridge.py:1852
↓ 14 callersFunction_batch
(seed=0, *, device="cpu", valid_density=0.9)
tests/test_grpo_loss.py:36
↓ 14 callersMethodforward_fp32
fp32 gold standard: cast inputs to fp32, then call `torch.matmul` once.
rl_engine/kernels/ops/pytorch/linear/matmul.py:26
↓ 14 callersMethodpublish
( self, model: torch.nn.Module, *, weight_version: int, metadata: Opti
rl_engine/executors/bridge.py:1713
↓ 13 callersMethodforward
(ctx, x)
tests/test_op_checks.py:161
↓ 13 callersMethodforward_fp32
Ground truth: concat cache+new, then attend in strict fp32. Delegates to ``NativeAttentionOp.forward_fp32`` so the fp32 golden path is
rl_engine/kernels/ops/pytorch/attention/kv_cache.py:100
↓ 13 callersMethodscore
(self, inputs: StatelessForwardInputs)
rl_engine/executors/stateless_executor.py:113
↓ 12 callersFunction_arg_int
(args: argparse.Namespace, name: str, default: int)
rl_engine/kernels/gtest/operator_inputs.py:255
↓ 12 callersFunction_batch_seq
(args: argparse.Namespace)
rl_engine/kernels/gtest/operator_inputs.py:235
↓ 12 callersFunction_make_inputs
( batch: int, seq: int, k: int, n: int, *, dtype: torch.dtype = torch.float32, see
tests/test_matmul.py:29
↓ 12 callersFunction_q
(batch, sq, *, seed, dtype=torch.float32)
tests/test_kv_cache_attention.py:86
↓ 12 callersMethodforward
Canonical entry: attend in the input dtype, output the input dtype. This is the dtype-behavior path used as the Axis-B accuracy candi
rl_engine/kernels/ops/pytorch/attention/standard_attn.py:59
↓ 11 callersFunction_environment
()
benchmarks/benchmark_weight_sync_bridge.py:57
↓ 11 callersFunction_inputs
()
tests/test_stateless_executor.py:80
↓ 11 callersFunction_qkv
Fixed-seed random q [B,Hq,Sq,D], k/v [B,Hkv,Skv,D] for reproducibility.
tests/test_attention.py:80
↓ 11 callersFunction_timestamp
()
benchmarks/benchmark_weight_sync_bridge.py:53
↓ 10 callersFunction_inputs
A full ratio/KL input set: (policy_logits, ref_logits, action_ids, mask, old_logps).
tests/test_ratio_kl.py:48
↓ 10 callersFunction_logit_pair
(policy_logits, ref_logits) for a batch.
tests/test_grpo_loss.py:54
↓ 10 callersFunction_make_inputs
( batch: int, seq: int, vocab: int, *, dtype: torch.dtype = torch.float32, seed: int =
tests/test_logp.py:14
↓ 10 callersFunction_rand_hidden
Fixed-seed random hidden-state tensor for reproducibility.
tests/test_lm_head.py:84
↓ 10 callersMethod_row
(key: str, value: Any)
scripts/run_perf.py:42
↓ 10 callersMethodforward_fp32
Ground-truth: native-dtype gather, then upcast the result to fp32.
rl_engine/kernels/ops/pytorch/linear/embedding.py:31
↓ 10 callersMethodfrom_tensor
(cls, name: str, tensor: torch.Tensor)
rl_engine/executors/bridge.py:1373
↓ 10 callersMethodgroup_advantages
Normalize raw per-sequence rewards within each generation group.
rl_engine/kernels/ops/pytorch/loss/grpo_loss.py:53
↓ 10 callersMethodupdate_status
(self, update_id: str)
rl_engine/executors/bridge.py:1625
↓ 9 callersMethod_fallback_op
(self)
rl_engine/kernels/ops/cuda/loss/logp.py:28
↓ 9 callersFunction_inputs
(seed, *, device, dtype=torch.float32, bias=True, lead=None)
tests/test_linear_logp.py:152
↓ 9 callersFunction_install_fake_deepspeed
(monkeypatch)
tests/test_deepspeed_training_worker.py:84
↓ 9 callersMethod_prepare_inputs
( self, logits: torch.Tensor, token_ids: torch.Tensor, )
rl_engine/kernels/ops/cuda/loss/logp.py:105
↓ 9 callersFunction_rand_weight
Fixed-seed random lm_head weight tensor for reproducibility.
tests/test_lm_head.py:90
↓ 9 callersMethodacknowledge
(self, update_id: str)
rl_engine/executors/bridge.py:1423
↓ 9 callersMethodforward_fp32
Ground-truth: compute in fp32 and force fp32 output.
rl_engine/kernels/ops/pytorch/activation/swiglu.py:30
↓ 9 callersMethodforward_fp32
Ground-truth: compute in fp32 and force fp32 output.
rl_engine/kernels/ops/pytorch/activation/swiglu.py:65
↓ 9 callersMethodremote
(self, actor_class)
tests/test_ray_actor_manager.py:85
↓ 9 callersFunctionrun_operator_suite
Run candidates against gold outputs and return a structured report.
rl_engine/kernels/gtest/op_checks.py:93
↓ 9 callersMethodtrain
(self, rollout: RolloutStageResult)
rl_engine/executors/deepspeed_trainer.py:128
↓ 8 callersFunction_fused_logp_op
(op_type: str = "logp")
tests/test_op_accuracy.py:13
↓ 8 callersFunction_rand_ids
(shape, *, seed, vocab=_VOCAB)
tests/test_embedding.py:35
↓ 8 callersFunction_rand_weight
(vocab=_VOCAB, hidden=_HIDDEN, *, seed, dtype=torch.float32)
tests/test_embedding.py:30
↓ 8 callersFunction_rollout
(iteration=2, weight_version=9)
tests/test_deepspeed_training_worker.py:102
↓ 8 callersFunction_safe_getattr
(value: Any, attr: str)
rl_engine/executors/vllm_sampler.py:279
↓ 8 callersMethodforward
Canonical entry: project in the input dtype, output the input dtype. This is the dtype-behavior path used as the Axis-B accuracy cand
rl_engine/kernels/ops/pytorch/linear/lm_head.py:47
↓ 8 callersMethodforward_fp32
Ground truth: upcast to fp32, accumulate in fp32, force fp32 output. The matmul is wrapped to disable autocast and TF32 so this stays a true
rl_engine/kernels/ops/pytorch/linear/lm_head.py:62
↓ 8 callersMethodrelease
(self, update_id: str)
rl_engine/executors/bridge.py:1427
↓ 8 callersFunctionsummarize_kernel_drift
Summarize candidate-vs-reference drift for benchmark/test output.
rl_engine/testing/reference_ops.py:109
↓ 7 callersFunction_cuda_define_from_env
(name: str, macro: str)
setup.py:24
↓ 7 callersFunction_logp_case
(name: str, dtype: torch.dtype, *, seed: int = 0)
tests/test_op_checks.py:12
↓ 7 callersFunction_reference_selected_logp
(logits: torch.Tensor, token_ids: torch.Tensor)
tests/test_op_accuracy.py:38
↓ 7 callersFunction_require_distributed_initialized
()
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:41
↓ 7 callersFunction_single_thread
Pin CPU GEMM to one thread so the matmul reduction order is stable.
tests/test_kv_cache_attention.py:76
↓ 7 callersFunction_sm90_inputs
(seed, *, bias=True, dtype=torch.bfloat16, lead=None)
tests/test_linear_logp.py:136
↓ 7 callersFunctioncompute_policy_ratio
Compute exp(current - old) with masked entries set to zero.
rl_engine/testing/reference_ops.py:82
↓ 7 callersFunctioncompute_reference_kl
Compute the common GRPO/PPO reference KL approximation.
rl_engine/testing/reference_ops.py:95
↓ 7 callersMethodforward
( self, policy_logits: torch.Tensor, ref_logits: torch.Tensor, action_ids: tor
rl_engine/kernels/ops/triton/loss/grpo_loss.py:216
↓ 7 callersMethodforward
( self, policy_logits: torch.Tensor, ref_logits: torch.Tensor, action_ids: tor
rl_engine/kernels/ops/pytorch/loss/grpo_loss.py:127
↓ 7 callersMethodinfo_once
(self, msg: str, *args: Any)
rl_engine/utils/logger.py:47
↓ 7 callersFunctionmake_weight_bridge
( transport: str = "local-clone", *, source_worker: str = "local-training", source_rank: int =
rl_engine/executors/bridge.py:2631
↓ 7 callersFunctionmasked_mean
Mean values while ignoring masked-out entries.
rl_engine/testing/reference_ops.py:71
↓ 7 callersMethodrelease
(self, update_id: str)
rl_engine/executors/bridge.py:1616
↓ 7 callersMethodrelease
(self, update_id: str)
rl_engine/executors/bridge.py:2532
↓ 7 callersMethodto_dict
(self)
benchmarks/profiler.py:75
↓ 7 callersMethodupdate_weights
Import a complete published weight manifest and switch active version. The underlying bridge owns transport-specific semantics. Same
rl_engine/executors/rollout.py:50
↓ 6 callersFunction_bool_mask
(mask: torch.Tensor, *, device: torch.device)
rl_engine/testing/reference_ops.py:11
↓ 6 callersFunction_model
(hidden_dim: int, layers: int)
benchmarks/benchmark_weight_sync_bridge.py:80
↓ 6 callersFunction_time_ms
(fn, warmup, iters)
benchmarks/benchmark_linear_logp.py:62
↓ 6 callersMethodapply_fp32
Backward-compatible alias for forward_fp32.
rl_engine/kernels/ops/pytorch/loss/logp.py:64
↓ 6 callersFunctionextract_logits
Extract logits from common model output shapes.
rl_engine/alignment/model_wrappers.py:40
↓ 6 callersMethodforward
Baseline selected-token log probability extraction using torch.gather.
rl_engine/kernels/ops/pytorch/loss/logp.py:52
↓ 6 callersMethodfrom_metadata
(cls, metadata: Optional[Mapping[str, Any]])
rl_engine/executors/bridge.py:53
↓ 6 callersFunctionfwd
(op, h=hidden, w=weight, t=target)
benchmarks/benchmark_linear_logp.py:108
↓ 6 callersMethodinstall
( self, manifest: WeightUpdateManifest, tensors: Mapping[str, torch.Tensor], )
rl_engine/executors/bridge.py:136
↓ 6 callersFunctionmake_operator_inputs
( op_name: str, args: argparse.Namespace, dtype: torch.dtype, device: torch.device, )
rl_engine/kernels/gtest/operator_inputs.py:22
↓ 6 callersFunctionrun_automated_suite
(args: argparse.Namespace)
benchmarks/profiler.py:741
↓ 6 callersMethodvalidate_supported
(self)
rl_engine/executors/bridge.py:84
↓ 5 callersFunction_bool_mask
(mask: torch.Tensor, *, device: torch.device)
rl_engine/executors/stateless_executor.py:565
↓ 5 callersFunction_extract_hidden_states
( model_output: Any, *, expected_hidden_dim: Optional[int] = None, )
rl_engine/executors/deepspeed_trainer.py:424
↓ 5 callersFunction_fmt
(value: Any, spec: str)
benchmarks/benchmark_ratio_kl.py:203
next →1–100 of 1,290, ranked by callers