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
↓ 2 callersFunctionceil_div
(a, b)
tutorials/04-deepseek-infer-all2all.py:78
↓ 2 callersFunctioncheck_allclose
Check if two tensors are close within a tolerance.
python/triton_dist/test/amd/test_tp_e2e.py:50
↓ 2 callersFunctioncheck_allclose
Checks if two tensors are close, with detailed logging on failure.
python/triton_dist/test/nvidia/test_tp_e2e.py:76
↓ 2 callersMethodcheck_input
(self, tensor)
python/triton_dist/kernels/nvidia/sp_ulysess_qkv_gemm_all2all.py:668
↓ 2 callersFunctioncheck_with_golden
(golden_results: tuple, received_out_data: torch.Tensor, received_out_splits: torch.Tensor = None,
python/triton_dist/test/nvidia/test_all_to_all_vdev_2d_offset.py:176
↓ 2 callersFunctioncompare_ld_reduce_precision
(symm_tensor: torch.Tensor, acc_dtype: torch.dtype)
python/triton_dist/test/nvidia/test_multimem_ld_reduce.py:119
↓ 2 callersFunctioncompute_launch_params
(M, N, K)
python/little_kernel/design/sm90_bf16_gemm.py:281
↓ 2 callersMethodconfigure_function
Configure function attributes for shared memory and cluster launch. Sets MaxDynamicSharedMemorySize and optionally NonPortableCluster
python/little_kernel/runtime/cuda_runtime.py:499
↓ 2 callersMethodconnect
Establish a connection to the server
python/triton_dist/mega_triton_kernel/test/models/chat.py:41
↓ 2 callersFunctionconst_fold
(tree: ast.AST, ctx: Dict[str, Any] = None)
python/little_kernel/core/passes/constfold.py:641
↓ 2 callersFunctionconsumer_all_reduce
(symm_input, symm_ar_out, ar_out, gemm_barrier, multi_st_barrier, BLOCK_SIZE_M=16, BLO
python/triton_dist/kernels/nvidia/gemm_allreduce.py:607
↓ 2 callersFunctioncopy_continuous_kernel
(src_ptr, dst_ptr, N, BLOCK_SIZE: tl.constexpr)
python/triton_dist/kernels/nvidia/reduce_scatter.py:195
↓ 2 callersFunctioncopy_kernel_producer_all_gather
( rank, num_ranks, local_tensor: torch.Tensor, remote_tensor_buffers: List[torch.Tensor],
python/triton_dist/kernels/amd/allgather_gemm.py:360
↓ 2 callersFunctioncopy_warp
( dst_ptr, src_ptr, nbytes, )
python/triton_dist/kernels/nvidia/memory_ops.py:250
↓ 2 callersFunctioncp_engine_producer_all_gather_full_mesh_push_multi_stream
( rank, num_ranks, local_tensor: torch.Tensor, remote_tensor_buffers: List[torch.Tensor],
python/triton_dist/kernels/amd/allgather_gemm.py:296
↓ 2 callersFunctioncreate_ag_group_gemm_context
Create context for allgather group gemm. Args: rank (int): current rank num_ranks (int): total number of ranks tensor_A (
python/triton_dist/kernels/nvidia/allgather_group_gemm.py:338
↓ 2 callersFunctioncreate_all_to_all_context
( max_m: int, hidden: int, rank: int, num_tot_experts: int, WORLD_SIZE: int, experts_p
python/triton_dist/kernels/nvidia/low_latency_all_to_all.py:176
↓ 2 callersFunctioncreate_context
( rank: int, world_size: int, ne: int, k: int, token_len_elem: int, token_dtype: torch
python/triton_dist/kernels/nvidia/all_to_all_vdev_2d_offset.py:528
↓ 2 callersFunctioncreate_context
( rank: int, world_size: int, ne: int, k: int, token_len_elem: int, token_dtype: torch
python/triton_dist/kernels/nvidia/all_to_all_vdev_2d_offset_inter_node.py:571
↓ 2 callersFunctioncreate_data
(rank_in_row)
python/triton_dist/test/nvidia/test_all_to_all_vdev_2d_offset.py:285
↓ 2 callersFunctioncreate_data
(rank_in_row)
python/triton_dist/test/nvidia/test_all_to_all_vdev_2d_offset_inter_node.py:284
↓ 2 callersFunctioncreate_descs
(A, B, C, K_aligned, N_stride)
python/little_kernel/design/sm90_bf16_gemm.py:292
↓ 2 callersFunctioncreate_ep_ll_a2a_ctx
(max_m, hidden, topk, num_experts, online_quant_fp8, fp8_gsize, dtype, world_size, rank)
python/triton_dist/kernels/amd/low_latency_all_to_all_v2.py:647
↓ 2 callersFunctioncreate_gemm_ar_context
(ar_stream: torch.cuda.Stream, rank, world_size, max_M, N, dtype, MIN_BLOCK_SIZE_M=64,
python/triton_dist/kernels/amd/gemm_allreduce.py:57
↓ 2 callersFunctioncreate_gemm_ar_context
(ar_stream: torch.cuda.Stream, rank, world_size, local_world_size, max_M, N, dtype,
python/triton_dist/kernels/nvidia/gemm_allreduce.py:103
↓ 2 callersFunctioncreate_moe_rs_context
Creates and initializes a context object for the MoE Reduce-Scatter operation. This context holds pre-allocated buffers (including symmetric
python/triton_dist/kernels/nvidia/moe_reduce_rs.py:88
↓ 2 callersFunctioncreate_symlink_rel
Make source path a symlink pointing to the target path. if both source/target under base_dir, create link with relative path. why it's tricky
python/build_helpers.py:45
↓ 2 callersMethodcustom_codegen
(node: ast.Call, emitter)
python/little_kernel/tests/unit/test_struct_stub.py:353
↓ 2 callersMethoddata
(self)
python/triton_dist/utils.py:1242
↓ 2 callersFunctiondecode_tag
Decode a profiler tag into (block_idx, group_idx, task_type, is_start). Tag layout: GLOBAL_ID | TASK TYPE | IS START
python/triton_dist/tools/profiler/viewer.py:40
↓ 2 callersFunctiondeinit_triton_dist_ep_op
(ep_implementation: str = "mega")
python/triton_dist/function/nvidia/common.py:285
↓ 2 callersFunctiondequant_fp8_bf16
(q_tensor: torch.Tensor, scales: torch.Tensor)
tutorials/04-deepseek-infer-all2all.py:423
↓ 2 callersMethodderef
(self)
python/little_kernel/core/expr.py:215
↓ 2 callersMethoddist_print
(self, *args, **kwargs)
python/triton_dist/autotuner.py:57
↓ 2 callersFunctiondot_2parts_k_const
( a_ptrs, b1_ptrs, b2_ptrs, c_ptrs, M, N, K1: tl.constexpr, K2: tl.constexpr,
python/triton_dist/kernels/nvidia/group_gemm.py:197
↓ 2 callersFunctiondot_k_const
A single (BLOCK_M, BLOCK_N) GEMM tile with a K reduction loop. Plain ``tl.dot`` accumulation in fp32.
python/triton_dist/kernels/amd/ep_all2all_fused.py:81
↓ 2 callersFunctiondot_k_const
( a_ptrs, b_ptrs, c_ptrs, M, N, K: tl.constexpr, stride_ak: tl.constexpr, stri
python/triton_dist/kernels/nvidia/ep_all2all_fused.py:561
↓ 2 callersFunctiondtype_size_in_bytes
(dtype: torch.dtype)
python/triton_dist/kernels/amd/low_latency_all_to_all.py:138
↓ 2 callersFunctiondtype_size_in_bytes
(dtype: torch.dtype)
python/triton_dist/kernels/nvidia/low_latency_all_to_all.py:121
↓ 2 callersFunctiondump_ast
Create a pass that dumps the AST with special attributes.
python/little_kernel/core/passes/dump.py:92
↓ 2 callersMethoddump_combine_trace
(self)
python/triton_dist/layers/amd/ep_ll_a2a_layer.py:179
↓ 2 callersMethoddump_dispatch_trace
(self)
python/triton_dist/layers/amd/ep_ll_a2a_layer.py:168
↓ 2 callersFunctioneval_return_type
Generated eval_return_type function.
python/little_kernel/language/simple_builtin.py:522
↓ 2 callersFunctionextract_empty_params
Extract shape, dtype, scope from ll.empty call (handles both positional and keyword args). Assumes ll.empty signature: empty(shape, dtype, sc
python/little_kernel/core/passes/utils/extract_empty_params.py:29
↓ 2 callersFunctionfence
( semantic: core.constexpr = core.constexpr("sc"), scope: core.constexpr = core.constexpr("gpu"), _sem
python/triton_dist/kernels/common_ops.py:124
↓ 2 callersFunctionffs
Find first set bit (1-indexed, 0 if no bits set).
python/little_kernel/language/intrin/simt.py:213
↓ 2 callersFunctionfill_random
(tensor: torch.Tensor)
python/triton_dist/test/nvidia/test_reduce_scatter.py:41
↓ 2 callersMethodfinalize
(self)
python/triton_dist/layers/amd/ep_a2a_fused_layer.py:147
↓ 2 callersFunctionfind_best_topk_fast
Greedy: pick `topk` rows to minimize a score of the column-wise minima over chosen rows. Score options: - "mean": minimize mean(col_m
python/triton_dist/tools/tune/find_topk.py:176
↓ 2 callersFunctionformat_bandwidth
Format bandwidth value for display.
python/little_kernel/benchmark/utils.py:222
↓ 2 callersMethodforward
return C = all_gather(A) @ B.T
python/triton_dist/test/amd/test_ag_gemm_intra_node.py:106
↓ 2 callersMethodforward
( self, x: torch.Tensor, w: torch.Tensor, chosen_experts: torch.Tensor,
python/triton_dist/test/nvidia/test_moe_reduce_ar.py:119
↓ 2 callersMethodforward
( self, x: torch.Tensor, w: torch.Tensor, chosen_experts: torch.Tensor,
python/triton_dist/test/nvidia/test_moe_reduce_rs.py:149
↓ 2 callersMethodforward
( self, input: torch.Tensor, # [M, local_K] weight: torch.Tensor, # [N, local_K]
python/triton_dist/layers/nvidia/gemm_allreduce_layer.py:107
↓ 2 callersMethodforward_push_2d
(self, symm_buffer: torch.Tensor)
python/triton_dist/layers/nvidia/low_latency_allgather_layer.py:66
↓ 2 callersMethodforward_push_2d_ll
(self, symm_buffer: torch.Tensor)
python/triton_dist/layers/nvidia/low_latency_allgather_layer.py:105
↓ 2 callersMethodfree
(self)
python/little_kernel/design/test_flashcomm_multi_gpu.py:147
↓ 2 callersFunctionfunc
()
python/triton_dist/test/nvidia/test_ag_small_msg.py:162
↓ 2 callersFunctionfwd
()
python/triton_dist/test/nvidia/test_ep_a2a.py:382
↓ 2 callersFunctiongemm_allreduce_op
(ctx: GemmARContext, A: torch.Tensor, B: torch.Tensor, gemm_config: triton.Config)
python/triton_dist/kernels/amd/gemm_allreduce.py:306
↓ 2 callersFunctiongemm_allreduce_torch
Reference torch implementation for gemm+allreduce
python/triton_dist/test/amd/test_gemm_ar_intra_node.py:37
↓ 2 callersFunctiongemm_rs_producer_persistent
(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, barrier: torch.Tensor, wor
python/triton_dist/kernels/nvidia/gemm_reduce_scatter.py:514
↓ 2 callersFunctiongen_routing
Top-k expert ids per token. ``random`` = distinct ids via top-k over random logits; ``skew_rank0`` = every token -> rank-0's first ``topk`` expert
python/triton_dist/test/amd/test_ep_moe_fused.py:144
↓ 2 callersFunctiongenerate_ntokens_per_rank_per_expert_random
(nexperts_per_rank, nexperts, TP_SIZE)
python/triton_dist/kernels/nvidia/threadblock_swizzle_ag_moe_triton.py:328
↓ 2 callersFunctiongenerate_response
Generate response from the model using token IDs tensor on GPU
python/triton_dist/mega_triton_kernel/test/models/model_server.py:61
↓ 2 callersFunctiongenerate_size_configs
Generate tensor configs for a fixed list of M values with K=4096.
python/triton_dist/test/amd/test_mori_shmem_bw.py:419
↓ 2 callersFunctiongetSharedMemAddress
lib/Conversion/TritonDistributedToLLVM/SIMTOpToLLVM.cpp:44
↓ 2 callersFunctionget_ag_splits_and_recv_offset_for_dispatch_intra_node
Compute recv_buf_offset_per_expert and num_recv_tokens_per_rank for intra-node dispatch. AMD implementation uses mori_shmem device API (putm
python/triton_dist/kernels/amd/ep_a2a_intra_node.py:513
↓ 2 callersFunctionget_auto_triton_config
(M, N, K, N_CHUNKS, dtype: torch.dtype)
python/triton_dist/kernels/nvidia/moe_reduce_ar.py:200
↓ 2 callersFunctionget_auto_triton_config
(M, N, K, topk, nexperts, N_CHUNKS, persistent: bool, dtype: torch.dtype)
python/triton_dist/kernels/nvidia/moe_reduce_rs.py:795
↓ 2 callersFunctionget_barrier_kernel
(gpu_id=0)
python/little_kernel/design/test_flashcomm_multi_gpu.py:353
↓ 2 callersFunctionget_cmake_dir
()
python/build_helpers.py:36
↓ 2 callersFunctionget_compute_offset_kernel
(num_experts, gpu_id=0)
python/little_kernel/design/test_flashcomm_multi_gpu.py:309
↓ 2 callersFunctionget_config_space
()
python/triton_dist/kernels/amd/gemm.py:543
↓ 2 callersMethodget_ctx
(self, input_ids: torch.LongTensor)
python/triton_dist/models/engine.py:107
↓ 2 callersFunctionget_dispatch_layout_kernel
(gpu_id=0)
python/little_kernel/design/test_flashcomm_multi_gpu.py:325
↓ 2 callersFunctionget_dispatch_send_reqs
(exp_indices, send_reqs_for_nodes, experts_per_rank, local_world_size, num_sms, use
python/triton_dist/kernels/nvidia/ep_a2a.py:725
↓ 2 callersFunctionget_dispatch_v1_kernel
(gpu_id=0)
python/little_kernel/design/test_flashcomm_multi_gpu.py:339
↓ 2 callersFunctionget_dram_gbps_by_device_name
(device_name: str)
python/triton_dist/kernels/nvidia/gemm_perf_model.py:194
↓ 2 callersFunctionget_enum
Get a registered Enum class by name.
python/little_kernel/codegen/registries/enum_registry.py:91
↓ 2 callersFunctionget_extra_packages
(extra_name)
python/setup.py:1035
↓ 2 callersMethodget_gemm_out_buf
(self, input)
python/triton_dist/kernels/nvidia/gemm_reduce_scatter.py:72
↓ 2 callersFunctionget_golden_results
(inp_splits: torch.Tensor, inp: torch.Tensor, world_size: int, ne: int, align: int, ran
python/triton_dist/test/nvidia/test_all_to_all_vdev_2d_offset.py:44
↓ 2 callersFunctionget_group_numa_world_size
allgather all ranks in the process group and get the NUMA world size
python/triton_dist/utils.py:536
↓ 2 callersFunctionget_intranode_fullmesh_world_size
(num_ranks)
python/triton_dist/kernels/metax/allgather_gemm.py:1006
↓ 2 callersMethodget_modifier
Get the LoopModifier for a Call node. Args: node: AST Call node representing the modifier call ctx:
python/little_kernel/core/passes/utils/registries/loop_modifier_registry.py:121
↓ 2 callersFunctionget_moe_optim_config
(use_mega: bool = False, is_forward: bool = True)
python/triton_dist/function/nvidia/common.py:420
↓ 2 callersMethodget_mxlink_adjacency_matrix
()
python/triton_dist/kernels/metax/utils.py:10
↓ 2 callersFunctionget_numa_node_count_in_group
(pg: torch.distributed.ProcessGroup)
python/triton_dist/utils.py:517
↓ 2 callersMethodget_nvshmem_build_from_src_home
()
python/triton_dist/nv_utils.py:450
↓ 2 callersFunctionget_nvshmem_hash
()
python/triton_dist/utils.py:632
↓ 2 callersMethodget_nvshmem_home
()
python/triton_dist/nv_utils.py:437
↓ 2 callersFunctionget_nvshmem_version
()
python/triton_dist/utils.py:599
↓ 2 callersFunctionget_rocshmem_hash
()
python/triton_dist/utils.py:658
↓ 2 callersFunctionget_rocshmem_version
()
python/triton_dist/utils.py:646
↓ 2 callersFunctionget_rt
()
python/little_kernel/design/test_ipc_minimal.py:37
↓ 2 callersMethodget_scatter_bufs_and_signal_for_each_node
(self, input, node_id)
python/triton_dist/kernels/nvidia/reduce_scatter.py:108
↓ 2 callersFunctionget_shmem_extern_lib
()
python/triton_dist/jit.py:104
↓ 2 callersFunctionget_source_from_ctx
Try to get source code from context (e.g., from LLKernel). Args: ctx: Context dictionary that may contain source information
python/little_kernel/core/passes/utils/error_report.py:134
↓ 2 callersMethodget_task_builder
(self, op_type: str)
python/triton_dist/mega_triton_kernel/models/model_builder.py:192
← previousnext →701–800 of 5,103, ranked by callers