diff --git a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh index a7f03548aa0..de4ba481ca7 100644 --- a/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh +++ b/csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh @@ -154,7 +154,7 @@ __global__ void buildMinLatencyActiveExpertMapsKernel( bool const smart_routing, int const cluster_rank, int const cluster_size, int const num_experts_smem) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif // Use one block to process the min latency case int tid = threadIdx.x; @@ -247,7 +247,7 @@ __global__ void buildMinLatencyActiveExpertMapsKernel( } } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -309,7 +309,7 @@ __global__ void fusedBuildExpertMapsSortFirstTokenKernel( // Wait PDL before reading token_selected_experts #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif // build expert map @@ -350,7 +350,7 @@ __global__ void fusedBuildExpertMapsSortFirstTokenKernel( // We are done with compute, launch the dependent kernels while the stores are in flight #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif // write to shared memory and global memory @@ -550,7 +550,7 @@ __global__ void blockExpertPrefixSumKernel(int const* token_selected_experts, int const token_id = block_id * kNumTokensPerBlock + threadIdx.x; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif int expanded_token_id = -1; @@ -579,7 +579,7 @@ __global__ void blockExpertPrefixSumKernel(int const* token_selected_experts, } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -633,7 +633,7 @@ __global__ void globalExpertPrefixSumLargeKernel(int const* blocked_expert_count int cnt = 0; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif // Note: Because of limited registers, cannot store thread-level prefix sum or enable #pragma @@ -662,7 +662,7 @@ __global__ void globalExpertPrefixSumLargeKernel(int const* blocked_expert_count } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -676,7 +676,7 @@ __global__ void globalExpertPrefixSumKernel(int const* blocked_expert_counts, __shared__ typename BlockScan::TempStorage temp_storage; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif int const cnt = threadIdx.x < num_experts_per_node * num_blocks_per_seq @@ -696,7 +696,7 @@ __global__ void globalExpertPrefixSumKernel(int const* blocked_expert_counts, } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -759,7 +759,7 @@ __global__ void mergeExpertPrefixSumKernel(int const* blocked_expert_counts, int const token_id = block_id * blockDim.x + threadIdx.x; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif int const cnt = blocked_expert_counts[target_expert_id * num_blocks_per_seq + block_id]; @@ -774,7 +774,7 @@ __global__ void mergeExpertPrefixSumKernel(int const* blocked_expert_counts, } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -1241,7 +1241,7 @@ __global__ void computeStridesTmaWarpSpecializedKernel( } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif // Both gemms use the same token offset @@ -1274,6 +1274,18 @@ __global__ void computeStridesTmaWarpSpecializedKernel( layout_info2.swap_ab ? gemm_m : gemm2_n, gemm2_k); } + // Skip expensive stride/pointer/SF setup for experts with no assigned tokens. + // All problem shapes (including int4_groupwise) are initialized above so CUTLASS + // can correctly traverse the problem list. The remaining work (alpha scales, + // block scaling factors, strides, pointers) is only needed for active experts. + // For decode (1 token, top_k=8, 128 experts), this skips ~120 of 128 experts. + if (gemm_m == 0) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + if (alpha_scale_flat1 && alpha_scale_flat2) { layout_info1.alpha_scale_ptr_array[expert] = alpha_scale_flat1 + expert; layout_info2.alpha_scale_ptr_array[expert] = alpha_scale_flat2 + expert; @@ -1319,7 +1331,7 @@ __global__ void computeStridesTmaWarpSpecializedKernel( quant_params.groupwise.fc2.weight_scales), bias2, gemm2_output, router_scales, permuted_row_to_unpermuted_row, expert); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -1386,7 +1398,7 @@ __global__ void expandInputRowsKernel( "of the expansion"); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif constexpr int VecSize = is_nvfp4 ? TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize @@ -1508,49 +1520,15 @@ __global__ void expandInputRowsKernel( } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif - // Pad zeros in the extra SFs along the N dimension, we do this to ensure there are no nan values - // in the padded SF atom - if constexpr (is_nvfp4 || is_mxfp8) { - int64_t const start_offset = threadIdx.x; - int64_t const stride = EXPAND_THREADS_PER_BLOCK; - // Use VecSize per thread since we are just writing out zeros so every thread can process a - // whole vector - int64_t const padded_num_elems_in_col = padded_hidden_size / VecSize; - assert(padded_hidden_size % VecSize == 0); - - constexpr int min_num_tokens_alignment = - is_nvfp4 ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4 - : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX; - static_assert((min_num_tokens_alignment & (min_num_tokens_alignment - 1)) == 0, - "Min num tokens alignment must be a power of two"); - // Since we don't know a priori how much padding is needed we assume the max per expert - // NOTE: we don't use (min_num_tokens_alignment-1) to be able to do power of two divisions - int64_t num_padding_tokens = min_num_tokens_alignment * num_experts_per_node; - - for (int64_t padding_token = blockIdx.x; padding_token < num_padding_tokens; - padding_token += gridDim.x) { - int64_t expert = padding_token / min_num_tokens_alignment; - int64_t num_tokens_before_expert = expert_first_token_offset[expert]; - int64_t num_tokens_after_expert = expert_first_token_offset[expert + 1]; - int64_t tokens_to_expert = num_tokens_after_expert - num_tokens_before_expert; - int64_t padding_to_expert = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( - tokens_to_expert, min_num_tokens_alignment) - - tokens_to_expert; - int64_t expert_pad_idx = padding_token % min_num_tokens_alignment; - if (expert_pad_idx < padding_to_expert) { - for (int64_t elem_index = start_offset; elem_index < padded_num_elems_in_col; - elem_index += stride) { - writeSF(num_tokens_before_expert, expert, /*source_row*/ -1, - num_tokens_after_expert + expert_pad_idx, elem_index, - padded_hidden_size, fc1_act_sf_flat, - /* input_sf */ nullptr); // Pass nulltpr input_sf so we write 0 - } - } - } - } + // N-dim SF padding (zeroing extra token rows beyond tokens_to_expert up to MinNDimAlignment) + // is intentionally omitted. The CUTLASS grouped GEMM sets gemm_m = tokens_to_expert per expert + // and never reads scale factors for rows beyond that. The N-dim padding rows don't correspond + // to any valid MMA tiles, so their content doesn't affect correctness. + // K-dim SF padding (above, inside the per-token loop) is still required because MMA tiles may + // straddle the inter_size boundary within valid rows. } template @@ -1568,18 +1546,13 @@ void expandInputRowsKernelLauncher( (std::is_same_v && fc1_act_sf_flat) || !use_per_expert_act_scale, "Per-expert act scale for FC1 is only supported for NVFP4 activations"); - constexpr int64_t min_num_tokens_alignment = - std::is_same_v - ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4 - : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX; - int64_t num_padding_tokens = min_num_tokens_alignment * num_experts_per_node; -#else - int64_t num_padding_tokens = 0; #endif static int64_t const smCount = tensorrt_llm::common::getMultiProcessorCount(); // Note: Launching 8 blocks per SM can fully leverage the memory bandwidth (tested on B200). - int64_t const blocks = std::min(smCount * 8, std::max(num_rows * k, num_padding_tokens)); + // N-dim SF padding has been removed (CUTLASS grouped GEMM never reads beyond + // tokens_to_expert), so the grid is driven purely by the expanded token count. + int64_t const blocks = std::min(smCount * 8, std::max(num_rows * k, int64_t{1})); int64_t const threads = EXPAND_THREADS_PER_BLOCK; auto func = [&]() { @@ -1710,7 +1683,7 @@ __global__ void finalizeMoeRoutingKernel( auto* reduced_row_ptr_v = reinterpret_cast(reduced_row_ptr); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif #pragma unroll @@ -1746,7 +1719,7 @@ __global__ void finalizeMoeRoutingKernel( reduced_row_ptr_v[elem_index] = output_elem; } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -1766,7 +1739,7 @@ __global__ void finalizeMoeRoutingNoFillingKernel( assert(unpadded_cols <= padded_cols); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif int64_t const num_valid_tokens = expert_first_token_offset[num_experts_per_node]; @@ -1849,7 +1822,7 @@ __global__ void finalizeMoeRoutingNoFillingKernel( } } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } @@ -2058,13 +2031,12 @@ void doGatedActivation(ActivationOutputType* output, GemmOutputType const* gemm_ template -__global__ void doActivationKernel(T* output, GemmOutputType const* gemm_result, - float const* fp8_quant, ScaleBiasType const* bias_ptr, - bool bias_is_broadcast, int64_t const* expert_first_token_offset, - int num_experts_per_node, int64_t inter_size, - float const* fc2_act_global_scale, bool use_per_expert_act_scale, - TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, - ActivationParams activation_params) { +__global__ __launch_bounds__(ACTIVATION_THREADS_PER_BLOCK) void doActivationKernel( + T* output, GemmOutputType const* gemm_result, float const* fp8_quant, + ScaleBiasType const* bias_ptr, bool bias_is_broadcast, int64_t const* expert_first_token_offset, + int num_experts_per_node, int64_t inter_size, float const* fc2_act_global_scale, + bool use_per_expert_act_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, + ActivationParams activation_params) { #ifdef ENABLE_FP4 constexpr bool IsNVFP4 = std::is_same_v && @@ -2101,7 +2073,7 @@ __global__ void doActivationKernel(T* output, GemmOutputType const* gemm_result, int64_t const num_valid_tokens = expert_first_token_offset[num_experts_per_node]; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif for (int64_t token = blockIdx.x; token < num_valid_tokens; token += gridDim.x) { size_t gemm_result_offset = token * inter_size * gated_size_mul; @@ -2216,52 +2188,15 @@ __global__ void doActivationKernel(T* output, GemmOutputType const* gemm_result, } #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif - // Pad zeros in the extra SFs along the N dimension, we do this to ensure there are no nan values - // in the padded SF atom - if constexpr (IsNVFP4 || IsMXFP8) { - int64_t const start_offset = threadIdx.x; - int64_t const stride = ACTIVATION_THREADS_PER_BLOCK; - // Use VecSize per thread since we are just writing out zeros so every thread can process a - // whole vector - int64_t const padded_num_elems_in_col = padded_inter_size / VecSize; - assert(padded_inter_size % VecSize == 0); - - constexpr int64_t min_num_tokens_alignment = - IsNVFP4 ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4 - : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX; - static_assert((min_num_tokens_alignment & (min_num_tokens_alignment - 1)) == 0, - "Min num tokens alignment must be a power of two"); - // Since we don't know a priori how much padding is needed we assume the max per expert - // NOTE: we don't (min_num_tokens_alignment-1) to have power of two divisions - int64_t num_padding_tokens = min_num_tokens_alignment * num_experts_per_node; - - for (int64_t padding_token = blockIdx.x; padding_token < num_padding_tokens; - padding_token += gridDim.x) { - int64_t expert = padding_token / min_num_tokens_alignment; - int64_t num_tokens_before_expert = expert_first_token_offset[expert]; - int64_t num_tokens_after_expert = expert_first_token_offset[expert + 1]; - int64_t tokens_to_expert = num_tokens_after_expert - num_tokens_before_expert; - int64_t padding_to_expert = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( - tokens_to_expert, min_num_tokens_alignment) - - tokens_to_expert; - int64_t expert_pad_idx = padding_token % min_num_tokens_alignment; - if (expert_pad_idx < padding_to_expert) { - for (int64_t elem_index = start_offset; elem_index < padded_num_elems_in_col; - elem_index += stride) { - // The SF buffer is padded to a multiple of MinNDimAlignment for each expert - // This means we can safely write to offset num_tokens_after_expert + padded_token, since - // the next expert will leave space for the padding - writeSF(num_tokens_before_expert, expert, /*source_row*/ -1, - num_tokens_after_expert + expert_pad_idx, elem_index, - padded_inter_size, fc2_act_sf_flat, - /* input_sf */ nullptr); // Pass nulltpr input_sf so we write 0 - } - } - } - } + // N-dim SF padding (zeroing extra token rows beyond tokens_to_expert up to MinNDimAlignment) + // is intentionally omitted. The CUTLASS grouped GEMM sets gemm_m = tokens_to_expert per expert + // and never reads scale factors for rows beyond that. The N-dim padding rows don't correspond + // to any valid MMA tiles, so their content doesn't affect correctness. + // K-dim SF padding (above, inside the per-token loop) is still required because MMA tiles may + // straddle the inter_size boundary within valid rows. } template @@ -2272,18 +2207,11 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 QuantParams const& quant_params, bool use_per_expert_act_scale, TmaWarpSpecializedGroupedGemmInput::ElementSF* fc2_act_sf_flat, bool enable_pdl, cudaStream_t stream) { -#ifdef ENABLE_FP4 - constexpr int64_t min_num_tokens_alignment = - std::is_same_v ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4 - : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX; - int64_t num_padding_tokens = min_num_tokens_alignment * num_experts_per_node; -#else - int64_t num_padding_tokens = 0; -#endif - static int64_t const smCount = tensorrt_llm::common::getMultiProcessorCount(); // Note: Launching 8 blocks per SM can fully leverage the memory bandwidth (tested on B200). - int64_t const blocks = std::min(smCount * 8, std::max(expanded_num_tokens, num_padding_tokens)); + // N-dim SF padding has been removed (CUTLASS grouped GEMM never reads beyond + // tokens_to_expert), so the grid is driven purely by the expanded token count. + int64_t const blocks = std::min(smCount * 8, std::max(expanded_num_tokens, int64_t{1})); int64_t const threads = ACTIVATION_THREADS_PER_BLOCK; auto fn = [&]() { @@ -3952,7 +3880,10 @@ CutlassMoeFCRunner padding. # Exercises the weight_scale_vec_size snap fix (issue #2847). (288, 128), + # Non-aligned hidden AND intermediate — exercises K-dim padding in BOTH + # expandInputRows (FC1 SFs, hidden→padded_hidden) and doActivation (FC2 SFs, + # inter→padded_inter). + (288, 192), # hidden 288→384, inter 192→256: K-dim padding in both kernels + (160, 192), # hidden 160→256, inter 192→256: K-dim padding in both kernels + (320, 160), # hidden 320→384, inter 160→256: K-dim padding in both kernels + # Aligned hidden, non-aligned intermediate — only doActivation K-dim padding + ( + 256, + 192, + ), # hidden 256→256 (no pad), inter 192→256: K-dim padding in doActivation only ], ) @pytest.mark.parametrize("num_experts", [2]) @@ -2091,5 +2102,303 @@ def round_up(x, y): # triggers the weight_scale_vec_size bug cannot occur. +# ============================================================================ +# Tests for N-dim SF padding removal safety +# ============================================================================ +# The N-dim SF padding (zeroing extra token rows beyond tokens_to_expert up to +# MinNDimAlignment) was removed because CUTLASS grouped GEMM sets gemm_m = +# tokens_to_expert per expert and never reads scale factors for rows beyond +# that. These tests exercise configurations where empty experts and uninitialized +# SF padding rows could cause incorrect results if the GEMM did read them. +# +# Key configurations that stress-test the removal: +# - num_experts >> top_k: many empty experts with uninitialized SF regions +# - Various batch sizes: different amounts of N-dim padding per expert +# - Large hidden/intermediate sizes: more SF buffer area at risk +# - Non-aligned intermediate sizes: confirms K-dim padding (still present) works + +NDIM_PADDING_BATCH_SIZES = [1, 4, 8] +NDIM_PADDING_HIDDEN_SIZES = [2048] +NDIM_PADDING_NUM_EXPERTS = [128] +NDIM_PADDING_TOP_K = [8] +NDIM_PADDING_INTERMEDIATE_SIZES = [768, 1024] + + +@pytest.mark.parametrize("batch_size", NDIM_PADDING_BATCH_SIZES) +@pytest.mark.parametrize("hidden_size", NDIM_PADDING_HIDDEN_SIZES) +@pytest.mark.parametrize("num_experts", NDIM_PADDING_NUM_EXPERTS) +@pytest.mark.parametrize("top_k", NDIM_PADDING_TOP_K) +@pytest.mark.parametrize("intermediate_size", NDIM_PADDING_INTERMEDIATE_SIZES) +@pytest.mark.parametrize("quantized_input", [False, True]) +@pytest.mark.skipif( + torch.cuda.get_device_capability()[0] not in [10, 11, 12], + reason="NVFP4 is only supported on SM100, SM110 and SM120/SM121", +) +def test_moe_nvfp4_ndim_padding_safety( + batch_size, + hidden_size, + num_experts, + top_k, + intermediate_size, + quantized_input, +): + """Test that N-dim SF padding removal is safe with many empty experts. + + With num_experts=128 and top_k=8, 120 experts have no tokens. Their SF + buffer regions contain uninitialized data. This test verifies the CUTLASS + grouped GEMM produces correct results despite those uninitialized regions. + """ + if top_k > num_experts: + pytest.skip( + f"top_k ({top_k}) cannot be greater than num_experts ({num_experts})" + ) + + torch.manual_seed(42) + otype = torch.bfloat16 + quant_blocksize = 16 + round_up = lambda x, y: (x + y - 1) // y * y + e = num_experts + m = batch_size + n = intermediate_size + k = hidden_size + + w1_n = 2 * n # Swiglu gated + w1 = torch.randn((e, w1_n, k), device="cuda", dtype=otype) / 10 + + sf_w1_2n = round_up(w1_n, 128) + sf_w1_k = round_up(k // quant_blocksize, 4) + w1_blockscale = torch.empty( + (e, sf_w1_2n, sf_w1_k), device="cuda", dtype=torch.float8_e4m3fn + ) + + w2 = torch.randn((e, k, n), device="cuda", dtype=otype) / 10 + sf_w2_k = round_up(k, 128) + sf_w2_n = round_up(n // quant_blocksize, 4) + w2_blockscale = torch.empty( + (e, sf_w2_k, sf_w2_n), device="cuda", dtype=torch.float8_e4m3fn + ) + w1_q = torch.empty((e, w1_n, k // 2), device="cuda", dtype=torch.uint8) + w2_q = torch.empty((e, k, n // 2), device="cuda", dtype=torch.uint8) + w1_gs = torch.empty((e,), device="cuda", dtype=torch.float32) + w2_gs = torch.empty((e,), device="cuda", dtype=torch.float32) + + for expert in range(e): + w1_amax = torch.abs(w1).max().to(torch.float32) + w2_amax = torch.abs(w2).max().to(torch.float32) + w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax + w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax + w1_q[expert], w1_blockscale[expert] = fp4_quantize(w1[expert], w1_gs[expert]) + w2_q[expert], w2_blockscale[expert] = fp4_quantize(w2[expert], w2_gs[expert]) + + x = torch.randn(m, k, dtype=otype).cuda() + a1_gs = torch.tensor(1.0, device="cuda", dtype=torch.float32) + a2_gs = torch.tensor(1.0, device="cuda", dtype=torch.float32) + router_logits = torch.randn(m, e, dtype=otype).cuda() + routing_weights, selected_experts = compute_routing(router_logits, top_k) + + flash_output = torch.zeros_like(x) + + quant_scales = [ + a1_gs, + w1_blockscale.view(torch.int32), + 1.0 / (a1_gs * w1_gs), + a2_gs, + w2_blockscale.view(torch.int32), + 1.0 / (a2_gs * w2_gs), + ] + hidden_states = x + input_sf = None + if quantized_input: + hidden_states, input_sf = fp4_quantize(x, a1_gs) + + _ = fused_moe.cutlass_fused_moe( + hidden_states, + selected_experts.to(torch.int), + routing_weights, + w1_q.contiguous().view(torch.long), + w2_q.contiguous().view(torch.long), + otype, + quant_scales=quant_scales, + input_sf=input_sf, + output=flash_output, + ) + + # Reference: dequantize and compute in high precision + a_fp4, a_scale_interleaved = fp4_quantize(x, a1_gs) + a_in_dtype = dequantize_nvfp4_to_dtype( + a_fp4, + a_scale_interleaved, + a1_gs, + dtype=otype, + device=x.device, + block_size=quant_blocksize, + ) + + w1_d = torch.empty((e, w1_n, k), device="cuda", dtype=otype) + w2_d = torch.empty((e, k, n), device="cuda", dtype=otype) + + for idx in range(0, e): + w1_d[idx] = dequantize_nvfp4_to_dtype( + w1_q[idx], + w1_blockscale[idx], + w1_gs[idx], + dtype=w1.dtype, + device=w1.device, + block_size=quant_blocksize, + ) + w2_d[idx] = dequantize_nvfp4_to_dtype( + w2_q[idx], + w2_blockscale[idx], + w2_gs[idx], + dtype=w2.dtype, + device=w2.device, + block_size=quant_blocksize, + ) + + ref_output = torch_moe_nvfp4( + a_in_dtype, + w1_d, + w2_d, + top_k, + routing_weights, + selected_experts, + ActivationType.Swiglu, + ) + # Two-tier tolerance for FP4 at larger K dimensions (2048 vs 128 in existing tests): + # 1. Tight: >=95% of elements within atol=0.5 (baseline on SM120 is ~98%+). + # If N-dim padding corruption occurs, this drops dramatically. + # 2. Relaxed: 100% within atol=2.0. Catches catastrophic NaN/corruption. + abs_diff = (ref_output - flash_output).abs() + tight_match_rate = (abs_diff <= 0.5).float().mean().item() + assert tight_match_rate >= 0.95, ( + f"Only {tight_match_rate * 100:.1f}% of elements within tight tolerance (0.5). " + f"Expected >=95%." + ) + assert abs_diff.max().item() <= 2.0, ( + f"Max absolute difference {abs_diff.max().item():.4f} exceeds relaxed tolerance (2.0)." + ) + + +@pytest.mark.parametrize("batch_size", NDIM_PADDING_BATCH_SIZES) +@pytest.mark.parametrize("hidden_size", NDIM_PADDING_HIDDEN_SIZES) +@pytest.mark.parametrize("num_experts", NDIM_PADDING_NUM_EXPERTS) +@pytest.mark.parametrize("top_k", NDIM_PADDING_TOP_K) +@pytest.mark.parametrize("intermediate_size", NDIM_PADDING_INTERMEDIATE_SIZES) +@pytest.mark.skipif( + torch.cuda.get_device_capability()[0] not in [10, 11, 12], + reason="MXFP8xMXFP4 is only supported on SM100, SM110 and SM120/SM121", +) +def test_moe_mxfp8_mxfp4_ndim_padding_safety( + batch_size, + hidden_size, + num_experts, + top_k, + intermediate_size, +): + """Test that N-dim SF padding removal is safe for MXFP8xMXFP4 with many empty experts. + + Same rationale as test_moe_nvfp4_ndim_padding_safety but for the MXFP8 activation + + MXFP4 weight path, which also had N-dim SF padding that was removed. + """ + if top_k > num_experts: + pytest.skip( + f"top_k ({top_k}) cannot be greater than num_experts ({num_experts})" + ) + + torch.manual_seed(42) + otype = torch.bfloat16 + e = num_experts + m = batch_size + n = intermediate_size + k = hidden_size + + x = torch.randn(m, k, dtype=otype).cuda() + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=otype) / 10 + w2 = torch.randn((e, k, n), device="cuda", dtype=otype) / 10 + + mxfp8_x, mxfp8_x_sf = mxfp8_quantize(x, True, 32) + + mxfp4_w1, mxfp4_w1_scale = quant_mxfp4_batches(w1, e) + mxfp4_w2, mxfp4_w2_scale = quant_mxfp4_batches(w2, e) + + router_logits = torch.randn(m, e, dtype=otype).cuda() + routing_weights, selected_experts = compute_routing(router_logits, top_k) + + fake_input_scale = torch.ones(e, device=x.device) + + quant_scales = [ + mxfp4_w1_scale.view(torch.int32), + fake_input_scale, + mxfp4_w2_scale.view(torch.int32), + fake_input_scale, + ] + + flash_output = torch.zeros_like(x) + + _ = fused_moe.cutlass_fused_moe( + mxfp8_x, + selected_experts.to(torch.int), + routing_weights, + mxfp4_w1.contiguous().view(torch.long), + mxfp4_w2.contiguous().view(torch.long), + otype, + quant_scales=quant_scales, + input_sf=mxfp8_x_sf, + use_mxfp8_act_scaling=True, + output=flash_output, + ) + + # Reference: dequantize and compute in high precision + dq_mxfp8_x = ( + mxfp8_dequantize_host( + mxfp8_x.cpu().view(torch.uint8), + mxfp8_x_sf.cpu().view(torch.uint8).reshape(-1), + True, + ) + .cuda() + .to(otype) + ) + + dq_mfxp4_w1 = ( + dequant_mxfp4_batches( + mxfp4_w1.cpu().view(torch.uint8), + mxfp4_w1_scale.cpu().view(torch.uint8).reshape(-1), + ) + .cuda() + .to(otype) + ) + + dq_mfxp4_w2 = ( + dequant_mxfp4_batches( + mxfp4_w2.cpu().view(torch.uint8), + mxfp4_w2_scale.cpu().view(torch.uint8).reshape(-1), + ) + .cuda() + .to(otype) + ) + + ref_output = compute_with_experts( + e, + dq_mxfp8_x, + dq_mfxp4_w1, + dq_mfxp4_w2, + selected_experts, + routing_weights, + ) + + # Two-tier tolerance — MXFP8×MXFP4 has significantly higher error than NVFP4 due to + # two levels of block scaling. Baseline on SM120 is ~76-85% at atol=0.5. + # 1. Tight: >=95% within atol=1.0. 2. Relaxed: 100% within atol=3.0. + abs_diff = (ref_output - flash_output).abs() + tight_match_rate = (abs_diff <= 1.0).float().mean().item() + assert tight_match_rate >= 0.95, ( + f"Only {tight_match_rate * 100:.1f}% of elements within tight tolerance (1.0). " + f"Expected >=95%." + ) + assert abs_diff.max().item() <= 3.0, ( + f"Max absolute difference {abs_diff.max().item():.4f} exceeds relaxed tolerance (3.0)." + ) + + if __name__ == "__main__": pytest.main([__file__, "-v"])