From e12bf66b1db004d2bc93c5df1b65160b88aa3425 Mon Sep 17 00:00:00 2001 From: MadeBy561 Date: Thu, 30 Jul 2026 15:34:46 -0400 Subject: [PATCH] moe: add isolated mixed K3/K4 Trellis path --- .../moe/_shared/kernels/w4a16/kernel.py | 1450 +++++++++++++++++ sparkinfer/moe/_trellis_moe/__init__.py | 100 ++ sparkinfer/moe/_trellis_moe/_impl.py | 778 +++++++++ sparkinfer/moe/_trellis_moe/_mixk.py | 462 ++++++ sparkinfer/moe/_trellis_moe/api.py | 52 + 5 files changed, 2842 insertions(+) create mode 100644 sparkinfer/moe/_trellis_moe/__init__.py create mode 100644 sparkinfer/moe/_trellis_moe/_impl.py create mode 100644 sparkinfer/moe/_trellis_moe/_mixk.py create mode 100644 sparkinfer/moe/_trellis_moe/api.py diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py index fc4ac543..d766761d 100644 --- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py +++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py @@ -5,6 +5,7 @@ import os from dataclasses import dataclass, replace from functools import partial +from types import SimpleNamespace from typing import NamedTuple import cuda.bindings.driver as cuda @@ -11405,14 +11406,1461 @@ def _weight_args(prepared) -> tuple[torch.Tensor, torch.Tensor]: return output +class W4A16FusedMoeFullRotationHybridKernel: + """One-grid mixed-bitrate Trellis MoE with the exact full-rotation contract. + + The existing hybrid kernel is the direct-topk/TC-decode path used by + packed NVFP4/NF3 mixtures. EXL3 Trellis is different: routes must be + packed once by GLOBAL expert, gate/up inputs and the SwiGLU intermediate + use the checkpoint's H128 rotations, and router weights are applied only + by the final FP32 top-k reduction. + + This kernel keeps that arithmetic byte-for-byte. A single global packed + route list drives both phases. For each packed expert block, + ``tier_local_map[global_expert]`` selects the K3 or K4 child decoder and + its tier-local expert index. Input rotation, activation rotation, and the + final reduction continue to index global-order rotation tables, so mixed + storage never changes the model's mathematical expert identity. + """ + + ABI_VERSION = 1 + + def __init__( + self, + *, + tier0: W4A16FusedMoeKernel, + tier1: W4A16FusedMoeKernel, + map_slots: int, + ): + for name, moe in (("tier0", tier0), ("tier1", tier1)): + if moe.direct_topk_routes: + raise ValueError( + f"full-rotation hybrid W4A16 {name} forbids direct_topk_routes" + ) + if moe.tc_decode_fused_sum: + raise ValueError( + f"full-rotation hybrid W4A16 {name} forbids TC fused sum" + ) + if not moe.activation_is_gated: + raise ValueError( + f"full-rotation hybrid W4A16 {name} requires gated activation" + ) + if moe.collect_activation_amax: + raise ValueError( + f"full-rotation hybrid W4A16 {name} forbids activation amax" + ) + if not (moe.intermediate_rotation and moe.dual_a and moe.full_rotation): + raise ValueError( + f"full-rotation hybrid W4A16 {name} requires the projection-" + "major Trellis rotation path" + ) + if moe.weight_layout != "trellis3_t256": + raise ValueError( + f"full-rotation hybrid W4A16 {name} requires trellis3_t256" + ) + if moe.element_dtype != "fp16": + raise ValueError( + f"full-rotation hybrid W4A16 {name} requires fp16 operands" + ) + if moe.apply_router_weight_on_input or moe.zero_fc2_output: + raise ValueError( + f"full-rotation hybrid W4A16 {name} has invalid reduction policy" + ) + for attr in ( + "size_m", + "hidden_size", + "intermediate_size", + "fc1_cols", + "top_k", + "moe_block_size", + "activation", + "activation_is_swigluoai", + "has_swiglu_limit", + "swiglu_limit", + "swiglu_alpha", + "swiglu_beta", + "element_dtype", + "is_fp16", + "fast_math", + "apply_router_weight_on_input", + "cta_threads", + "sms", + "blocks_per_sm", + "barrier_count_off", + "barrier_sense_off", + "schedule_whole_tiles", + "intermediate_rotation", + "dual_a", + "full_rotation", + "rotation_input_dtype", + ): + if getattr(tier0, attr) != getattr(tier1, attr): + raise ValueError( + f"full-rotation hybrid W4A16 tiers disagree on {attr}: " + f"{getattr(tier0, attr)!r} != {getattr(tier1, attr)!r}" + ) + for phase in ("fc1", "fc2"): + gemm0 = getattr(tier0, phase) + gemm1 = getattr(tier1, phase) + if ( + gemm0.n_tiles, + gemm0.k_tiles, + gemm0.tile_n, + gemm0.tile_k, + gemm0.top_k, + gemm0.mul_topk_weights, + gemm0.fused_topk_sum, + gemm0.moe_block_size, + gemm0.cta_threads, + gemm0.schedule_whole_tiles, + gemm0.route_major_a, + gemm0.dual_a, + ) != ( + gemm1.n_tiles, + gemm1.k_tiles, + gemm1.tile_n, + gemm1.tile_k, + gemm1.top_k, + gemm1.mul_topk_weights, + gemm1.fused_topk_sum, + gemm1.moe_block_size, + gemm1.cta_threads, + gemm1.schedule_whole_tiles, + gemm1.route_major_a, + gemm1.dual_a, + ): + raise ValueError( + f"full-rotation hybrid W4A16 tiers disagree on {phase} geometry" + ) + if int(map_slots) < tier0.num_experts + tier1.num_experts: + raise ValueError( + "full-rotation hybrid map_slots must cover both tier expert sets" + ) + if tier0.num_experts > 256 or tier1.num_experts > 256: + raise ValueError( + "full-rotation hybrid local expert ids exceed descriptor capacity" + ) + self.tier0 = tier0 + self.tier1 = tier1 + self.map_slots = int(map_slots) + self.size_m = tier0.size_m + self.hidden_size = tier0.hidden_size + self.intermediate_size = tier0.intermediate_size + self.top_k = tier0.top_k + self.element_dtype = tier0.element_dtype + self.rotation_input_dtype = tier0.rotation_input_dtype + self.cta_threads = tier0.cta_threads + self.sms = tier0.sms + self.blocks_per_sm = min(tier0.blocks_per_sm, tier1.blocks_per_sm) + self.shared_words = max(tier0.shared_words, tier1.shared_words) + + @property + def __cache_key__(self) -> tuple[object, ...]: + return ( + "w4a16_fused_moe_full_rotation_hybrid", + self.ABI_VERSION, + self.map_slots, + self.tier0.__cache_key__, + self.tier1.__cache_key__, + self.shared_words, + ) + + @cute.jit + def _emit_route_map_tile( + self, + is_fc1: cutlass.Constexpr, + a_fp16_flat: cute.Tensor, + a_alt_fp16_flat: cute.Tensor, + t0_b_i32_flat: cute.Tensor, + t0_scales_i32_flat: cute.Tensor, + t0_global_scale: cute.Tensor, + t1_b_i32_flat: cute.Tensor, + t1_scales_i32_flat: cute.Tensor, + t1_global_scale: cute.Tensor, + c_fp16_flat: cute.Tensor, + packed_route_indices: cute.Tensor, + block_expert_ids: cute.Tensor, + tier_local_map_i32_flat: cute.Tensor, + topk_weights_flat: cute.Tensor, + c_tmp_f32_flat: cute.Tensor, + locks_i32_flat: cute.Tensor, + smem_base: Int32, + tid: Int32, + active_size_m: Int32, + route_block_idx: Int32, + output_n_tile: Int32, + reduce_k_tile: Int32, + reduce_tile_count: Int32, + reduce_slice_count: Int32, + reduce_slice_idx: Int32, + lock_slot: Int32, + ): + gid = block_expert_ids[route_block_idx].to(Int32) + if gid >= Int32(0) and gid < Int32(self.map_slots): + descriptor = tier_local_map_i32_flat[gid].to(Int32) + if descriptor >= Int32(0): + tier = descriptor >> Int32(8) + local_expert = descriptor & Int32(0xFF) + if tier == Int32(0): + if local_expert < Int32(self.tier0.num_experts): + if cutlass.const_expr(is_fc1): + self.tier0.fc1._run_tile( + a_fp16_flat, + a_alt_fp16_flat, + t0_b_i32_flat, + c_fp16_flat, + t0_scales_i32_flat, + t0_global_scale, + packed_route_indices, + topk_weights_flat, + c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot, + active_size_m, + ) + else: + self.tier0.fc2._run_tile( + a_fp16_flat, + a_alt_fp16_flat, + t0_b_i32_flat, + c_fp16_flat, + t0_scales_i32_flat, + t0_global_scale, + packed_route_indices, + topk_weights_flat, + c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot, + active_size_m, + ) + elif tier == Int32(1): + if local_expert < Int32(self.tier1.num_experts): + if cutlass.const_expr(is_fc1): + self.tier1.fc1._run_tile( + a_fp16_flat, + a_alt_fp16_flat, + t1_b_i32_flat, + c_fp16_flat, + t1_scales_i32_flat, + t1_global_scale, + packed_route_indices, + topk_weights_flat, + c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot, + active_size_m, + ) + else: + self.tier1.fc2._run_tile( + a_fp16_flat, + a_alt_fp16_flat, + t1_b_i32_flat, + c_fp16_flat, + t1_scales_i32_flat, + t1_global_scale, + packed_route_indices, + topk_weights_flat, + c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + route_block_idx, + local_expert, + output_n_tile, + reduce_k_tile, + reduce_tile_count, + reduce_slice_count, + reduce_slice_idx, + lock_slot, + active_size_m, + ) + + @cute.jit + def __call__( + self, + a_gate_fp16_ptr: cute.Pointer, + a_up_fp16_ptr: cute.Pointer, + rotation_input_ptr: cute.Pointer, + t0_w13_i32_flat: cute.Tensor, + t0_w2_i32_flat: cute.Tensor, + t0_w13_scales_i32_flat: cute.Tensor, + t0_w2_scales_i32_flat: cute.Tensor, + t0_w13_global_scale: cute.Tensor, + t0_w2_global_scale: cute.Tensor, + t1_w13_i32_flat: cute.Tensor, + t1_w2_i32_flat: cute.Tensor, + t1_w13_scales_i32_flat: cute.Tensor, + t1_w2_scales_i32_flat: cute.Tensor, + t1_w13_global_scale: cute.Tensor, + t1_w2_global_scale: cute.Tensor, + packed_route_indices: cute.Tensor, + block_expert_ids: cute.Tensor, + packed_route_count: cute.Tensor, + tier_local_map_i32_flat: cute.Tensor, + fc1_fp16_flat: cute.Tensor, + activated_fp16_flat: cute.Tensor, + fc2_fp16_flat: cute.Tensor, + topk_weights_ptr: cute.Pointer, + fc1_c_tmp_f32_flat: cute.Tensor, + fc2_c_tmp_f32_flat: cute.Tensor, + locks_i32_flat: cute.Tensor, + rot_scales_flat: cute.Tensor, + suh_gate_flat: cute.Tensor, + suh_up_flat: cute.Tensor, + active_m: cutlass.Int32, + grid_x: cutlass.Int32, + stream: cuda.CUstream, + ): + routed_rows = active_m * Int32(self.top_k) + a_gate_fp16_flat = cute.make_tensor( + a_gate_fp16_ptr, + layout=cute.make_layout( + (routed_rows * Int32(self.hidden_size),), stride=(1,) + ), + ) + a_up_fp16_flat = cute.make_tensor( + a_up_fp16_ptr, + layout=cute.make_layout( + (routed_rows * Int32(self.hidden_size),), stride=(1,) + ), + ) + rotation_input_flat = cute.make_tensor( + rotation_input_ptr, + layout=cute.make_layout((active_m * Int32(self.hidden_size),), stride=(1,)), + ) + topk_weights_flat = cute.make_tensor( + topk_weights_ptr, + layout=cute.make_layout((routed_rows,), stride=(1,)), + ) + self.kernel( + a_gate_fp16_flat, + a_up_fp16_flat, + rotation_input_flat, + t0_w13_i32_flat, + t0_w2_i32_flat, + t0_w13_scales_i32_flat, + t0_w2_scales_i32_flat, + t0_w13_global_scale, + t0_w2_global_scale, + t1_w13_i32_flat, + t1_w2_i32_flat, + t1_w13_scales_i32_flat, + t1_w2_scales_i32_flat, + t1_w13_global_scale, + t1_w2_global_scale, + packed_route_indices, + block_expert_ids, + packed_route_count, + tier_local_map_i32_flat, + fc1_fp16_flat, + activated_fp16_flat, + fc2_fp16_flat, + topk_weights_flat, + fc1_c_tmp_f32_flat, + fc2_c_tmp_f32_flat, + locks_i32_flat, + rot_scales_flat, + suh_gate_flat, + suh_up_flat, + active_m, + ).launch( + grid=(grid_x, 1, 1), + block=[self.cta_threads, 1, 1], + min_blocks_per_mp=self.blocks_per_sm, + cooperative=True, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + a_gate_fp16_flat: cute.Tensor, + a_up_fp16_flat: cute.Tensor, + rotation_input_flat: cute.Tensor, + t0_w13_i32_flat: cute.Tensor, + t0_w2_i32_flat: cute.Tensor, + t0_w13_scales_i32_flat: cute.Tensor, + t0_w2_scales_i32_flat: cute.Tensor, + t0_w13_global_scale: cute.Tensor, + t0_w2_global_scale: cute.Tensor, + t1_w13_i32_flat: cute.Tensor, + t1_w2_i32_flat: cute.Tensor, + t1_w13_scales_i32_flat: cute.Tensor, + t1_w2_scales_i32_flat: cute.Tensor, + t1_w13_global_scale: cute.Tensor, + t1_w2_global_scale: cute.Tensor, + packed_route_indices: cute.Tensor, + block_expert_ids: cute.Tensor, + packed_route_count: cute.Tensor, + tier_local_map_i32_flat: cute.Tensor, + fc1_fp16_flat: cute.Tensor, + activated_fp16_flat: cute.Tensor, + fc2_fp16_flat: cute.Tensor, + topk_weights_flat: cute.Tensor, + fc1_c_tmp_f32_flat: cute.Tensor, + fc2_c_tmp_f32_flat: cute.Tensor, + locks_i32_flat: cute.Tensor, + rot_scales_flat: cute.Tensor, + suh_gate_flat: cute.Tensor, + suh_up_flat: cute.Tensor, + active_m: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + grid_x_raw, _, _ = cute.arch.grid_dim() + tid = Int32(tidx) + cta = Int32(bidx) + grid_x = Int32(grid_x_raw) + + smem = cutlass.utils.SmemAllocator() + + @cute.struct + class Storage: + words: cute.struct.Align[ + cute.struct.MemRange[cutlass.Uint32, self.shared_words], + 1024, + ] + + storage = smem.allocate(Storage) + smem_base = shared_ptr_to_u32(storage.words.data_ptr()) + + fc1_emit_tile = partial( + self._emit_route_map_tile, + True, + a_gate_fp16_flat, + a_up_fp16_flat, + t0_w13_i32_flat, + t0_w13_scales_i32_flat, + t0_w13_global_scale, + t1_w13_i32_flat, + t1_w13_scales_i32_flat, + t1_w13_global_scale, + fc1_fp16_flat, + packed_route_indices, + block_expert_ids, + tier_local_map_i32_flat, + topk_weights_flat, + fc1_c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + active_m, + ) + fc2_emit_tile = partial( + self._emit_route_map_tile, + False, + activated_fp16_flat, + activated_fp16_flat, + t0_w2_i32_flat, + t0_w2_scales_i32_flat, + t0_w2_global_scale, + t1_w2_i32_flat, + t1_w2_scales_i32_flat, + t1_w2_global_scale, + fc2_fp16_flat, + packed_route_indices, + block_expert_ids, + tier_local_map_i32_flat, + topk_weights_flat, + fc2_c_tmp_f32_flat, + locks_i32_flat, + smem_base, + tid, + active_m * Int32(self.top_k), + ) + # Tier 0 supplies the shared full-rotation phase machinery. The emit + # hooks above own only the bitrate/local-expert weight selection. + self.tier0._moe_body( + a_gate_fp16_flat, + a_up_fp16_flat, + rotation_input_flat, + t0_w13_i32_flat, + t0_w2_i32_flat, + fc1_fp16_flat, + activated_fp16_flat, + fc2_fp16_flat, + t0_w13_scales_i32_flat, + t0_w2_scales_i32_flat, + t0_w13_global_scale, + t0_w2_global_scale, + packed_route_indices, + block_expert_ids, + packed_route_count, + packed_route_count, + Int32(0), + topk_weights_flat, + fc1_c_tmp_f32_flat, + fc2_c_tmp_f32_flat, + locks_i32_flat, + rot_scales_flat, + suh_gate_flat, + suh_up_flat, + smem_base, + tid, + cta, + grid_x, + active_m, + fc1_emit_tile, + fc2_emit_tile, + ) + + +class W4A16FusedMoeFullRotationHybridCompileResult: + compiled: object + size_m: int + hidden_size: int + intermediate_size: int + top_k: int + activation: str + map_slots: int + tier0_num_experts: int + tier0_trellis_bits: int + tier1_num_experts: int + tier1_trellis_bits: int + fc1_tile_n: int + fc1_tile_k: int + fc2_tile_n: int + fc2_tile_k: int + moe_block_size: int + max_m_blocks: int + cta_threads: int + blocks_per_sm: int + shared_memory_bytes: int + rotation_input_dtype: str + registers_per_thread: int + local_memory_bytes: int + + +def compile_w4a16_fused_moe_full_rotation_hybrid( + *, + size_m: int, + route_capacity_m_blocks: int | None = None, + hidden_size: int, + intermediate_size: int, + tier0_num_experts: int, + tier0_trellis_bits: int, + tier1_num_experts: int, + tier1_trellis_bits: int, + top_k: int, + activation: str, + map_slots: int, + moe_block_size: int, + rotation_input_dtype: str, + fast_math: bool, + sms: int, + max_shared_mem: int, + force_tile_config: tuple[int, int, int, int], +) -> W4A16FusedMoeFullRotationHybridCompileResult: + """Compile the one-grid K3/K4 projection-major Trellis kernel. + + Both tier decoders retain their own compile-time bitrate and weight + strides. The route scheduler, H128 transforms, activation, and final route + buffers are shared. This is deliberately a separate ABI from the + direct-topk NVFP4/NF3 hybrid so neither path can silently select the wrong + reduction contract. + """ + + activation = normalize_moe_activation(activation) + if not validate_activation(activation): + raise ValueError("full-rotation hybrid W4A16 requires gated activation") + if activation != "silu": + raise ValueError("full-rotation Trellis hybrid requires silu") + rotation_input_dtype = str(rotation_input_dtype) + if rotation_input_dtype not in {"bf16", "fp16"}: + raise ValueError("rotation_input_dtype must be 'bf16' or 'fp16'") + tier0_trellis_bits = int(tier0_trellis_bits) + tier1_trellis_bits = int(tier1_trellis_bits) + if tier0_trellis_bits not in _TRELLIS256_BITS: + raise ValueError(f"invalid tier0 Trellis bits {tier0_trellis_bits}") + if tier1_trellis_bits not in _TRELLIS256_BITS: + raise ValueError(f"invalid tier1 Trellis bits {tier1_trellis_bits}") + if tier0_trellis_bits == tier1_trellis_bits: + raise ValueError("full-rotation hybrid tiers must use distinct bitrates") + fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = ( + int(value) for value in force_tile_config + ) + route_slots = max_packed_route_slots( + int(size_m) * int(top_k), + int(moe_block_size), + int(map_slots), + ) + live_m_blocks = (int(route_slots) + int(moe_block_size) - 1) // int(moe_block_size) + max_m_blocks = ( + live_m_blocks + if route_capacity_m_blocks is None + else int(route_capacity_m_blocks) + ) + if max_m_blocks < live_m_blocks: + raise ValueError( + "full-rotation hybrid route capacity under-covers the exact-M " + f"launch: capacity_blocks={max_m_blocks}, " + f"required_blocks={live_m_blocks}" + ) + + def _tier_kernel(num_experts: int, trellis_bits: int) -> W4A16FusedMoeKernel: + return W4A16FusedMoeKernel( + size_m=int(size_m), + hidden_size=int(hidden_size), + intermediate_size=int(intermediate_size), + num_experts=int(num_experts), + top_k=int(top_k), + activation=activation, + apply_router_weight_on_input=False, + zero_fc2_output=False, + fc1_tile_n=fc1_tile_n, + fc1_tile_k=fc1_tile_k, + fc2_tile_n=fc2_tile_n, + fc2_tile_k=fc2_tile_k, + moe_block_size=int(moe_block_size), + max_m_blocks=max_m_blocks, + element_dtype="fp16", + fast_math=bool(fast_math), + weight_layout="trellis3_t256", + scale_format="e4m3_k32", + w13_layout="trellis3_t256_proj", + trellis_bits=int(trellis_bits), + direct_topk_routes=False, + tc_decode_fused_sum=False, + collect_activation_amax=False, + schedule_whole_tiles=True, + intermediate_rotation=True, + full_rotation=True, + rotation_input_dtype=rotation_input_dtype, + ) + + kernel = W4A16FusedMoeFullRotationHybridKernel( + tier0=_tier_kernel(int(tier0_num_experts), tier0_trellis_bits), + tier1=_tier_kernel(int(tier1_num_experts), tier1_trellis_bits), + map_slots=int(map_slots), + ) + if kernel.shared_words * 4 > int(max_shared_mem) - 512: + raise ValueError( + "full-rotation hybrid W4A16 shared memory exceeds device limit: " + f"{kernel.shared_words * 4} > {int(max_shared_mem) - 512}" + ) + + device = int(torch.cuda.current_device()) if torch.cuda.is_available() else None + cache_key = ( + "w4a16_fused_moe_full_rotation_hybrid", + device, + kernel.__cache_key__, + ) + cached = _FUSED_CACHE.get(cache_key) + if cached is not None: + return replace(cached, size_m=int(size_m), max_m_blocks=max_m_blocks) + + compile_size_m = _fake_m_for_specialization(int(size_m)) + compile_routed_rows = int(compile_size_m) * int(top_k) + compile_route_blocks = max(1, min(max_m_blocks, compile_routed_rows)) + compile_route_slots = compile_route_blocks * int(moe_block_size) + fc1_cols = 2 * int(intermediate_size) + + def _trellis_weight_fake( + num_experts: int, + size_n: int, + size_k: int, + trellis_bits: int, + ): + return cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + ( + int(num_experts) + * (int(size_k) // 16) + * (int(size_n) // 16) + * (8 * int(trellis_bits)), + ), + assumed_align=16, + ) + + def _scales_fake(num_experts: int, size_n: int, size_k: int): + return cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + ( + _scale_fake_int32_elements( + num_experts=int(num_experts), + size_k=int(size_k), + size_n=int(size_n), + scale_format="e4m3_k32", + ), + ), + assumed_align=16, + ) + + def _global_fake(num_experts: int): + return cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (int(num_experts),), assumed_align=16 + ) + + fp16_ptr = make_ptr(cutlass.Float16, 16, cute.AddressSpace.gmem, assumed_align=16) + rotation_input_ptr = make_ptr( + _cutlass_element_dtype(rotation_input_dtype), + 16, + cute.AddressSpace.gmem, + assumed_align=16, + ) + topk_ptr = make_ptr(cutlass.Float32, 4, cute.AddressSpace.gmem, assumed_align=4) + scratch_elements = max( + fc1_cols * compile_route_slots, + int(hidden_size) * compile_route_slots, + 4 * 256 * int(moe_block_size) * 256, + ) + t0_ne = int(tier0_num_experts) + t1_ne = int(tier1_num_experts) + compile_args = ( + fp16_ptr, + fp16_ptr, + rotation_input_ptr, + _trellis_weight_fake(t0_ne, fc1_cols, hidden_size, tier0_trellis_bits), + _trellis_weight_fake(t0_ne, hidden_size, intermediate_size, tier0_trellis_bits), + _scales_fake(t0_ne, fc1_cols, hidden_size), + _scales_fake(t0_ne, hidden_size, intermediate_size), + _global_fake(t0_ne), + _global_fake(t0_ne), + _trellis_weight_fake(t1_ne, fc1_cols, hidden_size, tier1_trellis_bits), + _trellis_weight_fake(t1_ne, hidden_size, intermediate_size, tier1_trellis_bits), + _scales_fake(t1_ne, fc1_cols, hidden_size), + _scales_fake(t1_ne, hidden_size, intermediate_size), + _global_fake(t1_ne), + _global_fake(t1_ne), + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (compile_route_slots,), assumed_align=16 + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (compile_route_blocks,), assumed_align=16 + ), + cute.runtime.make_fake_compact_tensor(cutlass.Int32, (1,), assumed_align=4), + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (int(map_slots),), assumed_align=16 + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float16, + (compile_routed_rows * fc1_cols,), + assumed_align=16, + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float16, + (compile_routed_rows * int(intermediate_size),), + assumed_align=16, + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float16, + (compile_routed_rows * int(hidden_size),), + assumed_align=16, + ), + topk_ptr, + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (scratch_elements,), assumed_align=16 + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (scratch_elements,), assumed_align=16 + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (4 * 256 + 2,), assumed_align=16 + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float16, + (int(map_slots) * 3 * int(intermediate_size),), + assumed_align=16, + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float16, + (int(map_slots) * int(hidden_size),), + assumed_align=16, + ), + cute.runtime.make_fake_compact_tensor( + cutlass.Float16, + (int(map_slots) * int(hidden_size),), + assumed_align=16, + ), + 1, + 1, + current_cuda_stream(), + ) + + raise_if_kernel_resolution_frozen( + "cute.compile", target=kernel, cache_key=cache_key + ) + compiled = sparkinfer_compile( + kernel, + *compile_args, + compile_spec=KernelCompileSpec.from_key( + "moe.w4a16.fused_moe_full_rotation_hybrid", + W4A16FusedMoeFullRotationHybridKernel.ABI_VERSION, + cache_key, + ), + dsl_compile_options=OptLevel(2), + ) + registers_per_thread = -1 + local_memory_bytes = -1 + resources = _query_w4a16_kernel_resources(compiled) + if resources is not None: + _, registers_per_thread, local_memory_bytes = resources + if local_memory_bytes != 0: + raise RuntimeError( + "full-rotation hybrid W4A16 codegen spills to local memory " + f"({local_memory_bytes} bytes/thread); refusing admission" + ) + + result = W4A16FusedMoeFullRotationHybridCompileResult( + compiled=compiled, + size_m=int(size_m), + hidden_size=int(hidden_size), + intermediate_size=int(intermediate_size), + top_k=int(top_k), + activation=activation, + map_slots=int(map_slots), + tier0_num_experts=t0_ne, + tier0_trellis_bits=tier0_trellis_bits, + tier1_num_experts=t1_ne, + tier1_trellis_bits=tier1_trellis_bits, + fc1_tile_n=fc1_tile_n, + fc1_tile_k=fc1_tile_k, + fc2_tile_n=fc2_tile_n, + fc2_tile_k=fc2_tile_k, + moe_block_size=int(moe_block_size), + max_m_blocks=max_m_blocks, + cta_threads=kernel.cta_threads, + blocks_per_sm=kernel.blocks_per_sm, + shared_memory_bytes=kernel.shared_words * 4, + rotation_input_dtype=rotation_input_dtype, + registers_per_thread=registers_per_thread, + local_memory_bytes=local_memory_bytes, + ) + _FUSED_CACHE[cache_key] = result + return result + + +def _run_w4a16_moe_full_rotation_hybrid_eager( + a_input: torch.Tensor, + prepared_tier0, + prepared_tier1, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + tier_local_map: torch.Tensor, + *, + activation: str, + intermediate_cache13: torch.Tensor, + intermediate_cache2: torch.Tensor, + output: torch.Tensor, + fc1_c_tmp: torch.Tensor, + fc2_c_tmp: torch.Tensor, + packed_route_indices: torch.Tensor, + block_expert_ids: torch.Tensor, + packed_route_count: torch.Tensor, + expert_offsets: torch.Tensor, + expert_counts: torch.Tensor, + rotation_a_gate: torch.Tensor, + rotation_a_up: torch.Tensor, + global_intermediate_rotations: torch.Tensor, + global_suh_gate: torch.Tensor, + global_suh_up: torch.Tensor, + global_svh_down: torch.Tensor, + fused_launch: W4A16FusedMoeFullRotationHybridCompileResult, + topk_sum_launch: W4A16TopKSumCompileResult, + fast_math: bool = True, + stream: cuda.CUstream | None = None, +) -> torch.Tensor: + """Run one globally packed mixed-K full-rotation Trellis MoE layer.""" + + activation = normalize_moe_activation(activation) + if activation != "silu": + raise ValueError("mixed-K full-rotation Trellis requires silu") + if a_input.ndim != 2 or not a_input.is_cuda or not a_input.is_contiguous(): + raise ValueError("a_input must be a contiguous CUDA rank-2 tensor") + if a_input.dtype not in (torch.bfloat16, torch.float16): + raise TypeError("a_input must be bf16 or fp16") + m, hidden_size = (int(a_input.shape[0]), int(a_input.shape[1])) + if m < 1 or m > int(fused_launch.size_m): + raise ValueError( + f"mixed-K launch requires 1 <= m <= {int(fused_launch.size_m)}, got {m}" + ) + topk = int(topk_ids.shape[1]) + if tuple(topk_ids.shape) != (m, topk): + raise ValueError("topk_ids must be rank-2 [m, topk]") + if topk_ids.dtype not in (torch.int32, torch.int64): + raise TypeError("topk_ids must be int32 or int64") + if ( + tuple(topk_weights.shape) != (m, topk) + or topk_weights.dtype != torch.float32 + or not topk_weights.is_contiguous() + ): + raise ValueError("topk_weights must be contiguous float32 [m, topk]") + if not topk_ids.is_contiguous() or topk_ids.device != a_input.device: + raise ValueError("topk_ids must be contiguous on the input device") + if topk_weights.device != a_input.device: + raise ValueError("topk_weights must be on the input device") + if hidden_size != int(fused_launch.hidden_size) or topk != int(fused_launch.top_k): + raise ValueError("input/router geometry does not match fused launch") + if output.dtype != torch.float32 or tuple(output.shape) != (m, hidden_size): + raise ValueError("output must be contiguous float32 [m, hidden]") + if output.device != a_input.device or not output.is_contiguous(): + raise ValueError("output must be contiguous on the input device") + + map_slots = int(fused_launch.map_slots) + if ( + tier_local_map.dtype != torch.int32 + or tier_local_map.device != a_input.device + or tuple(tier_local_map.shape) != (map_slots,) + or not tier_local_map.is_contiguous() + ): + raise ValueError("tier_local_map must be contiguous int32 [map_slots]") + intermediate_size = int(fused_launch.intermediate_size) + expected_tiers = ( + ( + "tier0", + prepared_tier0, + int(fused_launch.tier0_num_experts), + int(fused_launch.tier0_trellis_bits), + ), + ( + "tier1", + prepared_tier1, + int(fused_launch.tier1_num_experts), + int(fused_launch.tier1_trellis_bits), + ), + ) + for name, prepared, experts, bits in expected_tiers: + actual = ( + getattr(prepared, "weight_layout", None), + getattr(prepared, "trellis_codebook", None), + int(getattr(prepared, "trellis_bits", -1)), + int(getattr(prepared, "num_experts", -1)), + int(getattr(prepared, "hidden_size", -1)), + int(getattr(prepared, "intermediate_size", -1)), + getattr(prepared, "params_dtype", None), + ) + wanted = ( + "trellis3_t256", + "mcg", + bits, + experts, + hidden_size, + intermediate_size, + torch.float16, + ) + if actual != wanted: + raise ValueError( + f"{name} prepared Trellis contract mismatch: {actual} != {wanted}" + ) + + capacity_m = int(fused_launch.size_m) + capacity_routed_rows = capacity_m * topk + fc1_cols = 2 * intermediate_size + cache13_need = capacity_routed_rows * max(fc1_cols, hidden_size) + cache2_need = capacity_routed_rows * intermediate_size + for name, tensor, need in ( + ("intermediate_cache13", intermediate_cache13, cache13_need), + ("intermediate_cache2", intermediate_cache2, cache2_need), + ("rotation_a_gate", rotation_a_gate, capacity_routed_rows * hidden_size), + ("rotation_a_up", rotation_a_up, capacity_routed_rows * hidden_size), + ): + if ( + tensor.dtype != torch.float16 + or tensor.device != a_input.device + or not tensor.is_contiguous() + or int(tensor.numel()) < int(need) + ): + raise ValueError( + f"{name} must be contiguous fp16 on {a_input.device} with " + f"at least {int(need)} elements" + ) + if rotation_a_gate.data_ptr() == rotation_a_up.data_ptr(): + raise ValueError("rotation gate/up scratch must not alias") + for name, table, shape in ( + ("global_suh_gate", global_suh_gate, (map_slots, hidden_size)), + ("global_suh_up", global_suh_up, (map_slots, hidden_size)), + ("global_svh_down", global_svh_down, (map_slots, hidden_size)), + ( + "global_intermediate_rotations", + global_intermediate_rotations, + (map_slots, 3 * intermediate_size), + ), + ): + if ( + table.dtype != torch.float16 + or table.device != a_input.device + or tuple(table.shape) != shape + or not table.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous fp16 {shape} on {a_input.device}" + ) + + block_size_m = int(fused_launch.moe_block_size) + route_slots = max_packed_route_slots(capacity_m * topk, block_size_m, map_slots) + route_blocks = (route_slots + block_size_m - 1) // block_size_m + for name, tensor, dtype, need in ( + ("packed_route_indices", packed_route_indices, torch.int32, route_slots), + ("block_expert_ids", block_expert_ids, torch.int32, route_blocks), + ("packed_route_count", packed_route_count, torch.int32, 1), + ("expert_offsets", expert_offsets, torch.int32, map_slots + 1), + ("expert_counts", expert_counts, torch.int32, map_slots), + ): + if ( + tensor.dtype != dtype + or tensor.device != a_input.device + or not tensor.is_contiguous() + or int(tensor.numel()) < int(need) + ): + raise ValueError(f"{name} does not cover planned route capacity") + packed_route_indices = packed_route_indices[:route_slots] + block_expert_ids = block_expert_ids[:route_blocks] + expert_offsets = expert_offsets[: map_slots + 1] + expert_counts = expert_counts[:map_slots] + + pack_topk_routes_by_expert( + topk_ids, + block_size_m, + map_slots, + packed_route_indices=packed_route_indices, + block_expert_ids=block_expert_ids, + packed_route_count=packed_route_count, + expert_offsets=expert_offsets, + expert_counts=expert_counts, + ) + + props = torch.cuda.get_device_properties(a_input.device) + sms = int(props.multi_processor_count) + fc1_scratch_elements = packed_gemm_scratch_elements( + size_n=fc1_cols, + route_slots=route_slots, + moe_block_size=block_size_m, + sms=sms, + ) + fc2_scratch_elements = packed_gemm_scratch_elements( + size_n=hidden_size, + route_slots=route_slots, + moe_block_size=block_size_m, + sms=sms, + ) + fc1_scratch = _get_c_tmp( + fc1_scratch_elements, device=a_input.device, scratch=fc1_c_tmp + ) + fc2_scratch = _get_c_tmp( + fc2_scratch_elements, device=a_input.device, scratch=fc2_c_tmp + ) + workspace = prepared_tier0.workspace + if int(workspace.numel()) < sms * 4 + 2: + raise ValueError("tier0 workspace is too small for cooperative barriers") + + cache13_flat = intermediate_cache13.view(-1) + cache2_flat = intermediate_cache2.view(-1) + fc1_out = cache13_flat[: capacity_routed_rows * fc1_cols] + activated = cache2_flat[: capacity_routed_rows * intermediate_size] + fc2_out = cache13_flat[: capacity_routed_rows * hidden_size] + stream = current_cuda_stream() if stream is None else stream + grid_x = _w4a16_fused_persistent_grid_x( + fused=fused_launch, + m=m, + topk=topk, + intermediate_size=intermediate_size, + activation=activation, + direct_topk_routes=False, + sms=sms, + ) + compiled = fused_launch.compiled + compiled( + make_ptr( + cutlass.Float16, + rotation_a_gate.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + cutlass.Float16, + rotation_a_up.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + _cutlass_element_dtype(_normalize_element_dtype(a_input.dtype)), + a_input.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + prepared_tier0.w13.view(torch.int32).view(-1), + prepared_tier0.w2.view(torch.int32).view(-1), + prepared_tier0.w13_scale.view(torch.uint8).view(torch.int32).view(-1), + prepared_tier0.w2_scale.view(torch.uint8).view(torch.int32).view(-1), + prepared_tier0.w13_global_scale, + prepared_tier0.w2_global_scale, + prepared_tier1.w13.view(torch.int32).view(-1), + prepared_tier1.w2.view(torch.int32).view(-1), + prepared_tier1.w13_scale.view(torch.uint8).view(torch.int32).view(-1), + prepared_tier1.w2_scale.view(torch.uint8).view(torch.int32).view(-1), + prepared_tier1.w13_global_scale, + prepared_tier1.w2_global_scale, + packed_route_indices, + block_expert_ids, + packed_route_count, + tier_local_map, + fc1_out, + activated, + fc2_out, + make_ptr( + cutlass.Float32, + topk_weights.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=4, + ), + fc1_scratch, + fc2_scratch, + workspace, + global_intermediate_rotations.view(-1), + global_suh_gate.view(-1), + global_suh_up.view(-1), + m, + grid_x, + cuda.CUstream(int(stream)), + ) + + expected_sum = ( + topk, + hidden_size, + True, + map_slots, + 0, + topk_ids.dtype, + False, + ) + actual_sum = ( + int(topk_sum_launch.topk), + int(topk_sum_launch.hidden_size), + bool(topk_sum_launch.full_rotation), + int(topk_sum_launch.num_experts), + int(topk_sum_launch.route_num_experts), + topk_sum_launch.route_ids_dtype, + bool(topk_sum_launch.use_expert_map), + ) + if actual_sum != expected_sum: + raise ValueError( + f"mixed-K top-k sum launch mismatch: {actual_sum} != {expected_sum}" + ) + _w4a16_topk_sum_launch_flat( + fc2_out, + output, + m, + topk, + hidden_size, + "fp16", + int(stream), + full_rotation=True, + num_experts=map_slots, + topk_weights=topk_weights, + route_expert_ids=topk_ids, + expert_map=None, + svh_table=global_svh_down, + ) + return output + + +_FULL_ROTATION_HYBRID_OP_HANDLES: dict[ + int, + tuple[ + W4A16FusedMoeFullRotationHybridCompileResult, + W4A16TopKSumCompileResult, + bool, + ], +] = {} + + +_FULL_ROTATION_HYBRID_OP_KEYS: dict[tuple[int, int, bool], int] = {} + + +_FULL_ROTATION_HYBRID_OP_NEXT_HANDLE = 1 + + +def _register_full_rotation_hybrid_op_handle( + fused_launch: W4A16FusedMoeFullRotationHybridCompileResult, + topk_sum_launch: W4A16TopKSumCompileResult, + fast_math: bool, +) -> int: + """Retain opaque launch objects behind a Dynamo-safe scalar handle.""" + + global _FULL_ROTATION_HYBRID_OP_NEXT_HANDLE + key = (id(fused_launch), id(topk_sum_launch), bool(fast_math)) + handle = _FULL_ROTATION_HYBRID_OP_KEYS.get(key) + if handle is not None: + return handle + handle = _FULL_ROTATION_HYBRID_OP_NEXT_HANDLE + _FULL_ROTATION_HYBRID_OP_NEXT_HANDLE += 1 + _FULL_ROTATION_HYBRID_OP_KEYS[key] = handle + _FULL_ROTATION_HYBRID_OP_HANDLES[handle] = ( + fused_launch, + topk_sum_launch, + bool(fast_math), + ) + return handle + + +def _w4a16_fused_moe_full_rotation_hybrid_launch_op( + a_input: torch.Tensor, + t0_w13: torch.Tensor, + t0_w2: torch.Tensor, + t0_w13_scale: torch.Tensor, + t0_w2_scale: torch.Tensor, + t0_w13_global: torch.Tensor, + t0_w2_global: torch.Tensor, + t1_w13: torch.Tensor, + t1_w2: torch.Tensor, + t1_w13_scale: torch.Tensor, + t1_w2_scale: torch.Tensor, + t1_w13_global: torch.Tensor, + t1_w2_global: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + tier_local_map: torch.Tensor, + intermediate_cache13: torch.Tensor, + intermediate_cache2: torch.Tensor, + output: torch.Tensor, + fc1_c_tmp: torch.Tensor, + fc2_c_tmp: torch.Tensor, + workspace: torch.Tensor, + packed_route_indices: torch.Tensor, + block_expert_ids: torch.Tensor, + packed_route_count: torch.Tensor, + expert_offsets: torch.Tensor, + expert_counts: torch.Tensor, + rotation_a_gate: torch.Tensor, + rotation_a_up: torch.Tensor, + global_intermediate_rotations: torch.Tensor, + global_suh_gate: torch.Tensor, + global_suh_up: torch.Tensor, + global_svh_down: torch.Tensor, + launch_handle: int, +) -> None: + try: + fused_launch, topk_sum_launch, fast_math = _FULL_ROTATION_HYBRID_OP_HANDLES[ + int(launch_handle) + ] + except KeyError as exc: + raise RuntimeError( + "full-rotation hybrid launch handle is not registered in this process" + ) from exc + + def prepared( + *, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_global: torch.Tensor, + w2_global: torch.Tensor, + num_experts: int, + trellis_bits: int, + ) -> SimpleNamespace: + return SimpleNamespace( + weight_layout="trellis3_t256", + trellis_codebook="mcg", + trellis_bits=int(trellis_bits), + num_experts=int(num_experts), + hidden_size=int(fused_launch.hidden_size), + intermediate_size=int(fused_launch.intermediate_size), + params_dtype=torch.float16, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + w13_global_scale=w13_global, + w2_global_scale=w2_global, + workspace=workspace, + ) + + tier0 = prepared( + w13=t0_w13, + w2=t0_w2, + w13_scale=t0_w13_scale, + w2_scale=t0_w2_scale, + w13_global=t0_w13_global, + w2_global=t0_w2_global, + num_experts=fused_launch.tier0_num_experts, + trellis_bits=fused_launch.tier0_trellis_bits, + ) + tier1 = prepared( + w13=t1_w13, + w2=t1_w2, + w13_scale=t1_w13_scale, + w2_scale=t1_w2_scale, + w13_global=t1_w13_global, + w2_global=t1_w2_global, + num_experts=fused_launch.tier1_num_experts, + trellis_bits=fused_launch.tier1_trellis_bits, + ) + _run_w4a16_moe_full_rotation_hybrid_eager( + a_input, + tier0, + tier1, + topk_weights, + topk_ids, + tier_local_map, + activation=fused_launch.activation, + intermediate_cache13=intermediate_cache13, + intermediate_cache2=intermediate_cache2, + output=output, + fc1_c_tmp=fc1_c_tmp, + fc2_c_tmp=fc2_c_tmp, + packed_route_indices=packed_route_indices, + block_expert_ids=block_expert_ids, + packed_route_count=packed_route_count, + expert_offsets=expert_offsets, + expert_counts=expert_counts, + rotation_a_gate=rotation_a_gate, + rotation_a_up=rotation_a_up, + global_intermediate_rotations=global_intermediate_rotations, + global_suh_gate=global_suh_gate, + global_suh_up=global_suh_up, + global_svh_down=global_svh_down, + fused_launch=fused_launch, + topk_sum_launch=topk_sum_launch, + fast_math=fast_math, + ) + + +def _w4a16_fused_moe_full_rotation_hybrid_launch_fake( + a_input: torch.Tensor, + t0_w13: torch.Tensor, + t0_w2: torch.Tensor, + t0_w13_scale: torch.Tensor, + t0_w2_scale: torch.Tensor, + t0_w13_global: torch.Tensor, + t0_w2_global: torch.Tensor, + t1_w13: torch.Tensor, + t1_w2: torch.Tensor, + t1_w13_scale: torch.Tensor, + t1_w2_scale: torch.Tensor, + t1_w13_global: torch.Tensor, + t1_w2_global: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + tier_local_map: torch.Tensor, + intermediate_cache13: torch.Tensor, + intermediate_cache2: torch.Tensor, + output: torch.Tensor, + fc1_c_tmp: torch.Tensor, + fc2_c_tmp: torch.Tensor, + workspace: torch.Tensor, + packed_route_indices: torch.Tensor, + block_expert_ids: torch.Tensor, + packed_route_count: torch.Tensor, + expert_offsets: torch.Tensor, + expert_counts: torch.Tensor, + rotation_a_gate: torch.Tensor, + rotation_a_up: torch.Tensor, + global_intermediate_rotations: torch.Tensor, + global_suh_gate: torch.Tensor, + global_suh_up: torch.Tensor, + global_svh_down: torch.Tensor, + launch_handle: int, +) -> None: + return None + + +def run_w4a16_moe_full_rotation_hybrid( + a_input: torch.Tensor, + prepared_tier0, + prepared_tier1, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + tier_local_map: torch.Tensor, + *, + activation: str, + intermediate_cache13: torch.Tensor, + intermediate_cache2: torch.Tensor, + output: torch.Tensor, + fc1_c_tmp: torch.Tensor, + fc2_c_tmp: torch.Tensor, + packed_route_indices: torch.Tensor, + block_expert_ids: torch.Tensor, + packed_route_count: torch.Tensor, + expert_offsets: torch.Tensor, + expert_counts: torch.Tensor, + rotation_a_gate: torch.Tensor, + rotation_a_up: torch.Tensor, + global_intermediate_rotations: torch.Tensor, + global_suh_gate: torch.Tensor, + global_suh_up: torch.Tensor, + global_svh_down: torch.Tensor, + fused_launch: W4A16FusedMoeFullRotationHybridCompileResult, + topk_sum_launch: W4A16TopKSumCompileResult, + fast_math: bool = True, + stream: cuda.CUstream | None = None, +) -> torch.Tensor: + """Run mixed full-rotation Trellis without an opaque custom-op boundary. + + The registered-op experiment changed full-model scheduling and materially + reduced sustained speculative acceptance. vLLM records the underlying + CuTe launches correctly in its outer CUDA graph, so retain the direct + launch path used by the last coherent baseline. + """ + + return _run_w4a16_moe_full_rotation_hybrid_eager( + a_input, + prepared_tier0, + prepared_tier1, + topk_weights, + topk_ids, + tier_local_map, + activation=activation, + intermediate_cache13=intermediate_cache13, + intermediate_cache2=intermediate_cache2, + output=output, + fc1_c_tmp=fc1_c_tmp, + fc2_c_tmp=fc2_c_tmp, + packed_route_indices=packed_route_indices, + block_expert_ids=block_expert_ids, + packed_route_count=packed_route_count, + expert_offsets=expert_offsets, + expert_counts=expert_counts, + rotation_a_gate=rotation_a_gate, + rotation_a_up=rotation_a_up, + global_intermediate_rotations=global_intermediate_rotations, + global_suh_gate=global_suh_gate, + global_suh_up=global_suh_up, + global_svh_down=global_svh_down, + fused_launch=fused_launch, + topk_sum_launch=topk_sum_launch, + fast_math=fast_math, + stream=stream, + ) + + __all__ = [ "W4A16ActivationCompileResult", "W4A16FusedMoeCompileResult", "W4A16FusedMoeHybridCompileResult", + "W4A16FusedMoeFullRotationHybridCompileResult", "W4A16GemmCompileResult", "W4A16TopKSumCompileResult", "W4A16FusedMoeKernel", "W4A16FusedMoeHybridKernel", + "W4A16FusedMoeFullRotationHybridKernel", "W4A16ActivationKernel", "W4A16GemmKernel", "W4A16TopKSumKernel", @@ -11421,10 +12869,12 @@ def _weight_args(prepared) -> tuple[torch.Tensor, torch.Tensor]: "compile_w4a16_activation", "compile_w4a16_fused_moe", "compile_w4a16_fused_moe_hybrid", + "compile_w4a16_fused_moe_full_rotation_hybrid", "compile_w4a16_gemm", "compile_w4a16_topk_sum", "pack_topk_routes_by_expert", "run_trellis256_dense", "run_w4a16_moe", "run_w4a16_moe_hybrid", + "run_w4a16_moe_full_rotation_hybrid", ] diff --git a/sparkinfer/moe/_trellis_moe/__init__.py b/sparkinfer/moe/_trellis_moe/__init__.py new file mode 100644 index 00000000..c69b7957 --- /dev/null +++ b/sparkinfer/moe/_trellis_moe/__init__.py @@ -0,0 +1,100 @@ +"""Private compatibility API for the mixed-K EXL3 comparison path. + +This module preserves the r12 ``trellis3_t256`` MCG lifecycle needed to +reproduce the mixed-K comparison. It is intentionally private: the production +uniform-K path remains :mod:`sparkinfer.moe.fused_moe`. + +Weight preparation +wraps projection-major native EXL3 tensors without repacking. Planning fixes the +token, route, tile, and exact M-block capacities and eagerly compiles both the +fused MoE launch and full-rotation FP32 top-k reductions. Binding maps one +caller-owned uint8 scratch arena into stable views; ``run`` performs no tensor +allocation and is CUDA-graph-capture safe after ordinary eager warmup. + +Example: + from sparkinfer.moe import _trellis_moe + + weights = _trellis_moe.prepare_weights( + w13, w2, + gate_suh=gate_suh, + up_suh=up_suh, + intermediate_rotations=intermediate_rotations, + down_svh=down_svh, + codebook="mcg", + ) + plan = _trellis_moe.plan(_trellis_moe.Caps( + max_tokens=32, + num_topk=8, + num_experts=weights.num_experts, + hidden_size=weights.hidden_size, + intermediate_size=weights.intermediate_size, + route_num_experts=256, + block_size_m=8, + input_dtype=torch.bfloat16, + device=weights.device, + )) + spec = plan.scratch_specs()[0] + scratch = torch.empty(spec.shape, dtype=spec.dtype, device=spec.device) + binding = _trellis_moe.bind( + plan, + scratch=scratch, + a=x, + weights=weights, + topk_weights=router_weights, + topk_ids=router_ids, + ) + output = _trellis_moe.run(binding=binding) +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ..._lib.meta import OpMeta, Provenance, install_lazy_api + +META = OpMeta( + name="_trellis_moe", + group="moe", + api_style="planned", + entry_points=( + "Caps", + "Plan", + "Weights", + "Binding", + "plan", + "prepare_weights", + "bind", + "run", + "is_supported", + "clear_caches", + ), + dtypes=("bf16", "fp16"), + recipes=("trellis3_t256_mcg",), + requires=("triton",), + provenance=Provenance( + repo="https://github.com/brandonmmusic-max/b12x", + commit="e611971", + paths=("b12x/moe/fused/w4a16/",), + ), + since="1.1.0", + notes=( + "Private r12 compatibility surface for the opt-in mixed-K comparison; " + "uniform-K production serving remains on moe.fused_moe." + ), +) + +if TYPE_CHECKING: # static analysis only; runtime resolution is lazy + from .api import ( # noqa: F401 + Binding, + Caps, + Plan, + Weights, + bind, + clear_caches, + is_supported, + plan, + prepare_weights, + run, + ) + +install_lazy_api(globals(), META) diff --git a/sparkinfer/moe/_trellis_moe/_impl.py b/sparkinfer/moe/_trellis_moe/_impl.py new file mode 100644 index 00000000..0d7493ec --- /dev/null +++ b/sparkinfer/moe/_trellis_moe/_impl.py @@ -0,0 +1,778 @@ +"""Implementation for the private full-rotation Trellis compatibility API.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field, replace +from math import prod + +import torch + +from ..._lib.scratch import ScratchBufferSpec, scratch_buffer_spec, scratch_tensor +from .._shared.kernels.w4a16.host import ( + W4A16BufferPlan, + max_packed_route_slots, + plan_w4a16_buffers, +) +from .._shared.kernels.w4a16.kernel import ( + W4A16FusedMoeCompileResult, + W4A16TopKSumCompileResult, + clear_w4a16_kernel_cache, + compile_w4a16_fused_moe, + compile_w4a16_topk_sum, + run_w4a16_moe, +) +from .._shared.kernels.w4a16.prepare import ( + PreparedNF3MoeWeights, + _normalize_trellis256_codebook, + prepare_trellis256_moe_weights, +) + + +_ALLOWED_BLOCK_M = (8, 16, 32, 48, 64) +_ALLOWED_BITS = (3, 4, 5, 6) +_DEFAULT_TILE_CONFIG = (64, 256, 64, 256) +_MCG_SENTINEL = 0xCBAC1FED +_ARENA_ALIGNMENT = 256 + + +def _align_up(value: int, alignment: int = _ARENA_ALIGNMENT) -> int: + return ((int(value) + int(alignment) - 1) // int(alignment)) * int(alignment) + + +def _dtype_nbytes(dtype: torch.dtype) -> int: + return torch.empty((), dtype=dtype).element_size() + + +def _normalize_tile_config(value: Sequence[int]) -> tuple[int, int, int, int]: + if len(value) != 4: + raise ValueError( + "tile_config must be (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n)" + ) + tile_config = tuple(int(item) for item in value) + if any(item <= 0 or item % 16 != 0 for item in tile_config): + raise ValueError( + "every tile_config value must be a positive multiple of 16, got " + f"{tile_config}" + ) + fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = tile_config + if fc1_tile_k * fc1_tile_n != fc2_tile_k * fc2_tile_n: + raise ValueError( + "FC1 and FC2 tile_config entries must select the same CTA thread count" + ) + return tile_config + + +def _normalize_input_dtype(dtype: torch.dtype) -> torch.dtype: + if dtype not in (torch.bfloat16, torch.float16): + raise TypeError( + "Trellis MoE input_dtype must be torch.bfloat16 or torch.float16, " + f"got {dtype}" + ) + return dtype + + +def _input_dtype_name(dtype: torch.dtype) -> str: + return "bf16" if dtype == torch.bfloat16 else "fp16" + + +def _resolve_cuda_device(device: torch.device | str | int) -> torch.device: + if isinstance(device, int): + result = torch.device("cuda", int(device)) + else: + result = torch.device(device) + if result.type != "cuda": + raise ValueError(f"Trellis MoE requires a CUDA device, got {result}") + if result.index is None and torch.cuda.is_available(): + result = torch.device("cuda", torch.cuda.current_device()) + return result + + +@dataclass(frozen=True, kw_only=True) +class TrellisMoECaps: + """Fixed serving capacity and compile policy for one Trellis MoE shape.""" + + max_tokens: int + num_topk: int + num_experts: int + hidden_size: int + intermediate_size: int + device: torch.device | str | int + input_dtype: torch.dtype = torch.bfloat16 + route_num_experts: int | None = None + block_size_m: int = 8 + trellis_bits: int = 3 + tile_config: tuple[int, int, int, int] = _DEFAULT_TILE_CONFIG + activation: str = "silu" + fast_math: bool = True + + def __post_init__(self) -> None: + for name in ( + "max_tokens", + "num_topk", + "num_experts", + "hidden_size", + "intermediate_size", + ): + value = int(getattr(self, name)) + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + object.__setattr__(self, name, value) + if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: + raise ValueError( + "full-rotation Trellis MoE requires hidden_size and " + "intermediate_size divisible by 128" + ) + object.__setattr__(self, "device", _resolve_cuda_device(self.device)) + object.__setattr__( + self, "input_dtype", _normalize_input_dtype(self.input_dtype) + ) + route_num_experts = ( + self.num_experts + if self.route_num_experts is None + else int(self.route_num_experts) + ) + if route_num_experts <= 0: + raise ValueError( + f"route_num_experts must be positive, got {route_num_experts}" + ) + if self.num_topk > route_num_experts: + raise ValueError( + f"num_topk={self.num_topk} exceeds route_num_experts={route_num_experts}" + ) + object.__setattr__(self, "route_num_experts", route_num_experts) + block_size_m = int(self.block_size_m) + if block_size_m not in _ALLOWED_BLOCK_M: + raise ValueError( + f"block_size_m must be one of {_ALLOWED_BLOCK_M}, got {block_size_m}" + ) + object.__setattr__(self, "block_size_m", block_size_m) + trellis_bits = int(self.trellis_bits) + if trellis_bits not in _ALLOWED_BITS: + raise ValueError( + f"trellis_bits must be one of {_ALLOWED_BITS}, got {trellis_bits}" + ) + object.__setattr__(self, "trellis_bits", trellis_bits) + tile_config = _normalize_tile_config(self.tile_config) + fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = tile_config + if self.hidden_size % fc1_tile_k != 0: + raise ValueError("hidden_size must be divisible by FC1 tile K") + if self.intermediate_size % fc1_tile_n != 0: + raise ValueError( + "projection-major intermediate_size must be divisible by FC1 tile N" + ) + if self.intermediate_size % fc2_tile_k != 0: + raise ValueError("intermediate_size must be divisible by FC2 tile K") + if self.hidden_size % fc2_tile_n != 0: + raise ValueError("hidden_size must be divisible by FC2 tile N") + object.__setattr__(self, "tile_config", tile_config) + if str(self.activation).strip().lower() != "silu": + raise ValueError("the validated full-rotation Trellis recipe requires silu") + object.__setattr__(self, "activation", "silu") + object.__setattr__(self, "fast_math", bool(self.fast_math)) + + @property + def is_gated(self) -> bool: + return True + + +@dataclass(frozen=True, eq=False) +class TrellisMoEWeights: + """Zero-copy native tensors and persistent full-rotation tables.""" + + w13: torch.Tensor + w2: torch.Tensor + gate_suh: torch.Tensor + up_suh: torch.Tensor + intermediate_rotations: torch.Tensor + down_svh: torch.Tensor + hidden_size: int + intermediate_size: int + num_experts: int + trellis_bits: int + tile_config: tuple[int, int, int, int] + device: torch.device + _prepared: PreparedNF3MoeWeights = field(repr=False) + + +@dataclass(frozen=True) +class _ArenaViewSpec: + name: str + offset_bytes: int + shape: tuple[int, ...] + dtype: torch.dtype + + @property + def nbytes(self) -> int: + return int(prod(self.shape)) * _dtype_nbytes(self.dtype) + + +@dataclass(frozen=True) +class _ArenaLayout: + nbytes: int + views: tuple[_ArenaViewSpec, ...] + + def materialize(self, scratch: torch.Tensor) -> dict[str, torch.Tensor]: + result: dict[str, torch.Tensor] = {} + for spec in self.views: + raw = scratch.narrow(0, spec.offset_bytes, spec.nbytes) + result[spec.name] = raw.view(spec.dtype).view(spec.shape) + return result + + +def _make_arena_layout( + caps: TrellisMoECaps, + buffers: W4A16BufferPlan, + *, + sms: int, +) -> _ArenaLayout: + assert caps.route_num_experts is not None + specs = ( + ( + "intermediate_cache13", + (buffers.intermediate_cache13_elements,), + torch.float16, + ), + ("intermediate_cache2", (buffers.intermediate_cache2_elements,), torch.float16), + ("output", (caps.max_tokens, caps.hidden_size), torch.float32), + ("fc1_c_tmp", (buffers.fc1_c_tmp_elements,), torch.float32), + ("fc2_c_tmp", (buffers.fc2_c_tmp_elements,), torch.float32), + ("packed_route_indices", (buffers.route_slots,), torch.int32), + ("block_expert_ids", (buffers.route_blocks,), torch.int32), + ("packed_route_count", (1,), torch.int32), + ("expert_offsets", (caps.route_num_experts + 1,), torch.int32), + ("expert_counts", (caps.route_num_experts,), torch.int32), + ("rotation_a_gate", (buffers.rotation_a_elements,), torch.float16), + ("rotation_a_up", (buffers.rotation_a_elements,), torch.float16), + ("kernel_workspace", (int(sms) * 4 + 2,), torch.int32), + ) + cursor = 0 + views: list[_ArenaViewSpec] = [] + for name, shape, dtype in specs: + cursor = _align_up(cursor, max(_ARENA_ALIGNMENT, _dtype_nbytes(dtype))) + spec = _ArenaViewSpec( + name=name, + offset_bytes=cursor, + shape=tuple(int(dim) for dim in shape), + dtype=dtype, + ) + views.append(spec) + cursor += spec.nbytes + return _ArenaLayout(nbytes=max(_align_up(cursor), 1), views=tuple(views)) + + +@dataclass(frozen=True) +class TrellisMoEPlan: + """Compiled launches and byte-arena layout for one fixed capacity.""" + + caps: TrellisMoECaps + buffer_plan: W4A16BufferPlan + fused_launch: W4A16FusedMoeCompileResult = field(repr=False) + identity_sums: tuple[W4A16TopKSumCompileResult, ...] = field(repr=False) + mapped_sums: tuple[W4A16TopKSumCompileResult, ...] = field(repr=False) + _arena_layout: _ArenaLayout = field(repr=False) + _scratch_specs: tuple[ScratchBufferSpec, ...] = field(repr=False) + + @property + def scratch_nbytes(self) -> int: + return self._arena_layout.nbytes + + def scratch_specs(self) -> tuple[ScratchBufferSpec, ...]: + return self._scratch_specs + + def shapes_and_dtypes(self) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]: + return tuple((spec.shape, spec.dtype) for spec in self._scratch_specs) + + def bind(self, **kwargs) -> "TrellisMoEBinding": + return bind_trellis_moe(self, **kwargs) + + def topk_sum_launch( + self, ids_dtype: torch.dtype, *, mapped: bool + ) -> W4A16TopKSumCompileResult: + launches = self.mapped_sums if mapped else self.identity_sums + for launch in launches: + if launch.route_ids_dtype == ids_dtype: + return launch + raise TypeError(f"Trellis MoE does not have a top-k sum for {ids_dtype}") + + +@dataclass(frozen=True, kw_only=True) +class TrellisMoEBinding: + """Stable tensor views for one launch; safe to retain across graph replay.""" + + plan: TrellisMoEPlan + weights: TrellisMoEWeights + a: torch.Tensor + topk_weights: torch.Tensor + topk_ids: torch.Tensor + output: torch.Tensor + route_expert_map: torch.Tensor | None + output_expert_map: torch.Tensor | None + prepared: PreparedNF3MoeWeights = field(repr=False) + intermediate_cache13: torch.Tensor = field(repr=False) + intermediate_cache2: torch.Tensor = field(repr=False) + fc1_c_tmp: torch.Tensor = field(repr=False) + fc2_c_tmp: torch.Tensor = field(repr=False) + packed_route_indices: torch.Tensor = field(repr=False) + block_expert_ids: torch.Tensor = field(repr=False) + packed_route_count: torch.Tensor = field(repr=False) + expert_offsets: torch.Tensor = field(repr=False) + expert_counts: torch.Tensor = field(repr=False) + rotation_a_gate: torch.Tensor = field(repr=False) + rotation_a_up: torch.Tensor = field(repr=False) + topk_sum_launch: W4A16TopKSumCompileResult = field(repr=False) + + def run(self) -> torch.Tensor: + return run_trellis_moe(binding=self) + + +def _validate_rotation_table( + name: str, + tensor: torch.Tensor, + *, + shape: tuple[int, ...], + device: torch.device, +) -> None: + if tensor.dtype != torch.float16: + raise TypeError(f"{name} must be torch.float16, got {tensor.dtype}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _validate_mcg(codebook: str | int, mcg: torch.Tensor | int | None) -> None: + normalized = _normalize_trellis256_codebook(codebook) + if normalized != "mcg": + raise NotImplementedError( + "the production Trellis MoE decoder accepts only the MCG codebook" + ) + if mcg is None: + return + if isinstance(mcg, torch.Tensor): + if mcg.numel() != 1 or mcg.dtype not in (torch.int32, torch.uint32): + raise ValueError("mcg must be a scalar int32/uint32 tensor") + marker = int(mcg.item()) & 0xFFFFFFFF + else: + marker = int(mcg) & 0xFFFFFFFF + if marker != _MCG_SENTINEL: + raise ValueError( + f"unexpected MCG marker {marker:#010x}; expected {_MCG_SENTINEL:#010x}" + ) + + +def prepare_trellis_moe_weights( + w13: torch.Tensor, + w2: torch.Tensor, + *, + gate_suh: torch.Tensor, + up_suh: torch.Tensor, + intermediate_rotations: torch.Tensor, + down_svh: torch.Tensor, + codebook: str | int = "mcg", + mcg: torch.Tensor | int | None = None, + tile_config: tuple[int, int, int, int] = _DEFAULT_TILE_CONFIG, + dummy_scale: torch.Tensor | None = None, +) -> TrellisMoEWeights: + """Validate and wrap projection-major native EXL3 tensors without copying.""" + _validate_mcg(codebook, mcg) + tile_config = _normalize_tile_config(tile_config) + if w13.ndim != 5 or int(w13.shape[0]) != 2: + raise ValueError( + "w13 must be projection-major [2,E,H/16,I/16,16*bits] int16 " + "or the byte-identical [...,8*bits] int32 view" + ) + if w2.ndim != 4: + raise ValueError( + "w2 must be [E,I/16,H/16,16*bits] int16 or the byte-identical " + "[...,8*bits] int32 view" + ) + num_experts = int(w13.shape[1]) + hidden_size = int(w13.shape[2]) * 16 + intermediate_size = int(w13.shape[3]) * 16 + if tuple(w2.shape[:3]) != ( + num_experts, + intermediate_size // 16, + hidden_size // 16, + ): + raise ValueError( + "w2 geometry does not match projection-major w13: " + f"got {tuple(w2.shape[:3])}" + ) + if hidden_size % 128 != 0 or intermediate_size % 128 != 0: + raise ValueError( + "full-rotation Trellis weights require hidden and intermediate " + "dimensions divisible by 128" + ) + fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = tile_config + if ( + hidden_size % fc1_tile_k != 0 + or intermediate_size % fc1_tile_n != 0 + or intermediate_size % fc2_tile_k != 0 + or hidden_size % fc2_tile_n != 0 + ): + raise ValueError( + f"tile_config={tile_config} does not divide H={hidden_size}, " + f"I={intermediate_size}" + ) + device = w13.device + if w2.device != device: + raise ValueError("w13 and w2 must be on the same device") + _validate_rotation_table( + "gate_suh", gate_suh, shape=(num_experts, hidden_size), device=device + ) + _validate_rotation_table( + "up_suh", up_suh, shape=(num_experts, hidden_size), device=device + ) + _validate_rotation_table( + "intermediate_rotations", + intermediate_rotations, + shape=(num_experts, 3 * intermediate_size), + device=device, + ) + _validate_rotation_table( + "down_svh", down_svh, shape=(num_experts, hidden_size), device=device + ) + prepared = prepare_trellis256_moe_weights( + w13, + w2, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + activation="silu", + params_dtype=torch.float16, + fc1_tile_n=fc1_tile_n, + fc2_tile_n=fc2_tile_n, + w13_layout="trellis3_t256_proj", + codebook=codebook, + dummy_scale=dummy_scale, + gate_suh=gate_suh, + up_suh=up_suh, + intermediate_rotations=intermediate_rotations, + down_svh=down_svh, + tile_config=tile_config, + ) + if prepared.trellis_codebook != "mcg": + raise RuntimeError("Trellis preparation did not preserve the MCG contract") + if ( + prepared.w13.data_ptr() != w13.data_ptr() + or prepared.w2.data_ptr() != w2.data_ptr() + ): + raise RuntimeError("Trellis preparation unexpectedly copied native weights") + return TrellisMoEWeights( + w13=w13, + w2=w2, + gate_suh=gate_suh, + up_suh=up_suh, + intermediate_rotations=intermediate_rotations, + down_svh=down_svh, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + trellis_bits=int(prepared.trellis_bits), + tile_config=tile_config, + device=device, + _prepared=prepared, + ) + + +def plan_trellis_moe(caps: TrellisMoECaps) -> TrellisMoEPlan: + """Compile every launch and produce the fixed caller-scratch layout.""" + if not isinstance(caps, TrellisMoECaps): + raise TypeError("caps must be a TrellisMoECaps") + device = _resolve_cuda_device(caps.device) + if device.index is None: + raise RuntimeError("CUDA must be available before planning Trellis MoE") + if device != caps.device: + caps = replace(caps, device=device) + with torch.cuda.device(device): + props = torch.cuda.get_device_properties(device) + sms = int(props.multi_processor_count) + max_shared_mem = int(getattr(props, "shared_memory_per_block_optin", 101_376)) + buffer_plan = plan_w4a16_buffers( + caps, + m=caps.max_tokens, + topk=caps.num_topk, + route_num_experts=caps.route_num_experts, + sms=sms, + full_rotation=True, + block_size_m=caps.block_size_m, + ) + route_slots = max_packed_route_slots( + caps.max_tokens * caps.num_topk, + caps.block_size_m, + caps.route_num_experts, + ) + max_m_blocks = (route_slots + caps.block_size_m - 1) // caps.block_size_m + fused_launch = compile_w4a16_fused_moe( + size_m=caps.max_tokens, + hidden_size=caps.hidden_size, + intermediate_size=caps.intermediate_size, + num_experts=caps.num_experts, + top_k=caps.num_topk, + activation=caps.activation, + apply_router_weight_on_input=False, + zero_fc2_output=False, + moe_block_size=caps.block_size_m, + max_m_blocks=max_m_blocks, + element_dtype="fp16", + fast_math=caps.fast_math, + sms=sms, + max_shared_mem=max_shared_mem, + weight_layout="trellis3_t256", + scale_format="e4m3_k32", + w13_layout="trellis3_t256_proj", + trellis_bits=caps.trellis_bits, + force_tile_config=caps.tile_config, + intermediate_rotation=True, + full_rotation=True, + rotation_input_dtype=_input_dtype_name(caps.input_dtype), + ) + identity_sums = tuple( + compile_w4a16_topk_sum( + m=caps.max_tokens, + topk=caps.num_topk, + hidden_size=caps.hidden_size, + element_dtype="fp16", + full_rotation=True, + num_experts=caps.num_experts, + route_num_experts=0, + route_ids_dtype=ids_dtype, + use_expert_map=False, + ) + for ids_dtype in (torch.int32, torch.int64) + ) + mapped_sums = tuple( + compile_w4a16_topk_sum( + m=caps.max_tokens, + topk=caps.num_topk, + hidden_size=caps.hidden_size, + element_dtype="fp16", + full_rotation=True, + num_experts=caps.num_experts, + route_num_experts=caps.route_num_experts, + route_ids_dtype=ids_dtype, + use_expert_map=True, + ) + for ids_dtype in (torch.int32, torch.int64) + ) + arena_layout = _make_arena_layout(caps, buffer_plan, sms=sms) + specs = ( + scratch_buffer_spec( + "trellis_moe", + nbytes=arena_layout.nbytes, + device=device, + ), + ) + return TrellisMoEPlan( + caps=caps, + buffer_plan=buffer_plan, + fused_launch=fused_launch, + identity_sums=identity_sums, + mapped_sums=mapped_sums, + _arena_layout=arena_layout, + _scratch_specs=specs, + ) + + +def _validate_runtime_tensor( + name: str, + tensor: torch.Tensor, + *, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> None: + if tensor.dtype != dtype: + raise TypeError(f"{name} must be {dtype}, got {tensor.dtype}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _validate_expert_map( + name: str, + tensor: torch.Tensor | None, + *, + caps: TrellisMoECaps, +) -> None: + if tensor is None: + return + assert caps.route_num_experts is not None + _validate_runtime_tensor( + name, + tensor, + shape=(caps.route_num_experts,), + dtype=torch.int32, + device=caps.device, + ) + + +def bind_trellis_moe( + plan: TrellisMoEPlan, + *, + scratch: torch.Tensor | Mapping[str, torch.Tensor] | Sequence[torch.Tensor], + a: torch.Tensor, + weights: TrellisMoEWeights, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + route_expert_map: torch.Tensor | None = None, + output_expert_map: torch.Tensor | None = None, + output: torch.Tensor | None = None, +) -> TrellisMoEBinding: + """Bind runtime tensors by carving views from ``scratch`` only.""" + if not isinstance(plan, TrellisMoEPlan): + raise TypeError("plan must come from trellis_moe.plan") + if not isinstance(weights, TrellisMoEWeights): + raise TypeError("weights must come from trellis_moe.prepare_weights") + caps = plan.caps + if ( + weights.hidden_size != caps.hidden_size + or weights.intermediate_size != caps.intermediate_size + or weights.num_experts != caps.num_experts + or weights.trellis_bits != caps.trellis_bits + or weights.tile_config != caps.tile_config + or weights.device != caps.device + ): + raise ValueError("weights do not match the Trellis MoE plan") + if a.ndim != 2: + raise ValueError(f"a must be rank 2, got shape {tuple(a.shape)}") + tokens = int(a.shape[0]) + if tokens < 1 or tokens > caps.max_tokens: + raise ValueError( + f"input tokens must be in [1, {caps.max_tokens}], got {tokens}" + ) + _validate_runtime_tensor( + "a", + a, + shape=(tokens, caps.hidden_size), + dtype=caps.input_dtype, + device=caps.device, + ) + _validate_runtime_tensor( + "topk_weights", + topk_weights, + shape=(tokens, caps.num_topk), + dtype=torch.float32, + device=caps.device, + ) + if topk_ids.dtype not in (torch.int32, torch.int64): + raise TypeError("topk_ids must be torch.int32 or torch.int64") + _validate_runtime_tensor( + "topk_ids", + topk_ids, + shape=(tokens, caps.num_topk), + dtype=topk_ids.dtype, + device=caps.device, + ) + _validate_expert_map("route_expert_map", route_expert_map, caps=caps) + _validate_expert_map("output_expert_map", output_expert_map, caps=caps) + + scratch_storage = scratch_tensor(scratch, plan._scratch_specs, owner="Trellis MoE") + if int(scratch_storage.data_ptr()) % _ARENA_ALIGNMENT != 0: + raise ValueError(f"Trellis MoE scratch must be {_ARENA_ALIGNMENT}-byte aligned") + views = plan._arena_layout.materialize(scratch_storage) + views["kernel_workspace"].zero_() + if output is None: + output_view = views["output"][:tokens] + else: + if output.ndim != 2 or tuple(output.shape) not in ( + (tokens, caps.hidden_size), + (caps.max_tokens, caps.hidden_size), + ): + raise ValueError( + "output must be the live or capacity FP32 view: expected " + f"{(tokens, caps.hidden_size)} or {(caps.max_tokens, caps.hidden_size)}, " + f"got {tuple(output.shape)}" + ) + if output.dtype != torch.float32: + raise TypeError(f"output must be torch.float32, got {output.dtype}") + if output.device != caps.device or not output.is_contiguous(): + raise ValueError("output must be contiguous on the planned CUDA device") + output_view = output[:tokens] + prepared = replace(weights._prepared, workspace=views["kernel_workspace"]) + return TrellisMoEBinding( + plan=plan, + weights=weights, + a=a, + topk_weights=topk_weights, + topk_ids=topk_ids, + output=output_view, + route_expert_map=route_expert_map, + output_expert_map=output_expert_map, + prepared=prepared, + intermediate_cache13=views["intermediate_cache13"], + intermediate_cache2=views["intermediate_cache2"], + fc1_c_tmp=views["fc1_c_tmp"], + fc2_c_tmp=views["fc2_c_tmp"], + packed_route_indices=views["packed_route_indices"], + block_expert_ids=views["block_expert_ids"], + packed_route_count=views["packed_route_count"], + expert_offsets=views["expert_offsets"], + expert_counts=views["expert_counts"], + rotation_a_gate=views["rotation_a_gate"], + rotation_a_up=views["rotation_a_up"], + topk_sum_launch=plan.topk_sum_launch( + topk_ids.dtype, mapped=output_expert_map is not None + ), + ) + + +def run_trellis_moe(*, binding: TrellisMoEBinding) -> torch.Tensor: + """Run the preplanned full-rotation path into ``binding.output``.""" + if not isinstance(binding, TrellisMoEBinding): + raise TypeError("binding must come from trellis_moe.bind") + caps = binding.plan.caps + return run_w4a16_moe( + binding.a, + binding.prepared, + binding.topk_weights, + binding.topk_ids, + activation=caps.activation, + intermediate_cache13=binding.intermediate_cache13, + intermediate_cache2=binding.intermediate_cache2, + output=binding.output, + fc1_c_tmp=binding.fc1_c_tmp, + fc2_c_tmp=binding.fc2_c_tmp, + packed_route_indices=binding.packed_route_indices, + block_expert_ids=binding.block_expert_ids, + packed_route_count=binding.packed_route_count, + expert_offsets=binding.expert_offsets, + expert_counts=binding.expert_counts, + expert_map=binding.route_expert_map, + output_expert_map=binding.output_expert_map, + apply_router_weight_on_input=False, + fast_math=caps.fast_math, + fused_launch=binding.plan.fused_launch, + topk_sum_launch=binding.topk_sum_launch, + intermediate_rotation_scales=binding.weights.intermediate_rotations, + full_rotation=True, + suh_gate_table=binding.weights.gate_suh, + suh_up_table=binding.weights.up_suh, + svh_table=binding.weights.down_svh, + rotation_a_gate=binding.rotation_a_gate, + rotation_a_up=binding.rotation_a_up, + ) + + +def clear_trellis_moe_caches() -> None: + """Clear the shared W4A16 compile caches used by this planned op.""" + clear_w4a16_kernel_cache() + + +__all__ = [ + "TrellisMoEBinding", + "TrellisMoECaps", + "TrellisMoEPlan", + "TrellisMoEWeights", + "bind_trellis_moe", + "clear_trellis_moe_caches", + "plan_trellis_moe", + "prepare_trellis_moe_weights", + "run_trellis_moe", +] diff --git a/sparkinfer/moe/_trellis_moe/_mixk.py b/sparkinfer/moe/_trellis_moe/_mixk.py new file mode 100644 index 00000000..44185c9c --- /dev/null +++ b/sparkinfer/moe/_trellis_moe/_mixk.py @@ -0,0 +1,462 @@ +"""Private one-grid mixed-K full-rotation Trellis MoE comparison path. + +This module composes the existing production Trellis weight preparation and +scratch contract with the full-rotation K3/K4 hybrid kernel. It is kept +separate from the stable uniform API so a deployment must explicitly opt in +and can retain the two-pass implementation as a rollback oracle. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field, replace + +import torch + +from ..._lib.scratch import ScratchBufferSpec, scratch_buffer_spec, scratch_tensor +from .._shared.kernels.w4a16.host import ( + W4A16BufferPlan, + max_packed_route_slots, + plan_w4a16_buffers, +) +from .._shared.kernels.w4a16.kernel import ( + W4A16FusedMoeFullRotationHybridCompileResult, + W4A16TopKSumCompileResult, + compile_w4a16_fused_moe_full_rotation_hybrid, + compile_w4a16_topk_sum, + run_w4a16_moe_full_rotation_hybrid, +) +from ._impl import ( + TrellisMoECaps, + TrellisMoEWeights, + _ARENA_ALIGNMENT, + _input_dtype_name, + _make_arena_layout, + _resolve_cuda_device, + _validate_runtime_tensor, +) + + +@dataclass(frozen=True) +class MixedTrellisMoEPlan: + caps: TrellisMoECaps + tier0_num_experts: int + tier0_trellis_bits: int + tier1_num_experts: int + tier1_trellis_bits: int + buffer_plan: W4A16BufferPlan + fused_launches: tuple[ + tuple[int, W4A16FusedMoeFullRotationHybridCompileResult], ... + ] = field(repr=False) + identity_sums: tuple[W4A16TopKSumCompileResult, ...] = field(repr=False) + _arena_layout: object = field(repr=False) + _scratch_specs: tuple[ScratchBufferSpec, ...] = field(repr=False) + + @property + def scratch_nbytes(self) -> int: + return int(self._arena_layout.nbytes) + + def scratch_specs(self) -> tuple[ScratchBufferSpec, ...]: + return self._scratch_specs + + def topk_sum_launch(self, ids_dtype: torch.dtype) -> W4A16TopKSumCompileResult: + for launch in self.identity_sums: + if launch.route_ids_dtype == ids_dtype: + return launch + raise TypeError(f"mixed Trellis plan has no top-k sum for {ids_dtype}") + + def fused_launch(self, tokens: int) -> W4A16FusedMoeFullRotationHybridCompileResult: + tokens = int(tokens) + for planned_tokens, launch in self.fused_launches: + if planned_tokens == tokens: + return launch + # Small decode plans contain an exact launch for every admitted M. + # Large prefill plans intentionally retain one capacity launch, just + # like r12's stock Trellis scratch plan. That launch accepts any live + # M covered by its caller-owned arena. + for planned_tokens, launch in self.fused_launches: + if planned_tokens >= tokens: + return launch + raise ValueError( + f"mixed Trellis plan has no launch covering {tokens} tokens; " + f"planned={[m for m, _ in self.fused_launches]}" + ) + + +@dataclass(frozen=True, kw_only=True) +class MixedTrellisMoEBinding: + plan: MixedTrellisMoEPlan + tier0_weights: TrellisMoEWeights + tier1_weights: TrellisMoEWeights + a: torch.Tensor + topk_weights: torch.Tensor + topk_ids: torch.Tensor + tier_local_map: torch.Tensor + global_suh_gate: torch.Tensor + global_suh_up: torch.Tensor + global_intermediate_rotations: torch.Tensor + global_svh_down: torch.Tensor + output: torch.Tensor + prepared_tier0: object = field(repr=False) + prepared_tier1: object = field(repr=False) + intermediate_cache13: torch.Tensor = field(repr=False) + intermediate_cache2: torch.Tensor = field(repr=False) + fc1_c_tmp: torch.Tensor = field(repr=False) + fc2_c_tmp: torch.Tensor = field(repr=False) + packed_route_indices: torch.Tensor = field(repr=False) + block_expert_ids: torch.Tensor = field(repr=False) + packed_route_count: torch.Tensor = field(repr=False) + expert_offsets: torch.Tensor = field(repr=False) + expert_counts: torch.Tensor = field(repr=False) + rotation_a_gate: torch.Tensor = field(repr=False) + rotation_a_up: torch.Tensor = field(repr=False) + fused_launch: W4A16FusedMoeFullRotationHybridCompileResult = field(repr=False) + topk_sum_launch: W4A16TopKSumCompileResult = field(repr=False) + + +def plan_mixed_trellis_moe( + caps: TrellisMoECaps, + *, + tier0_weights: TrellisMoEWeights, + tier1_weights: TrellisMoEWeights, +) -> MixedTrellisMoEPlan: + """Compile the fused K3/K4 launch and one global full-rotation sum.""" + + if not isinstance(caps, TrellisMoECaps): + raise TypeError("caps must be TrellisMoECaps") + device = _resolve_cuda_device(caps.device) + if device != caps.device: + caps = replace(caps, device=device) + tiers = (tier0_weights, tier1_weights) + if any(not isinstance(weights, TrellisMoEWeights) for weights in tiers): + raise TypeError("tier weights must come from trellis_moe.prepare_weights") + if tier0_weights.device != device or tier1_weights.device != device: + raise ValueError("tier weights must be on the planned CUDA device") + if ( + tier0_weights.hidden_size != caps.hidden_size + or tier1_weights.hidden_size != caps.hidden_size + or tier0_weights.intermediate_size != caps.intermediate_size + or tier1_weights.intermediate_size != caps.intermediate_size + ): + raise ValueError("mixed tier dimensions do not match caps") + if tier0_weights.tile_config != tier1_weights.tile_config: + raise ValueError("mixed tiers require an identical tile_config") + if tier0_weights.tile_config != caps.tile_config: + raise ValueError("mixed tier tile_config does not match caps") + if tier0_weights.trellis_bits == tier1_weights.trellis_bits: + raise ValueError("mixed tiers must use distinct Trellis bitrates") + if tier0_weights.num_experts + tier1_weights.num_experts != caps.num_experts: + raise ValueError("mixed tier expert counts must exactly cover caps.num_experts") + assert caps.route_num_experts is not None + if caps.route_num_experts != caps.num_experts: + raise ValueError( + "fused mixed-K routes must use the global expert count directly" + ) + + with torch.cuda.device(device): + props = torch.cuda.get_device_properties(device) + sms = int(props.multi_processor_count) + max_shared_mem = int(getattr(props, "shared_memory_per_block_optin", 101_376)) + buffer_plan = plan_w4a16_buffers( + caps, + m=caps.max_tokens, + topk=caps.num_topk, + route_num_experts=caps.route_num_experts, + sms=sms, + full_rotation=True, + block_size_m=caps.block_size_m, + ) + route_slots = max_packed_route_slots( + caps.max_tokens * caps.num_topk, + caps.block_size_m, + caps.route_num_experts, + ) + route_blocks = (route_slots + caps.block_size_m - 1) // caps.block_size_m + # Match r12's consolidated full-rotation contract exactly: small decode + # gets an exact live-M specialization, while every specialization keeps + # the full caller-owned route-arena capacity. The retired compatibility + # wrapper instead reused one M=32 kernel (v8) or shrank max_m_blocks with + # each exact M (v6); neither is the production r12 launch contract. + fused_token_counts = ( + tuple(range(1, int(caps.max_tokens) + 1)) + if int(caps.max_tokens) <= 32 + else (int(caps.max_tokens),) + ) + fused_launches = tuple( + ( + token_count, + compile_w4a16_fused_moe_full_rotation_hybrid( + size_m=token_count, + route_capacity_m_blocks=route_blocks, + hidden_size=caps.hidden_size, + intermediate_size=caps.intermediate_size, + tier0_num_experts=tier0_weights.num_experts, + tier0_trellis_bits=tier0_weights.trellis_bits, + tier1_num_experts=tier1_weights.num_experts, + tier1_trellis_bits=tier1_weights.trellis_bits, + top_k=caps.num_topk, + activation=caps.activation, + map_slots=caps.num_experts, + moe_block_size=caps.block_size_m, + rotation_input_dtype=_input_dtype_name(caps.input_dtype), + fast_math=caps.fast_math, + sms=sms, + max_shared_mem=max_shared_mem, + force_tile_config=caps.tile_config, + ), + ) + for token_count in fused_token_counts + ) + if any( + int(launch.max_m_blocks) < int(route_blocks) for _, launch in fused_launches + ): + raise RuntimeError("compiled mixed-K launch under-covers route capacity") + identity_sums = tuple( + compile_w4a16_topk_sum( + m=caps.max_tokens, + topk=caps.num_topk, + hidden_size=caps.hidden_size, + element_dtype="fp16", + full_rotation=True, + num_experts=caps.num_experts, + route_num_experts=0, + route_ids_dtype=ids_dtype, + use_expert_map=False, + ) + for ids_dtype in (torch.int32, torch.int64) + ) + + arena_layout = _make_arena_layout(caps, buffer_plan, sms=sms) + specs = ( + scratch_buffer_spec( + "mixed_trellis_moe", + nbytes=arena_layout.nbytes, + device=device, + ), + ) + return MixedTrellisMoEPlan( + caps=caps, + tier0_num_experts=tier0_weights.num_experts, + tier0_trellis_bits=tier0_weights.trellis_bits, + tier1_num_experts=tier1_weights.num_experts, + tier1_trellis_bits=tier1_weights.trellis_bits, + buffer_plan=buffer_plan, + fused_launches=fused_launches, + identity_sums=identity_sums, + _arena_layout=arena_layout, + _scratch_specs=specs, + ) + + +def _validate_global_table( + name: str, + tensor: torch.Tensor, + *, + shape: tuple[int, ...], + device: torch.device, +) -> None: + if ( + tensor.dtype != torch.float16 + or tensor.device != device + or tuple(tensor.shape) != shape + or not tensor.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous fp16 {shape} on {device}; got " + f"{tuple(tensor.shape)}/{tensor.dtype}/{tensor.device}" + ) + + +def bind_mixed_trellis_moe( + plan: MixedTrellisMoEPlan, + *, + scratch: torch.Tensor | Mapping[str, torch.Tensor] | Sequence[torch.Tensor], + a: torch.Tensor, + tier0_weights: TrellisMoEWeights, + tier1_weights: TrellisMoEWeights, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + tier_local_map: torch.Tensor, + global_suh_gate: torch.Tensor, + global_suh_up: torch.Tensor, + global_intermediate_rotations: torch.Tensor, + global_svh_down: torch.Tensor, + output: torch.Tensor | None = None, +) -> MixedTrellisMoEBinding: + if not isinstance(plan, MixedTrellisMoEPlan): + raise TypeError("plan must come from plan_mixed_trellis_moe") + caps = plan.caps + tokens = int(a.shape[0]) + _validate_runtime_tensor( + "a", + a, + shape=(tokens, caps.hidden_size), + dtype=caps.input_dtype, + device=caps.device, + ) + if tokens < 1 or tokens > caps.max_tokens: + raise ValueError("input token count exceeds mixed Trellis plan") + _validate_runtime_tensor( + "topk_weights", + topk_weights, + shape=(tokens, caps.num_topk), + dtype=torch.float32, + device=caps.device, + ) + if topk_ids.dtype not in (torch.int32, torch.int64): + raise TypeError("topk_ids must be int32 or int64") + _validate_runtime_tensor( + "topk_ids", + topk_ids, + shape=(tokens, caps.num_topk), + dtype=topk_ids.dtype, + device=caps.device, + ) + if ( + tier_local_map.dtype != torch.int32 + or tier_local_map.device != caps.device + or tuple(tier_local_map.shape) != (caps.num_experts,) + or not tier_local_map.is_contiguous() + ): + raise ValueError("tier_local_map must be contiguous int32 [global_num_experts]") + _validate_global_table( + "global_suh_gate", + global_suh_gate, + shape=(caps.num_experts, caps.hidden_size), + device=caps.device, + ) + _validate_global_table( + "global_suh_up", + global_suh_up, + shape=(caps.num_experts, caps.hidden_size), + device=caps.device, + ) + _validate_global_table( + "global_svh_down", + global_svh_down, + shape=(caps.num_experts, caps.hidden_size), + device=caps.device, + ) + _validate_global_table( + "global_intermediate_rotations", + global_intermediate_rotations, + shape=(caps.num_experts, 3 * caps.intermediate_size), + device=caps.device, + ) + for name, weights, expected_experts, expected_bits in ( + ( + "tier0_weights", + tier0_weights, + plan.tier0_num_experts, + plan.tier0_trellis_bits, + ), + ( + "tier1_weights", + tier1_weights, + plan.tier1_num_experts, + plan.tier1_trellis_bits, + ), + ): + if ( + weights.num_experts != expected_experts + or weights.trellis_bits != expected_bits + or weights.device != caps.device + ): + raise ValueError(f"{name} does not match the mixed Trellis plan") + + storage = scratch_tensor(scratch, plan._scratch_specs, owner="mixed Trellis MoE") + if int(storage.data_ptr()) % _ARENA_ALIGNMENT != 0: + raise ValueError( + f"mixed Trellis scratch must be {_ARENA_ALIGNMENT}-byte aligned" + ) + views = plan._arena_layout.materialize(storage) + views["kernel_workspace"].zero_() + if output is None: + output_view = views["output"][:tokens] + else: + if ( + output.dtype != torch.float32 + or output.device != caps.device + or not output.is_contiguous() + or tuple(output.shape) + not in ( + (tokens, caps.hidden_size), + (caps.max_tokens, caps.hidden_size), + ) + ): + raise ValueError("output must be a live/capacity contiguous FP32 view") + output_view = output[:tokens] + prepared0 = replace(tier0_weights._prepared, workspace=views["kernel_workspace"]) + prepared1 = replace(tier1_weights._prepared, workspace=views["kernel_workspace"]) + return MixedTrellisMoEBinding( + plan=plan, + tier0_weights=tier0_weights, + tier1_weights=tier1_weights, + a=a, + topk_weights=topk_weights, + topk_ids=topk_ids, + tier_local_map=tier_local_map, + global_suh_gate=global_suh_gate, + global_suh_up=global_suh_up, + global_intermediate_rotations=global_intermediate_rotations, + global_svh_down=global_svh_down, + output=output_view, + prepared_tier0=prepared0, + prepared_tier1=prepared1, + intermediate_cache13=views["intermediate_cache13"], + intermediate_cache2=views["intermediate_cache2"], + fc1_c_tmp=views["fc1_c_tmp"], + fc2_c_tmp=views["fc2_c_tmp"], + packed_route_indices=views["packed_route_indices"], + block_expert_ids=views["block_expert_ids"], + packed_route_count=views["packed_route_count"], + expert_offsets=views["expert_offsets"], + expert_counts=views["expert_counts"], + rotation_a_gate=views["rotation_a_gate"], + rotation_a_up=views["rotation_a_up"], + fused_launch=plan.fused_launch(tokens), + topk_sum_launch=plan.topk_sum_launch(topk_ids.dtype), + ) + + +def run_mixed_trellis_moe(*, binding: MixedTrellisMoEBinding) -> torch.Tensor: + if not isinstance(binding, MixedTrellisMoEBinding): + raise TypeError("binding must come from bind_mixed_trellis_moe") + caps = binding.plan.caps + return run_w4a16_moe_full_rotation_hybrid( + binding.a, + binding.prepared_tier0, + binding.prepared_tier1, + binding.topk_weights, + binding.topk_ids, + binding.tier_local_map, + activation=caps.activation, + intermediate_cache13=binding.intermediate_cache13, + intermediate_cache2=binding.intermediate_cache2, + output=binding.output, + fc1_c_tmp=binding.fc1_c_tmp, + fc2_c_tmp=binding.fc2_c_tmp, + packed_route_indices=binding.packed_route_indices, + block_expert_ids=binding.block_expert_ids, + packed_route_count=binding.packed_route_count, + expert_offsets=binding.expert_offsets, + expert_counts=binding.expert_counts, + rotation_a_gate=binding.rotation_a_gate, + rotation_a_up=binding.rotation_a_up, + global_intermediate_rotations=binding.global_intermediate_rotations, + global_suh_gate=binding.global_suh_gate, + global_suh_up=binding.global_suh_up, + global_svh_down=binding.global_svh_down, + fused_launch=binding.fused_launch, + topk_sum_launch=binding.topk_sum_launch, + fast_math=caps.fast_math, + ) + + +__all__ = [ + "MixedTrellisMoEBinding", + "MixedTrellisMoEPlan", + "bind_mixed_trellis_moe", + "plan_mixed_trellis_moe", + "run_mixed_trellis_moe", +] diff --git a/sparkinfer/moe/_trellis_moe/api.py b/sparkinfer/moe/_trellis_moe/api.py new file mode 100644 index 00000000..c82ad6f9 --- /dev/null +++ b/sparkinfer/moe/_trellis_moe/api.py @@ -0,0 +1,52 @@ +"""Compatibility surface for :mod:`sparkinfer.moe._trellis_moe`.""" + +from __future__ import annotations + +from ..._lib.gating import default_is_supported +from . import META +from ._impl import ( + TrellisMoEBinding as Binding, +) +from ._impl import ( + TrellisMoECaps as Caps, +) +from ._impl import ( + TrellisMoEPlan as Plan, +) +from ._impl import ( + TrellisMoEWeights as Weights, +) +from ._impl import ( + bind_trellis_moe as bind, +) +from ._impl import ( + clear_trellis_moe_caches as clear_caches, +) +from ._impl import ( + plan_trellis_moe as plan, +) +from ._impl import ( + prepare_trellis_moe_weights as prepare_weights, +) +from ._impl import ( + run_trellis_moe as run, +) + + +def is_supported(device=None) -> bool: + """Return whether ``device`` supports the SM12x Trellis kernel stack.""" + return default_is_supported(device, requires=META.requires) + + +__all__ = [ + "Binding", + "Caps", + "Plan", + "Weights", + "bind", + "clear_caches", + "is_supported", + "plan", + "prepare_weights", + "run", +]