diff --git a/sonicmoe/ernie_compat/deepep_metadata.py b/sonicmoe/ernie_compat/deepep_metadata.py index 22f65e6..428cdd7 100644 --- a/sonicmoe/ernie_compat/deepep_metadata.py +++ b/sonicmoe/ernie_compat/deepep_metadata.py @@ -44,6 +44,7 @@ from .deepep_topk_metadata_cuda import ( deepep_topk_metadata_cuda, deepep_topk_metadata_cuda_with_scales, + deepep_topk_metadata_cuda_with_scales_and_gated_outputs, deepep_topk_metadata_cuda_with_scales_rowpack, deepep_topk_metadata_cuda_with_scales_scatterpack, ) @@ -123,11 +124,14 @@ def get_or_alloc( needed = math.prod(shape) if shape else 1 t = self._bufs.get(name) + # In Paddle compatibility mode Tensor.numel() is a device scalar. Shape + # metadata keeps this allocator path free of accidental GPU->host syncs. + capacity = math.prod(t.shape) if t is not None else 0 - if t is not None and t.dtype == dtype and str(t.device) == dev_s and t.numel() >= needed: + if t is not None and t.dtype == dtype and str(t.device) == dev_s and capacity >= needed: view = t.view(-1)[:needed].view(shape) else: - alloc = max(needed, int(t.numel() * 1.25) if t is not None else needed) + alloc = max(needed, int(capacity * 1.25) if t is not None else needed) t = torch.empty(alloc, dtype=dtype, device=device) self._bufs[name] = t view = t[:needed].view(shape) @@ -342,12 +346,16 @@ def deepep_topk_to_sonic_metadata_with_scales( dispatched_probs: torch.Tensor, # [N_recv, topk] float32 tokens_per_expert: Sequence[int] | torch.Tensor, # [E] E: int, - raw_scales: torch.Tensor, # [N_recv, ceil(cols/32)] int32/uint8 + raw_scales: torch.Tensor, # raw bytes or compact int32 DeepEP carrier cols: int, device: str | torch.device = "cuda", block: int = 128, pack_scales_in_scatter: bool = False, pack_scales_rowmajor: bool = False, + gated_output_prototype: torch.Tensor | None = None, + gated_n: int | None = None, + gated_preact_bf16: bool = False, + gated_allocate_z_scale: bool = True, ): """Topk metadata conversion plus optional Sonic FP8 scale packing. @@ -381,11 +389,42 @@ def deepep_topk_to_sonic_metadata_with_scales( f"raw_scales row mismatch: expected {N_recv}, got {raw_scales_2d.shape[0]}" ) expected_scale_cols = (int(cols) + 31) // 32 - if raw_scales_2d.shape[1] != expected_scale_cols: + compact_scale_cols = ( + expected_scale_cols // 4 if expected_scale_cols % 4 == 0 else None + ) + is_compact_word_carrier = ( + "int32" in str(raw_scales_2d.dtype) + and compact_scale_cols is not None + and raw_scales_2d.shape[1] == compact_scale_cols + ) + if ( + raw_scales_2d.shape[1] != expected_scale_cols + and not is_compact_word_carrier + ): + compact_hint = ( + f" or compact int32 [{N_recv}, {compact_scale_cols}]" + if compact_scale_cols is not None + else "" + ) + raise ValueError( + f"raw_scales shape mismatch: expected [{N_recv}, " + f"{expected_scale_cols}]{compact_hint} for cols={cols}, got " + f"{tuple(raw_scales_2d.shape)}/{raw_scales_2d.dtype}" + ) + if (gated_output_prototype is None) != (gated_n is None): raise ValueError( - f"raw_scales col mismatch: expected {expected_scale_cols} for cols={cols}, " - f"got {raw_scales_2d.shape[1]}" + "gated_output_prototype and gated_n must be provided together" ) + if gated_output_prototype is not None: + if "float8_e4m3" not in str(gated_output_prototype.dtype): + raise ValueError( + "gated_output_prototype must use float8_e4m3, got " + f"{gated_output_prototype.dtype}" + ) + if int(gated_n) <= 0 or int(gated_n) % 256 != 0: + raise ValueError( + f"gated_n must be positive and divisible by 256, got {gated_n}" + ) if _HAS_TOPK_CUDA_KERNEL and _HAS_TOPK_CUDA_SCALES_KERNEL: return _deepep_topk_to_sonic_metadata_cuda( @@ -399,6 +438,10 @@ def deepep_topk_to_sonic_metadata_with_scales( cols=int(cols), pack_scales_in_scatter=pack_scales_in_scatter, pack_scales_rowmajor=pack_scales_rowmajor, + gated_output_prototype=gated_output_prototype, + gated_n=gated_n, + gated_preact_bf16=bool(gated_preact_bf16), + gated_allocate_z_scale=bool(gated_allocate_z_scale), ) meta = deepep_topk_to_sonic_metadata( @@ -538,6 +581,10 @@ def _deepep_topk_to_sonic_metadata_cuda( cols: int | None = None, pack_scales_in_scatter: bool = False, pack_scales_rowmajor: bool = False, + gated_output_prototype: torch.Tensor | None = None, + gated_n: int | None = None, + gated_preact_bf16: bool = False, + gated_allocate_z_scale: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]: """CUDA fused topk metadata conversion (warp-ballot, zero argsort). @@ -566,16 +613,13 @@ def _deepep_topk_to_sonic_metadata_cuda( "Each token's topk slots must be unique experts." ) - # Compute TK and TK_padded on host (need tokens_per_expert list) + # DeepEP already returns the counts on host, so use them only to size the + # outputs. The CUDA kernels derive their own expert counts from block_hist; + # no redundant pinned H2D copy is needed on the launch-critical path. if isinstance(tokens_per_expert, torch.Tensor): tpe_list = tokens_per_expert.tolist() - tpe_dev = _TOPK_CACHE.get_or_alloc( - "tpe_dev", (E,), torch.int32, device, zero=False, - ) - tpe_dev.copy_(tokens_per_expert.to(device=device, dtype=torch.int32)) else: tpe_list = list(tokens_per_expert) - tpe_dev = _copy_tpe_h2d_async(tokens_per_expert, device) TK = sum(tpe_list) # Compute TK_padded (padded sum) @@ -592,6 +636,8 @@ def _deepep_topk_to_sonic_metadata_cuda( naept = torch.zeros(N_recv + 1, dtype=torch.int32, device=device) base = (efo, empty_i, empty_i, empty_i, naept, empty_f, 0, 0, N_recv, None) if raw_scales is not None: + if gated_output_prototype is not None: + return (*base, None, None, None, None, None) return (*base, None) return base @@ -615,17 +661,35 @@ def _deepep_topk_to_sonic_metadata_cuda( if raw_scales_2d.shape[0] != 1: raise ValueError(f"raw_scales: expected batch=1, got shape {tuple(raw_scales_2d.shape)}") raw_scales_2d = raw_scales_2d.squeeze(0) - raw_scales_2d = raw_scales_2d.contiguous() - if pack_scales_in_scatter: + if not raw_scales_2d.is_contiguous(): + raw_scales_2d = raw_scales_2d.contiguous() + if gated_output_prototype is not None: + if pack_scales_in_scatter or pack_scales_rowmajor: + raise ValueError( + "preallocated gated outputs require the default scale packer" + ) + metadata_with_scales = ( + deepep_topk_metadata_cuda_with_scales_and_gated_outputs + ) + elif pack_scales_in_scatter: metadata_with_scales = deepep_topk_metadata_cuda_with_scales_scatterpack elif pack_scales_rowmajor: metadata_with_scales = deepep_topk_metadata_cuda_with_scales_rowpack else: metadata_with_scales = deepep_topk_metadata_cuda_with_scales - outputs = metadata_with_scales( - dispatched_indices.contiguous(), - dispatched_probs.contiguous(), - tpe_dev, + indices_2d = ( + dispatched_indices + if dispatched_indices.is_contiguous() + else dispatched_indices.contiguous() + ) + probs_2d = ( + dispatched_probs + if dispatched_probs.is_contiguous() + else dispatched_probs.contiguous() + ) + args = [ + indices_2d, + probs_2d, N_recv, E, topk, @@ -634,13 +698,32 @@ def _deepep_topk_to_sonic_metadata_cuda( block, raw_scales_2d, int(cols), - stream, - ) + ] + if gated_output_prototype is not None: + args.extend( + [ + gated_output_prototype, + int(gated_n), + bool(gated_preact_bf16), + bool(gated_allocate_z_scale), + ] + ) + args.append(stream) + outputs = metadata_with_scales(*args) else: + indices_2d = ( + dispatched_indices + if dispatched_indices.is_contiguous() + else dispatched_indices.contiguous() + ) + probs_2d = ( + dispatched_probs + if dispatched_probs.is_contiguous() + else dispatched_probs.contiguous() + ) outputs = deepep_topk_metadata_cuda( - dispatched_indices.contiguous(), - dispatched_probs.contiguous(), - tpe_dev, + indices_2d, + probs_2d, N_recv, E, topk, @@ -660,6 +743,7 @@ def _deepep_topk_to_sonic_metadata_cuda( score_src_idx, ) = outputs[:7] packed_scales = outputs[7] if len(outputs) > 7 else None + gated_outputs = tuple(outputs[8:12]) if len(outputs) >= 12 else () # topk_scores is TK_padded-sized with zeros at pad positions. # For the topk path, _DownProjection uses T=N_recv and _router_forward @@ -680,7 +764,7 @@ def _deepep_topk_to_sonic_metadata_cuda( score_src_idx, ) if raw_scales is not None: - return (*base, packed_scales) + return (*base, packed_scales, *gated_outputs) return base diff --git a/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/__init__.py b/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/__init__.py index 80dcc5f..01233cd 100644 --- a/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/__init__.py +++ b/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/__init__.py @@ -29,7 +29,6 @@ def deepep_topk_metadata_cuda( dispatched_indices: torch.Tensor, dispatched_probs: torch.Tensor, - tokens_per_expert: torch.Tensor, N_recv: int, E: int, topk: int, @@ -48,7 +47,6 @@ def deepep_topk_metadata_cuda( def deepep_topk_metadata_cuda_with_scales( dispatched_indices: torch.Tensor, dispatched_probs: torch.Tensor, - tokens_per_expert: torch.Tensor, N_recv: int, E: int, topk: int, @@ -61,6 +59,30 @@ def deepep_topk_metadata_cuda_with_scales( ) -> list[torch.Tensor]: ... +@torch.library.custom_op( + f"{LIBRARY_NAME}::deepep_topk_metadata_cuda_with_scales_and_gated_outputs", + mutates_args=(), +) +@cpp_jit() +def deepep_topk_metadata_cuda_with_scales_and_gated_outputs( + dispatched_indices: torch.Tensor, + dispatched_probs: torch.Tensor, + N_recv: int, + E: int, + topk: int, + TK: int, + TK_padded: int, + alignment: int, + raw_scales: torch.Tensor, + cols: int, + gated_output_prototype: torch.Tensor, + gated_n: int, + gated_preact_bf16: bool, + gated_allocate_z_scale: bool, + stream: int, +) -> list[torch.Tensor]: ... + + @torch.library.custom_op( f"{LIBRARY_NAME}::deepep_topk_metadata_cuda_with_scales_scatterpack", mutates_args=(), @@ -69,7 +91,6 @@ def deepep_topk_metadata_cuda_with_scales( def deepep_topk_metadata_cuda_with_scales_scatterpack( dispatched_indices: torch.Tensor, dispatched_probs: torch.Tensor, - tokens_per_expert: torch.Tensor, N_recv: int, E: int, topk: int, @@ -90,7 +111,6 @@ def deepep_topk_metadata_cuda_with_scales_scatterpack( def deepep_topk_metadata_cuda_with_scales_rowpack( dispatched_indices: torch.Tensor, dispatched_probs: torch.Tensor, - tokens_per_expert: torch.Tensor, N_recv: int, E: int, topk: int, diff --git a/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/kernel.cu b/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/kernel.cu index a23be1e..0fd19cf 100644 --- a/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/kernel.cu +++ b/sonicmoe/ernie_compat/deepep_topk_metadata_cuda/kernel.cu @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -302,11 +303,12 @@ void block_offset_scan_kernel( // ============================================================================ __global__ __launch_bounds__(BLOCK_DIM) void prefix_sums_kernel( - const int* __restrict__ tokens_per_expert, // [E] + const int* __restrict__ block_hist, // [E * scatter_blocks] + const int* __restrict__ block_offset, // [E * scatter_blocks] const int* __restrict__ block_naept_sum, // [scatter_blocks] in int* __restrict__ expert_offsets, // [E+1] out int* __restrict__ seg_starts, // [E] out - int* __restrict__ real_bases, // [E] out + int* __restrict__ expert_counts, // [E] out int* __restrict__ block_naept_base, // [scatter_blocks] out: exclusive prefix int* __restrict__ naept, // [N_recv+1] (only naept[0] and naept[N_recv] written here) int N_recv, @@ -320,15 +322,16 @@ void prefix_sums_kernel( // --- B.1: Expert offsets (thread 0, O(E) — typically 8..128) --- if (threadIdx.x == 0) { - int padded_cum = 0, real_cum = 0; + int padded_cum = 0; expert_offsets[0] = 0; for (int e = 0; e < num_experts; e++) { - int count = tokens_per_expert[e]; + const int last = scatter_blocks - 1; + const int idx = e * scatter_blocks + last; + const int count = block_offset[idx] + block_hist[idx]; int padded = (count > 0) ? ((count + alignment - 1) / alignment * alignment) : 0; seg_starts[e] = padded_cum; - real_bases[e] = real_cum; + expert_counts[e] = count; padded_cum += padded; - real_cum += count; expert_offsets[e + 1] = padded_cum; } naept[0] = 0; @@ -406,7 +409,7 @@ void scatter_and_fixup_kernel( const int* __restrict__ block_naept_base, // [scatter_blocks] (Fusion v2) const int* __restrict__ expert_offsets, // [E+1] (padded cumsum) const int* __restrict__ seg_starts, // [E] - const int* __restrict__ tokens_per_expert, // [E] + const int* __restrict__ expert_counts, // [E] int* __restrict__ naept, // [N_recv+1] OUT (we materialize naept[global_row]) int* __restrict__ x_gather_idx, // [TK_padded] output int* __restrict__ s_scatter_idx, // [TK_padded] output @@ -586,7 +589,7 @@ void scatter_and_fixup_kernel( else hi = mid; } const int seg_start = expert_offsets[lo]; - const int real_count = tokens_per_expert[lo]; + const int real_count = expert_counts[lo]; const int local_pos = pos - seg_start; if (local_pos >= real_count) { @@ -601,6 +604,13 @@ void scatter_and_fixup_kernel( } } } + + // Real scores densely cover [0, TK); clear only the padded tail here. + for (int score_pos = TK + global_tid; + score_pos < TK_padded; + score_pos += total_threads) { + topk_scores[score_pos] = 0.0f; + } } // ============================================================================ @@ -764,7 +774,6 @@ void pack_raw_scales_rowmajor_kernel( std::vector deepep_topk_metadata_cuda_impl( torch::Tensor& dispatched_indices, torch::Tensor& dispatched_probs, - torch::Tensor& tokens_per_expert, int64_t N_recv, int64_t E, int64_t topk, @@ -775,46 +784,45 @@ std::vector deepep_topk_metadata_cuda_impl( torch::Tensor* raw_scales, int64_t cols, bool pack_scales_in_scatter, - bool pack_scales_rowmajor) + bool pack_scales_rowmajor, + torch::Tensor* gated_output_prototype, + int64_t gated_n, + bool gated_preact_bf16, + bool gated_allocate_z_scale) { cudaStream_t stream = reinterpret_cast(stream_ptr); auto opt_i = torch::dtype(torch::kInt32).device(torch::kCUDA); auto opt_f = torch::dtype(torch::kFloat32).device(torch::kCUDA); - // ── Output tensors (saved on autograd ctx) — independent storage/call ──── - // Mirror the exact init the previous Python path used: x_gather_idx and - // topk_scores are zero-initialised (pad slots must read 0), the rest are - // fully written by the kernels below so torch::empty is sufficient. + const int scatter_blocks = (N_recv + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK; + const int64_t workspace_elements = + 2 * static_cast(scatter_blocks) * (E + 1) + 2 * E; + + // Custom-op outputs must own independent storage. Undeclared output + // aliasing can let live PP/VPP contexts observe later invocations. torch::Tensor expert_offsets = torch::empty({E + 1}, opt_i); - torch::Tensor x_gather_idx = torch::zeros({TK_padded}, opt_i); + torch::Tensor x_gather_idx = torch::empty({TK_padded}, opt_i); torch::Tensor s_scatter_idx = torch::empty({TK_padded}, opt_i); torch::Tensor s_reverse_scatter_idx = torch::empty({TK}, opt_i); - torch::Tensor topk_scores = torch::zeros({TK_padded}, opt_f); torch::Tensor naept = torch::empty({N_recv + 1}, opt_i); torch::Tensor score_src_idx = torch::empty({TK}, opt_i); - - const int scatter_blocks = (N_recv + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK; - - // ── Pure scratch (NOT returned, NOT ctx-saved) ─────────────────────────── - torch::Tensor seg_starts = torch::empty({E}, opt_i); - torch::Tensor real_bases = torch::empty({E}, opt_i); - // Workspace layout (must match the launcher partition below): - // block_hist[scatter_blocks*E] + block_offset[scatter_blocks*E] - // + block_naept_sum[scatter_blocks] + block_naept_base[scatter_blocks] - // = 2*scatter_blocks*(E+1) int32 slots. torch::Tensor cumsum_workspace = - torch::empty({2 * scatter_blocks * (static_cast(E) + 1)}, opt_i); + torch::empty({workspace_elements}, opt_i); + torch::Tensor topk_scores = torch::empty({TK_padded}, opt_f); // Workspace layout within cumsum_workspace: // [0 .. scatter_blocks*E-1]: block_hist // [scatter_blocks*E .. 2*scatter_blocks*E-1]: block_offset // [2*scatter_blocks*E ..]: block_naept_sum / base + // [2*scatter_blocks*(E+1) ..]: seg_starts / expert_counts int* workspace = cumsum_workspace.data_ptr(); int* block_hist = workspace; int* block_offset = workspace + scatter_blocks * E; int* block_naept_sum = workspace + 2 * scatter_blocks * E; int* block_naept_base = workspace + 2 * scatter_blocks * E + scatter_blocks; + int* seg_starts = workspace + 2 * scatter_blocks * (E + 1); + int* expert_counts = seg_starts + E; // Shared memory for histogram_kernel: expert_bitmask[E] int smem_hist = static_cast(E * sizeof(uint32_t)); @@ -847,35 +855,47 @@ std::vector deepep_topk_metadata_cuda_impl( // ── Kernel 1c: Tail prefix sums (B.1 expert_offsets + scan over per-block sums) ── prefix_sums_kernel<<>>( - tokens_per_expert.data_ptr(), + block_hist, + block_offset, block_naept_sum, expert_offsets.data_ptr(), - seg_starts.data_ptr(), - real_bases.data_ptr(), + seg_starts, + expert_counts, block_naept_base, naept.data_ptr(), static_cast(N_recv), scatter_blocks, static_cast(E), static_cast(alignment)); torch::Tensor packed_scales; + uint8_t* packed_scale_output_bytes = nullptr; int64_t scale_cols = 0; int64_t stride_row = 0; int64_t stride_col = 0; int64_t k_tiles = 0; at::ScalarType raw_dtype = at::ScalarType::Byte; + bool compact_scale_words = false; + const uint8_t* compact_scale_bytes = nullptr; if (raw_scales != nullptr) { TORCH_CHECK(cols > 0, "cols must be positive when raw_scales is provided"); TORCH_CHECK(raw_scales->is_cuda(), "raw_scales must be a CUDA tensor"); - TORCH_CHECK(raw_scales->dim() == 2, "raw_scales must be rank-2 [N_recv, ceil(cols/32)]"); + TORCH_CHECK(raw_scales->dim() == 2, + "raw_scales must be a rank-2 raw scale matrix or compact int32 carrier"); TORCH_CHECK(raw_scales->size(0) == N_recv, "raw_scales row count mismatch: expected ", N_recv, ", got ", raw_scales->size(0)); const int64_t expected_scale_cols = div_up_i64(cols, SF_VEC_SIZE); - scale_cols = raw_scales->size(1); - TORCH_CHECK(scale_cols == expected_scale_cols, + const int64_t storage_scale_cols = raw_scales->size(1); + raw_dtype = raw_scales->scalar_type(); + compact_scale_words = + raw_dtype == at::ScalarType::Int && + expected_scale_cols % 4 == 0 && + storage_scale_cols == expected_scale_cols / 4; + TORCH_CHECK(storage_scale_cols == expected_scale_cols || compact_scale_words, "raw_scales column count mismatch: expected ", expected_scale_cols, - " for cols=", cols, ", got ", scale_cols); + " raw scale elements or ", expected_scale_cols / 4, + " compact int32 words for cols=", cols, + ", got ", storage_scale_cols); auto opt_u8 = torch::dtype(torch::kUInt8).device(torch::kCUDA); const int64_t per_batch_storage = scale_storage_per_batch_i64(TK_padded, cols); @@ -883,16 +903,64 @@ std::vector deepep_topk_metadata_cuda_impl( (TK_padded % SF_TILE_M == 0 && cols % SF_TILE_K == 0) ? torch::empty({1, per_batch_storage}, opt_u8) : torch::full({1, per_batch_storage}, 1, opt_u8); + packed_scale_output_bytes = packed_scales.data_ptr(); k_tiles = div_up_i64(cols, SF_TILE_K); - stride_row = raw_scales->stride(0); - stride_col = raw_scales->stride(1); - raw_dtype = raw_scales->scalar_type(); TORCH_CHECK(raw_dtype == at::ScalarType::Int || raw_dtype == at::ScalarType::Byte || raw_dtype == at::ScalarType::Char, "raw_scales dtype must be int32, uint8, or int8, got ", raw_scales->scalar_type()); + scale_cols = expected_scale_cols; + if (compact_scale_words) { + // DeepEP transports four opaque E8M0 bytes in each int32 word. + stride_row = raw_scales->stride(0) * sizeof(int); + stride_col = 1; + compact_scale_bytes = reinterpret_cast( + raw_scales->data_ptr()); + } else { + stride_row = raw_scales->stride(0); + stride_col = raw_scales->stride(1); + } + } + + // Sink the four hot-path output allocations into this native bridge. + // Each invocation still receives independent storage for autograd safety. + torch::Tensor gated_preact; + torch::Tensor gated_postact; + torch::Tensor gated_z_scales; + torch::Tensor gated_postact_scales; + if (gated_output_prototype != nullptr) { + TORCH_CHECK(raw_scales != nullptr, + "gated output allocation requires the with-scales path"); + TORCH_CHECK(gated_output_prototype->is_cuda(), + "gated_output_prototype must be a CUDA tensor"); + TORCH_CHECK(gated_output_prototype->element_size() == 1, + "gated_output_prototype must use a one-byte FP8 dtype"); + TORCH_CHECK(gated_n > 0 && gated_n % 256 == 0, + "gated_n must be positive and divisible by 256, got ", gated_n); + TORCH_CHECK(TK_padded % SF_TILE_M == 0, + "TK_padded must be divisible by 128, got ", TK_padded); + TORCH_CHECK( + TK_padded <= std::numeric_limits::max() / gated_n, + "gated output shape overflows int64: TK_padded=", TK_padded, + ", gated_n=", gated_n); + + auto fp8_options = gated_output_prototype->options(); + auto preact_options = gated_preact_bf16 + ? torch::dtype(torch::kBFloat16).device(gated_output_prototype->device()) + : fp8_options; + auto u8_options = torch::dtype(torch::kUInt8).device( + gated_output_prototype->device()); + const int64_t postact_n = gated_n / 2; + gated_preact = torch::empty({TK_padded, gated_n}, preact_options); + gated_postact = torch::empty({TK_padded, postact_n}, fp8_options); + gated_z_scales = gated_allocate_z_scale + ? torch::empty({TK_padded, gated_n / SF_VEC_SIZE}, u8_options) + : torch::empty({0}, u8_options); + gated_postact_scales = torch::empty( + {TK_padded / SF_TILE_M, postact_n / SF_TILE_K, SF_TILE_STORAGE}, + u8_options); } // ── Kernel 2: Scatter + fixup ──────────────────────────────────────────── @@ -918,8 +986,8 @@ std::vector deepep_topk_metadata_cuda_impl( block_offset, \ block_naept_base, \ expert_offsets.data_ptr(), \ - seg_starts.data_ptr(), \ - tokens_per_expert.data_ptr(), \ + seg_starts, \ + expert_counts, \ naept.data_ptr(), \ x_gather_idx.data_ptr(), \ s_scatter_idx.data_ptr(), \ @@ -930,15 +998,15 @@ std::vector deepep_topk_metadata_cuda_impl( static_cast(topk), static_cast(TK), \ static_cast(TK_padded), scatter_blocks); - #define LAUNCH_K2_PACK(TV, RAW_T) \ + #define LAUNCH_K2_PACK(TV, RAW_T, RAW_PTR) \ scatter_and_fixup_kernel<<>>( \ dispatched_indices.data_ptr(), \ dispatched_probs.data_ptr(), \ block_offset, \ block_naept_base, \ expert_offsets.data_ptr(), \ - seg_starts.data_ptr(), \ - tokens_per_expert.data_ptr(), \ + seg_starts, \ + expert_counts, \ naept.data_ptr(), \ x_gather_idx.data_ptr(), \ s_scatter_idx.data_ptr(), \ @@ -948,23 +1016,27 @@ std::vector deepep_topk_metadata_cuda_impl( static_cast(N_recv), static_cast(E), \ static_cast(topk), static_cast(TK), \ static_cast(TK_padded), scatter_blocks, \ - raw_scales->data_ptr(), \ - packed_scales.data_ptr(), \ + RAW_PTR, \ + packed_scale_output_bytes, \ scale_cols, stride_row, stride_col, k_tiles); if (raw_scales != nullptr && pack_scales_in_scatter) { - if (raw_dtype == at::ScalarType::Int) { - if (topk <= 4) { LAUNCH_K2_PACK(4, int); } - else if (topk <= 8) { LAUNCH_K2_PACK(8, int); } - else { LAUNCH_K2_PACK(16, int); } + if (compact_scale_words) { + if (topk <= 4) { LAUNCH_K2_PACK(4, uint8_t, compact_scale_bytes); } + else if (topk <= 8) { LAUNCH_K2_PACK(8, uint8_t, compact_scale_bytes); } + else { LAUNCH_K2_PACK(16, uint8_t, compact_scale_bytes); } + } else if (raw_dtype == at::ScalarType::Int) { + if (topk <= 4) { LAUNCH_K2_PACK(4, int, raw_scales->data_ptr()); } + else if (topk <= 8) { LAUNCH_K2_PACK(8, int, raw_scales->data_ptr()); } + else { LAUNCH_K2_PACK(16, int, raw_scales->data_ptr()); } } else if (raw_dtype == at::ScalarType::Byte) { - if (topk <= 4) { LAUNCH_K2_PACK(4, uint8_t); } - else if (topk <= 8) { LAUNCH_K2_PACK(8, uint8_t); } - else { LAUNCH_K2_PACK(16, uint8_t); } + if (topk <= 4) { LAUNCH_K2_PACK(4, uint8_t, raw_scales->data_ptr()); } + else if (topk <= 8) { LAUNCH_K2_PACK(8, uint8_t, raw_scales->data_ptr()); } + else { LAUNCH_K2_PACK(16, uint8_t, raw_scales->data_ptr()); } } else { - if (topk <= 4) { LAUNCH_K2_PACK(4, int8_t); } - else if (topk <= 8) { LAUNCH_K2_PACK(8, int8_t); } - else { LAUNCH_K2_PACK(16, int8_t); } + if (topk <= 4) { LAUNCH_K2_PACK(4, int8_t, raw_scales->data_ptr()); } + else if (topk <= 8) { LAUNCH_K2_PACK(8, int8_t, raw_scales->data_ptr()); } + else { LAUNCH_K2_PACK(16, int8_t, raw_scales->data_ptr()); } } } else { if (topk <= 4) { LAUNCH_K2(4); } @@ -996,51 +1068,69 @@ std::vector deepep_topk_metadata_cuda_impl( if (use_rowmajor_pack) { dim3 grid_scale(static_cast(row_tiles)); size_t smem_scale = static_cast(SF_TILE_M * scale_cols); - if (raw_dtype == at::ScalarType::Int) { + if (compact_scale_words) { + pack_raw_scales_rowmajor_kernel<<>>( + compact_scale_bytes, + x_gather_idx.data_ptr(), + packed_scale_output_bytes, + TK_padded, scale_cols, stride_row, stride_col, k_tiles); + } else if (raw_dtype == at::ScalarType::Int) { pack_raw_scales_rowmajor_kernel<<>>( raw_scales->data_ptr(), x_gather_idx.data_ptr(), - packed_scales.data_ptr(), + packed_scale_output_bytes, TK_padded, scale_cols, stride_row, stride_col, k_tiles); } else if (raw_dtype == at::ScalarType::Byte) { pack_raw_scales_rowmajor_kernel<<>>( raw_scales->data_ptr(), x_gather_idx.data_ptr(), - packed_scales.data_ptr(), + packed_scale_output_bytes, TK_padded, scale_cols, stride_row, stride_col, k_tiles); } else { pack_raw_scales_rowmajor_kernel<<>>( raw_scales->data_ptr(), x_gather_idx.data_ptr(), - packed_scales.data_ptr(), + packed_scale_output_bytes, TK_padded, scale_cols, stride_row, stride_col, k_tiles); } } else { dim3 grid_scale(static_cast(div_up_i64(row_tiles, SCALE_PACK_WARPS_PER_CTA)), static_cast(k_tiles)); - if (raw_dtype == at::ScalarType::Int) { + if (compact_scale_words) { + pack_raw_scales_from_gather_kernel<<>>( + compact_scale_bytes, + x_gather_idx.data_ptr(), + packed_scale_output_bytes, + TK_padded, scale_cols, stride_row, stride_col, k_tiles); + } else if (raw_dtype == at::ScalarType::Int) { pack_raw_scales_from_gather_kernel<<>>( raw_scales->data_ptr(), x_gather_idx.data_ptr(), - packed_scales.data_ptr(), + packed_scale_output_bytes, TK_padded, scale_cols, stride_row, stride_col, k_tiles); } else if (raw_dtype == at::ScalarType::Byte) { pack_raw_scales_from_gather_kernel<<>>( raw_scales->data_ptr(), x_gather_idx.data_ptr(), - packed_scales.data_ptr(), + packed_scale_output_bytes, TK_padded, scale_cols, stride_row, stride_col, k_tiles); } else { pack_raw_scales_from_gather_kernel<<>>( raw_scales->data_ptr(), x_gather_idx.data_ptr(), - packed_scales.data_ptr(), + packed_scale_output_bytes, TK_padded, scale_cols, stride_row, stride_col, k_tiles); } } } out.push_back(packed_scales); } + if (gated_output_prototype != nullptr) { + out.push_back(gated_preact); + out.push_back(gated_postact); + out.push_back(gated_z_scales); + out.push_back(gated_postact_scales); + } return out; } @@ -1048,7 +1138,6 @@ std::vector deepep_topk_metadata_cuda_impl( std::vector deepep_topk_metadata_cuda( torch::Tensor& dispatched_indices, torch::Tensor& dispatched_probs, - torch::Tensor& tokens_per_expert, int64_t N_recv, int64_t E, int64_t topk, @@ -1058,15 +1147,14 @@ std::vector deepep_topk_metadata_cuda( int64_t stream_ptr) { return deepep_topk_metadata_cuda_impl( - dispatched_indices, dispatched_probs, tokens_per_expert, + dispatched_indices, dispatched_probs, N_recv, E, topk, TK, TK_padded, alignment, stream_ptr, - nullptr, 0, false, false); + nullptr, 0, false, false, nullptr, 0, false, false); } std::vector deepep_topk_metadata_cuda_with_scales( torch::Tensor& dispatched_indices, torch::Tensor& dispatched_probs, - torch::Tensor& tokens_per_expert, int64_t N_recv, int64_t E, int64_t topk, @@ -1078,15 +1166,39 @@ std::vector deepep_topk_metadata_cuda_with_scales( int64_t stream_ptr) { return deepep_topk_metadata_cuda_impl( - dispatched_indices, dispatched_probs, tokens_per_expert, + dispatched_indices, dispatched_probs, + N_recv, E, topk, TK, TK_padded, alignment, stream_ptr, + &raw_scales, cols, false, false, nullptr, 0, false, false); +} + +std::vector +deepep_topk_metadata_cuda_with_scales_and_gated_outputs( + torch::Tensor& dispatched_indices, + torch::Tensor& dispatched_probs, + int64_t N_recv, + int64_t E, + int64_t topk, + int64_t TK, + int64_t TK_padded, + int64_t alignment, + torch::Tensor& raw_scales, + int64_t cols, + torch::Tensor& gated_output_prototype, + int64_t gated_n, + bool gated_preact_bf16, + bool gated_allocate_z_scale, + int64_t stream_ptr) +{ + return deepep_topk_metadata_cuda_impl( + dispatched_indices, dispatched_probs, N_recv, E, topk, TK, TK_padded, alignment, stream_ptr, - &raw_scales, cols, false, false); + &raw_scales, cols, false, false, &gated_output_prototype, gated_n, + gated_preact_bf16, gated_allocate_z_scale); } std::vector deepep_topk_metadata_cuda_with_scales_rowpack( torch::Tensor& dispatched_indices, torch::Tensor& dispatched_probs, - torch::Tensor& tokens_per_expert, int64_t N_recv, int64_t E, int64_t topk, @@ -1098,15 +1210,14 @@ std::vector deepep_topk_metadata_cuda_with_scales_rowpack( int64_t stream_ptr) { return deepep_topk_metadata_cuda_impl( - dispatched_indices, dispatched_probs, tokens_per_expert, + dispatched_indices, dispatched_probs, N_recv, E, topk, TK, TK_padded, alignment, stream_ptr, - &raw_scales, cols, false, true); + &raw_scales, cols, false, true, nullptr, 0, false, false); } std::vector deepep_topk_metadata_cuda_with_scales_scatterpack( torch::Tensor& dispatched_indices, torch::Tensor& dispatched_probs, - torch::Tensor& tokens_per_expert, int64_t N_recv, int64_t E, int64_t topk, @@ -1118,9 +1229,9 @@ std::vector deepep_topk_metadata_cuda_with_scales_scatterpack( int64_t stream_ptr) { return deepep_topk_metadata_cuda_impl( - dispatched_indices, dispatched_probs, tokens_per_expert, + dispatched_indices, dispatched_probs, N_recv, E, topk, TK, TK_padded, alignment, stream_ptr, - &raw_scales, cols, true, false); + &raw_scales, cols, true, false, nullptr, 0, false, false); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -1129,7 +1240,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "DeepEP topk metadata: allocates + returns metadata tensors (CUDA)", py::arg("dispatched_indices"), py::arg("dispatched_probs"), - py::arg("tokens_per_expert"), py::arg("N_recv"), py::arg("E"), py::arg("topk"), @@ -1142,7 +1252,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "DeepEP topk metadata: returns metadata tensors plus packed Sonic FP8 scales (CUDA)", py::arg("dispatched_indices"), py::arg("dispatched_probs"), - py::arg("tokens_per_expert"), py::arg("N_recv"), py::arg("E"), py::arg("topk"), @@ -1152,12 +1261,29 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("raw_scales"), py::arg("cols"), py::arg("stream")); + m.def("deepep_topk_metadata_cuda_with_scales_and_gated_outputs", + &deepep_topk_metadata_cuda_with_scales_and_gated_outputs, + "DeepEP metadata + packed scales + preallocated FP8 gated outputs (CUDA)", + py::arg("dispatched_indices"), + py::arg("dispatched_probs"), + py::arg("N_recv"), + py::arg("E"), + py::arg("topk"), + py::arg("TK"), + py::arg("TK_padded"), + py::arg("alignment"), + py::arg("raw_scales"), + py::arg("cols"), + py::arg("gated_output_prototype"), + py::arg("gated_n"), + py::arg("gated_preact_bf16"), + py::arg("gated_allocate_z_scale"), + py::arg("stream")); m.def("deepep_topk_metadata_cuda_with_scales_scatterpack", &deepep_topk_metadata_cuda_with_scales_scatterpack, "DeepEP topk metadata: packs Sonic FP8 scales inside scatter/fixup (CUDA)", py::arg("dispatched_indices"), py::arg("dispatched_probs"), - py::arg("tokens_per_expert"), py::arg("N_recv"), py::arg("E"), py::arg("topk"), @@ -1172,7 +1298,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "DeepEP topk metadata: row-major-load shared-transpose Sonic FP8 scale pack (CUDA)", py::arg("dispatched_indices"), py::arg("dispatched_probs"), - py::arg("tokens_per_expert"), py::arg("N_recv"), py::arg("E"), py::arg("topk"), diff --git a/sonicmoe/functional/__init__.py b/sonicmoe/functional/__init__.py index d3b3788..ea2dd44 100644 --- a/sonicmoe/functional/__init__.py +++ b/sonicmoe/functional/__init__.py @@ -33,6 +33,7 @@ from quack.gemm_interface import gemm from ..quack_utils.gemm_dgated import gemm_dgated as gemm_dgated_kernel from ..quack_utils.fp8_quack_patch import apply_fp8_quack_patch +from ..quack_utils.gemm_gated import try_gemm_gated_tvm_ffi_sonic_prepared apply_fp8_quack_patch() @@ -75,6 +76,7 @@ _run_cutlass_blockscaled_gemm_varlen_k, _run_cutlass_blockscaled_gemm_varlen_k_accumulate, _run_cutlass_blockscaled_gemm_varlen_k_tma_add, + _scatter_router_scores_i32, colwise_quantize_and_pack, fused_z_save_y1_quant, pack_blockscaled_1x32_scales_fast, @@ -95,6 +97,7 @@ ) _E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", torch.uint8) +_GATED_OUTPUT_PAYLOAD_ATTR = "_sonic_preallocated_gated_outputs" # --------------------------------------------------------------------------- @@ -169,15 +172,28 @@ def _swiglu_backward_clamp_reference( def _is_raw_1x32_scale_layout(scales: torch.Tensor, rows: int, cols: int) -> bool: return scales.ndim == 2 and tuple(scales.shape) == (rows, _div_up(cols, _SF_VEC_SIZE)) + +def _is_compact_1x32_scale_carrier( + scales: torch.Tensor, rows: int, cols: int +) -> bool: + groups = _div_up(cols, _SF_VEC_SIZE) + return ( + groups % 4 == 0 + and str(scales.dtype) in ("torch.int32", "paddle.int32", "int32") + and scales.ndim == 2 + and tuple(scales.shape) == (rows, groups // 4) + ) + + def _ensure_isa_1x32_scales(scales: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + if _is_compact_1x32_scale_carrier(scales, rows, cols): + scales = scales.view(torch.uint8).reshape( + rows, _div_up(cols, _SF_VEC_SIZE) + ) if _is_raw_1x32_scale_layout(scales, rows, cols): return pack_blockscaled_1x32_scales_fast(scales, cols).view(_E8M0_DTYPE) return scales - -def _is_raw_1x32_scale_layout(scales: torch.Tensor, rows: int, cols: int) -> bool: - return scales.ndim == 2 and tuple(scales.shape) == (rows, _div_up(cols, _SF_VEC_SIZE)) - def _gather_1x32_scales_to_isa( scales: torch.Tensor, gather_idx: torch.Tensor, @@ -186,6 +202,10 @@ def _gather_1x32_scales_to_isa( *, fill_value: int = 127, ) -> torch.Tensor: + if _is_compact_1x32_scale_carrier(scales, rows, cols): + scales = scales.view(torch.uint8).reshape( + rows, _div_up(cols, _SF_VEC_SIZE) + ) if _is_raw_1x32_scale_layout(scales, rows, cols): return gather_raw_blockscaled_1x32_scales_to_isa( scales, gather_idx, cols @@ -211,10 +231,20 @@ def _gather_1x32_scales_to_isa( def _is_fp8_e4m3_dtype(dtype) -> bool: return dtype == torch.float8_e4m3fn or str(dtype) == "paddle.float8_e4m3fn" -def _raw_1x32_scale_bytes(scales: torch.Tensor) -> torch.Tensor: +def _raw_1x32_scale_bytes( + scales: torch.Tensor, rows: int | None = None, cols: int | None = None +) -> torch.Tensor: if str(scales.dtype) in ("torch.uint8", "paddle.uint8", "uint8"): return scales if str(scales.dtype) in ("torch.int32", "paddle.int32", "int32"): + if ( + rows is not None + and cols is not None + and _is_compact_1x32_scale_carrier(scales, rows, cols) + ): + return scales.view(torch.uint8).reshape( + rows, _div_up(cols, _SF_VEC_SIZE) + ) return scales.to(torch.uint8) return scales.view(torch.uint8) @@ -256,11 +286,22 @@ def _fused_blockscaled_gated_forward( # Step 1: Quantize at T-size (NOT TK) x_scales_tk_pre = None + gated_output_payload = None if x_fp8_pre is not None: - if len(x_fp8_pre) == 3: - x_fp8, x_scales_t, x_scales_tk_pre = x_fp8_pre - else: - x_fp8, x_scales_t = x_fp8_pre + if len(x_fp8_pre) not in (2, 3, 7): + raise ValueError( + "prequant activation payload must contain 2, 3, or 7 tensors, " + f"got {len(x_fp8_pre)}" + ) + x_fp8, x_scales_t = x_fp8_pre[:2] + if len(x_fp8_pre) >= 3: + x_scales_tk_pre = x_fp8_pre[2] + if len(x_fp8_pre) == 7: + gated_output_payload = x_fp8_pre[3:7] + elif x_scales_tk_pre is not None: + gated_output_payload = getattr( + x_scales_tk_pre, _GATED_OUTPUT_PAYLOAD_ATTR, None + ) _PREQUANT_HIT_COUNT["activation_fwd"] += 1 else: x_fp8, x_scales_t = quantize_and_pack_activation(x) @@ -269,7 +310,11 @@ def _fused_blockscaled_gated_forward( TK = x_gather_idx.shape[0] K = x.shape[1] if x_scales_tk_pre is not None: - x_scales_tk_e8m0 = x_scales_tk_pre.view(_E8M0_DTYPE) + x_scales_tk_e8m0 = ( + x_scales_tk_pre + if x_scales_tk_pre.dtype == _E8M0_DTYPE + else x_scales_tk_pre.view(_E8M0_DTYPE) + ) else: x_scales_tk_e8m0 = _gather_1x32_scales_to_isa( x_scales_t, x_gather_idx, int(x_fp8.shape[0]), K @@ -285,9 +330,16 @@ def _fused_blockscaled_gated_forward( # and halves D bandwidth (192MB fp8 vs 384MB bf16). cfg = fp8_config if fp8_config is not None else _get_fp8_config() epilogue_quant = cfg.epilogue_quant and cfg.save_z_fp8 + N = w1.shape[1] # pretransposed weight is (E, 2I, H) if epilogue_quant: - N = w1.shape[1] # (2I, H, E) -> w1.shape[0] = 2I - z_scale_out = torch.empty(TK, N // 32, dtype=torch.uint8, device=x.device) + z_scale_out = ( + gated_output_payload[2] + if gated_output_payload is not None + # Paddle's Tensor.numel() returns a device scalar; use shape metadata + # here so output selection cannot synchronize the launch stream. + and all(dim != 0 for dim in gated_output_payload[2].shape) + else torch.empty(TK, N // 32, dtype=torch.uint8, device=x.device) + ) else: z_scale_out = None @@ -308,7 +360,7 @@ def _fused_blockscaled_gated_forward( ) postact_scale_out = torch.empty( (TK // 128, I_dim // 128, 512), dtype=torch.uint8, device=x.device - ) + ) if gated_output_payload is None else gated_output_payload[3] postact_dtype = torch.float8_e4m3fn else: postact_scale_out = None @@ -322,26 +374,85 @@ def _fused_blockscaled_gated_forward( # tensors in the autograd chain which cause illegal memory access in # backward at large shapes. z_out_dtype = torch.float8_e4m3fn if epilogue_quant else torch.bfloat16 + preact_out = None + postact_out = None + if gated_output_payload is not None: + if not fuse_y1: + raise ValueError( + "preallocated gated outputs require fused y1 epilogue quantization" + ) + preact_out, postact_out = gated_output_payload[:2] + expected = ( + (preact_out, (TK, N), z_out_dtype, "preact"), + (postact_out, (TK, I_dim), postact_dtype, "postact"), + ( + gated_output_payload[2], + (TK, N // 32) if epilogue_quant else (0,), + torch.uint8, + "z_scale", + ), + ( + postact_scale_out, + (TK // 128, I_dim // 128, 512), + torch.uint8, + "postact_scale", + ), + ) + for tensor, shape, dtype, name in expected: + if tuple(tensor.shape) != shape or tensor.dtype != dtype: + raise ValueError( + f"invalid preallocated {name}: shape={tuple(tensor.shape)}, " + f"dtype={tensor.dtype}; expected shape={shape}, dtype={dtype}" + ) - z, y1 = gemm_gated( - x_fp8, w1_fp8, - activation="swiglu", - out_dtype=z_out_dtype, - postact_dtype=postact_dtype, - cu_seqlens_m=expert_frequency_offset, - A_idx=x_gather_idx, - a_scales=x_scales_tk_e8m0, - b_scales=w1_scales, - store_preact=True, - dynamic_scheduler=False, - tuned=False, - z_scale_out=z_scale_out, - postact_scale_out=postact_scale_out, - swiglu_clamp_value=cfg.swiglu_clamp_value, - postact_bf16_trunc=cfg.fuse_y1_bf16_trunc, - current_stream=current_stream, - b_pretransposed=w1_pretransposed, + launched_prepared = ( + w1_pretransposed + and preact_out is not None + and postact_out is not None + and try_gemm_gated_tvm_ffi_sonic_prepared( + x_fp8, + w1_fp8, + preact_out, + postact_out, + activation="swiglu", + cu_seqlens_m=expert_frequency_offset, + A_idx=x_gather_idx, + a_scales=x_scales_tk_e8m0, + b_scales=w1_scales, + z_scale_out=z_scale_out, + postact_scale_out=postact_scale_out, + swiglu_clamp_value=cfg.swiglu_clamp_value, + postact_bf16_trunc=cfg.fuse_y1_bf16_trunc, + current_stream=current_stream, + ) ) + if launched_prepared: + z, y1 = preact_out, postact_out + else: + # The first call validates and compiles through the generic interface; + # subsequent calls reuse only its pointer-stable launch plan. + z, y1 = gemm_gated( + x_fp8, + w1_fp8, + activation="swiglu", + preact_out=preact_out, + postact_out=postact_out, + out_dtype=z_out_dtype, + postact_dtype=postact_dtype, + cu_seqlens_m=expert_frequency_offset, + A_idx=x_gather_idx, + a_scales=x_scales_tk_e8m0, + b_scales=w1_scales, + store_preact=True, + dynamic_scheduler=False, + tuned=False, + z_scale_out=z_scale_out, + postact_scale_out=postact_scale_out, + swiglu_clamp_value=cfg.swiglu_clamp_value, + postact_bf16_trunc=cfg.fuse_y1_bf16_trunc, + current_stream=current_stream, + b_pretransposed=w1_pretransposed, + ) del x_fp8, x_scales_tk_e8m0 if epilogue_quant: @@ -732,6 +843,26 @@ def _matches_prequant_tensor(lhs: torch.Tensor | None, rhs: torch.Tensor | None) ) +def attach_preallocated_gated_outputs( + packed_scales: torch.Tensor, + outputs: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], +) -> torch.Tensor: + """Carry per-call gated outputs without making them explicit PyLayer inputs.""" + if len(outputs) != 4: + raise ValueError(f"expected four gated outputs, got {len(outputs)}") + setattr(packed_scales, _GATED_OUTPUT_PAYLOAD_ATTR, tuple(outputs)) + return packed_scales + + +def _alias_preallocated_pylayer_output(tensor: torch.Tensor) -> torch.Tensor: + """Return independent autograd metadata over the same temporary storage.""" + alias = torch.empty(1, dtype=tensor.dtype, device=tensor.device).as_strided( + tensor.shape, tensor.stride() + ) + tensor._share_buffer_to(alias) + return alias + + def _get_cu_seqlens_cpu(cu_seqlens: torch.Tensor) -> tuple: """Return cu_seqlens values as a Python tuple, cached on the tensor object. @@ -1305,7 +1436,9 @@ def forward( ctx._has_num_activated = num_activated_expert_per_token_offset is not None ctx._prequant_activation_payload = prequant_activation_payload is not None ctx._prequant_activation_payload_len = ( - len(prequant_activation_payload) if prequant_activation_payload is not None else 0 + len(prequant_activation_payload) + if prequant_activation_payload is not None + else 0 ) ctx._w1_original = w1 ctx._wgrad_w1_accumulator = _main_grad_accumulator(w1) @@ -1363,6 +1496,16 @@ def forward( num_activated_expert_per_token_offset, ) + if ( + prequant_activation_payload is not None + and len(prequant_activation_payload) == 7 + ): + # Paddle forbids returning a differentiable PyLayer input directly. + # Keep the temporary storage but return fresh autograd metadata. + preallocated_z = prequant_activation_payload[3] + if _matches_prequant_tensor(preallocated_z, z): + z = _alias_preallocated_pylayer_output(preallocated_z) + ctx.mark_non_differentiable(y1) ctx.set_materialize_grads(False) @@ -1530,7 +1673,14 @@ def backward(ctx, _: None, dz: torch.Tensor): bwd_col = _PREQUANTIZED_SCALES.pop("bwd_col", None) if ctx._prequant_activation_payload: assert x_fp8_pre is not None and x_scales_pre is not None, "Pre-quantized input is None." - x_bf16 = dequantize_blockscaled_fp8(x_fp8_pre, _raw_1x32_scale_bytes(x_scales_pre)) + x_bf16 = dequantize_blockscaled_fp8( + x_fp8_pre, + _raw_1x32_scale_bytes( + x_scales_pre, + int(x_fp8_pre.shape[0]), + int(x_fp8_pre.shape[1]), + ), + ) else: x_bf16 = x _wgrad_accum = getattr(ctx, '_wgrad_w1_accumulator', None) @@ -1739,6 +1889,8 @@ def forward( fp8_protocol: FP8Protocol | None, fp8_combine_grad_handle=None, fp8_config: _FP8Config | None = None, + router_score_source: torch.Tensor | None = None, + router_score_src_idx: torch.Tensor | None = None, ) -> torch.Tensor: # w2 = w2_original.permute([1, 2, 0]) # w2.fp8 = getattr(w2_original, "fp8", None) @@ -1746,6 +1898,24 @@ def forward( TK = y1.size(0) E, H, I = w2.shape + if (router_score_source is None) != (router_score_src_idx is None): + raise ValueError("router_score_source and router_score_src_idx must be provided together") + has_router_score_source = router_score_source is not None + if has_router_score_source: + if router_score_source.ndim != 2: + raise ValueError("router_score_source must be rank-2, got " f"shape={router_score_source.shape}") + if router_score_src_idx.ndim != 1 or str(router_score_src_idx.dtype) not in ( + "torch.int32", + "paddle.int32", + "int32", + ): + raise ValueError( + "router_score_src_idx must be rank-1 int32, got " + f"shape={router_score_src_idx.shape}, dtype={router_score_src_idx.dtype}" + ) + if router_score_src_idx.shape[0] > topk_scores.shape[0]: + raise ValueError("router_score_src_idx cannot exceed the metadata score buffer") + use_quack_gemm = is_using_quack_gemm() if use_quack_gemm: @@ -1882,6 +2052,12 @@ def forward( ctx._has_b2 = b2 is not None ctx._has_num_activated = num_activated_expert_per_token_offset is not None ctx._fp8_combine_grad_handle = fp8_combine_grad_handle + ctx._has_router_score_source = has_router_score_source + if ctx._has_router_score_source: + # The source tensor is an autograd edge only. Forward keeps reading the + # compact metadata scores; backward scatters ds to this original shape. + ctx._router_score_source_shape = tuple(router_score_source.shape) + ctx._router_score_source_numel = int(router_score_source.shape[0]) * int(router_score_source.shape[1]) ctx._w2_original = w2 ctx._wgrad_w2_accumulator = _main_grad_accumulator(w2) # Always compute ds (topk_scores gradient) — needed for router training. @@ -1978,6 +2154,7 @@ def forward( x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) else: ctx.save_for_backward( @@ -1990,6 +2167,7 @@ def forward( x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) else: if _w2_decouple: @@ -2007,6 +2185,7 @@ def forward( x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) else: ctx.save_for_backward( @@ -2018,6 +2197,7 @@ def forward( x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) # Keep w2 FP8 cache — backward hits cache (~38µs savings) at ~37MB memory cost. @@ -2050,6 +2230,7 @@ def backward(ctx, dout: torch.Tensor): x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) = ctx.saved_tensor() w2_shape = ctx._w2_shape # (H, I, E) w2_dtype = ctx._w2_dtype @@ -2065,6 +2246,7 @@ def backward(ctx, dout: torch.Tensor): x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) = ctx.saved_tensor() w2_shape = w2.shape w2_dtype = w2.dtype @@ -2087,6 +2269,7 @@ def backward(ctx, dout: torch.Tensor): x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) = ctx.saved_tensor() w2_shape = ctx._w2_shape w2_dtype = ctx._w2_dtype @@ -2101,6 +2284,7 @@ def backward(ctx, dout: torch.Tensor): x_gather_idx, s_scatter_idx, s_reverse_scatter_idx, + router_score_src_idx, ) = ctx.saved_tensor() w2_shape = w2.shape w2_dtype = w2.dtype @@ -2190,7 +2374,11 @@ def backward(ctx, dout: torch.Tensor): dout_comm_fp8 = combine_grad_data dout_fp8 = dout_comm_fp8 if not ctx._fp8_cfg.fp8_wgrad: - dout_dequant_scales = _raw_1x32_scale_bytes(dout_raw_scales_t) + dout_dequant_scales = _raw_1x32_scale_bytes( + dout_raw_scales_t, + int(dout_comm_fp8.shape[0]), + int(dout_comm_fp8.shape[1]), + ) dout = dequantize_blockscaled_fp8(dout_comm_fp8, dout_dequant_scales) else: dout_fp8, dout_raw_scales_t = quantize_activation_blockscaled_fast(dout) @@ -2346,9 +2534,10 @@ def backward(ctx, dout: torch.Tensor): _wgrad_accum_w2 = getattr(ctx, '_wgrad_w2_accumulator', None) if _wgrad_accum_w2 is not None: # BF16 wgrad + fp32 accumulate. - # _wgrad_accum_w2 stays in grouped layout [E, I, H]; - # GEMM out is [E, H, I], so accumulate via permute(0,2,1) - # view (pure transpose, no gate/up interleave for w2). + # Fleet converts both w2 and main_grad to Sonic layout + # [E, H, I] before forward. The GEMM output has that same + # logical shape, so writing through a transposed view would + # silently corrupt the persistent optimizer buffer. y1s_wgrad = ( y1s if y1s.dtype == torch.bfloat16 else y1s.to(torch.bfloat16) ) @@ -2358,7 +2547,7 @@ def backward(ctx, dout: torch.Tensor): y1s_wgrad, expert_frequency_offset, x_gather_idx, - accumulator=_wgrad_accum_w2.permute(0, 2, 1), + accumulator=_wgrad_accum_w2, M=dout.shape[1], N=w2_shape[2], total_K=TK_wgrad, @@ -2371,7 +2560,7 @@ def backward(ctx, dout: torch.Tensor): y1s_wgrad, expert_frequency_offset, x_gather_idx, - accumulator=_wgrad_accum_w2.permute(0, 2, 1), + accumulator=_wgrad_accum_w2, M=dout.shape[1], N=w2_shape[2], total_K=TK_wgrad, @@ -2521,7 +2710,7 @@ def backward(ctx, dout: torch.Tensor): y1s_wgrad, expert_frequency_offset, x_gather_idx, - accumulator=_wgrad_accum_w2.permute(0, 2, 1), + accumulator=_wgrad_accum_w2, M=dout.shape[1], N=w2.shape[2], total_K=x_gather_idx.shape[0], @@ -2534,7 +2723,7 @@ def backward(ctx, dout: torch.Tensor): y1s_wgrad, expert_frequency_offset, x_gather_idx, - accumulator=_wgrad_accum_w2.permute(0, 2, 1), + accumulator=_wgrad_accum_w2, M=dout.shape[1], N=w2.shape[2], total_K=x_gather_idx.shape[0], @@ -2576,6 +2765,17 @@ def backward(ctx, dout: torch.Tensor): elif not is_varlen_K: ds = ds.view(T, K) + router_score_source_grad = None + if ctx._has_router_score_source: + # score_src_idx covers only real routes; the metadata score buffer may + # include aligned padding. The scatter helper zero-fills source slots + # not represented by a received route, matching the former carrier. + router_score_source_grad = _scatter_router_scores_i32( + ds.contiguous(), + router_score_src_idx, + ctx._router_score_source_numel, + ).reshape(ctx._router_score_source_shape) + if fp8_combine_grad_handle is not None: fp8_combine_grad_handle.pop("data", None) fp8_combine_grad_handle.pop("scale", None) @@ -2589,10 +2789,15 @@ def backward(ctx, dout: torch.Tensor): grads.append(None) if ctx._has_b2: grads.append(db2) - grads.extend([ds if ctx._topk_scores_needs_grad else None, None, None]) # topk_scores, selected_experts, expert_frequency_offset + topk_scores_grad = None if ctx._has_router_score_source else ds if ctx._topk_scores_needs_grad else None + grads.extend([topk_scores_grad, None, None]) # topk_scores, selected_experts, expert_frequency_offset grads.extend([None, None, None]) # x_gather_idx, s_scatter_idx, s_reverse_scatter_idx if ctx._has_num_activated: grads.append(None) + if ctx._has_router_score_source: + # Optional tensor inputs are returned only when present; fp8_config is + # non-tensor and therefore has no corresponding Paddle PyLayer slot. + grads.extend([router_score_source_grad, None]) if ctx._wgrad_w2_accumulator is not None: _apply_weight_backward_hook(ctx._w2_original) return tuple(grads) diff --git a/sonicmoe/quack_utils/bf16_wgrad_gemm.py b/sonicmoe/quack_utils/bf16_wgrad_gemm.py index a9c68fc..fe8e123 100644 --- a/sonicmoe/quack_utils/bf16_wgrad_gemm.py +++ b/sonicmoe/quack_utils/bf16_wgrad_gemm.py @@ -156,6 +156,19 @@ def _run_bf16_wgrad_varlen_k( epi_args, config: Optional[GemmConfig] = None, ) -> torch.Tensor: + expected_out_shape = (int(num_experts), int(M), int(N)) + if tuple(out.shape) != expected_out_shape: + raise ValueError( + "BF16 varlen-K wgrad output shape mismatch: expected " + f"{expected_out_shape}, got {tuple(out.shape)}. " + "A non-contiguous output view is supported, but its logical shape " + "must remain [num_experts, M, N]." + ) + if C is not None and tuple(C.shape) != expected_out_shape: + raise ValueError( + "BF16 varlen-K wgrad accumulator shape mismatch: expected " + f"{expected_out_shape}, got {tuple(C.shape)}." + ) if config is None: config = _bf16_wgrad_default_config(device, M=M, N=N) fast_key = ( diff --git a/sonicmoe/quack_utils/blockscaled_fp8_gemm.py b/sonicmoe/quack_utils/blockscaled_fp8_gemm.py index b03976f..315c1f0 100644 --- a/sonicmoe/quack_utils/blockscaled_fp8_gemm.py +++ b/sonicmoe/quack_utils/blockscaled_fp8_gemm.py @@ -16,6 +16,7 @@ import torch _E8M0_DTYPE = getattr(torch, "float8_e8m0fnu", torch.uint8) +_GATED_FFI_B_VIEW_ATTR = "_sonic_gated_tvm_ffi_b_view" import triton import triton.language as tl @@ -455,6 +456,7 @@ def _quantize_flat_v2_kernel( TILE_ROWS: tl.constexpr, # 128 TILE_COLS: tl.constexpr, # 256 SCALE_OUT_INT32: tl.constexpr = False, + SCALE_PACKED_WORDS: tl.constexpr = False, SAFE_INT64: tl.constexpr = False, ): """High-BW blockscaled quantize: large tiles, vectorized, pipelined. @@ -506,10 +508,15 @@ def _quantize_flat_v2_kernel( tl.store(dst_ptrs, quantized, mask=mask) group_id = (col_base // GROUP_SIZE) + g - scale_ptrs = dst_scale_ptr + row_offs * scale_stride_row + group_id * scale_stride_col - if SCALE_OUT_INT32: + if SCALE_PACKED_WORDS: + scale_byte_ptr = dst_scale_ptr.to(tl.pointer_type(tl.uint8)) + scale_ptrs = scale_byte_ptr + row_offs * scale_stride_row + group_id + tl.store(scale_ptrs, e8m0_byte, mask=row_mask) + elif SCALE_OUT_INT32: + scale_ptrs = dst_scale_ptr + row_offs * scale_stride_row + group_id * scale_stride_col tl.store(scale_ptrs, e8m0_byte.to(tl.int32), mask=row_mask) else: + scale_ptrs = dst_scale_ptr + row_offs * scale_stride_row + group_id * scale_stride_col tl.store(scale_ptrs, e8m0_byte, mask=row_mask) @@ -517,6 +524,7 @@ def quantize_activation_blockscaled_fast( x: torch.Tensor, group_size: int = 32, scale_dtype=None, + pack_scale_words: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Fast fused 1×group_size blockscaled quantization using a single Triton kernel. @@ -529,7 +537,31 @@ def quantize_activation_blockscaled_fast( fp8_out = torch.empty(M, K, dtype=torch.float8_e4m3fn, device=x.device) scale_out_int32 = str(scale_dtype) in ("torch.int32", "paddle.int32", "int32") - scale_out = torch.empty(M, num_groups, dtype=(torch.int32 if scale_out_int32 else torch.uint8), device=x.device) + if pack_scale_words: + if scale_out_int32: + raise ValueError( + "pack_scale_words and an explicit int32 scale_dtype are mutually exclusive" + ) + if num_groups % 4 != 0: + raise ValueError( + "pack_scale_words requires the number of 1x32 groups to be " + f"divisible by 4, got {num_groups} for K={K}" + ) + # DeepEP transports int32 words as four opaque E8M0 scale bytes. + scale_out = torch.empty( + M, num_groups // 4, dtype=torch.int32, device=x.device + ) + scale_stride_row = scale_out.stride(0) * 4 + scale_stride_col = 1 + else: + scale_out = torch.empty( + M, + num_groups, + dtype=(torch.int32 if scale_out_int32 else torch.uint8), + device=x.device, + ) + scale_stride_row = scale_out.stride(0) + scale_stride_col = scale_out.stride(1) TILE_ROWS = 128 # Larger tile: fewer CTAs, better wave occupancy TILE_COLS = min(K, 256) @@ -547,13 +579,14 @@ def quantize_activation_blockscaled_fast( x.stride(1), fp8_out.stride(0), fp8_out.stride(1), - scale_out.stride(0), - scale_out.stride(1), + scale_stride_row, + scale_stride_col, fp8_max=_FP8_E4M3_MAX, GROUP_SIZE=group_size, TILE_ROWS=TILE_ROWS, TILE_COLS=TILE_COLS, SCALE_OUT_INT32=scale_out_int32, + SCALE_PACKED_WORDS=pack_scale_words, SAFE_INT64=_needs_int64, ) return fp8_out, scale_out @@ -4233,6 +4266,9 @@ def _quantize_weight_pair_3d_triton( # block. Reshape (E, per_batch) → (1, E*per_batch) view-only. A_scales_flat = A_scales.view(1, E * A_per_batch) B_scales_flat = B_scales.view(1, E * B_per_batch) + # Cache the metadata-only TVM-FFI descriptor while the weight is prepared. + # Rebuilding this permute view on every microbatch is pure host overhead. + setattr(A_fp8, _GATED_FFI_B_VIEW_ATTR, A_fp8.permute(1, 2, 0)) return A_fp8, A_scales_flat.view(_E8M0_DTYPE), B_fp8, B_scales_flat.view(_E8M0_DTYPE) diff --git a/sonicmoe/quack_utils/gemm_gated.py b/sonicmoe/quack_utils/gemm_gated.py index 9fa76b4..3f7c33b 100644 --- a/sonicmoe/quack_utils/gemm_gated.py +++ b/sonicmoe/quack_utils/gemm_gated.py @@ -3,7 +3,7 @@ # ******************************************************************************** import os -from functools import partial +from functools import lru_cache, partial from pathlib import Path from typing import Callable, NamedTuple, Optional, Tuple @@ -93,6 +93,50 @@ def _stream_matches_current(stream) -> bool: return stream_id is None or stream_id == _current_raw_stream_id() +_GATED_FFI_B_VIEW_ATTR = "_sonic_gated_tvm_ffi_b_view" +_GATED_FFI_LAST_PLAN_ATTR = "_sonic_gated_tvm_ffi_last_plan" + + +class _PreparedGatedTvmFfiPlan(NamedTuple): + """Pointer-stable weight state for the Sonic gated FFI launch.""" + + b_view: Tensor + b_data_ptr: int + b_view_data_ptr: int + compiled_fn: Callable + gemm_cls: type + scheduler_args: object + activation: str + epilogue_quant: bool + postact_quant: bool + swiglu_clamp_value: float + postact_bf16_trunc: bool + + +def prepare_gated_tvm_ffi_weight(B: Tensor) -> Tensor: + """Attach the storage-free TVM-FFI B view to one quantized weight.""" + if B.ndim == 3 and B.is_cuda and B.is_contiguous(): + setattr(B, _GATED_FFI_B_VIEW_ATTR, B.permute(1, 2, 0)) + return B + + +def _prepared_gated_tvm_ffi_weight(B: Tensor) -> Optional[Tensor]: + view = getattr(B, _GATED_FFI_B_VIEW_ATTR, None) + if view is None or B.ndim != 3: + return None + expected_shape = (B.shape[1], B.shape[2], B.shape[0]) + expected_stride = (B.stride(1), B.stride(2), B.stride(0)) + if ( + view.dtype != B.dtype + or view.device != B.device + or tuple(view.shape) != expected_shape + or tuple(view.stride()) != expected_stride + or view.data_ptr() != B.data_ptr() + ): + return None + return view + + class GemmGatedSm90(GemmGatedMixin, GemmSm90): pass @@ -134,6 +178,7 @@ def _fake_tensor_rank(dtype, ndim: int, leading_dim: int): return fake_tensor(dtype, shape, leading_dim=leading_dim, divisibility=div_for_dtype(dtype)) +@lru_cache(maxsize=32) @jit_cache def _compile_gemm_gated_tvm_ffi( gemm_cls_name, @@ -357,6 +402,333 @@ def _run_gemm_gated_tvm_ffi( ) +@lru_cache(maxsize=16) +def _static_gated_scheduler_args( + cluster_m: int, + cluster_n: int, + max_swizzle_size: int, + device_index: int, +): + """Cache scheduler metadata that never contains a tensor pointer.""" + del device_index # Part of the cache key because cluster capacity is device-local. + max_active_clusters = get_max_active_clusters(cluster_m * cluster_n) + return make_scheduler_args(max_active_clusters, max_swizzle_size, None) + + +def _launch_prepared_gated_tvm_ffi( + plan: _PreparedGatedTvmFfiPlan, + A: Tensor, + D: Tensor, + PostAct: Tensor, + cu_seqlens_m: Tensor, + A_idx: Tensor, + a_scales: Tensor, + b_scales: Tensor, + z_scale_out: Optional[Tensor], + postact_scale_out: Tensor, +) -> None: + """Build only pointer-bearing arguments, then issue one prepared FFI call.""" + epi_kwargs = {"mPostActScaleIsa": postact_scale_out} + if plan.epilogue_quant: + epi_kwargs["mZScale"] = z_scale_out + epilogue_args = plan.gemm_cls.EpilogueArguments( + PostAct, + None, + swiglu_clamp_value=None, + postact_bf16_trunc=None, + mRowVecBroadcast=None, + mColVecBroadcast=None, + rounding_mode=None, + sr_seed=None, + **epi_kwargs, + ) + # These structures contain current-microbatch pointers and must not be + # cached across PP/VPP contexts. + varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx) + plan.compiled_fn( + A, + plan.b_view, + D, + None, + epilogue_args, + plan.scheduler_args, + varlen_args, + a_scales, + b_scales, + ) + + +def try_gemm_gated_tvm_ffi_sonic_prepared( + A: Tensor, + B: Tensor, + D: Optional[Tensor], + PostAct: Optional[Tensor], + *, + activation: str, + cu_seqlens_m: Optional[Tensor], + A_idx: Optional[Tensor], + a_scales: Optional[Tensor], + b_scales: Optional[Tensor], + z_scale_out: Optional[Tensor], + postact_scale_out: Optional[Tensor], + swiglu_clamp_value: float, + postact_bf16_trunc: bool, + current_stream=None, +) -> bool: + """Issue a previously validated plan from Sonic's preallocated fast path. + + The caller has already validated dynamic output shapes and dtypes. The plan + owns only the persistent quantized-weight view and pointer-free compiled + state; all microbatch-dependent pointer structures are rebuilt above. + """ + if os.getenv("SONIC_MOE_DISABLE_GATED_TVM_FFI", "0").lower() in { + "1", + "true", + "yes", + "on", + }: + return False + plan = getattr(B, _GATED_FFI_LAST_PLAN_ATTR, None) + if not isinstance(plan, _PreparedGatedTvmFfiPlan): + return False + epilogue_quant = z_scale_out is not None + if ( + D is None + or PostAct is None + or cu_seqlens_m is None + or A_idx is None + or a_scales is None + or b_scales is None + or postact_scale_out is None + or plan.activation != activation + or plan.epilogue_quant != epilogue_quant + or not plan.postact_quant + or plan.swiglu_clamp_value != float(swiglu_clamp_value) + or plan.postact_bf16_trunc != bool(postact_bf16_trunc) + or not _stream_matches_current(current_stream) + ): + return False + # Quantization normally replaces B wholesale. Detect uncommon in-place + # storage rebinding so a stale prepared view can never launch. + if int(B.data_ptr()) != plan.b_data_ptr or int(plan.b_view.data_ptr()) != plan.b_view_data_ptr: + setattr(B, _GATED_FFI_LAST_PLAN_ATTR, None) + return False + _launch_prepared_gated_tvm_ffi( + plan, + A, + D, + PostAct, + cu_seqlens_m, + A_idx, + a_scales, + b_scales, + z_scale_out, + postact_scale_out, + ) + return True + + +def try_gemm_gated_tvm_ffi_sonic_fast( + A: Tensor, + B: Tensor, + D: Optional[Tensor], + C: Optional[Tensor], + PostAct: Optional[Tensor], + *, + activation: str, + tile_M: int, + tile_N: int, + cluster_M: int, + cluster_N: int, + pingpong: bool, + max_swizzle_size: int, + bias: Optional[Tensor], + cu_seqlens_m: Optional[Tensor], + A_idx: Optional[Tensor], + dynamic_scheduler: bool, + a_scales: Optional[Tensor], + b_scales: Optional[Tensor], + z_scale_out: Optional[Tensor], + postact_scale_out: Optional[Tensor], + swiglu_clamp_value: float, + postact_bf16_trunc: bool, + current_stream=None, +) -> bool: + """Launch the exact Sonic forward contract without generic CuTe setup. + + Only pointer-free scheduler and compilation state are cached. Epilogue and + varlen arguments are reconstructed for every call so PP/VPP microbatches + cannot observe stale tensor pointers. + """ + if os.getenv("SONIC_MOE_DISABLE_GATED_TVM_FFI", "0").lower() in { + "1", + "true", + "yes", + "on", + }: + return False + if ( + activation != "swiglu" + or D is None + or PostAct is None + or C is not None + or bias is not None + or dynamic_scheduler + or cu_seqlens_m is None + or A_idx is None + or a_scales is None + or b_scales is None + or postact_scale_out is None + or cluster_N != 1 + or not _stream_matches_current(current_stream) + ): + return False + if ( + A.ndim != 2 + or B.ndim != 3 + or D.ndim != 2 + or PostAct.ndim != 2 + or cu_seqlens_m.ndim != 1 + or A_idx.ndim != 1 + or not all( + tensor.is_cuda + for tensor in (A, B, D, PostAct, cu_seqlens_m, A_idx) + ) + ): + return False + + E, N, K = B.shape + total_m = A_idx.shape[0] + epilogue_quant = z_scale_out is not None + postact_quant = postact_scale_out is not None + expected_d_dtype = ( + torch.float8_e4m3fn if epilogue_quant else torch.bfloat16 + ) + if ( + A.shape[1] != K + or tuple(D.shape) != (total_m, N) + or tuple(PostAct.shape) != (total_m, N // 2) + or tuple(cu_seqlens_m.shape) != (E + 1,) + or N % 128 != 0 + or total_m % 128 != 0 + or A.stride(-1) != 1 + or D.stride(-1) != 1 + or PostAct.stride(-1) != 1 + or A.dtype != torch.float8_e4m3fn + or B.dtype != torch.float8_e4m3fn + or D.dtype != expected_d_dtype + or PostAct.dtype != torch.float8_e4m3fn + or cu_seqlens_m.dtype != torch.int32 + or A_idx.dtype != torch.int32 + or postact_scale_out.dtype != torch.uint8 + or tuple(postact_scale_out.shape) + != (total_m // 128, (N // 2) // 128, 512) + ): + return False + if epilogue_quant and ( + z_scale_out.dtype != torch.uint8 + or tuple(z_scale_out.shape) != (total_m, N // 32) + ): + return False + + B_ffi = _prepared_gated_tvm_ffi_weight(B) + if B_ffi is None: + return False + device_capacity = get_device_capacity(A.device) + if device_capacity[0] <= 9: + return False + scale_tensors = tuple( + tensor + for tensor in (a_scales, b_scales, z_scale_out, postact_scale_out) + if tensor is not None + ) + if any(tensor.dtype not in _TORCH_TO_CUTLASS_DTYPE for tensor in scale_tensors): + return False + + GemmCls = ( + GemmGatedSm100ZeroMatQuantPostActQuant + if epilogue_quant + else GemmGatedSm100ZeroMatPostActQuant + ) + compiled_fn = _compile_gemm_gated_tvm_ffi( + GemmCls.__name__, + _TORCH_TO_CUTLASS_DTYPE[A.dtype], + _TORCH_TO_CUTLASS_DTYPE[B_ffi.dtype], + _TORCH_TO_CUTLASS_DTYPE[D.dtype], + None, + _TORCH_TO_CUTLASS_DTYPE[PostAct.dtype], + "k", + "k", + "n", + None, + "n", + (tile_M, tile_N), + (cluster_M, cluster_N, 1), + pingpong, + True, + activation, + device_capacity, + max_swizzle_size, + True, + True, + True, + epilogue_quant, + True, + float(swiglu_clamp_value), + bool(postact_bf16_trunc), + _TORCH_TO_CUTLASS_DTYPE[a_scales.dtype], + a_scales.ndim, + _TORCH_TO_CUTLASS_DTYPE[b_scales.dtype], + b_scales.ndim, + ( + _TORCH_TO_CUTLASS_DTYPE[z_scale_out.dtype] + if z_scale_out is not None + else None + ), + z_scale_out.ndim if z_scale_out is not None else 0, + _TORCH_TO_CUTLASS_DTYPE[postact_scale_out.dtype], + postact_scale_out.ndim, + ) + from quack.cache_utils import COMPILE_ONLY as _COMPILE_ONLY + + if _COMPILE_ONLY: + return True + + scheduler_args = _static_gated_scheduler_args( + cluster_M, + cluster_N, + max_swizzle_size, + int(torch.cuda.current_device()), + ) + plan = _PreparedGatedTvmFfiPlan( + b_view=B_ffi, + b_data_ptr=int(B.data_ptr()), + b_view_data_ptr=int(B_ffi.data_ptr()), + compiled_fn=compiled_fn, + gemm_cls=GemmCls, + scheduler_args=scheduler_args, + activation=activation, + epilogue_quant=epilogue_quant, + postact_quant=postact_quant, + swiglu_clamp_value=float(swiglu_clamp_value), + postact_bf16_trunc=bool(postact_bf16_trunc), + ) + setattr(B, _GATED_FFI_LAST_PLAN_ATTR, plan) + _launch_prepared_gated_tvm_ffi( + plan, + A, + D, + PostAct, + cu_seqlens_m, + A_idx, + a_scales, + b_scales, + z_scale_out, + postact_scale_out, + ) + return True + + def gemm_gated( A: Tensor, # (l, m, k) or (total_m, k) if varlen_m or (whatever, k) if gather_A with varlen_m B: Tensor, # (l, n, k) diff --git a/sonicmoe/quack_utils/gemm_interface.py b/sonicmoe/quack_utils/gemm_interface.py index ccca51a..a620693 100644 --- a/sonicmoe/quack_utils/gemm_interface.py +++ b/sonicmoe/quack_utils/gemm_interface.py @@ -33,7 +33,10 @@ def decorator(fn): return decorator from .gemm_dgated import gemm_dgated as gemm_dgated_sm90_sm100 -from .gemm_gated import gemm_gated as gemm_gated_sm90_sm100 +from .gemm_gated import ( + gemm_gated as gemm_gated_sm90_sm100, + try_gemm_gated_tvm_ffi_sonic_fast, +) default_device_capacity = get_device_capacity(paddle.device("cuda")) @@ -154,6 +157,32 @@ def gemm_gated_tuned( varlen_m = cu_seqlens_m is not None if varlen_m: assert not config.swap_ab, "Variable-length sequences not supported with swap_ab" + if b_pretransposed and try_gemm_gated_tvm_ffi_sonic_fast( + A, + B, + preact_out, + C, + postact_out, + activation=activation, + tile_M=config.tile_m, + tile_N=config.tile_n, + cluster_M=config.cluster_m, + cluster_N=config.cluster_n, + pingpong=config.pingpong, + max_swizzle_size=config.max_swizzle_size, + bias=bias, + cu_seqlens_m=cu_seqlens_m, + A_idx=A_idx, + dynamic_scheduler=dynamic_scheduler, + a_scales=a_scales, + b_scales=b_scales, + z_scale_out=z_scale_out, + postact_scale_out=postact_scale_out, + swiglu_clamp_value=swiglu_clamp_value, + postact_bf16_trunc=postact_bf16_trunc, + current_stream=current_stream, + ): + return if A.ndim == 2 and not varlen_m: A = A.unsqueeze(0) # (1, M, K) if not b_pretransposed: diff --git a/tests/ops/test_bf16_wgrad_varlen_k.py b/tests/ops/test_bf16_wgrad_varlen_k.py index ba0d08d..1aca28d 100644 --- a/tests/ops/test_bf16_wgrad_varlen_k.py +++ b/tests/ops/test_bf16_wgrad_varlen_k.py @@ -320,6 +320,63 @@ def test_bf16_wgrad_default_correctness_and_determinism( torch.testing.assert_close(first.float(), _gold(case).float(), atol=3e-2, rtol=2e-2) +@pytest.mark.parametrize("variant", ["beta", "tma_add"]) +def test_bf16_wgrad_rejects_transposed_w2_main_grad(variant: str) -> None: + """Down-projection main_grad is already Sonic [E,H,I], not [E,I,H].""" + from sonicmoe.quack_utils.bf16_wgrad_gemm import ( + bf16_wgrad_gemm_varlen_k_accumulate, + bf16_wgrad_gemm_varlen_k_tma_add, + ) + + case = _make_case( + M=256, + N=128, + E=2, + tokens_per_expert=128, + source_tokens=256, + layout="contiguous", + ) + fn = ( + bf16_wgrad_gemm_varlen_k_accumulate + if variant == "beta" + else bf16_wgrad_gemm_varlen_k_tma_add + ) + main_grad = torch.zeros( + (case["E"], case["M"], case["N"]), + dtype=torch.float32, + device=case["device"], + ) + + fn( + case["A"], + case["B"], + case["cu"], + case["A_idx"], + accumulator=main_grad, + M=case["M"], + N=case["N"], + total_K=case["total_K"], + num_experts=case["E"], + device=case["device"], + ) + torch.cuda.synchronize() + assert float(main_grad.norm()) > 0.0 + + with pytest.raises(ValueError, match="output shape mismatch"): + fn( + case["A"], + case["B"], + case["cu"], + case["A_idx"], + accumulator=main_grad.permute(0, 2, 1), + M=case["M"], + N=case["N"], + total_K=case["total_K"], + num_experts=case["E"], + device=case["device"], + ) + + def test_bf16_wgrad_auto_large_shape_matches_reference_config() -> None: case = _make_case( M=4096, diff --git a/tests/ops/test_deepep_topk_metadata.py b/tests/ops/test_deepep_topk_metadata.py index 28b4a3f..da71f15 100644 --- a/tests/ops/test_deepep_topk_metadata.py +++ b/tests/ops/test_deepep_topk_metadata.py @@ -34,6 +34,12 @@ ) from sonicmoe.quack_utils.blockscaled_fp8_gemm import ( gather_raw_blockscaled_1x32_scales_to_isa, + quantize_activation_blockscaled_fast, +) +from sonicmoe.functional import ( + _ensure_isa_1x32_scales, + _gather_1x32_scales_to_isa, + _raw_1x32_scale_bytes, ) @@ -745,6 +751,25 @@ def test_element_wise_consistency(self, N_recv, topk, E): class TestCudaScalePacking: + @pytest.mark.parametrize("rows,cols", [(7, 256), (129, 4096)]) + def test_quantizer_compact_scale_words_are_byte_exact(self, rows, cols): + torch.manual_seed(314159) + source = torch.randn( + (rows, cols), dtype=torch.bfloat16, device="cuda" + ) + fp8_raw, raw_scales = quantize_activation_blockscaled_fast(source) + fp8_compact, compact_scales = quantize_activation_blockscaled_fast( + source, pack_scale_words=True + ) + + assert compact_scales.dtype == torch.int32 + assert compact_scales.shape == (rows, cols // 128) + assert _t_eq(fp8_raw.view(torch.uint8), fp8_compact.view(torch.uint8)) + assert _t_eq( + raw_scales, + compact_scales.view(torch.uint8).reshape(rows, cols // 32), + ) + """Validate metadata-side raw scale packing against the existing reference.""" @pytest.mark.parametrize("N_recv,topk,E,cols,raw_dtype", [ @@ -790,6 +815,223 @@ def test_with_scales_matches_raw_gather_reference( "metadata-side packed scales differ from raw gather reference" ) + @pytest.mark.parametrize("N_recv,topk,E,cols", [ + (257, 4, 8, 256), + (513, 8, 16, 4096), + (129, 4, 24, 2048), + ]) + @pytest.mark.parametrize("pack_mode", ["default", "scatter", "rowmajor"]) + def test_compact_int32_carrier_is_byte_exact( + self, N_recv, topk, E, cols, pack_mode + ): + if not _HAS_TOPK_CUDA_SCALES_KERNEL: + pytest.skip("CUDA scale-packing metadata kernel not compiled") + + torch.manual_seed(314159) + indices, probs, tpe = fabricate_dispatch_result(N_recv, topk, E) + scale_cols = (cols + 31) // 32 + assert scale_cols % 4 == 0 + raw_scales = ( + torch.arange(N_recv * scale_cols, dtype=torch.int32, device="cuda") + .reshape(N_recv, scale_cols) + % 251 + ).to(torch.uint8) + compact_carrier = raw_scales.view(torch.int32).reshape( + N_recv, scale_cols // 4 + ) + assert compact_carrier.data_ptr() == raw_scales.data_ptr() + + kwargs = { + "block": 128, + "pack_scales_in_scatter": pack_mode == "scatter", + "pack_scales_rowmajor": pack_mode == "rowmajor", + } + raw_result = deepep_topk_to_sonic_metadata_with_scales( + indices, probs, tpe, E, raw_scales, cols, **kwargs + ) + compact_result = deepep_topk_to_sonic_metadata_with_scales( + indices, probs, tpe, E, compact_carrier, cols, **kwargs + ) + + for index in (0, 1, 2, 3, 4, 5, 9, 10): + assert _t_eq(raw_result[index], compact_result[index]), ( + f"output {index} differs for compact carrier in {pack_mode} mode" + ) + assert raw_result[6:9] == compact_result[6:9] + + def test_invalid_compact_carrier_shape_is_rejected(self): + N_recv, topk, E, cols = 32, 4, 8, 256 + indices, probs, tpe = fabricate_dispatch_result(N_recv, topk, E) + malformed = torch.empty((N_recv, 3), dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="raw_scales shape mismatch"): + deepep_topk_to_sonic_metadata_with_scales( + indices, probs, tpe, E, malformed, cols + ) + + @pytest.mark.parametrize("rows,cols", [(7, 256), (31, 4096), (129, 2048)]) + def test_compact_carrier_fallbacks_are_byte_exact(self, rows, cols): + scale_cols = cols // 32 + raw = ( + torch.arange(rows * scale_cols, dtype=torch.int32, device="cuda") + .reshape(rows, scale_cols) + % 251 + ).to(torch.uint8) + compact = raw.view(torch.int32).reshape(rows, scale_cols // 4) + gather_idx = torch.arange(128, dtype=torch.int32, device="cuda") % rows + + restored = _raw_1x32_scale_bytes(compact, rows, cols) + assert restored.data_ptr() == compact.data_ptr() + assert _t_eq(restored, raw) + + raw_gathered = _gather_1x32_scales_to_isa( + raw, gather_idx, rows, cols + ) + compact_gathered = _gather_1x32_scales_to_isa( + compact, gather_idx, rows, cols + ) + assert _t_eq( + raw_gathered.view(torch.uint8), + compact_gathered.view(torch.uint8), + ) + + raw_packed = _ensure_isa_1x32_scales(raw, rows, cols) + compact_packed = _ensure_isa_1x32_scales(compact, rows, cols) + assert _t_eq( + raw_packed.view(torch.uint8), compact_packed.view(torch.uint8) + ) + + @pytest.mark.parametrize( + "preact_bf16,allocate_z_scale", + [(False, True), (True, False)], + ids=["combined_quant", "postact_quant"], + ) + def test_gated_output_allocation_contract( + self, preact_bf16, allocate_z_scale + ): + if not _HAS_TOPK_CUDA_SCALES_KERNEL: + pytest.skip("CUDA scale-packing metadata kernel not compiled") + + N_recv, topk, E, cols, gated_n = 256, 4, 8, 256, 256 + indices, probs, tpe = fabricate_dispatch_result(N_recv, topk, E) + raw_scales = torch.ones( + (N_recv, cols // 32), dtype=torch.uint8, device="cuda" + ) + prototype = torch.empty( + (N_recv, cols), dtype=torch.float8_e4m3fn, device="cuda" + ) + + def run(): + return deepep_topk_to_sonic_metadata_with_scales( + indices, + probs, + tpe, + E, + raw_scales=raw_scales, + cols=cols, + block=128, + gated_output_prototype=prototype, + gated_n=gated_n, + gated_preact_bf16=preact_bf16, + gated_allocate_z_scale=allocate_z_scale, + ) + + first = run() + second = run() + assert len(first) == 15 + TK_padded = first[6] + for index in (0, 1, 2, 3, 4, 9): + assert first[index].data_ptr() % 16 == 0, ( + f"metadata output {index} must satisfy TVM-FFI 16-byte alignment" + ) + preact, postact, z_scale, postact_scale = first[11:15] + assert preact.shape == (TK_padded, gated_n) + assert postact.shape == (TK_padded, gated_n // 2) + assert z_scale.shape == ( + (TK_padded, gated_n // 32) if allocate_z_scale else (0,) + ) + assert postact_scale.shape == ( + TK_padded // 128, + gated_n // 256, + 512, + ) + assert preact.dtype == ( + torch.bfloat16 if preact_bf16 else torch.float8_e4m3fn + ) + assert postact.dtype == torch.float8_e4m3fn + assert z_scale.dtype == torch.uint8 + assert postact_scale.dtype == torch.uint8 + assert all(t.is_contiguous() for t in first[11:15]) + if not allocate_z_scale: + assert first[13].numel() == 0 and second[13].numel() == 0 + assert all( + lhs.data_ptr() != rhs.data_ptr() + for lhs, rhs in zip(first[11:15], second[11:15]) + if lhs.numel() > 0 + ), "non-empty ctx-saved gated outputs must have independent storage per microbatch" + + def test_large_live_outputs_stay_byte_exact_across_calls(self): + """PP/VPP-live metadata must not alias later custom-op invocations.""" + if not _HAS_TOPK_CUDA_SCALES_KERNEL: + pytest.skip("CUDA scale-packing metadata kernel not compiled") + + N_recv, TK, topk, E, cols = 16385, 20001, 4, 24, 4096 + rows = torch.arange(N_recv, dtype=torch.int32, device="cuda") + indices = torch.full( + (N_recv, topk), -1, dtype=torch.int32, device="cuda" + ) + probs = torch.zeros( + (N_recv, topk), dtype=torch.float32, device="cuda" + ) + indices[:, 0] = rows % E + probs[:, 0] = 1.0 + extra = TK - N_recv + indices[:extra, 1] = (rows[:extra] + 1) % E + probs[:extra, :2] = 0.5 + valid_experts = indices[indices >= 0].long() + tokens_per_expert = torch.bincount( + valid_experts, minlength=E + ).tolist() + compact_scales = torch.zeros( + (N_recv, cols // 128), dtype=torch.int32, device="cuda" + ) + + def run(): + return deepep_topk_to_sonic_metadata_with_scales( + indices, + probs, + tokens_per_expert, + E, + compact_scales, + cols, + block=128, + ) + + output_indices = (0, 1, 2, 3, 4, 5, 9, 10) + first = run() + snapshots = { + index: first[index].clone() for index in output_indices + } + torch.cuda.synchronize() + + live_outputs = [first] + for _ in range(4): + current = run() + live_outputs.append(current) + torch.cuda.synchronize() + for index in output_indices: + assert _t_eq(current[index], snapshots[index]), ( + f"metadata output {index} changed across live invocations" + ) + assert _t_eq(first[index], snapshots[index]), ( + f"live metadata output {index} was overwritten by a later call" + ) + + for index in output_indices: + pointers = [result[index].data_ptr() for result in live_outputs] + assert len(set(pointers)) == len(pointers), ( + f"metadata output {index} must own per-call storage" + ) + # ── Performance Benchmark ─────────────────────────────────────────────────── class TestPerformance: diff --git a/tests/ops/test_gemm_tvm_ffi_bitexact.py b/tests/ops/test_gemm_tvm_ffi_bitexact.py index 01ea9ee..8143fb4 100644 --- a/tests/ops/test_gemm_tvm_ffi_bitexact.py +++ b/tests/ops/test_gemm_tvm_ffi_bitexact.py @@ -2,6 +2,7 @@ from contextlib import contextmanager, nullcontext from importlib import import_module +import os import numpy as np import pytest @@ -25,6 +26,35 @@ def _force_executor(module_name: str, predicate_name: str): setattr(module, predicate_name, original) +@contextmanager +def _disable_gated_tvm_ffi(): + name = "SONIC_MOE_DISABLE_GATED_TVM_FFI" + original = os.environ.get(name) + os.environ[name] = "1" + try: + yield + finally: + if original is None: + os.environ.pop(name, None) + else: + os.environ[name] = original + + +@contextmanager +def _forbid_gated_executor_fallback(): + module = import_module("sonicmoe.quack_utils.gemm_interface") + original = module.gemm_gated_sm90_sm100 + + def fail(*args, **kwargs): + raise AssertionError("native-layout TVM-FFI path fell back to executor") + + module.gemm_gated_sm90_sm100 = fail + try: + yield + finally: + module.gemm_gated_sm90_sm100 = original + + def _assert_byte_exact(actual, expected, label): a = actual.contiguous().view(torch.uint8) b = expected.contiguous().view(torch.uint8) @@ -39,12 +69,23 @@ def _current_raw_stream(): return int(stream.cuda_stream) -def test_gemm_gated_tvm_ffi_matches_executor_combined_quant(): +@pytest.mark.parametrize( + "b_pretransposed,preallocated", + [(False, False), (True, False), (True, True)], + ids=["legacy_b", "native_enk", "native_enk_preallocated"], +) +@pytest.mark.parametrize( + "epilogue_mode", ["postact", "combined"], ids=["postact", "combined"] +) +def test_gemm_gated_tvm_ffi_matches_executor_quantized_postact( + b_pretransposed, preallocated, epilogue_mode +): from sonicmoe.quack_utils.blockscaled_fp8_gemm import ( precompute_weight_fp8_for_fused_gated, quantize_and_pack_activation, ) from sonicmoe.quack_utils.gemm_interface import gemm_gated + from sonicmoe.quack_utils.gemm_gated import prepare_gated_tvm_ffi_weight T, H, I, E, topk = _SHAPE TK = T * topk // E @@ -60,21 +101,61 @@ def test_gemm_gated_tvm_ffi_matches_executor_combined_quant(): a_idx = torch.arange(total_m, dtype=torch.int32, device="cuda") x_fp8, a_scales = quantize_and_pack_activation(x) w_fp8, b_scales = precompute_weight_fp8_for_fused_gated(w1) + b = w_fp8.mT if b_pretransposed else w_fp8 + if b_pretransposed: + prepare_gated_tvm_ffi_weight(b) def run(enabled): - context = ( - nullcontext() - if enabled - else _force_executor("sonicmoe.quack_utils.gemm_gated", "_can_use_gated_tvm_ffi") - ) + if enabled and b_pretransposed: + context = _forbid_gated_executor_fallback() + elif enabled: + context = nullcontext() + else: + context = _disable_gated_tvm_ffi() with context: - z_scale = torch.empty((total_m, 2 * I // 32), dtype=torch.uint8, device="cuda") + z_scale = ( + torch.empty( + (total_m, 2 * I // 32), + dtype=torch.uint8, + device="cuda", + ) + if epilogue_mode == "combined" + else None + ) y1_scale = torch.empty((total_m // 128, I // 128, 512), dtype=torch.uint8, device="cuda") + preact = ( + torch.empty( + (total_m, 2 * I), + dtype=( + torch.float8_e4m3fn + if epilogue_mode == "combined" + else torch.bfloat16 + ), + device="cuda", + ) + if preallocated + else None + ) + postact = ( + torch.empty( + (total_m, I), + dtype=torch.float8_e4m3fn, + device="cuda", + ) + if preallocated + else None + ) z, y1 = gemm_gated( x_fp8, - w_fp8, + b, activation="swiglu", - out_dtype=torch.float8_e4m3fn, + preact_out=preact, + postact_out=postact, + out_dtype=( + torch.float8_e4m3fn + if epilogue_mode == "combined" + else torch.bfloat16 + ), postact_dtype=torch.float8_e4m3fn, cu_seqlens_m=cu_seqlens, A_idx=a_idx, @@ -84,16 +165,141 @@ def run(enabled): postact_scale_out=y1_scale, tuned=False, current_stream=_current_raw_stream(), + b_pretransposed=b_pretransposed, ) torch.cuda.synchronize() - return z, y1, z_scale, y1_scale + outputs = [z, y1, y1_scale] + if z_scale is not None: + outputs.append(z_scale) + return outputs executor = run(False) tvm_ffi = run(True) - for label, lhs, rhs in zip(("z", "y1", "z_scale", "y1_scale"), tvm_ffi, executor): + labels = ["z", "y1", "y1_scale"] + if epilogue_mode == "combined": + labels.append("z_scale") + for label, lhs, rhs in zip(labels, tvm_ffi, executor): _assert_byte_exact(lhs, rhs, label) +@pytest.mark.parametrize("epilogue_mode", ["postact", "combined"], ids=["postact", "combined"]) +def test_gemm_gated_prepared_plan_rebuilds_microbatch_pointers(epilogue_mode): + """A prepared plan never retains activation, routing, scale, or output pointers.""" + from sonicmoe.quack_utils.blockscaled_fp8_gemm import ( + precompute_weight_fp8_for_fused_gated, + quantize_and_pack_activation, + ) + from sonicmoe.quack_utils.gemm_gated import prepare_gated_tvm_ffi_weight, try_gemm_gated_tvm_ffi_sonic_prepared + from sonicmoe.quack_utils.gemm_interface import gemm_gated + + T, H, I, E, topk = _SHAPE + TK = T * topk // E + total_m = TK * E + rng = np.random.RandomState(314159) + w1 = torch.from_numpy((rng.randn(2 * I, H, E).astype(np.float32) * 0.02)).to("cuda", dtype=torch.bfloat16) + w_fp8, b_scales = precompute_weight_fp8_for_fused_gated(w1) + b = prepare_gated_tvm_ffi_weight(w_fp8.mT) + + def make_inputs(seed): + local_rng = np.random.RandomState(seed) + x = torch.from_numpy((local_rng.randn(total_m, H).astype(np.float32) * 0.02)).to("cuda", dtype=torch.bfloat16) + x_fp8, a_scales = quantize_and_pack_activation(x) + a_idx = torch.arange(total_m - 1, -1, -1, dtype=torch.int32, device="cuda") + cu_seqlens = torch.arange(0, (E + 1) * TK, TK, dtype=torch.int32, device="cuda") + return x_fp8, a_scales, a_idx, cu_seqlens + + def allocate_outputs(): + z_scale = ( + torch.empty((total_m, 2 * I // 32), dtype=torch.uint8, device="cuda") + if epilogue_mode == "combined" + else None + ) + return ( + torch.empty( + (total_m, 2 * I), + dtype=(torch.float8_e4m3fn if epilogue_mode == "combined" else torch.bfloat16), + device="cuda", + ), + torch.empty((total_m, I), dtype=torch.float8_e4m3fn, device="cuda"), + z_scale, + torch.empty( + (total_m // 128, I // 128, 512), + dtype=torch.uint8, + device="cuda", + ), + ) + + # Populate the static plan through the fully validated path. + x1, scales1, idx1, offsets1 = make_inputs(1) + z1, y1, z_scale1, y1_scale1 = allocate_outputs() + gemm_gated( + x1, + b, + activation="swiglu", + preact_out=z1, + postact_out=y1, + out_dtype=z1.dtype, + postact_dtype=y1.dtype, + cu_seqlens_m=offsets1, + A_idx=idx1, + a_scales=scales1, + b_scales=b_scales, + z_scale_out=z_scale1, + postact_scale_out=y1_scale1, + tuned=False, + current_stream=_current_raw_stream(), + b_pretransposed=True, + ) + torch.cuda.synchronize() + + # Use distinct tensors for every dynamic argument. The executor provides + # the expected bytes; the prepared launch must write only the new outputs. + x2, scales2, idx2, offsets2 = make_inputs(2) + expected = allocate_outputs() + with _disable_gated_tvm_ffi(): + gemm_gated( + x2, + b, + activation="swiglu", + preact_out=expected[0], + postact_out=expected[1], + out_dtype=expected[0].dtype, + postact_dtype=expected[1].dtype, + cu_seqlens_m=offsets2, + A_idx=idx2, + a_scales=scales2, + b_scales=b_scales, + z_scale_out=expected[2], + postact_scale_out=expected[3], + tuned=False, + current_stream=_current_raw_stream(), + b_pretransposed=True, + ) + actual = allocate_outputs() + launched = try_gemm_gated_tvm_ffi_sonic_prepared( + x2, + b, + actual[0], + actual[1], + activation="swiglu", + cu_seqlens_m=offsets2, + A_idx=idx2, + a_scales=scales2, + b_scales=b_scales, + z_scale_out=actual[2], + postact_scale_out=actual[3], + swiglu_clamp_value=0.0, + postact_bf16_trunc=False, + current_stream=_current_raw_stream(), + ) + assert launched + torch.cuda.synchronize() + labels = ("z", "y1", "z_scale", "y1_scale") + for label, lhs, rhs in zip(labels, actual, expected): + if lhs is not None: + _assert_byte_exact(lhs, rhs, label) + + def test_gemm_dgated_tvm_ffi_matches_executor_fp8_preact_reduce(): from sonicmoe.quack_utils.blockscaled_fp8_gemm import ( precompute_weight_fp8_for_direct_fused_dgated, diff --git a/tests/ops/test_moe_module.py b/tests/ops/test_moe_module.py index 8938f23..e04c5e0 100644 --- a/tests/ops/test_moe_module.py +++ b/tests/ops/test_moe_module.py @@ -1391,6 +1391,159 @@ def test_sonicmoe_bf16_backward_vs_gold(T, H, I, E, K, seed): assert r_dw2 < BF16_DW_RRMSE, f"dw2 RRMSE {r_dw2:.6f} >= {BF16_DW_RRMSE}" +def test_down_projection_router_score_bridge_bit_exact(): + """The fused DeepEP score edge matches the standalone carrier exactly.""" + from sonicmoe.enums import ActivationType + from sonicmoe.functional import _DownProjection, _UpProjection + from sonicmoe.functional import utils as _fp8_utils + from sonicmoe.functional.triton_kernels import TC_topk_router_metadata_triton + from sonicmoe.quack_utils.blockscaled_fp8_gemm import _gather_router_scores_i32, _scatter_router_scores_i32 + + T, H, I, E, K = 256, 768, 384, 8, 2 + TK = T * K + N_PAD = 128 + x, w1_split, w2, topk_indices, source_values = _make_test_data(T, H, I, E, K, seed=314159) + w1_param, w2_param = _convert_weights_for_sonicmoe(w1_split, w2) + # The current functional API receives expert-major Sonic weights directly. + w1_func = w1_param + w2_func = w2_param + w2_func.stop_gradient = False + + s_scatter_idx = torch.empty(TK, dtype=torch.int32, device="cuda") + s_reverse_scatter_idx = torch.empty(TK, dtype=torch.int32, device="cuda") + expert_frequency = torch.empty(E, dtype=torch.int32, device="cuda") + expert_frequency_offset = torch.empty(E + 1, dtype=torch.int32, device="cuda") + x_gather_idx = torch.empty(TK, dtype=torch.int32, device="cuda") + TC_topk_router_metadata_triton( + topk_indices, + E, + expert_frequency, + expert_frequency_offset, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + ) + token_offsets = torch.arange(0, TK + 1, K, dtype=torch.int32, device="cuda") + + old_env = os.environ.get("SONIC_MOE_FP8_MODE") + old_fp8_active = _fp8_utils._IS_FP8_ACTIVE + os.environ["SONIC_MOE_FP8_MODE"] = "off" + _fp8_utils._IS_FP8_ACTIVE = False + try: + y1, z = _UpProjection.apply( + x, + w1_func, + None, + expert_frequency_offset, + TK, + K, + 0, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + None, + False, + ActivationType.SWIGLU, + False, + False, + ) + y1 = y1.detach() + z = z.detach() + z.stop_gradient = False + + # DeepEP may reorder received rows. Reverse order makes an accidental + # identity mapping visible while keeping every source index unique. + score_src_idx = torch.arange(TK - 1, -1, -1, dtype=torch.int32, device="cuda") + metadata_values = _gather_router_scores_i32(source_values.detach().reshape([-1]), score_src_idx) + # Real DeepEP metadata is block-aligned and can contain score padding. + # Keep source indices route-sized while exercising a larger score buffer. + metadata_values = torch.cat( + [ + metadata_values, + torch.zeros(N_PAD, dtype=metadata_values.dtype, device="cuda"), + ] + ) + grad_out = torch.randn(T, H, dtype=torch.bfloat16, device="cuda") + + metadata_reference = metadata_values.detach().clone() + metadata_reference.stop_gradient = False + z_reference = z.detach().clone() + z_reference.stop_gradient = False + w2_reference = w2_func.detach().clone() + w2_reference.stop_gradient = False + out_reference = _DownProjection.apply( + y1, + z_reference, + w2_reference, + None, + metadata_reference, + s_scatter_idx, + expert_frequency_offset, + T, + K, + 0, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + token_offsets, + True, + ActivationType.SWIGLU, + None, + ) + out_reference.backward(grad_out) + expected_source_grad = _scatter_router_scores_i32( + metadata_reference.grad.contiguous(), score_src_idx, TK + ).reshape([T, K]) + expected_z_grad = z_reference.grad.detach().clone() + expected_w2_grad = w2_reference.grad.detach().clone() + + source_fused = source_values.detach().clone() + source_fused.stop_gradient = False + metadata_fused = metadata_values.detach().clone() + z_fused = z.detach().clone() + z_fused.stop_gradient = False + w2_fused = w2_func.detach().clone() + w2_fused.stop_gradient = False + out_fused = _DownProjection.apply( + y1, + z_fused, + w2_fused, + None, + metadata_fused, + s_scatter_idx, + expert_frequency_offset, + T, + K, + 0, + x_gather_idx, + s_scatter_idx, + s_reverse_scatter_idx, + token_offsets, + True, + ActivationType.SWIGLU, + None, + None, + None, + source_fused, + score_src_idx, + ) + out_fused.backward(grad_out) + + assert bool(torch.all(out_fused.view(torch.uint8) == out_reference.view(torch.uint8)).item()) + assert bool(torch.all(source_fused.grad.view(torch.uint8) == expected_source_grad.view(torch.uint8)).item()) + assert bool(torch.all(z_fused.grad.view(torch.uint8) == expected_z_grad.view(torch.uint8)).item()) + assert bool(torch.all(w2_fused.grad.view(torch.uint8) == expected_w2_grad.view(torch.uint8)).item()) + assert source_fused.stop_gradient is False + assert metadata_fused.stop_gradient is True + assert metadata_fused.grad is None + finally: + _fp8_utils._IS_FP8_ACTIVE = old_fp8_active + if old_env is None: + os.environ.pop("SONIC_MOE_FP8_MODE", None) + else: + os.environ["SONIC_MOE_FP8_MODE"] = old_env + + # ═══════════════════════════════════════════════════════════════════════════ # Test 14: Gold reference handles all-same-expert routing # ═══════════════════════════════════════════════════════════════════════════