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
↓ 2 callersMethod_silu
(x: torch.Tensor, *, output_dtype: torch.dtype)
rl_engine/kernels/ops/pytorch/activation/swiglu.py:38
↓ 2 callersFunction_single_thread
Pin CPU GEMM to one thread so the K reduction order is M-independent.
tests/test_lm_head.py:73
↓ 2 callersFunction_sm90_supported
Whether the bf16 TMA+MMA forward can run these inputs directly.
rl_engine/kernels/ops/cuda/loss/linear_logp.py:30
↓ 2 callersFunction_squeeze_reward_tensor
(value: torch.Tensor)
rl_engine/executors/stateless_executor.py:557
↓ 2 callersMethod_swiglu
( gate: torch.Tensor, up: torch.Tensor, *, output_dtype: torch.dtype, )
rl_engine/kernels/ops/pytorch/activation/swiglu.py:73
↓ 2 callersFunction_sync_torch_cuda_home
(path: Path)
rl_engine/executors/deepspeed_trainer.py:375
↓ 2 callersFunction_synchronize
(device: torch.device)
tests/linear_logp_tp.py:185
↓ 2 callersFunction_temporarily_configure_stateless_model
( model: torch.nn.Module, config: StatelessForwardConfig, )
rl_engine/executors/stateless_executor.py:370
↓ 2 callersMethod_tensor_from_address
( self, address: int, shape: tuple[int, ...], stride: tuple[int, ...],
rl_engine/executors/bridge.py:2217
↓ 2 callersFunction_time_ms
(fn, device: torch.device, *, warmup: int = 3, repeat: int = 10)
benchmarks/benchmark_rl_kernels.py:96
↓ 2 callersFunction_time_ms
(fn, device: torch.device, *, warmup: int = 3, repeat: int = 10)
benchmarks/benchmark_ratio_kl.py:91
↓ 2 callersFunction_tiny_gpt2_reference_model
()
tests/test_stateless_hf_integration.py:55
↓ 2 callersMethod_validate_decode_alignment
Enforce the contract that ``q`` holds exactly the newly appended positions: ``Sq == S_new``. q's rows are the queries for ``k_new``,
rl_engine/kernels/ops/pytorch/attention/kv_cache.py:126
↓ 2 callersFunction_validate_global_targets
( target_1d: torch.Tensor, global_vocab_size: int, tp_group: Any = None, )
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:132
↓ 2 callersFunction_validate_inputs
(inputs: StatelessForwardInputs, config: PagedKVScoringConfig)
rl_engine/executors/paged_kv_baseline.py:312
↓ 2 callersMethod_validate_output_shape
(self, output: torch.Tensor, logits: torch.Tensor)
rl_engine/kernels/ops/pytorch/loss/logp.py:45
↓ 2 callersFunction_validate_tp_vocab_partition
( *, tp_group: Any, device: torch.device, vocab_start_index: int, local_vocab_size: int,
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:86
↓ 2 callersMethod_wait_for_cuda_vmm_publish_event
(self, manifest: WeightUpdateManifest)
rl_engine/executors/bridge.py:2197
↓ 2 callersMethodapply
Selected-token log-prob ``z[t] - logsumexp(z)``, returned in float32.
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:476
↓ 2 callersMethodapply
Compute ``(loss, policy_loss, kl)`` from logits + per-sequence advantages. ``sample_advantages`` is per-sequence and is broadcast to per-toke
rl_engine/kernels/ops/pytorch/loss/grpo_loss.py:96
↓ 2 callersMethodapply
( self, hidden: torch.Tensor, lm_head_weight: torch.Tensor, target_ids: torch.
rl_engine/kernels/ops/cuda/loss/linear_logp.py:289
↓ 2 callersMethodapply_model
(self, func)
tests/test_weight_sync_bridge.py:900
↓ 2 callersMethodbackward
(ctx, do)
rl_engine/kernels/ops/triton/triton_attn.py:396
↓ 2 callersMethodcreate_training_worker
(self, spec: RayWorkerSpec)
rl_engine/executors/ray_actor_manager.py:120
↓ 2 callersFunctiondescribe_attention_diff
(actual, expected, *, dtype, atol, rtol, causal, softmax_scale)
tests/test_attention_correctness.py:96
↓ 2 callersMethodexport_model_handles
(self, model: torch.nn.Module)
rl_engine/executors/bridge.py:2499
↓ 2 callersFunctionextract_kv_cache_outputs
Extract common cache-bearing output fields from model outputs.
rl_engine/executors/stateless_executor.py:382
↓ 2 callersFunctionextract_rollout_candidate_groups
Extract generated token ids while preserving prompt-level candidate groups.
rl_engine/executors/training_contract.py:451
↓ 2 callersFunctionextract_rollout_reference_logp_groups
Extract per-candidate reference logprobs from rollout payloads when present.
rl_engine/executors/training_contract.py:491
↓ 2 callersMethodforward_logits
(self, input_ids: torch.Tensor, **model_kwargs: Any)
rl_engine/alignment/model_wrappers.py:87
↓ 2 callersMethodgenerate
(self, prompts, sampling_params)
tests/test_vllm_rollout_sampler.py:36
↓ 2 callersMethodgenerate_candidates
Generate GRPO rollout candidates through vLLM with shared prefix caching.
rl_engine/executors/rollout.py:155
↓ 2 callersMethodget_metrics
(self)
benchmarks/profiler.py:385
↓ 2 callersMethodget_target_info
(device_index: int = 0)
benchmarks/profiler.py:85
↓ 2 callersMethodhealth_check
(self)
rl_engine/executors/ray_actor_manager.py:124
↓ 2 callersMethodindexed_fp32
( self, logits: torch.Tensor, token_ids: torch.Tensor, row_indices: torch.Tensor )
rl_engine/kernels/ops/pytorch/loss/logp.py:81
↓ 2 callersMethodinit
(self, **kwargs)
tests/test_ray_actor_manager.py:81
↓ 2 callersFunctionmasked_sum
Sum values while ignoring masked-out entries.
rl_engine/testing/reference_ops.py:50
↓ 2 callersFunctionmeasure_extra_vram
Measures PEAK extra VRAM allocated by fn(*args), excluding the memory already occupied by the input tensors. Returns (extra_vram_gb, avg_
benchmarks/benchmark_grpo_op.py:15
↓ 2 callersFunctionnative_sampling
Simulates standard PyTorch sampling logic (Top-K -> Top-P -> Softmax -> Multinomial)
benchmarks/benchmark_sampling.py:15
↓ 2 callersFunctionnormalize_grouped_outputs
( grouped_outputs: Sequence[Sequence[Any]], )
rl_engine/executors/vllm_sampler.py:203
↓ 2 callersMethodonline_indexed_out
( self, logits: torch.Tensor, token_ids: torch.Tensor, row_indices: torch.Tens
rl_engine/kernels/ops/pytorch/loss/logp.py:95
↓ 2 callersMethodonline_out
( self, logits: torch.Tensor, token_ids: torch.Tensor, output: torch.Tensor )
rl_engine/kernels/ops/pytorch/loss/logp.py:87
↓ 2 callersFunctionoperator_shape_name
(op_name: str, args: argparse.Namespace)
rl_engine/kernels/gtest/operator_inputs.py:47
↓ 2 callersMethodprint_summary
(self)
benchmarks/profiler.py:439
↓ 2 callersMethodprofile_logp
Profile a log-probability operator and compute end-to-end metrics. Args: candidate_fn: Callable that performs the fused/
benchmarks/profiler.py:247
↓ 2 callersFunctionpytorch_sdpa_reference
Compute a PyTorch SDPA math-backend reference for FlashAttention-layout inputs. The reference runs in fp32 and validates each low-precision
tests/test_attention_correctness.py:57
↓ 2 callersMethodrelease
(self, update_id)
tests/test_weight_sync_bridge.py:894
↓ 2 callersMethodrelease_weights
(self)
tests/test_ray_actor_manager.py:172
↓ 2 callersFunctionrocm_flash_attention_availability
()
tests/test_attention_correctness.py:153
↓ 2 callersMethodrollout
(self, spec)
tests/test_ray_actor_manager.py:104
↓ 2 callersFunctionrun
(op)
benchmarks/benchmark_ratio_kl.py:154
↓ 2 callersMethodscore
(self, inputs: StatelessForwardInputs)
rl_engine/executors/paged_kv_baseline.py:108
↓ 2 callersMethodscore
(self, rollout: RolloutStageResult)
rl_engine/executors/training_contract.py:327
↓ 2 callersMethodstate_dict
(self, *args: Any, **kwargs: Any)
benchmarks/benchmark_weight_sync_bridge.py:654
↓ 2 callersMethodstep
(self)
tests/test_deepspeed_training_worker.py:43
↓ 2 callersMethodtrain
(self, rollout)
tests/test_ray_actor_manager.py:120
↓ 1 callersMethodCUmemAccessDesc
(self)
rl_engine/executors/bridge.py:1131
↓ 1 callersMethodCUmemAllocationProp
(self)
rl_engine/executors/bridge.py:1111
↓ 1 callersMethod__init__
(self, state_dict: Mapping[str, torch.Tensor])
rl_engine/executors/deepspeed_trainer.py:645
↓ 1 callersMethod__init__
(self)
rl_engine/kernels/ops/pytorch/activation/swiglu.py:20
↓ 1 callersMethod__init__
(self, model: torch.nn.Module)
rl_engine/alignment/model_wrappers.py:80
↓ 1 callersFunction_add_windows_dll_dirs
()
tests/conftest.py:9
↓ 1 callersMethod_adjust_priority_for_hardware
Adjust CUDA priorities for hardware-gated experimental and production kernels.
rl_engine/kernels/registry.py:186
↓ 1 callersMethod_adjust_priority_from_env
(self)
rl_engine/kernels/registry.py:166
↓ 1 callersFunction_adv_tokens
(batch)
tests/test_grpo_loss.py:82
↓ 1 callersFunction_append_row
(path: Path, row: dict[str, Any])
benchmarks/benchmark_stateless_executor.py:149
↓ 1 callersMethod_attach_cuda_ipc_manifest
(self, manifest: WeightUpdateManifest)
rl_engine/executors/bridge.py:2537
↓ 1 callersMethod_attach_cuda_vmm_manifest
(self, manifest: WeightUpdateManifest)
rl_engine/executors/bridge.py:2281
↓ 1 callersFunction_attach_reference_logps_to_grouped_candidates
( payload: dict[str, Any], reference_logps: torch.Tensor, completion_mask: torch.Tensor, )
rl_engine/executors/training_contract.py:573
↓ 1 callersFunction_attach_reference_logps_to_grouped_outputs
( grouped_outputs: Sequence[Any], logp_rows: Sequence[Sequence[float]], )
rl_engine/executors/training_contract.py:722
↓ 1 callersFunction_attach_reward_tensor_to_grouped_candidates
( payload: dict[str, Any], rewards: torch.Tensor, )
rl_engine/executors/training_contract.py:546
↓ 1 callersFunction_attach_rewards_to_grouped_outputs
( grouped_outputs: Sequence[Any], rewards: Sequence[float], )
rl_engine/executors/training_contract.py:661
↓ 1 callersFunction_attach_shared_memory
(name: str)
rl_engine/executors/bridge.py:754
↓ 1 callersMethod_attach_shared_memory_manifest
(self, manifest: WeightUpdateManifest)
rl_engine/executors/bridge.py:1884
↓ 1 callersFunction_batch
(seed=0, *, device="cpu", valid_density=0.9)
tests/test_ratio_kl.py:30
↓ 1 callersFunction_broadcast_sequence_advantages
( sequence_advantages: torch.Tensor, completion_mask: torch.Tensor, )
rl_engine/executors/training_contract.py:994
↓ 1 callersMethod_build_bounds
Return CSR-style group offsets (int32) and the max group size.
rl_engine/kernels/ops/triton/loss/grpo_loss.py:102
↓ 1 callersMethod_build_engine
(self)
rl_engine/executors/vllm_sampler.py:126
↓ 1 callersMethod_build_sampling_params
(self, overrides: Optional[Mapping[str, Any]])
rl_engine/executors/vllm_sampler.py:143
↓ 1 callersFunction_call_accepts_keyword
(model: torch.nn.Module, keyword: str)
rl_engine/executors/stateless_executor.py:516
↓ 1 callersFunction_call_accepts_keyword
(model: torch.nn.Module, keyword: str)
rl_engine/executors/paged_kv_baseline.py:366
↓ 1 callersMethod_can_use_sm90
(self, logits: torch.Tensor)
rl_engine/kernels/ops/cuda/loss/logp.py:33
↓ 1 callersFunction_candidate_prompt_token_ids
(candidate: Any)
rl_engine/executors/training_contract.py:808
↓ 1 callersFunction_candidate_records_from_payload
(payload: Any)
rl_engine/executors/training_contract.py:511
↓ 1 callersFunction_candidate_reference_logp_values
(candidate: Any)
rl_engine/executors/training_contract.py:851
↓ 1 callersFunction_candidate_reference_logps
( reference_logp_groups: Sequence[Sequence[Sequence[float]]], group_index: int, candidate_index: i
rl_engine/executors/training_contract.py:937
↓ 1 callersFunction_candidate_reward
(candidate: Any)
rl_engine/executors/training_contract.py:824
↓ 1 callersFunction_candidate_reward_value
( token_ids: Sequence[int], reward_groups: Optional[Sequence[Sequence[float]]], group_index: int,
rl_engine/executors/training_contract.py:921
↓ 1 callersFunction_chunked_local_linear_logp_stats
( hidden_2d: torch.Tensor, weight: torch.Tensor, target_1d: torch.Tensor, bias_t: torch.Tensor
rl_engine/kernels/ops/pytorch/loss/linear_logp.py:167
↓ 1 callersMethod_clone_state_dict
(self, model: torch.nn.Module)
rl_engine/executors/bridge.py:1640
↓ 1 callersFunction_completion_mask_starts_at_position_zero
( completion_mask: torch.Tensor, *, device: torch.device, )
rl_engine/executors/paged_kv_baseline.py:398
↓ 1 callersFunction_compute_group_relative_advantages
( rewards: torch.Tensor, *, group_ids: torch.Tensor, num_groups: int, eps: float, )
rl_engine/executors/training_contract.py:963
↓ 1 callersFunction_configure_cuda_home_from_python_packages
()
rl_engine/executors/deepspeed_trainer.py:332
↓ 1 callersMethod_configure_signatures
(self)
rl_engine/executors/bridge.py:1030
↓ 1 callersFunction_cpu_fp16_matmul_supported
Probe whether this CPU backend implements float16 matmul.
tests/test_lm_head.py:49
↓ 1 callersFunction_cpu_fp16_matmul_supported
Probe whether this CPU backend implements float16 matmul.
tests/test_kv_cache_attention.py:56
↓ 1 callersFunction_cpu_fp16_matmul_supported
Probe whether this CPU backend implements float16 matmul.
tests/test_attention.py:46
↓ 1 callersFunction_create_shared_memory
(size: int)
rl_engine/executors/bridge.py:750
↓ 1 callersMethod_current_gpu_uuid
(self)
rl_engine/executors/bridge.py:318
← previousnext →301–400 of 1,290, ranked by callers