"""Host-side helpers for the CuTeDSL W4A16 MoE path.""" from __future__ import annotations from dataclasses import dataclass import torch from b12x.moe.fused.activations import ( SUPPORTED_MOE_ACTIVATIONS, is_gated_moe_activation, normalize_moe_activation, ) _W4A16_ALLOWED_ROUTED_SIZES = (8, 16, 32, 48, 64) _ROUTED_SIZE_TARGET_FILL = 0.9 _SUPPORTED_ACTIVATIONS = SUPPORTED_MOE_ACTIVATIONS @dataclass(frozen=True) class W4A16PackedShape: num_experts: int hidden_size: int intermediate_size: int w13_rows: int is_gated: bool @dataclass(frozen=True) class W4A16PackedBuffers: intermediate_cache13: torch.Tensor intermediate_cache2: torch.Tensor output: torch.Tensor fc1_c_tmp: torch.Tensor | None = None fc2_c_tmp: torch.Tensor | None = None packed_route_indices: torch.Tensor | None = None block_expert_ids: torch.Tensor | None = None packed_route_count: torch.Tensor | None = None expert_offsets: torch.Tensor | None = None @dataclass(frozen=True) class W4A16BufferPlan: routed_rows: int fc1_cols: int route_slots: int route_blocks: int fc1_c_tmp_elements: int fc2_c_tmp_elements: int intermediate_cache13_elements: int intermediate_cache2_elements: int def validate_activation(activation: str) -> bool: activation = normalize_moe_activation(activation) return is_gated_moe_activation(activation) def validate_w4a16_packed_inputs( w13_fp4: torch.Tensor, w13_global_scale: torch.Tensor, w2_fp4: torch.Tensor, w2_global_scale: torch.Tensor, *, activation: str, ) -> W4A16PackedShape: is_gated = validate_activation(activation) if w13_fp4.dtype != torch.uint8 or w2_fp4.dtype != torch.uint8: raise TypeError("packed FP4 weights must be torch.uint8") if ( w13_global_scale.dtype != torch.float32 or w2_global_scale.dtype != torch.float32 ): raise TypeError("global scales must be torch.float32") num_experts = int(w13_fp4.shape[0]) hidden_size = int(w2_fp4.shape[1]) intermediate_size = int(w2_fp4.shape[2] * 2) w13_rows = intermediate_size * (2 if is_gated else 1) if tuple(w13_fp4.shape) != (num_experts, w13_rows, hidden_size // 2): raise ValueError( f"expected w13_fp4 shape {(num_experts, w13_rows, hidden_size // 2)}, " f"got {tuple(w13_fp4.shape)}" ) if tuple(w2_fp4.shape) != (num_experts, hidden_size, intermediate_size // 2): raise ValueError( f"expected w2_fp4 shape {(num_experts, hidden_size, intermediate_size // 2)}, " f"got {tuple(w2_fp4.shape)}" ) return W4A16PackedShape( num_experts=num_experts, hidden_size=hidden_size, intermediate_size=intermediate_size, w13_rows=w13_rows, is_gated=is_gated, ) def validate_nf3_moe_inputs( w13_codes: torch.Tensor, w13_scale: torch.Tensor, w2_codes: torch.Tensor, w2_scale: torch.Tensor, *, activation: str, ) -> W4A16PackedShape: """Validate NF3 ("nf3_2p1") code/scale inputs to the packer. Codes are integer tensors of values 0..7 in kernel-native output order: w13_codes [E, 2*I, hidden] w2_codes [E, hidden, I] Scales are per-K/32-group ``t_s`` floats: w13_scale [E, 2*I, hidden//32] w2_scale [E, hidden, I//32] The NF3 runtime contract requires K % 32 == 0 and N % 64 == 0 for both GEMMs (the scale byte encoding and the 64-N tile packing). """ is_gated = validate_activation(activation) _int_dtypes = (torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64) if w13_codes.dtype not in _int_dtypes or w2_codes.dtype not in _int_dtypes: raise TypeError("NF3 codes must be integer tensors of values 0..7") num_experts = int(w2_codes.shape[0]) hidden_size = int(w2_codes.shape[1]) intermediate_size = int(w2_codes.shape[2]) w13_rows = intermediate_size * (2 if is_gated else 1) if tuple(w13_codes.shape) != (num_experts, w13_rows, hidden_size): raise ValueError( f"expected w13_codes shape {(num_experts, w13_rows, hidden_size)}, " f"got {tuple(w13_codes.shape)}" ) if tuple(w2_codes.shape) != (num_experts, hidden_size, intermediate_size): raise ValueError( f"expected w2_codes shape {(num_experts, hidden_size, intermediate_size)}, " f"got {tuple(w2_codes.shape)}" ) for name, size_k, size_n in ( ("w13", hidden_size, w13_rows), ("w2", intermediate_size, hidden_size), ): if int(size_k) % 32 != 0: raise ValueError(f"NF3 {name} requires K % 32 == 0, got {size_k}") if int(size_n) % 64 != 0: raise ValueError(f"NF3 {name} requires N % 64 == 0, got {size_n}") if tuple(w13_scale.shape) != (num_experts, w13_rows, hidden_size // 32): raise ValueError( f"expected w13_scale shape {(num_experts, w13_rows, hidden_size // 32)}, " f"got {tuple(w13_scale.shape)}" ) if tuple(w2_scale.shape) != (num_experts, hidden_size, intermediate_size // 32): raise ValueError( "expected w2_scale shape " f"{(num_experts, hidden_size, intermediate_size // 32)}, " f"got {tuple(w2_scale.shape)}" ) return W4A16PackedShape( num_experts=num_experts, hidden_size=hidden_size, intermediate_size=intermediate_size, w13_rows=w13_rows, is_gated=is_gated, ) def unswizzle_block_scale( swizzled_scale: torch.Tensor, rows: int, cols_blocks: int ) -> torch.Tensor: cols_padded = ((cols_blocks + 3) // 4) * 4 rows_padded = ((rows + 127) // 128) * 128 unswizzled = swizzled_scale.view(torch.float8_e4m3fn).reshape( rows_padded // 128, cols_padded // 4, 32, 4, 4, ) unswizzled = unswizzled.permute(0, 3, 2, 1, 4).contiguous() unswizzled = unswizzled.reshape(rows_padded, cols_padded) return unswizzled[:rows, :cols_blocks].to(torch.float32) def unswizzle_expert_scales( swizzled: torch.Tensor, *, rows: int, cols: int, ) -> torch.Tensor: if swizzled.dtype != torch.float8_e4m3fn: swizzled = swizzled.view(torch.float8_e4m3fn) scales = [ unswizzle_block_scale(swizzled[e], rows, cols // 16).to(torch.float8_e4m3fn) for e in range(swizzled.shape[0]) ] return torch.stack(scales, dim=0).contiguous() def reorder_w13_to_gate_up( w13: torch.Tensor, w13_scale: torch.Tensor, *, intermediate_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: half = int(intermediate_size) return ( torch.cat([w13[:, half:], w13[:, :half]], dim=1).contiguous(), torch.cat([w13_scale[:, half:], w13_scale[:, :half]], dim=1).contiguous(), ) def select_route_block_size_m(m: int, topk: int, num_experts: int) -> int: avg_routes_per_expert = (int(m) * int(topk)) / int(num_experts) for routed_size in _W4A16_ALLOWED_ROUTED_SIZES: if avg_routes_per_expert < _ROUTED_SIZE_TARGET_FILL * routed_size: return routed_size return _W4A16_ALLOWED_ROUTED_SIZES[-1] def max_packed_route_slots(numel: int, block_size: int, num_experts: int) -> int: max_packed_routes = int(numel) + int(num_experts) * (int(block_size) - 1) if int(numel) < int(num_experts): max_packed_routes = min( int(numel) * int(block_size), max_packed_routes, ) return max_packed_routes def route_pack_numel_capacity(numel: int, topk: int = 1) -> int: topk = max(int(topk), 1) tokens = (max(int(numel), 1) + topk - 1) // topk return route_pack_token_capacity(tokens, topk) * topk def route_pack_token_capacity(tokens: int, topk: int) -> int: del topk return 1 << (max(int(tokens), 1) - 1).bit_length() def max_w4a16_route_capacity(routed_rows: int, num_experts: int) -> tuple[int, int]: route_slots = 0 route_blocks = 0 for block_size in _W4A16_ALLOWED_ROUTED_SIZES: slots = max_packed_route_slots( int(routed_rows), int(block_size), int(num_experts) ) route_slots = max(route_slots, slots) route_blocks = max( route_blocks, (slots + int(block_size) - 1) // int(block_size) ) return max(route_slots, 1), max(route_blocks, 1) def packed_gemm_scratch_elements( *, size_n: int, route_slots: int, moe_block_size: int, sms: int, ) -> int: elements = min( int(size_n) * int(route_slots), int(sms) * 4 * int(moe_block_size) * 256, ) if moe_block_size == 8: elements *= 2 return max(elements, 1) def plan_w4a16_buffers( prepared, *, m: int, topk: int, route_num_experts: int | None = None, sms: int, ) -> W4A16BufferPlan: routed_rows = int(m) * int(topk) route_num_experts = ( int(prepared.num_experts) if route_num_experts is None else int(route_num_experts) ) intermediate_size = int(prepared.intermediate_size) hidden_size = int(prepared.hidden_size) fc1_cols = (2 if prepared.is_gated else 1) * intermediate_size block_size_m = select_route_block_size_m(m, topk, route_num_experts) route_slots = max_packed_route_slots(routed_rows, block_size_m, route_num_experts) route_blocks = (route_slots + block_size_m - 1) // block_size_m scratch_sms = int(sms) return W4A16BufferPlan( routed_rows=routed_rows, fc1_cols=fc1_cols, route_slots=route_slots, route_blocks=route_blocks, fc1_c_tmp_elements=packed_gemm_scratch_elements( size_n=fc1_cols, route_slots=route_slots, moe_block_size=block_size_m, sms=scratch_sms, ), fc2_c_tmp_elements=packed_gemm_scratch_elements( size_n=hidden_size, route_slots=route_slots, moe_block_size=block_size_m, sms=scratch_sms, ), intermediate_cache13_elements=routed_rows * max(fc1_cols, hidden_size), intermediate_cache2_elements=routed_rows * intermediate_size, ) def make_w4a16_packed_buffers( prepared, *, m: int, topk: int, dtype: torch.dtype, device: torch.device, route_num_experts: int | None = None, ) -> W4A16PackedBuffers: route_num_experts = ( int(prepared.num_experts) if route_num_experts is None else int(route_num_experts) ) sms = int(torch.cuda.get_device_properties(device).multi_processor_count) plan = plan_w4a16_buffers( prepared, m=m, topk=topk, route_num_experts=route_num_experts, sms=sms, ) fc1_c_tmp = torch.empty( (plan.fc1_c_tmp_elements,), dtype=torch.float32, device=device, ) fc2_c_tmp = torch.empty( (plan.fc2_c_tmp_elements,), dtype=torch.float32, device=device, ) return W4A16PackedBuffers( intermediate_cache13=torch.empty( (plan.intermediate_cache13_elements,), dtype=dtype, device=device, ), intermediate_cache2=torch.empty( (plan.routed_rows, int(prepared.intermediate_size)), dtype=dtype, device=device, ), output=torch.empty((m, prepared.hidden_size), dtype=dtype, device=device), fc1_c_tmp=fc1_c_tmp, fc2_c_tmp=fc2_c_tmp, packed_route_indices=torch.empty( (plan.route_slots,), dtype=torch.int32, device=device ), block_expert_ids=torch.empty( (plan.route_blocks,), dtype=torch.int32, device=device ), packed_route_count=torch.empty((1,), dtype=torch.int32, device=device), expert_offsets=torch.empty( (route_num_experts + 1,), dtype=torch.int32, device=device ), ) __all__ = [ "W4A16BufferPlan", "W4A16PackedBuffers", "W4A16PackedShape", "make_w4a16_packed_buffers", "max_w4a16_route_capacity", "packed_gemm_scratch_elements", "max_packed_route_slots", "plan_w4a16_buffers", "reorder_w13_to_gate_up", "route_pack_numel_capacity", "route_pack_token_capacity", "select_route_block_size_m", "unswizzle_block_scale", "unswizzle_expert_scales", "validate_activation", "validate_nf3_moe_inputs", "validate_w4a16_packed_inputs", ]