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
↓ 5 callersFunction_inputs
(device: torch.device | str = "cpu")
tests/test_stateless_hf_integration.py:24
↓ 5 callersFunction_logits_for
(inputs: StatelessForwardInputs, vocab_size: int = 8)
tests/test_stateless_executor.py:108
↓ 5 callersFunction_minimal_rl_loss
( current_logps, old_logps, ref_logps, advantages, completion_mask, eps=0.2, beta=
tests/test_rl_kernel_loss_step.py:18
↓ 5 callersFunction_print_rank0
(rank: int, message: str)
tests/linear_logp_tp.py:92
↓ 5 callersFunction_ref_softmax_attn
Independent naive-softmax reference mirroring the contract (GQA, masks). Dtype-preserving: pass fp32 tensors for a bitwise fp32 check, or .double
tests/test_attention.py:89
↓ 5 callersMethod_require_record
(self, update_id: str)
rl_engine/executors/bridge.py:1648
↓ 5 callersMethod_sync
(self)
benchmarks/profiler.py:173
↓ 5 callersFunction_tensor_nbytes
(tensor: torch.Tensor)
rl_engine/executors/paged_kv_baseline.py:412
↓ 5 callersMethod_validate_manifest
( self, record: _LocalUpdateRecord, manifest: WeightUpdateManifest, )
rl_engine/executors/bridge.py:1654
↓ 5 callersMethodacknowledge
(self, update_id: str)
rl_engine/executors/bridge.py:1592
↓ 5 callersFunctionactive_token_count
Return the number of active tokens as an fp32 scalar tensor.
rl_engine/testing/reference_ops.py:59
↓ 5 callersFunctionavailable_workloads
()
benchmarks/profiler.py:728
↓ 5 callersMethodensure_runtime
(self)
rl_engine/executors/ray_actor_manager.py:95
↓ 5 callersMethodindexed_out
( self, logits: torch.Tensor, token_ids: torch.Tensor, row_indices: torch.Tens
rl_engine/kernels/ops/pytorch/loss/logp.py:68
↓ 5 callersMethodis_initialized
(self)
tests/test_ray_actor_manager.py:78
↓ 5 callersMethodpublish
( self, model: torch.nn.Module, *, weight_version: int, metadata: Opti
rl_engine/executors/bridge.py:1409
↓ 5 callersMethodpublish
( self, model: torch.nn.Module, *, weight_version: int, metadata: Opti
rl_engine/executors/bridge.py:2023
↓ 5 callersMethodpublish
( self, model: torch.nn.Module, *, weight_version: int, metadata: Opti
rl_engine/executors/bridge.py:2397
↓ 5 callersMethodpublish_weights
Publish the current DeepSpeed model state as a complete weight manifest. ZeRO-3 partitions parameters across ranks, so publication f
rl_engine/executors/deepspeed_trainer.py:204
↓ 5 callersMethodrelease
(self, update_id: str)
rl_engine/executors/bridge.py:2143
↓ 5 callersMethodrelease_weights
Release the active weight update held by this rollout worker.
rl_engine/executors/rollout.py:85
↓ 5 callersFunctionscore_reference_logprobs
Compute causal next-token selected logprobs aligned to ``[B, S]`` masks.
rl_engine/executors/stateless_executor.py:212
↓ 5 callersFunctionscore_rewards
Run a reward adapter and validate that it returns one scalar per sequence.
rl_engine/executors/stateless_executor.py:256
↓ 5 callersMethodselected_logprobs
( self, input_ids: torch.Tensor, token_ids: torch.Tensor, *, mask: Opt
rl_engine/alignment/model_wrappers.py:90
↓ 5 callersMethodzero_grad
(self, *args, **kwargs)
tests/test_deepspeed_training_worker.py:35
↓ 4 callersFunction_args
(**overrides)
tests/test_operator_inputs.py:14
↓ 4 callersMethod_batch_from_rollout_or_synthetic
( self, rollout: RolloutStageResult, )
rl_engine/executors/training_contract.py:79
↓ 4 callersFunction_bool_mask
(mask: torch.Tensor, *, device: torch.device)
rl_engine/executors/paged_kv_baseline.py:394
↓ 4 callersFunction_candidate_token_ids
(candidate: Any)
rl_engine/executors/training_contract.py:783
↓ 4 callersFunction_check_metric
( *, name: str, actual: torch.Tensor, expected: torch.Tensor, atol: float, rtol: float
tests/linear_logp_tp.py:210
↓ 4 callersFunction_copy_int_list
(value: Any)
rl_engine/executors/training_contract.py:1004
↓ 4 callersFunction_copy_int_list
(value: Any)
rl_engine/executors/vllm_sampler.py:291
↓ 4 callersMethod_engine_world_size
(self)
rl_engine/executors/deepspeed_trainer.py:288
↓ 4 callersFunction_fail
(message: str)
scripts/check_rocm_env.py:10
↓ 4 callersFunction_flatten_tensors
(value: Any)
rl_engine/kernels/gtest/op_checks.py:392
↓ 4 callersFunction_generator
(device: torch.device, seed: int)
tests/linear_logp_tp.py:97
↓ 4 callersMethod_get_backend
(self)
rl_engine/executors/bridge.py:2189
↓ 4 callersFunction_linear_logits
( hidden_2d: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], *, use_fp32
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:25
↓ 4 callersFunction_normalized_dim
(args: argparse.Namespace)
rl_engine/kernels/gtest/operator_inputs.py:239
↓ 4 callersMethod_prepare_indices
(self, row_indices: torch.Tensor, logits: torch.Tensor)
rl_engine/kernels/ops/cuda/loss/logp.py:123
↓ 4 callersMethod_prepare_output
(self, output: torch.Tensor, orig_shape: torch.Size)
rl_engine/kernels/ops/cuda/loss/logp.py:115
↓ 4 callersFunction_reduce_max
(value: float, device: torch.device)
tests/linear_logp_tp.py:173
↓ 4 callersFunction_safe_token_ids
(token_ids: torch.Tensor, mask: Optional[torch.Tensor] = None)
rl_engine/executors/deepspeed_trainer.py:559
↓ 4 callersMethod_selected_logps
( self, logits: torch.Tensor, token_ids: torch.Tensor, *, output_dtype
rl_engine/kernels/ops/pytorch/loss/logp.py:20
↓ 4 callersFunction_single_thread
Pin CPU GEMM to one thread so the matmul reduction order is batch-independent.
tests/test_attention.py:69
↓ 4 callersFunction_sync
(device: torch.device)
benchmarks/benchmark_rl_kernels.py:91
↓ 4 callersFunction_sync
(device: torch.device)
benchmarks/benchmark_ratio_kl.py:86
↓ 4 callersFunction_time_ms
(fn, warmup, iters)
benchmarks/benchmark_grpo_loss.py:46
↓ 4 callersFunction_use_fp32_matmul
(*tensors: torch.Tensor)
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:17
↓ 4 callersFunction_validated_manifest_metadata
(metadata: Optional[Mapping[str, Any]])
rl_engine/executors/bridge.py:1351
↓ 4 callersMethodapply
Evaluate the loss from logits + per-sequence advantages.
rl_engine/kernels/ops/triton/loss/grpo_loss.py:185
↓ 4 callersFunctionbuild_stateless_inputs_from_rollout_payload
Build dense no-cache scoring inputs from grouped rollout token payloads.
rl_engine/executors/training_contract.py:377
↓ 4 callersMethodforward
Canonical entry: concat cache+new, then attend in the input dtype. Delegates to ``NativeAttentionOp.forward`` (the Axis-B dtype path)
rl_engine/kernels/ops/pytorch/attention/kv_cache.py:78
↓ 4 callersMethodforward
Canonical entry: accumulate in fp32, cast the result back to x.dtype. This is the dtype-behavior path used as the Axis-B accuracy can
rl_engine/kernels/ops/pytorch/norm/rms_norm.py:27
↓ 4 callersMethodforward
(self, input_ids, attention_mask=None, use_cache=None)
benchmarks/benchmark_stateless_executor.py:108
↓ 4 callersFunctiongenerate
(llm: Any)
benchmarks/benchmark_weight_sync_bridge.py:672
↓ 4 callersMethodgenerate
Generate grouped candidates while keeping each prompt prefix byte-identical.
rl_engine/executors/vllm_sampler.py:99
↓ 4 callersMethodimport_update
(self, manifest: WeightUpdateManifest)
rl_engine/executors/bridge.py:2482
↓ 4 callersMethodinstall
( self, manifest: WeightUpdateManifest, tensors: Mapping[str, torch.Tensor], )
rl_engine/executors/bridge.py:356
↓ 4 callersMethodinstall
( self, manifest: WeightUpdateManifest, tensors: Mapping[str, torch.Tensor], )
rl_engine/executors/bridge.py:647
↓ 4 callersFunctionload_contract
Load the dtype/operator-class tolerance contract.
rl_engine/kernels/gtest/tolerance.py:14
↓ 4 callersFunctionmake_qkv
(batch, seqlen, nheads, headdim, device, dtype, nheads_k=None)
tests/test_attention_correctness.py:86
↓ 4 callersMethodrelease
(self, update_id: str)
rl_engine/executors/bridge.py:277
↓ 4 callersMethodrelease
(self, update_id: str)
rl_engine/executors/bridge.py:685
↓ 4 callersFunctionrun
(op, h, w, b)
tests/test_linear_logp.py:457
↓ 4 callersMethodsample
Unified sampling interface.
rl_engine/kernels/sampling.py:42
↓ 4 callersMethodsave_csv
(self, path: Path | str)
benchmarks/profiler.py:404
↓ 4 callersMethodsave_json
(self, path: Path | str)
benchmarks/profiler.py:391
↓ 4 callersMethodshutdown
(self)
rl_engine/executors/ray_actor_manager.py:128
↓ 3 callersMethod__init__
( self, *, address: int, shape: tuple[int, ...], stride: tuple[int, ..
rl_engine/executors/bridge.py:942
↓ 3 callersMethod__init__
(self)
tests/test_stateless_executor.py:221
↓ 3 callersFunction_arg_float
(args: argparse.Namespace, name: str, default: float)
rl_engine/kernels/gtest/operator_inputs.py:251
↓ 3 callersFunction_detach_optional
(tensor: Optional[torch.Tensor])
rl_engine/executors/stateless_executor.py:579
↓ 3 callersFunction_detach_optional
(tensor: Optional[torch.Tensor])
rl_engine/executors/paged_kv_baseline.py:408
↓ 3 callersFunction_environment
()
benchmarks/benchmark_stateless_executor.py:141
↓ 3 callersFunction_extract_logps
( model_output: Any, model: torch.nn.Module, token_ids: torch.Tensor, completion_mask: Optiona
rl_engine/executors/deepspeed_trainer.py:618
↓ 3 callersFunction_fake_cuda_tensors_for_manifest
(manifest)
tests/test_weight_sync_bridge.py:78
↓ 3 callersFunction_inputs
()
tests/test_paged_kv_baseline.py:39
↓ 3 callersFunction_install_fake_deepspeed_with_gather
(monkeypatch)
tests/test_deepspeed_training_worker.py:90
↓ 3 callersFunction_looks_like_hidden_tensor
( tensor: torch.Tensor, expected_hidden_dim: Optional[int], )
rl_engine/executors/deepspeed_trainer.py:533
↓ 3 callersFunction_make_loss_inputs
(device="cpu", dtype=torch.float32)
tests/test_rl_kernel_loss_step.py:42
↓ 3 callersFunction_matmul_operand
(tensor: torch.Tensor, use_fp32: bool)
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:21
↓ 3 callersMethod_measure_peak_memory_once
Measure peak allocation for one invocation after the configured warmup.
benchmarks/profiler.py:216
↓ 3 callersFunction_parse_int_list
(value: str)
benchmarks/benchmark_rl_kernels.py:72
↓ 3 callersFunction_parse_int_list
(value: str)
benchmarks/profiler.py:484
↓ 3 callersFunction_parse_int_list
(value: str)
benchmarks/benchmark_ratio_kl.py:67
↓ 3 callersMethod_peak_memory_gb
(self)
benchmarks/profiler.py:182
↓ 3 callersFunction_peak_vram_gb
(fn, warmup=3, iters=5)
benchmarks/benchmark_linear_logp.py:76
↓ 3 callersFunction_recover_canonical_rows
(layout_values, order)
tests/test_linear_logp.py:182
↓ 3 callersFunction_reduce_min_int
(value: bool, device: torch.device)
tests/linear_logp_tp.py:179
↓ 3 callersFunction_rollout
(payload)
tests/test_stateless_training_contract.py:28
↓ 3 callersFunction_scoring_mode
(args: argparse.Namespace)
benchmarks/benchmark_stateless_executor.py:381
↓ 3 callersMethod_time_kernel
Run fn with warmup + repeat, return (last_result, median_ms, std_ms).
benchmarks/profiler.py:187
↓ 3 callersFunction_timestamp
()
benchmarks/benchmark_stateless_executor.py:137
↓ 3 callersFunction_tiny_past_key_values
(hidden: torch.Tensor, use_cache: bool | None)
benchmarks/benchmark_stateless_executor.py:120
↓ 3 callersFunction_token_ids
( shape: tuple[int, ...], vocab: int, args: argparse.Namespace, device: torch.device, )
rl_engine/kernels/gtest/operator_inputs.py:215
↓ 3 callersFunction_unix_stream_socket
()
rl_engine/executors/bridge.py:860
↓ 3 callersMethodacknowledge
(self, update_id)
tests/test_weight_sync_bridge.py:891
↓ 3 callersFunctionassert_flash_attention_matches_sdpa
( op, dtype, atol, rtol, causal, use_explicit_scale, batch, seqlen, nheads
tests/test_attention_correctness.py:183
↓ 3 callersMethodbenchmark_metadata
(self)
rl_engine/testing/rl_batch.py:70
← previousnext →101–200 of 1,290, ranked by callers