/* * Portions of this file are adapted from DeepEP (https://github.com/deepseek-ai/DeepEP). * Copyright (c) 2025 DeepSeek. Licensed under the MIT License. * SPDX-License-Identifier: MIT */ /* * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * See LICENSE.txt for more license information. */ #pragma once #include #include #include #include "nccl_device.h" namespace nccl_ep { namespace hybridep { // ============================================================================ // Runtime constants (from ep_group, not hardcoded): // lsa_team_size - ncclTeamLsa(comm).nRanks (up to 72 for MNNVL/NVL72) // rdma_ranks - nRanks / lsa_team_size // num_local_experts, hidden // // Use HYBRIDEP_SWITCH_* macros to instantiate templates with runtime values. // LSA team size is dispatched in steps of 4 up to 72; each instantiation // sizes its own register-file arrays via the LSA_TEAM_SIZE template param. // ============================================================================ // ============================================================================ // Kernel: Convert sparse topk_idx to dense routing map // ============================================================================ // NCCL API uses sparse format: topk_idx[token][k] = expert_id // HT uses bitmap format: routing_bitmap[token][expert/8] has bit (expert%8) set void convert_topk_to_routing_map( const int64_t* topk_idx, uint8_t* routing_bitmap, int num_tokens, int num_topk, int num_experts_packed, // = ceil(num_experts / 8) cudaStream_t stream); // ============================================================================ // Kernel: Convert sparse topk_weights to dense prob (for dispatch input) // ============================================================================ // NCCL API uses sparse format: topk_weights[token][k] = weight, topk_idx[token][k] = expert_id // HT uses dense format: prob[token][expert] = weight (0 for non-selected) void sparse_to_dense_prob( const int64_t* topk_idx, // [num_tokens, topk] - which experts const float* topk_weights, // [num_tokens, topk] - weights for those experts float* dense_prob, // [num_tokens, num_experts] - output dense prob (pre-zeroed) int num_tokens, int num_topk, int num_experts, cudaStream_t stream); // ============================================================================ // Kernel: Convert sparse topk_weights to dense prob (for combine input) // ============================================================================ // Used for combine backward pass. Unlike dispatch which has explicit topk_idx, // combine uses local_expert_routing_map to determine expert positions. // The order matches dispatch output (dense_to_sparse_prob): experts scanned in order 0,1,2... // NOTE: No topk_idx needed - expert mapping comes from local_expert_routing_map void sparse_to_dense_prob_combine( const float* topk_weights, // [num_tokens, topk] - sparse weights from combine input const bool* local_expert_routing_map, // [num_tokens, experts_per_rank] - from handle float* dense_prob, // [num_tokens, experts_per_node] - output (pre-zeroed) int num_tokens, int num_topk, int experts_per_rank, int experts_per_node, // = experts_per_rank * ranks_per_node int local_rank, // rank within node cudaStream_t stream); // ============================================================================ // Kernel: Convert dense prob output to sparse format (for dispatch output) // ============================================================================ // HT outputs: dense_prob[recv_token][experts_per_node] // HT expects: recv_topk_weights[recv_token][topk], recv_topk_idx[recv_token][topk] // Uses local_expert_routing_map to know which experts each token is routed to. // NOTE: recv_topk_idx contains LOCAL expert indices (0-based within this rank) // to match NCCL API expectations. void dense_to_sparse_prob( const float* dense_prob, // [num_recv_tokens, experts_per_node] - from IPC buffer const bool* local_expert_routing_map, // [num_recv_tokens, experts_per_rank] - from preprocessing float* recv_topk_weights, // [num_recv_tokens, topk] - output sparse weights int64_t* recv_topk_idx, // [num_recv_tokens, topk] - output LOCAL expert indices (0-based) int num_recv_tokens, int topk, int experts_per_rank, int experts_per_node, // = experts_per_rank * ranks_per_node int local_rank, // rank within node cudaStream_t stream); // ============================================================================ // Kernel: Convert dense prob output to sparse format (for combine output) // ============================================================================ // Used for combine backward pass output. Converts kernel's dense output back to // sparse format with GLOBAL expert indices (matching original dispatch input format). // Uses local_routing_map (original tokens → global experts) for the conversion. void dense_to_sparse_prob_combine( const float* dense_prob, // [num_tokens, num_experts] - from kernel output const uint8_t* routing_bitmap, // [num_tokens, ceil(num_experts / 8)] - bitmap routing map float* combined_topk_weights, // [num_tokens, topk] - output sparse weights int64_t* combined_topk_idx, // [num_tokens, topk] - output GLOBAL expert indices (optional, can be nullptr) int num_tokens, int topk, int num_experts, cudaStream_t stream); // ============================================================================ // Switch macros for compile-time template instantiation // Uses constants from configs.cuh where applicable // ============================================================================ // Switch on LSA team size (multiples of 4, up to 32). // Instantiates templates with LSA_TEAM_SIZE, sizing register-file arrays exactly. // MNNVL configurations (NVL72, LSA_TEAM_SIZE > 32) are not yet supported: // the scan kernel in metadata_preprocessing uses warp-reduction and requires // LSA_TEAM_SIZE <= 32. Extend this macro (and the scan kernel) when adding MNNVL support. // // Compile-time filtering: set _NCCL_EP_LSA_TEAM_SIZE_MIN / _NCCL_EP_LSA_TEAM_SIZE_MAX // (via make _NCCL_EP_LSA_TEAM_SIZE_MIN=N or cmake -D_NCCL_EP_LSA_TEAM_SIZE_MIN=N) to // restrict which instantiations are compiled. Reduces build time during development. // When unset, the full set (4–32, multiples of 4) is compiled. // Per-size case helpers — expand to the switch case when N is within the configured // [MIN, MAX] range, empty otherwise. #if must live at file scope (not inside a macro // body). MIN/MAX are always defined by the build system (Makefile/?= 4/32, CMake default 4/32). #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 4 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 4 #define _NCCL_EP_LSA_CASE_4(...) case 4: { constexpr int LSA_TEAM_SIZE = 4; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_4(...) #endif #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 8 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 8 #define _NCCL_EP_LSA_CASE_8(...) case 8: { constexpr int LSA_TEAM_SIZE = 8; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_8(...) #endif #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 12 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 12 #define _NCCL_EP_LSA_CASE_12(...) case 12: { constexpr int LSA_TEAM_SIZE = 12; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_12(...) #endif #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 16 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 16 #define _NCCL_EP_LSA_CASE_16(...) case 16: { constexpr int LSA_TEAM_SIZE = 16; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_16(...) #endif #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 20 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 20 #define _NCCL_EP_LSA_CASE_20(...) case 20: { constexpr int LSA_TEAM_SIZE = 20; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_20(...) #endif #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 24 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 24 #define _NCCL_EP_LSA_CASE_24(...) case 24: { constexpr int LSA_TEAM_SIZE = 24; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_24(...) #endif #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 28 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 28 #define _NCCL_EP_LSA_CASE_28(...) case 28: { constexpr int LSA_TEAM_SIZE = 28; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_28(...) #endif #if _NCCL_EP_LSA_TEAM_SIZE_MIN <= 32 && _NCCL_EP_LSA_TEAM_SIZE_MAX >= 32 #define _NCCL_EP_LSA_CASE_32(...) case 32: { constexpr int LSA_TEAM_SIZE = 32; __VA_ARGS__; } break; #else #define _NCCL_EP_LSA_CASE_32(...) #endif #define HYBRIDEP_SWITCH_LSA_TEAM_SIZE(lsa_val, ...) \ do { switch (lsa_val) { \ _NCCL_EP_LSA_CASE_4(__VA_ARGS__) \ _NCCL_EP_LSA_CASE_8(__VA_ARGS__) \ _NCCL_EP_LSA_CASE_12(__VA_ARGS__) \ _NCCL_EP_LSA_CASE_16(__VA_ARGS__) \ _NCCL_EP_LSA_CASE_20(__VA_ARGS__) \ _NCCL_EP_LSA_CASE_24(__VA_ARGS__) \ _NCCL_EP_LSA_CASE_28(__VA_ARGS__) \ _NCCL_EP_LSA_CASE_32(__VA_ARGS__) \ default: assert(false && "Unsupported LSA team size (must be multiple of 4, " \ "in [_NCCL_EP_LSA_TEAM_SIZE_MIN, _NCCL_EP_LSA_TEAM_SIZE_MAX])"); \ } } while(0) // Switch on number of LSA domains (RDMA peers = nRanks / lsa_team_size). // Each LSA domain is one NVLink/MNNVL clique; domains communicate via RDMA. // // Compile-time filtering: define _NCCL_EP_NUM_LSA_TEAMS_N=1 for each N to include. // Set via make (e.g. _NCCL_EP_NUM_LSA_TEAMS_LIST="1 2") or cmake. // Undefined flags default to 0 (excluded). Full set when unset: 1 2 3 4 8. #if defined(_NCCL_EP_NUM_LSA_TEAMS_1) && _NCCL_EP_NUM_LSA_TEAMS_1 #define _NCCL_EP_NLT_CASE_1(...) case 1: { constexpr int NUM_LSA_TEAMS = 1; __VA_ARGS__; } break; #else #define _NCCL_EP_NLT_CASE_1(...) #endif #if defined(_NCCL_EP_NUM_LSA_TEAMS_2) && _NCCL_EP_NUM_LSA_TEAMS_2 #define _NCCL_EP_NLT_CASE_2(...) case 2: { constexpr int NUM_LSA_TEAMS = 2; __VA_ARGS__; } break; #else #define _NCCL_EP_NLT_CASE_2(...) #endif #if defined(_NCCL_EP_NUM_LSA_TEAMS_3) && _NCCL_EP_NUM_LSA_TEAMS_3 #define _NCCL_EP_NLT_CASE_3(...) case 3: { constexpr int NUM_LSA_TEAMS = 3; __VA_ARGS__; } break; #else #define _NCCL_EP_NLT_CASE_3(...) #endif #if defined(_NCCL_EP_NUM_LSA_TEAMS_4) && _NCCL_EP_NUM_LSA_TEAMS_4 #define _NCCL_EP_NLT_CASE_4(...) case 4: { constexpr int NUM_LSA_TEAMS = 4; __VA_ARGS__; } break; #else #define _NCCL_EP_NLT_CASE_4(...) #endif #if defined(_NCCL_EP_NUM_LSA_TEAMS_8) && _NCCL_EP_NUM_LSA_TEAMS_8 #define _NCCL_EP_NLT_CASE_8(...) case 8: { constexpr int NUM_LSA_TEAMS = 8; __VA_ARGS__; } break; #else #define _NCCL_EP_NLT_CASE_8(...) #endif #define HYBRIDEP_SWITCH_NUM_LSA_TEAMS(num_lsa_domains_val, ...) \ do { switch (num_lsa_domains_val) { \ _NCCL_EP_NLT_CASE_1(__VA_ARGS__) \ _NCCL_EP_NLT_CASE_2(__VA_ARGS__) \ _NCCL_EP_NLT_CASE_3(__VA_ARGS__) \ _NCCL_EP_NLT_CASE_4(__VA_ARGS__) \ _NCCL_EP_NLT_CASE_8(__VA_ARGS__) \ default: assert(false && "Unsupported LSA domain count (must be in " \ "_NCCL_EP_NUM_LSA_TEAMS_LIST, default 1 2 3 4 8)"); \ } } while(0) // Switch on token data type (BF16 = uint16_t, FP8 = uint8_t) // use_fp8: false = BF16 (uint16_t), true = FP8 (uint8_t) #define HYBRIDEP_SWITCH_DATATYPE(use_fp8, ...) \ do { \ if (use_fp8) { \ using TOKEN_DATA_TYPE = uint8_t; \ __VA_ARGS__; \ } else { \ using TOKEN_DATA_TYPE = uint16_t; \ __VA_ARGS__; \ } \ } while(0) // ============================================================================ // Preprocessing wrapper with template parameter resolution // ============================================================================ // Helper to call metadata_preprocessing with resolved template parameters // Note: MAX_NUM_OF_TOKENS_PER_RANK is not used by metadata_preprocessing, // so we use a fixed dummy value for class instantiation. // Note: Caller must provide a pre-allocated scan temp buffer (see get_preprocessing_scan_tmp_size). // Also computes per-expert token counts for NCCL API compatibility when requested. void call_metadata_preprocessing( const uint8_t* global_routing_map, // Already allgathered bitmap routing map int32_t* sparse_to_dense_map, // Output: token→rank→position mapping bool* rdma_to_attn_map, // Output: which tokens come from RDMA bool* attn_to_rdma_map, // Output: which tokens go to RDMA uint8_t* token_rank_mask, // Scratch: cached per-token rank mask [tokens * ranks_per_node * nodes] int32_t* num_tokens_for_experts, // Output: total tokens for local experts bool* local_expert_routing_map, // Output: per-expert routing for local tokens int32_t* per_expert_token_counts, // Optional output: per-expert counts (nullptr to skip) void* scan_tmp, // Pre-allocated temp buffer (from get_preprocessing_scan_tmp_size) int node_rank, // This node's rank (0 to num_nodes-1) int local_rank, // Rank within node (0 to num_ranks_per_node-1) int num_tokens_per_rank, // Actual tokens per rank this iteration (runtime) int hidden_dim, // Model hidden dimension int num_nodes, // Number of nodes (RDMA domain size) int num_ranks_per_node, // Ranks per node (NVLink domain size, 1-8) int experts_per_rank, // Experts per GPU cudaStream_t stream); // Returns required size in bytes for the scan temp buffer used by call_metadata_preprocessing. // Caller must allocate at least this many bytes and pass the pointer to call_metadata_preprocessing. size_t get_preprocessing_scan_tmp_size(int num_ranks_per_node); // ============================================================================ // Memory region info structs for GIN // All buffers are part of a single large gin_base_ptr buffer // Offsets are relative to gin_base_ptr (stored as size_t for offset calculation) // ============================================================================ struct dispatch_memory_region_info_t { // Offsets relative to gin_base_ptr for RDMA operations size_t attn_input_token_offset; // Offset of token staging buffer from gin_base_ptr size_t attn_input_prob_offset; // Offset of prob staging buffer from gin_base_ptr size_t attn_input_scaling_factor_offset; // Offset of scaling factor staging buffer // Batched RDMA staging (packed layout: token+prob+sf per entry) size_t rdma_send_staging_offset; // Offset of per-destination staging buffer size_t rdma_inter_node_group_packed_offset;// Offset of packed receive buffer (token+prob+sf per entry) size_t bytes_per_entry; // Size of packed entry (token + prob + sf) size_t max_tokens_per_dest; // Max tokens that can be staged per destination // Streaming RDMA signals unsigned signals_tail_base; // Base signal ID for tail tracking (sender -> receiver) // Streaming buffer configuration int num_max_rdma_chunked_send_tokens; // Batch size per RDMA put (default: 6) }; struct combine_memory_region_info_t { // Offsets relative to gin_base_ptr for RDMA operations size_t rdma_intra_node_red_token_offset; // Offset of intra-node reduced token buffer size_t combine_rdma_inter_node_group_token_offset; // Offset of combine rdma token buffer size_t rdma_intra_node_red_prob_offset; // Offset of intra-node reduced prob buffer size_t combine_rdma_inter_node_group_prob_offset; // Offset of combine rdma prob buffer }; // ============================================================================ // Dispatch wrapper with template parameter resolution // ============================================================================ // All parameters needed for dispatch kernel, needed redifine because those in hybrid_ep.cuh are with template on data type. struct DispatchParams { // User inputs int hidden_dim; // Model hidden dimension int experts_per_rank; // Experts per GPU int num_ranks_per_node; // Ranks per node (NVLink domain size, 1-8) const void* attn_input_token; const float* attn_input_prob; // Forward dispatch only const float* attn_input_scaling_factor; // FP8 only // IPC-mapped output buffers (from ep_group) // Pointer tables are expected to be device-resident (d_* arrays). void* const* expert_output_token_ptrs; // Array[num_ranks_per_node] float* const* expert_output_prob_ptrs; // Forward only float* const* expert_output_scaling_factor_ptrs; // FP8 only // Metadata (from handle->hybridep preprocessing outputs) const bool* rdma_to_attn_map; const bool* attn_to_rdma_map; const int32_t* sparse_to_dense_map; // Sync state (group-level monotonic counters, computed on host) uint64_t expected_rdma_flag_value; uint32_t expected_intra_node_flag_value; uint64_t* rdma_inter_node_group_flags; uint32_t* intra_node_write_completion_flags; // Grid barrier counter for fused device_sync in dispatch tail uint32_t* dispatch_grid_barrier_counter; // GIN context (from ep_group, multi-node only) ncclDevComm_t* dcomms; // Device communicators array ncclWindow_t nccl_window; // Single registered window handle (by value) int num_gin_comms; // Number of GIN communicators int num_ctx_per_comm; // Number of contexts per communicator void* gin_base_ptr; // Base pointer for offset calculations unsigned signals_base; // Base signal ID dispatch_memory_region_info_t mr_info; // Runtime config int local_rank; int node_rank; int num_tokens_per_rank; }; // Call dispatch kernel with runtime template parameter resolution void call_dispatch( const DispatchParams& params, int max_tokens_per_rank, // Max tokens for buffer sizing int num_nodes, // Number of nodes (RDMA domain size) bool use_fp8, // false = BF16 (uint16_t), true = FP8 (uint8_t) bool forward_dispatch, // True for forward, false for backward cudaStream_t stream); // ============================================================================ // Combine wrapper with template parameter resolution // ============================================================================ // All parameters needed for combine kernel struct CombineParams { // IPC-mapped input buffers (expert outputs from MLP) // NOTE: Pass HOST arrays containing device pointers (not device arrays). // These pointers are copied into the kernel param struct for fast __grid_constant__ access. int hidden_dim; // Model hidden dimension int experts_per_rank; // Experts per GPU int num_ranks_per_node; // Ranks per node (NVLink domain size, 1-8) uint16_t* const* expert_input_token_ptrs; // HOST array[num_ranks_per_node], BF16 only float* const* expert_input_prob_ptrs; // HOST array, backward only // User output buffers void* attn_output_token; float* attn_output_prob; // Backward only (dense format for kernel) // RDMA buffers (multi-node only, from ep_group) uint16_t* rdma_intra_node_red_token; float* rdma_intra_node_red_prob; // Backward only const uint16_t* combine_rdma_inter_node_group_token; const float* combine_rdma_inter_node_group_prob; // Backward only // Metadata (from handle->hybridep preprocessing outputs) const int32_t* sparse_to_dense_map; const bool* rdma_to_attn_map; const bool* attn_to_rdma_map; // For multi-node RDMA routing const bool* local_expert_routing_map; // For backward gradient routing // Sync state (group-level monotonic counters, computed on host) uint64_t combine_expected_rdma_flag_value; uint32_t combine_expected_intra_node_flag_value; uint64_t* combine_rdma_inter_node_group_flags; uint32_t* combine_intra_node_write_completion_flags; // GIN context (multi-node only) ncclDevComm_t* dcomms; // Device communicators array ncclWindow_t nccl_window; // Single registered window handle (by value) int num_gin_comms; // Number of GIN communicators int num_ctx_per_comm; // Number of contexts per communicator void* gin_base_ptr; // Base pointer for offset calculations unsigned signals_base; // Base signal ID unsigned combine_signal_offset; // Signal offset for combine operations combine_memory_region_info_t mr_info; // Runtime config int local_rank; int node_rank; int num_tokens_per_rank; // Original token count from dispatch int num_recv_tokens; // Actual received tokens this rank }; // Call combine kernel with runtime template parameter resolution // Note: HT combine doesn't support FP8, only BF16 void call_combine( const CombineParams& params, int max_tokens_per_rank, // Max tokens for buffer sizing int num_nodes, // Number of nodes (RDMA domain size) bool backward_combine, // True for backward (training), false for forward cudaStream_t stream); } // namespace hybridep } // namespace nccl_ep