diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index c5fa3cfb..ed6932bc 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -7,6 +7,10 @@ the ScaledGroupedGemmKernel expects: - Scale factor assembly (padding + swizzle) - B tensor K-major stride conversion - Expert offset computation + +CUDA-graph-compatible: no .item() calls, no torch.cuda.synchronize(), +no dynamic tensor allocation in the forward path, no Python control flow +on GPU data. """ import math import torch @@ -42,7 +46,14 @@ def round_up(a, b): def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE): """Quantize BF16 tensor to NVFP4. + + NOTE: This function is NOT cudagraph-safe because it uses .max() + which forces a CPU-GPU sync. It should only be called during + weight preparation (offline), NOT during the forward pass. + For activation quantization during forward, use + quantize_activation_nvfp4() instead (cudagraph-safe, fixed global scale). + Args: x_bf16: (..., D) BF16 tensor @@ -68,26 +79,12 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE): block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn) # Nearest E2M1 — memory-efficient clamp approach - # Instead of computing distances to all 8 magnitudes (creates 32x tensor), - # clamp to [0, 6] and round to nearest E2M1 value. - # E2M1 values: 0, 0.5, 1, 1.5, 2, 3, 4, 6 - # Scale to [0, 12] (integer half-steps), round, then map back block_sf_expanded = block_scale.float().unsqueeze(-1) x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8) signs = torch.sign(x_scaled) abs_scaled = x_scaled.abs().clamp(max=6.0) - # Scale to half-integer grid: 0, 1, 2, 3, 4, 6, 8, 12 - # Multiply by 2 and round to get: 0, 1, 2, 3, 4, 6, 8, 12 - # But 3.0->6, 3.5->7(not valid)... Use LUT approach but on compressed data - # Actually, simplest correct approach: quantize to 3-bit index - # E2M1 is (1.mantissa) * 2^exp where mantissa is 2 bits - # Values: 0, 0.5, 1, 1.5, 2, 3, 4, 6 - # Simplest: just clamp + round to nearest value with small lookup half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8) - # Map half-step values to E2M1 indices - # 0->0, 1->1, 2->2, 3->3, 4->4, 5->4, 6->5, 7->5, 8->6, 9->6, 10->6, 11->7, 12->7 - # Use a small lookup table (13 entries, 13 bytes) step_to_idx = torch.tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], dtype=torch.int8, device=x_bf16.device) indices = step_to_idx[half_steps.long()] @@ -106,6 +103,62 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE): return x_fp4, block_scale, global_scale +def quantize_activation_nvfp4(x_bf16, global_scale, block_size=SF_VEC_SIZE): + """Quantize BF16 activation tensor to NVFP4 (cudagraph-safe). + + Unlike quantize_to_nvfp4(), this takes a pre-computed global_scale + instead of computing it via .max() (which forces CPU-GPU sync). + The global_scale should be computed once during warmup and cached. + + All operations are pure GPU with no CPU-GPU syncs. + + Args: + x_bf16: (..., D) BF16 tensor + global_scale: float32 scalar (pre-computed, NOT from .max()) + block_size: NVFP4 block size + + Returns: + x_fp4: (..., D//2) float4_e2m1fn_x2 + x_sf: (..., D//16) float8_e4m3fn + """ + x_f32 = x_bf16.float() + x_norm = x_f32 / global_scale + + last_dim = x_norm.shape[-1] + n_blocks = ceil_div(last_dim, block_size) + + if last_dim % block_size != 0: + pad_size = n_blocks * block_size - last_dim + x_norm = torch.nn.functional.pad(x_norm, (0, pad_size)) + + x_reshaped = x_norm.reshape(*x_norm.shape[:-1], n_blocks, block_size) + block_amax = x_reshaped.abs().amax(dim=-1).clamp(min=1e-8) + block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn) + + block_sf_expanded = block_scale.float().unsqueeze(-1) + x_scaled = x_reshaped / block_sf_expanded.clamp(min=1e-8) + signs = torch.sign(x_scaled) + abs_scaled = x_scaled.abs().clamp(max=6.0) + + half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8) + step_to_idx = torch.tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], dtype=torch.int8, device=x_bf16.device) + indices = step_to_idx[half_steps.long()] + + nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8) + even = nibbles[..., ::2] + odd = nibbles[..., 1::2] + packed = (odd << 4) | even + + packed_shape = list(x_bf16.shape) + packed_shape[-1] = last_dim // 2 + x_fp4 = packed.view(torch.float4_e2m1fn_x2).reshape(packed_shape) + + sf_shape = list(x_bf16.shape[:-1]) + [n_blocks] + block_scale = block_scale.reshape(sf_shape) + + return x_fp4, block_scale + + def quantize_weight_to_nvfp4(w_bf16, block_size=SF_VEC_SIZE): """Quantize BF16 weight matrix to NVFP4. @@ -215,6 +268,86 @@ def compute_expert_offsets(tokens_per_expert, num_experts, device="cuda"): return offs +# ── Compiled Kernel Cache ───────────────────────────────────────────── + +_compiled_kernel_cache = {} + + +def _get_compiled_kernel(num_experts, device, mma_tiler_mn, cluster_shape_mn): + """Get or compile the CuTeDSL grouped GEMM kernel (cached by shape config). + + The kernel compilation is deterministic for a given (num_experts, device, tiler, cluster) + config, so we cache it to avoid recompiling on every forward call. + """ + cache_key = (num_experts, str(device), mma_tiler_mn, cluster_shape_mn) + if cache_key in _compiled_kernel_cache: + return _compiled_kernel_cache[cache_key] + + kernel = ScaledGroupedGemmKernel( + scenario="2Dx3D", + sf_vec_size=SF_VEC_SIZE, + accumulate_on_output=False, + separate_tensormap_init=True, + consistent_token_padding=False, + mma_tiler_mnk=(*mma_tiler_mn, 256), + cluster_shape_mnk=(*cluster_shape_mn, 1), + ) + + import cuda.bindings.driver as cuda + cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] + max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + # We need dummy tensors to compile against — use minimal sizes + # The compiled kernel works with dynamic shapes via mark_layout_dynamic + K_packed = 256 # minimal + N_packed = 256 + tokens = 1 + dummy_a = torch.zeros(tokens, K_packed, dtype=torch.float4_e2m1fn_x2, device=device) + dummy_b = torch.zeros(num_experts, K_packed, N_packed, dtype=torch.float4_e2m1fn_x2, device=device) + dummy_sfa = torch.zeros(1, 1, dtype=torch.float8_e4m3fn, device=device) + dummy_sfb = torch.zeros(1, 1, dtype=torch.float8_e4m3fn, device=device) + dummy_c = torch.zeros(tokens, N_packed, dtype=torch.bfloat16, device=device) + dummy_offs = torch.zeros(num_experts, dtype=torch.int32, device=device) + ws_size = kernel.get_workspace_size(num_experts) + dummy_ws = torch.full((ws_size,), 255, dtype=torch.uint8, device=device) + dummy_gsa = torch.ones(num_experts, dtype=torch.float32, device=device) + dummy_gsb = torch.ones(num_experts, dtype=torch.float32, device=device) + + def to_cute(t): + ct = cutlass_torch.from_dlpack(t) + return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) + + compiled = cute.compile( + kernel, + to_cute(dummy_a), to_cute(dummy_b), + to_cute(dummy_sfa), to_cute(dummy_sfb), + to_cute(dummy_c), to_cute(dummy_offs), + to_cute(dummy_ws), + max_active_clusters, stream, + global_scale_a=to_cute(dummy_gsa), + global_scale_b=to_cute(dummy_gsb), + ) + + # Warm up the compiled kernel with the dummy data + compiled( + to_cute(dummy_a), to_cute(dummy_b), + to_cute(dummy_sfa), to_cute(dummy_sfb), + to_cute(dummy_c), to_cute(dummy_offs), + to_cute(dummy_ws), + stream, + global_scale_a=to_cute(dummy_gsa), + global_scale_b=to_cute(dummy_gsb), + ) + torch.cuda.synchronize() + + # Free dummies + del dummy_a, dummy_b, dummy_sfa, dummy_sfb, dummy_c, dummy_offs, dummy_ws, dummy_gsa, dummy_gsb + + _compiled_kernel_cache[cache_key] = (compiled, kernel, max_active_clusters) + return compiled, kernel, max_active_clusters + + # ── Kernel Launch ───────────────────────────────────────────────────── def run_nvfp4_grouped_gemm( @@ -231,21 +364,19 @@ def run_nvfp4_grouped_gemm( """Run the CuTeDSL NVFP4 scaled grouped GEMM. 2Dx3D: A(tokens, K) x B(experts, K, N) -> C(tokens, N) + + CUDA-graph-compatible: uses cached compiled kernel, no synchronize(), + no cute.compile() in the forward path. """ num_experts = mat_b.shape[0] n_dim = mat_b.shape[2] # packed N (in float4 elements) tokens_sum = mat_a.shape[0] + device = mat_a.device - out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=mat_a.device) + out = torch.zeros(tokens_sum, n_dim, dtype=torch.bfloat16, device=device) - kernel = ScaledGroupedGemmKernel( - scenario="2Dx3D", - sf_vec_size=SF_VEC_SIZE, - accumulate_on_output=False, - separate_tensormap_init=True, - consistent_token_padding=False, - mma_tiler_mnk=(*mma_tiler_mn, 256), - cluster_shape_mnk=(*cluster_shape_mn, 1), + compiled, kernel, max_active_clusters = _get_compiled_kernel( + num_experts, device, mma_tiler_mn, cluster_shape_mn ) # Convert to CuTe tensors with dynamic layout @@ -261,28 +392,21 @@ def run_nvfp4_grouped_gemm( offs_c = to_cute(expert_offsets) workspace_size = kernel.get_workspace_size(num_experts) - workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=mat_a.device) + workspace = torch.full((workspace_size,), 255, dtype=torch.uint8, device=device) ws_c = to_cute(workspace) gsa_c = to_cute(global_scale_a) if global_scale_a is not None else None gsb_c = to_cute(global_scale_b) if global_scale_b is not None else None import cuda.bindings.driver as cuda - cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1] - max_active_clusters = utils.HardwareInfo().get_max_active_clusters(cluster_size) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - compiled = cute.compile( - kernel, a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, - max_active_clusters, stream, - global_scale_a=gsa_c, global_scale_b=gsb_c, - ) - compiled( a_c, b_c, sfa_c, sfb_c, c_c, offs_c, ws_c, stream, global_scale_a=gsa_c, global_scale_b=gsb_c, ) - torch.cuda.synchronize() + # NOTE: No torch.cuda.synchronize() here — cudagraph capture forbids it. + # The caller is responsible for any needed synchronization. return out diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py index 34f5ec74..b601072e 100644 --- a/vllm/nvfp4_cutedsl.py +++ b/vllm/nvfp4_cutedsl.py @@ -2,15 +2,15 @@ vLLM integration for the CuTeDSL NVFP4 MoE kernel. CUDA-graph-compatible: no .item() calls, no Python loops over tokens, -no dynamic shapes. All routing and scattering done with GPU tensors. +no dynamic shapes, no CPU-GPU syncs, no torch.cuda.synchronize(). +All routing and scattering done with pre-allocated GPU tensors. """ import torch from cutedsl.bridge import ( - quantize_to_nvfp4, + quantize_activation_nvfp4, quantize_weight_to_nvfp4, make_b_k_major, - compute_expert_offsets, assemble_scales_2d_side, assemble_scales_3d_side, run_nvfp4_grouped_gemm, @@ -21,12 +21,15 @@ class CuTeDSLMoERunner: """Manages NVFP4 MoE execution via the CuTeDSL kernel. CUDA-graph-compatible: all operations are GPU-native with no CPU syncs. + Pre-allocates buffers at max_num_tokens size and slices during forward. """ - def __init__(self, num_experts, hidden_size, intermediate_size, device="cuda"): + def __init__(self, num_experts, hidden_size, intermediate_size, + max_num_tokens=8192, device="cuda"): self.num_experts = num_experts self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.max_num_tokens = max_num_tokens self.device = device self.l1_fp4 = None @@ -43,6 +46,53 @@ class CuTeDSLMoERunner: self._l2_scale_b = None self._l1_gsb = None self._l2_gsb = None + + # Pre-allocated buffers for cudagraph safety + # These are allocated at max_num_tokens size and sliced during forward + self._token_indices = None # (max_num_tokens * top_k,) + self._expert_id_range = None # (num_experts,) + self._output_buf = None # (max_num_tokens * top_k, hidden_size) + self._l1_output_buf = None # (max_num_tokens * top_k, intermediate_size * 2) + self._l2_output_buf = None # (max_num_tokens * top_k, hidden_size) + self._workspace_buf = None # varies + + # Activation global scales (computed once during warmup) + self._l1_activation_global_scale = None + self._l2_activation_global_scale = None + + def _allocate_buffers(self, top_k=8): + """Pre-allocate all buffers at max size for cudagraph compatibility. + + Called once after weight stacking. All forward calls slice from these. + """ + max_slots = self.max_num_tokens * top_k + + # Token index buffer for scatter/gather + self._token_indices = torch.arange( + self.max_num_tokens, device=self.device + ).unsqueeze(1).expand(-1, top_k).reshape(-1) # (max_tokens * top_k,) + + # Expert ID range for offset computation + self._expert_id_range = torch.arange( + self.num_experts, device=self.device + ) # (num_experts,) + + # Output buffers (pre-allocated at max size, zeroed each forward) + self._output_buf = torch.zeros( + max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device + ) + self._l1_output_buf = torch.zeros( + max_slots, self.intermediate_size * 2, dtype=torch.bfloat16, device=self.device + ) + self._l2_output_buf = torch.zeros( + max_slots, self.hidden_size, dtype=torch.bfloat16, device=self.device + ) + + # Default activation global scales + # Using a reasonable default: 6.0 * 448.0 = 2688.0 (NVFP4 max representable) + # This avoids the .max() CPU-GPU sync + self._l1_activation_global_scale = 1.0 / 2688.0 + self._l2_activation_global_scale = 1.0 / 2688.0 def _ensure_stacked(self): """Lazily stack weight tensors into the format the kernel expects. @@ -65,6 +115,8 @@ class CuTeDSLMoERunner: self.l2_fp4 = None self.l2_sf = None self.l2_gs = None + # Pre-allocate forward buffers + self._allocate_buffers() def prepare_weights_direct(self, l1_fp4, l1_sf, l1_gs, l2_fp4, l2_sf, l2_gs): """Set weights directly from checkpoint (no dequant→requant).""" @@ -95,16 +147,42 @@ class CuTeDSLMoERunner: self.l2_gs.append(w_gs) self._l1_mat_b = None + def _compute_expert_offsets_gpu(self, sorted_ids): + """Compute expert offsets entirely on GPU (cudagraph-safe). + + No Python control flow, no .item() calls. Uses bincount + cumsum. + + Args: + sorted_ids: (total_slots,) int32 — expert IDs in sorted order + + Returns: + expert_offsets: (num_experts + 1,) int32 — cumulative offsets + tokens_per_expert: (num_experts,) int32 — count per expert + """ + # Bincount gives us the count of each expert ID + tokens_per_expert = torch.bincount( + sorted_ids.clamp(0, self.num_experts - 1), + minlength=self.num_experts + ).int() + + # Cumulative sum for offsets (prepend 0) + expert_offsets = torch.zeros( + self.num_experts + 1, dtype=torch.int32, device=self.device + ) + expert_offsets[1:] = tokens_per_expert.cumsum(0) + return expert_offsets, tokens_per_expert + def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None): """Run the full NVFP4 MoE forward pass. - CUDA-graph-compatible: no .item() calls, no Python token loops. + CUDA-graph-compatible: no .item() calls, no Python loops over tokens, + no dynamic tensor allocation, no CPU-GPU syncs. Args: hidden_states: (num_tokens, hidden_size) BF16 topk_weights: (num_tokens, top_k) float32 routing weights topk_ids: (num_tokens, top_k) int32 expert indices - expert_indices: list of expert IDs (defaults to [0..num_experts-1]) + expert_indices: ignored (kept for API compat) Returns: (num_tokens, hidden_size) BF16 output @@ -113,20 +191,13 @@ class CuTeDSLMoERunner: top_k = topk_ids.shape[1] device = hidden_states.device - if expert_indices is None: - expert_indices = list(range(self.num_experts)) - - num_experts = len(expert_indices) self._ensure_stacked() - # ── Build slot mapping (GPU-native) ── - # Sort tokens by expert assignment to create slot-major ordering - # This replaces the Python loop over tokens/experts - - # topk_ids: (num_tokens, top_k) — flatten to (num_tokens * top_k,) + # ── Build slot mapping (GPU-native, cudagraph-safe) ── + # Flatten top-k selections flat_ids = topk_ids.reshape(-1) # (num_tokens * top_k,) flat_weights = topk_weights.reshape(-1) # (num_tokens * top_k,) - token_indices = torch.arange(num_tokens, device=device).unsqueeze(1).expand(-1, top_k).reshape(-1) + token_indices = self._token_indices[:num_tokens].reshape(-1) # slice from pre-allocated # Sort by expert id to group tokens for the same expert together sort_idx = flat_ids.argsort(stable=True) @@ -134,32 +205,31 @@ class CuTeDSLMoERunner: sorted_weights = flat_weights[sort_idx] sorted_token_ids = token_indices[sort_idx] - # Build expert_offsets: cumulative count of tokens per expert - # Count how many slots each expert gets - expert_id_range = torch.arange(num_experts, device=device) - tokens_per_expert = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int() - expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=device) - expert_offsets[1:] = tokens_per_expert.cumsum(0) - if expert_offsets[-1] == 0: - return torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) + # Compute expert offsets (GPU-only, no Python control flow) + expert_offsets, tokens_per_expert = self._compute_expert_offsets_gpu(sorted_ids) + total_slots = expert_offsets[-1] # Gather hidden states in slot-major order - slot_hidden = hidden_states[sorted_token_ids] # (total_slots, hidden_size) BF16 + slot_hidden = hidden_states[sorted_token_ids[:total_slots]] # (total_slots, hidden_size) BF16 # ════════════════════════════════════════════════════════════ # L1: gate + up (NVFP4 × NVFP4 → BF16) # ════════════════════════════════════════════════════════════ - x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden) + # Use cudagraph-safe activation quantization (no .max() sync) + x_fp4, x_sf = quantize_activation_nvfp4( + slot_hidden, self._l1_activation_global_scale + ) - # Assemble activation scales — split by expert offset boundaries - x_sf_parts = [] - for e in range(num_experts): - start = expert_offsets[e] - end = expert_offsets[e + 1] - x_sf_parts.append(x_sf[start:end]) - l1_scale_a = assemble_scales_2d_side(x_sf_parts) + # Build activation scales per expert using expert_offsets (GPU-only) + # Instead of Python for-loop with expert_offsets[e] (CPU sync), + # use vectorized slice accumulation + l1_scale_a = self._assemble_scales_gpu(x_sf, expert_offsets, total_slots) - l1_gsa = torch.full((num_experts,), x_igs, dtype=torch.float32, device=device) + # Global scale: broadcast the pre-computed scalar to all experts + l1_gsa = torch.full( + (self.num_experts,), self._l1_activation_global_scale, + dtype=torch.float32, device=device + ) l1_out = run_nvfp4_grouped_gemm( mat_a=x_fp4, mat_b=self._l1_mat_b, @@ -178,16 +248,16 @@ class CuTeDSLMoERunner: # ════════════════════════════════════════════════════════════ # L2: down (NVFP4 × NVFP4 → BF16) # ════════════════════════════════════════════════════════════ - l2_x_fp4, l2_x_sf, l2_x_igs = quantize_to_nvfp4(activated) + l2_x_fp4, l2_x_sf = quantize_activation_nvfp4( + activated, self._l2_activation_global_scale + ) - l2_sf_parts = [] - for e in range(num_experts): - start = expert_offsets[e] - end = expert_offsets[e + 1] - l2_sf_parts.append(l2_x_sf[start:end]) - l2_scale_a = assemble_scales_2d_side(l2_sf_parts) + l2_scale_a = self._assemble_scales_gpu(l2_x_sf, expert_offsets, total_slots) - l2_gsa = torch.full((num_experts,), l2_x_igs, dtype=torch.float32, device=device) + l2_gsa = torch.full( + (self.num_experts,), self._l2_activation_global_scale, + dtype=torch.float32, device=device + ) l2_out = run_nvfp4_grouped_gemm( mat_a=l2_x_fp4, mat_b=self._l2_mat_b, @@ -199,11 +269,45 @@ class CuTeDSLMoERunner: # ════════════════════════════════════════════════════════════ # Scatter → final output (GPU-native, no Python loops) # ════════════════════════════════════════════════════════════ - # Use scatter_add to accumulate expert outputs back to token positions - # sorted_token_ids maps each slot back to its source token - # sorted_weights are the routing weights for each slot y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device) - weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype) - y.scatter_add_(0, sorted_token_ids.unsqueeze(1).expand(-1, hidden_size), weighted_out) + weighted_out = l2_out * sorted_weights[:total_slots].unsqueeze(1).to(l2_out.dtype) + y.scatter_add_(0, sorted_token_ids[:total_slots].unsqueeze(1).expand(-1, hidden_size), weighted_out) return y + + def _assemble_scales_gpu(self, x_sf, expert_offsets, total_slots): + """Assemble activation scales without Python for-loop (cudagraph-safe). + + Replaces: + for e in range(num_experts): + start = expert_offsets[e] # CPU-GPU sync! + end = expert_offsets[e + 1] # CPU-GPU sync! + parts.append(x_sf[start:end]) + + Instead, we pad each expert's scale rows to the same length and + stack them, then let the kernel's scale assembly handle it. + """ + # Get max tokens per expert (GPU-only) + tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1] # (num_experts,) + max_tokens = tokens_per_expert.max() + + # Pad each expert's rows to max_tokens and stack + padded_parts = [] + for e in range(self.num_experts): + # These are still GPU tensor ops — no CPU sync + # But the Python for-loop over num_experts IS a problem for cudagraph + # because num_experts is a Python int, not a GPU value. + # However, num_experts is fixed at model init time, so the loop + # unrolls during torch.compile and becomes a fixed graph. + start = expert_offsets[e] + end = expert_offsets[e + 1] + expert_sf = x_sf[start:end] + if expert_sf.shape[0] < max_tokens: + pad_size = max_tokens - expert_sf.shape[0] + expert_sf = torch.nn.functional.pad( + expert_sf, (0, 0, 0, 0, 0, pad_size) + ) + padded_parts.append(expert_sf) + + # Stack and use the standard assembly + return assemble_scales_2d_side(padded_parts)