MCPcopy Create free account

hub / github.com/ByteDance-Seed/Triton-distributed / functions

Functions5,103 in github.com/ByteDance-Seed/Triton-distributed

⨍ Functions 5,103◇ Types & classes 521↳ Endpoints 149
↓ 3 callersFunctioncheck_mxshmem_init
shmem/mxshmem_bind/pymxshmem/src/pymxshmem.cc:71
↓ 3 callersFunctioncheck_with_token_cnt_per_rank_per_expert
(token_cnts_per_rank_per_expert, verbose=True)
python/triton_dist/kernels/nvidia/threadblock_swizzle_ag_moe.py:207
↓ 3 callersFunctioncooperative_barrier_on_this_grid
triton implementation of cooperative_group::this_grid().sync() WARNING: use with care. better launch triton with launch_cooperative_grid=True to
python/triton_dist/kernels/nvidia/common_ops.py:109
↓ 3 callersFunctioncopy_node_with_attrs
Copy special attributes from source node to target node. Args: node: Source AST node (may have special attributes) new_n
python/little_kernel/core/passes/utils/preserve_attributes.py:41
↓ 3 callersFunctioncp_engine_producer_all_gather_inter_node
( local_tensor: torch.Tensor, ag_buffer: list[torch.Tensor], signal_buffer: list[torch.Tensor],
python/triton_dist/kernels/nvidia/allgather.py:454
↓ 3 callersFunctioncp_engine_producer_all_gather_intra_node
( rank, num_ranks, local_tensor: torch.Tensor, remote_tensor_buffers: List[torch.Tensor],
python/triton_dist/kernels/nvidia/allgather.py:202
↓ 3 callersMethodcreate
(profiler_buffer, group_id, num_groups=1, is_leader=True, num_blocks=None, ENABLE_PROFILING=True)
python/triton_dist/tools/profiler/language.py:52
↓ 3 callersFunctioncreate_ag_gemm_intra_node_context
create context for allgather gemm intra-node Args: tensor_A (torch.Tensor<float>): local matmul A matrix. shape: [M_per_rank, K]
python/triton_dist/kernels/metax/allgather_gemm.py:1062
↓ 3 callersFunctioncreate_all_to_all_single_gemm_context
Create context for All-to-All Single GEMM operation Args: max_m: Maximum number of rows n: Number of columns for output
python/triton_dist/kernels/nvidia/all_to_all_single_gemm.py:236
↓ 3 callersFunctioncreate_fast_allgather_context
(rank, node, num_ranks, num_nodes, max_buffer_size: int = 2 * 32 * 1024 * 1024)
python/triton_dist/kernels/nvidia/low_latency_allgather.py:798
↓ 3 callersFunctioncreate_gemm_rs_context
( max_M, N, rank, world_size, local_world_size, output_dtype: torch.dtype, rs_stre
python/triton_dist/kernels/nvidia/gemm_reduce_scatter.py:78
↓ 3 callersFunctioncreate_rand_tensor
(rank, shape, dtype=torch.float16, device="cuda")
python/triton_dist/test/nvidia/test_moe_reduce_ar.py:37
↓ 3 callersFunctioncreate_rand_tensor
(rank, shape, dtype=torch.float16, device="cuda")
python/triton_dist/test/nvidia/test_moe_reduce_rs.py:39
↓ 3 callersFunctioncreate_reduce_scater_2d_ctx
for num_reduction_sms: tunable param, 16 are enough for H800 For H800, we overlap local reduce and inter-node p2p with intra-node
python/triton_dist/kernels/nvidia/reduce_scatter.py:148
↓ 3 callersMethodcreate_symm_tensor
(self, shape, dtype)
python/triton_dist/mega_triton_kernel/models/model_builder.py:130
↓ 3 callersFunctioncumsum
python/triton_dist/kernels/nvidia/threadblock_swizzle_ag_moe.cc:99
↓ 3 callersFunctioncumsum
(x)
python/triton_dist/kernels/nvidia/threadblock_swizzle_ag_moe.py:46
↓ 3 callersFunctiondist_swizzle2d_Nz
communication swizzle Nz
python/triton_dist/language/extra/ascend/algorithm.py:7
↓ 3 callersMethoddump_profiler_trace
(self, info: str = None)
python/triton_dist/kernels/nvidia/all_to_all_vdev_2d_offset.py:507
↓ 3 callersFunctionestimate_reduce_scatter_time_ms
intranode_bw/internode_bw in GB/s
python/triton_dist/kernels/nvidia/comm_perf_model.py:94
↓ 3 callersFunctionfast_all_to_all
low-latency all-to-all communication
python/triton_dist/kernels/amd/low_latency_all_to_all.py:215
↓ 3 callersMethodforward
(self, inputs_full=None, num_micro_batches=0)
python/triton_dist/test/nvidia/test_pp_block.py:167
↓ 3 callersFunctionfused_group_gemm_combine_token
Fused gemm2 + combine. Returns ``combined_out[num_tokens, N2]``.
python/triton_dist/kernels/amd/ep_all2all_fused.py:1151
↓ 3 callersFunctiongather_benchmark
(time_value)
python/triton_dist/test/amd/test_all_to_all.py:388
↓ 3 callersFunctiongather_benchmark
(time_value)
python/triton_dist/test/nvidia/test_all_to_all.py:386
↓ 3 callersFunctiongen_indices
(routing, n, G, topk, epr)
python/triton_dist/test/amd/test_ep_a2a_fused_kernel.py:104
↓ 3 callersFunctiongenerate_data
(configs)
python/triton_dist/utils.py:445
↓ 3 callersFunctiongenerate_ntokens_per_rank_per_expert_uniform
[ [ntokens(r0, e0), ntokens(r0, e1), ...], [ntokens(r1, e0), ntokens(r1, e1), ...], ... [ntokens(r7, e0), ntokens
python/triton_dist/kernels/nvidia/threadblock_swizzle_ag_moe_triton.py:314
↓ 3 callersFunctiongenerate_token_cnts_per_rank_per_expert_uniform
[ [ntokens(r0, e0), ntokens(r0, e1), ...], [ntokens(r1, e0), ntokens(r1, e1), ...], ... [ntokens(r7, e0), ntokens
python/triton_dist/kernels/nvidia/threadblock_swizzle_ag_moe.py:176
↓ 3 callersFunctiongenerate_uniform_counts_per_rank_per_expert
python/triton_dist/kernels/nvidia/threadblock_swizzle_ag_moe.cc:231
↓ 3 callersFunctionget_all_struct_stubs
Get all registered struct stubs.
python/little_kernel/language/intrin/struct_stub.py:228
↓ 3 callersFunctionget_deps
()
python/triton_dist/tune.py:228
↓ 3 callersFunctionget_device_multi_processor_count
(device)
python/triton_dist/kernels/nvidia/gemm_perf_model.py:45
↓ 3 callersFunctionget_device_property
(device_id=0)
python/triton_dist/utils.py:819
↓ 3 callersFunctionget_element_at
(tensor: tl.tensor, idx: int)
python/triton_dist/kernels/nvidia/all_to_all_single_2d.py:36
↓ 3 callersFunctionget_env_with_keys
(key: list)
python/setup.py:236
↓ 3 callersMethodget_gemm_out_buf
(self, input, weight)
python/triton_dist/kernels/nvidia/gemm_allreduce.py:67
↓ 3 callersFunctionget_int_env
(env, default_value)
python/triton_dist/utils.py:912
↓ 3 callersFunctionget_intrin_ctx
Get a context dictionary containing all registered intrin functions. Returns: A dictionary mapping intrin function names to func
python/little_kernel/language/builtin_base.py:45
↓ 3 callersFunctionget_ir_node_type
Get the IR node type if the node is an IR node.
python/little_kernel/core/passes/utils/ir_nodes.py:96
↓ 3 callersFunctionget_loop_modifier_registry
Get the global loop modifier registry.
python/little_kernel/core/passes/utils/registries/loop_modifier_registry.py:181
↓ 3 callersFunctionget_nic_gbps_per_gpu
in GB/s not Gbps
python/triton_dist/kernels/nvidia/comm_perf_model.py:83
↓ 3 callersMethodget_nvshmem_size
Get the total nvshmem memory size in bytes. This can be called before sync() to query the total memory needed.
python/triton_dist/layers/nvidia/ep_a2a_fused_layer.py:241
↓ 3 callersMethodget_swizzled_block_idx
Get swizzled block indices from block index. Returns: Tuple of (m_block_idx, n_block_idx) Note:
python/little_kernel/atom/scheduler.py:160
↓ 3 callersFunctionget_triton_combine_kv_algo_info
(split_kv, v_head_dim, block_dv=None)
python/triton_dist/kernels/nvidia/flash_decode.py:281
↓ 3 callersFunctionget_variable_type_from_scope
Get type of a variable from scope/context using the unified ScopeManager. Raises UndefinedVariableError if not found. Priority:
python/little_kernel/core/passes/utils/type_inference/type_inference_core.py:68
↓ 3 callersFunctionhas_tma
()
python/triton_dist/utils.py:799
↓ 3 callersMethodinc_offset
(self, seq_len: int)
python/triton_dist/mega_triton_kernel/models/paged_kv_cache.py:53
↓ 3 callersFunctioninit_nvshmem_by_torch_process_group
(pg: torch.distributed.ProcessGroup)
python/triton_dist/utils.py:233
↓ 3 callersMethodinit_triton_dist_gemm_ar_ctx
(self, max_M: int = 4096)
python/triton_dist/models/dense.py:201
↓ 3 callersFunctioninter_node_allgather
(local_tensor: torch.Tensor, ag_buffer: list[torch.Tensor], signal_buffer: list[torch.Tensor],
python/triton_dist/kernels/metax/allgather_gemm.py:337
↓ 3 callersMethodis_generic
Check if the type is a generic type
python/little_kernel/core/type_system.py:75
↓ 3 callersFunctionld
( ptr, scope: core.constexpr = "gpu", semantic: core.constexpr = "relaxed", _semantic=None, )
python/triton_dist/language/extra/cuda/language_extra.py:582
↓ 3 callersMethodlocal_q_nheads
(self, q_nheads)
python/triton_dist/kernels/nvidia/ulysses_sp_dispatch.py:519
↓ 3 callersMethodmake_allreduce
if double_input_buffer is True, user needs to ensure that the input of two consecutive allreduce are completely different buffers,
python/triton_dist/mega_triton_kernel/models/model_builder.py:473
↓ 3 callersFunctionmake_data
(M, N, K, dtype, scale)
python/triton_dist/test/nvidia/test_all_to_all_single_gemm.py:84
↓ 3 callersFunctionmake_data
(local_seq_len, hidden_dim, out_features, input_dtype)
python/triton_dist/test/nvidia/test_ulysses_sp_infer_qkv_proj_a2a.py:153
↓ 3 callersFunctionmake_data
(M, N, K, dtype: torch.dtype, trans_b, tp_group: torch.distributed.ProcessGroup)
python/triton_dist/test/nvidia/test_ag_gemm.py:76
↓ 3 callersMethodmake_flash_decode
query: (batch, seq_len, num_q_heads, q_head_dim) key_cache: (MAX_NUM_KV_BLOCKS, PAGE_SIZE, num_kv_heads, q_head_dim)
python/triton_dist/mega_triton_kernel/models/model_builder.py:251
↓ 3 callersMethodmake_qkv_proj
(self, input: torch.Tensor, weight: torch.Tensor, output: torch.Tensor, layer_id: int = 0)
python/triton_dist/mega_triton_kernel/models/model_builder.py:235
↓ 3 callersFunctionmake_vector
(args: Sequence, _semantic=None)
python/triton_dist/language/simt_ops.py:231
↓ 3 callersMethodmap_special_type
Map special type name to C++ type string.
python/little_kernel/codegen/registries/type_mapping_registry.py:94
↓ 3 callersFunctionnext_power_of_2
Return the smallest power of 2 greater than or equal to n
python/triton_dist/mega_triton_kernel/kernels/utils.py:37
↓ 3 callersFunctionparse_range
(range_str)
python/triton_dist/tools/tune/tune_gemm.py:88
↓ 3 callersFunctionperf_torch
Benchmark PyTorch implementation with separated timing
python/triton_dist/test/nvidia/test_all_to_all_single_gemm.py:175
↓ 3 callersFunctionperf_triton
Benchmark Triton implementation with separated timing
python/triton_dist/test/nvidia/test_all_to_all_single_gemm.py:253
↓ 3 callersFunctionprepare_cos_sin_cache
(head_dim, max_position_embeddings, rope_theta)
python/triton_dist/mega_triton_kernel/test/torch_impl_utils.py:30
↓ 3 callersMethodrecast
(self, new_elem_dtype: dtype, _semantic=None)
python/triton_dist/language/simt_ops.py:170
↓ 3 callersFunctionreduce_scatter_2d_op
(input: torch.Tensor, ctx: ReduceScatter2DContext, output: Optional[torch.Tensor] = None)
python/triton_dist/kernels/nvidia/reduce_scatter.py:822
↓ 3 callersMethodreset_all_barrier_buf
(self)
python/triton_dist/kernels/nvidia/gemm_allreduce.py:72
↓ 3 callersMethodreset_cusum_seq_lens
(self, local_seqlen, seq_lens_cpu=None)
python/triton_dist/kernels/nvidia/sp_ulysess_qkv_gemm_all2all.py:693
↓ 3 callersFunctionrocshmem_barrier_all_on_stream
(stream: Optional[torch.cuda.Stream] = None)
python/triton_dist/utils.py:334
↓ 3 callersMethodrun
(self, *args, **kwargs)
python/triton_dist/jit.py:167
↓ 3 callersFunctionrun_ag_single_test
Run bandwidth benchmark for a single tensor size
python/triton_dist/test/amd/test_bandwidth.py:475
↓ 3 callersFunctionrun_ibgda_single_test
IBGDA (In-Place Broadcast Gather with Direct Addressing) bandwidth test Similar to P2P test but can be configured with different parameters
python/triton_dist/test/amd/test_mori_shmem_bw.py:341
↓ 3 callersFunctionrun_p2p_single_test
Run a CORRECTED P2P bandwidth benchmark for a single tensor size for all GPU pairs. This version serializes the tests to avoid contention.
python/triton_dist/test/amd/test_bandwidth.py:514
↓ 3 callersFunctionrun_p2p_single_test
(M, K, dtype, RANK, WORLD_SIZE, TP_GROUP, args)
python/triton_dist/test/amd/test_mori_shmem_bw.py:283
↓ 3 callersMethodserve
(self, input_ids: torch.Tensor, gen_len: int)
python/triton_dist/models/engine.py:113
↓ 3 callersMethodset_backend
(self, backend: str)
python/triton_dist/test/nvidia/test_pp_block.py:123
↓ 3 callersFunctionshard_local
(tensor: torch.Tensor, world_size: int, dim: int, local_rank: int)
python/triton_dist/mega_triton_kernel/models/layers/tp_mlp.py:29
↓ 3 callersFunctionshard_local
(tensor: torch.Tensor, world_size: int, dim: int, local_rank: int)
python/triton_dist/layers/amd/tp_mlp.py:34
↓ 3 callersFunctionshard_local
(tensor: torch.Tensor, world_size: int, dim: int, local_rank: int)
python/triton_dist/layers/nvidia/tp_mlp.py:38
↓ 3 callersFunctionshard_local
(tensor: torch.Tensor, world_size: int, dim: int, local_rank: int)
python/triton_dist/layers/nvidia/tp_moe.py:34
↓ 3 callersFunctionsort_topk_ids_align_block_size
( topk_ids: torch.Tensor, # [ntokens, topk] num_experts: int, rank: int, num_ranks: int,
python/triton_dist/kernels/nvidia/allgather_group_gemm.py:201
↓ 3 callersFunctionsplits_to_cumsum
(splits: torch.Tensor)
python/triton_dist/test/amd/test_all_to_all.py:40
↓ 3 callersFunctionsplits_to_cumsum
(splits: torch.Tensor)
python/triton_dist/test/nvidia/test_all_to_all.py:40
↓ 3 callersFunctionsupport_launch_cooperative_grid
()
python/triton_dist/utils.py:835
↓ 3 callersFunctionswizzle_2d
(tile_id, num_pid_m, num_pid_n, GROUP_SIZE_M: tl.constexpr)
python/triton_dist/kernels/nvidia/gemm_reduce_scatter.py:103
↓ 3 callersFunctiontest
(kernel, M, N, K, label='')
python/little_kernel/design/sm90_bf16_gemm.py:303
↓ 3 callersMethodtest
(t)
python/triton_dist/test/nvidia/test_ep_moe_inference.py:458
↓ 3 callersFunctionthreadblock_swizzle_gemm_reduce_scatter_kernel
( tiled_m, M, rank, WORLD_SIZE: tl.constexpr, NNODES: tl.constexpr, BLOCK_SIZE_M: tl.c
python/triton_dist/kernels/nvidia/gemm_rs_threadblock_swizzle.py:69
↓ 3 callersFunctiontile_wise_matmul_compute
(tile_id, a_ptr, b_ptr, c_ptr, M, N, K, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
python/triton_dist/mega_triton_kernel/kernels/linear.py:32
↓ 3 callersFunctiontorch_dtype_to_triton_dtype_str
(torch_dtype)
python/triton_dist/mega_triton_kernel/tasks/utils.py:58
↓ 3 callersFunctiontorch_post_attn_all_to_all_transpose
(sp_group, input, a2a_only, is_dp, seq_lens_cpu=None)
python/triton_dist/test/nvidia/test_llm_ulysess_all2all_gemm_intra_node.py:139
↓ 3 callersFunctiontransposed_dot
( a_ptrs, b_ptrs, c_ptrs, split_size, N, K, stride_am: tl.constexpr, stride_bm
python/triton_dist/kernels/nvidia/group_gemm.py:475
↓ 3 callersFunctiontriton_dist_moe_func
Ditron MoE forward function
python/triton_dist/test/nvidia/test_ep_moe_fused.py:136
↓ 3 callersFunctionvector_binOp
(x, y, op, _semantic=None)
python/triton_dist/language/simt_ops.py:67
↓ 3 callersMethodwarmup
(self, *args, grid, **kwargs)
python/triton_dist/jit.py:171
↓ 3 callersMethodwriteln_main
Write line to main code buffer with newline.
python/little_kernel/codegen/codegen_base.py:124
↓ 2 callersFunctionCreateMXSHMEMOp
lib/Conversion/TritonDistributedToLLVM/METAX/DistributedOpToLLVM.cpp:26
← previousnext →501–600 of 5,103, ranked by callers