"""Local NVFP4/BF16 weight preparation for the CuTeDSL W4A16 path.""" from __future__ import annotations import math from dataclasses import dataclass import torch from b12x.moe.fused.w4a16.host import ( W4A16PackedBuffers, make_w4a16_packed_buffers as _make_w4a16_packed_buffers, unswizzle_expert_scales, validate_nf3_moe_inputs, validate_w4a16_packed_inputs, ) _PACKED_TILE_SIZE = 16 _PACKED_TILE_N_SIZE = 64 _PACK_FACTOR_4BIT = 8 _MODEL_OPT_W13_LAYOUTS = {"w13", "w31"} _SOURCE_FORMATS = { "modelopt_nvfp4": "modelopt_nvfp4", "fp4_e8m0_k32": "fp4_e8m0_k32", "compressed_tensors": "compressed_tensors", } _E8M0_K32_BF16_MAX_SCALE_BYTE = 247 _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT = 64 # Canonical W13 layout names are "w13"/"w31"; accept the physical FC1-half # spellings as aliases. Logical checkpoint order "w13" arrives up/gate and # needs a swap before the kernel's SwiGLU; "w31" is already kernel-native # gate/up order. _W13_LAYOUTS = { "w13": "w13", "w31": "w31", "up_gate": "w13", "gate_up": "w31", } _MODEL_OPT_NVFP4_FORMATS = {"modelopt_nvfp4"} # NF3 ("nf3_2p1") 3-bit codebook. Must match kernel._NF3_CODEBOOK. The scale # convention (see NF3_KERNEL_MODSPEC.md): weights decode to the FULL-precision # bf16 codebook value, the K/32 scale byte encodes t_s * 2**-4 (e4m3-style, # E4=0 arm), and the per-tensor global_scale carries the matching 2**116. _NF3_CODEBOOK = (-1.0, -0.6047, -0.3563, -0.1275, 0.1275, 0.3563, 0.6047, 1.0) _NF3_SCALE_GLOBAL = 2.0**116 _NF3_SCALE_FLOOR = 2.0 ** -10 # f16(t_s*2^-4) must stay f16-NORMAL (>=2^-14) -> t_s >= 2^-10; real early-layer scales dip below the old 2^-7 @dataclass(frozen=True) class W4A16PackedWeights: w13: torch.Tensor w13_scale: torch.Tensor w13_global_scale: torch.Tensor w2: torch.Tensor w2_scale: torch.Tensor w2_global_scale: torch.Tensor workspace: torch.Tensor hidden_size: int intermediate_size: int num_experts: int is_gated: bool params_dtype: torch.dtype source_format: str = "modelopt_nvfp4" w13_layout: str = "w13" weight_layout: str = "packed" scale_format: str = "e4m3_k16" @dataclass(frozen=True) class W4A16ModelOptWeights: w13: torch.Tensor w13_scale: torch.Tensor w13_global_scale: torch.Tensor w2: torch.Tensor w2_scale: torch.Tensor w2_global_scale: torch.Tensor workspace: torch.Tensor hidden_size: int intermediate_size: int num_experts: int is_gated: bool params_dtype: torch.dtype source_format: str = "modelopt_nvfp4" weight_layout: str = "modelopt" scale_format: str = "e4m3_k16" micro_w13_scale: torch.Tensor | None = None micro_w13_global_scale: torch.Tensor | None = None micro_w2_scale: torch.Tensor | None = None micro_w2_global_scale: torch.Tensor | None = None # Physical order of the two fused FC1 halves in source W13. "w13" (logical, # == "up_gate" physical) needs a row rotation before W4A16 SwiGLU; "w31" # (== "gate_up") is already in the kernel-native order. w13_layout: str = "w13" @dataclass(frozen=True) class PreparedNF3MoeWeights: """Runtime NF3 ("nf3_2p1") hybrid-expert W4A16 weights. Mirrors W4A16PackedWeights so run_w4a16_moe consumes it identically. The packed planes are int32 (3 words per 32-code unit); the K/32 scales are e4m3-style bytes viewed uint8; the global scales carry the 2**116 the NF3 scale convention defers out of the codebook. fc1_tile_n / fc2_tile_n record the CTA N-tile the flat-span packing was built for -- the kernel MUST be compiled/launched with the SAME tile_n (read them back from the compiled W4A16FusedMoeCompileResult). """ w13: torch.Tensor w13_scale: torch.Tensor w13_global_scale: torch.Tensor w2: torch.Tensor w2_scale: torch.Tensor w2_global_scale: torch.Tensor workspace: torch.Tensor hidden_size: int intermediate_size: int num_experts: int is_gated: bool params_dtype: torch.dtype fc1_tile_n: int fc2_tile_n: int source_format: str = "nf3_2p1" w13_layout: str = "packed" weight_layout: str = "nf3_2p1" scale_format: str = "e4m3_k32" def _make_workspace( device: torch.device, *, max_blocks_per_sm: int = 4 ) -> torch.Tensor: props = torch.cuda.get_device_properties(device) sms = int(props.multi_processor_count) return torch.zeros( (sms * int(max_blocks_per_sm) + 2,), dtype=torch.int32, device=device, ) def _scale_perms() -> tuple[list[int], list[int]]: scale_perm: list[int] = [] for i in range(8): scale_perm.extend([i + 8 * j for j in range(8)]) scale_perm_single: list[int] = [] for i in range(4): scale_perm_single.extend([2 * i + j for j in [0, 1, 8, 9, 16, 17, 24, 25]]) return scale_perm, scale_perm_single def _e8m0_logical_tail_scale_n(size_n: int) -> int: return ( (int(size_n) + _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT - 1) // _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT ) * _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT def _permute_packed_scales( scales: torch.Tensor, *, size_k: int, size_n: int, group_size: int, output_size_n: int | None = None, ) -> torch.Tensor: scale_perm, scale_perm_single = _scale_perms() if group_size < size_k and group_size != -1: block = len(scale_perm) if output_size_n is not None or int(size_n) % block != 0: padded_n = ( ((int(size_n) + block - 1) // block) * block if output_size_n is None else int(output_size_n) ) if padded_n < int(size_n) or padded_n % block != 0: raise ValueError( f"output_size_n must be a multiple of {block} and >= size_n" ) padded = scales.new_zeros((int(scales.shape[0]), padded_n)) padded[:, : int(size_n)] = scales scales = padded.reshape((int(scales.shape[0]), -1, block)) scales = scales[:, :, scale_perm].reshape((int(scales.shape[0]), padded_n)) if output_size_n is None: scales = scales[:, : int(size_n)] return scales.contiguous() scales = scales.reshape((-1, block))[:, scale_perm] else: block = len(scale_perm_single) if output_size_n is not None or int(size_n) % block != 0: padded_n = ( ((int(size_n) + block - 1) // block) * block if output_size_n is None else int(output_size_n) ) if padded_n < int(size_n) or padded_n % block != 0: raise ValueError( f"output_size_n must be a multiple of {block} and >= size_n" ) padded = scales.new_zeros((int(scales.shape[0]), padded_n)) padded[:, : int(size_n)] = scales scales = padded.reshape((int(scales.shape[0]), -1, block)) scales = scales[:, :, scale_perm_single].reshape( (int(scales.shape[0]), padded_n) ) if output_size_n is None: scales = scales[:, : int(size_n)] return scales.contiguous() scales = scales.reshape((-1, block))[:, scale_perm_single] return scales.reshape((-1, size_n)).contiguous() def _nvfp4_compute_scale_factor( packed_scales: torch.Tensor, a_dtype: torch.dtype, ) -> float: if a_dtype == torch.float16: return 1.0 max_scalar = 0.0 for expert in range(int(packed_scales.shape[0])): ws_float = packed_scales[expert].float() * (2**7) nonzero_mask = ws_float > 0 if bool(nonzero_mask.any().item()): max_scalar = max(max_scalar, float(ws_float[nonzero_mask].max().item())) if max_scalar > 0.0 and max_scalar < 448 * (2**7): return float(2 ** math.floor(math.log2((448 * (2**7)) / max_scalar))) return 1.0 def _process_nvfp4_packed_scales( packed_scales: torch.Tensor, *, scale_factor: float, ) -> torch.Tensor: packed_scales = packed_scales.to(torch.float16) packed_scales = packed_scales.view(-1, 4)[:, [0, 2, 1, 3]].view( packed_scales.size(0), -1, ) if scale_factor > 1.0: packed_scales = (packed_scales.float() * scale_factor).to(torch.float16) packed_scales = packed_scales * (2**7) packed_scales[packed_scales < 2] = 0 packed_scales = packed_scales.view(torch.int16) << 1 packed_scales = packed_scales.view(torch.float8_e4m3fn) return packed_scales[:, 1::2].contiguous() def _process_nvfp4_packed_global_scale( global_scale: torch.Tensor, *, a_dtype: torch.dtype, ) -> torch.Tensor: if a_dtype == torch.float16: target_exponent = 5 elif a_dtype == torch.bfloat16: target_exponent = 8 else: raise TypeError(f"unsupported W4A16 activation dtype {a_dtype}") fp4_exponent = 2 exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp4_exponent - 1) return global_scale * (2.0 ** (exponent_bias - 7)) def _process_nvfp4_micro_global_scale_from_packed( packed_global_scale: torch.Tensor, *, a_dtype: torch.dtype, ) -> torch.Tensor: if a_dtype == torch.float16: target_exponent = 5 elif a_dtype == torch.bfloat16: target_exponent = 8 else: raise TypeError(f"unsupported W4A16 activation dtype {a_dtype}") fp4_exponent = 2 exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp4_exponent - 1) return (packed_global_scale * (2.0 ** (-exponent_bias))).to(torch.float32).contiguous() def _normalize_source_format(source_format: str) -> str: if source_format.lower() == "mxfp4_native": raise ValueError( "source_format='mxfp4_native' has been removed; use " "source_format='fp4_e8m0_k32' for E8M0 K/32 scales, or add " "a real MXFP4 source contract" ) try: return _SOURCE_FORMATS[source_format.lower()] except KeyError as exc: raise ValueError( "source_format must be one of 'modelopt_nvfp4', " "'fp4_e8m0_k32', or 'compressed_tensors', " f"got {source_format!r}" ) from exc def _normalize_w13_layout(w13_layout: str) -> str: try: return _W13_LAYOUTS[w13_layout.lower()] except KeyError as exc: raise ValueError( "w13_layout must be one of 'w13'/'w31' (or the 'up_gate'/'gate_up' " f"aliases), got {w13_layout!r}" ) from exc def _source_global_scale( global_scale: torch.Tensor, *, source_format: str ) -> torch.Tensor: if source_format == "compressed_tensors": return (1.0 / global_scale).to(torch.float32).contiguous() return global_scale.contiguous() def _validate_e8m0_k32_scales( scales: torch.Tensor, *, rows: int, cols: int, name: str, allow_k_tail: bool = False, ) -> torch.Tensor: """Validate source E8M0 K/32 scale tensor shape and dtype.""" if scales.ndim != 3: raise ValueError( f"{name} must be [E, N, ceil(K/32)], got {tuple(scales.shape)}" ) if allow_k_tail: if int(cols) % 8 != 0: raise ValueError( f"{name} compact E8M0 K-tail requires K divisible by 8, got " f"{int(cols)}" ) expected_cols = (int(cols) + 31) // 32 elif int(cols) % 32 != 0: raise ValueError(f"{name} requires K divisible by 32, got {int(cols)}") else: expected_cols = int(cols) // 32 if tuple(scales.shape[1:]) != (int(rows), expected_cols): raise ValueError( f"{name} must have shape [E, {int(rows)}, {expected_cols}] for " f"E8M0 K/32 scales, got {tuple(scales.shape)}" ) e8m0_dtype = getattr(torch, "float8_e8m0fnu", None) if scales.dtype == torch.uint8: return scales.view(e8m0_dtype) if e8m0_dtype is not None else scales if e8m0_dtype is not None and scales.dtype == e8m0_dtype: return scales raise TypeError(f"{name} must be torch.uint8 or torch.float8_e8m0fnu") def _pack_e8m0_k32_scales( scales: torch.Tensor, *, size_k: int, size_n: int, row_rotation: int | None = None, reuse_input_storage: bool = False, allow_k_tail: bool = False, padded_size_n: int | None = None, ) -> torch.Tensor: if allow_k_tail: if int(size_k) % 8 != 0: raise ValueError( f"compact E8M0 K-tail requires K divisible by 8, got {size_k}" ) scale_cols = (int(size_k) + 31) // 32 elif int(size_k) % 32 != 0: raise ValueError(f"E8M0 K/32 scales require K divisible by 32, got {size_k}") else: scale_cols = int(size_k) // 32 if tuple(scales.shape[1:]) != (int(size_n), scale_cols): raise ValueError( f"expected E8M0 scale shape [E, {int(size_n)}, {scale_cols}], " f"got {tuple(scales.shape)}" ) output_size_n = int(size_n) if padded_size_n is None else int(padded_size_n) if output_size_n < int(size_n): raise ValueError( f"padded_size_n must be >= size_n, got {output_size_n} < {int(size_n)}" ) source = scales.view(torch.uint8) if reuse_input_storage: if output_size_n != int(size_n): raise ValueError("reuse_input_storage requires unpadded E8M0 scales") if allow_k_tail: raise ValueError( "reuse_input_storage is not supported for compact E8M0 K-tail scales" ) if not source.is_contiguous(): raise ValueError("reuse_input_storage requires contiguous E8M0 scales") source.clamp_(max=_E8M0_K32_BF16_MAX_SCALE_BYTE) packed = source.reshape( int(source.shape[0]), scale_cols, output_size_n, ) else: source = source.clamp(max=_E8M0_K32_BF16_MAX_SCALE_BYTE) packed = torch.empty( (int(source.shape[0]), scale_cols, output_size_n), dtype=torch.uint8, device=scales.device, ) for expert in range(int(source.shape[0])): expert_source = source[expert] if row_rotation is not None: expert_source = torch.cat( [expert_source[row_rotation:], expert_source[:row_rotation]], dim=0, ) expert_packed = _permute_packed_scales( expert_source.T.contiguous(), size_k=size_k, size_n=size_n, group_size=32, output_size_n=output_size_n, ) expert_packed = ( expert_packed.view(-1, 4)[:, [0, 2, 1, 3]] .reshape_as(expert_packed) .contiguous() ) packed[expert].copy_(expert_packed) return packed.view(scales.dtype) if scales.dtype != torch.uint8 else packed def _repack_4bit_no_perm( qweight_i32: torch.Tensor, *, size_k: int, size_n: int, out: torch.Tensor | None = None, flat_scratch: torch.Tensor | None = None, gather_scratch: torch.Tensor | None = None, ) -> torch.Tensor: """Pack 4-bit weights into the W4A16 A16 kernel layout.""" if qweight_i32.dtype != torch.int32: raise TypeError("qweight_i32 must be torch.int32") if tuple(qweight_i32.shape) != (size_k // _PACK_FACTOR_4BIT, size_n): raise ValueError( f"expected qweight shape {(size_k // _PACK_FACTOR_4BIT, size_n)}, " f"got {tuple(qweight_i32.shape)}" ) if size_k % _PACKED_TILE_SIZE != 0 or size_n % _PACKED_TILE_N_SIZE != 0: raise ValueError( f"W4A16 repack requires K,N multiples of 16,64; got {size_k},{size_n}" ) k_tiles = size_k // _PACKED_TILE_SIZE n_tiles = size_n // _PACKED_TILE_N_SIZE packed_shape = (k_tiles, n_tiles, 128) if out is not None and ( out.dtype != torch.int32 or tuple(out.shape) != packed_shape ): raise ValueError( f"out must be int32 with shape {packed_shape}, got " f"{out.dtype} {tuple(out.shape)}" ) if flat_scratch is not None and ( flat_scratch.dtype != torch.int32 or tuple(flat_scratch.shape) != packed_shape ): raise ValueError( f"flat_scratch must be int32 with shape {packed_shape}, got " f"{flat_scratch.dtype} {tuple(flat_scratch.shape)}" ) if gather_scratch is not None and ( gather_scratch.dtype != torch.int32 or tuple(gather_scratch.shape) != packed_shape ): raise ValueError( f"gather_scratch must be int32 with shape {packed_shape}, got " f"{gather_scratch.dtype} {tuple(gather_scratch.shape)}" ) tiles = qweight_i32.view( k_tiles, 2, n_tiles, _PACKED_TILE_N_SIZE, ) if flat_scratch is None: flat = tiles.permute(0, 2, 1, 3).reshape( k_tiles, n_tiles, 2 * _PACKED_TILE_N_SIZE, ) else: flat_scratch.view(k_tiles, n_tiles, 2, _PACKED_TILE_N_SIZE).copy_( tiles.permute(0, 2, 1, 3) ) flat = flat_scratch device = qweight_i32.device out_pos = torch.arange(128, device=device, dtype=torch.long) th_id = out_pos // 4 warp_id = out_pos % 4 tc_col = th_id // 4 tc_row = (th_id % 4) * 2 offsets = torch.tensor([0, 1, 8, 9], device=device, dtype=torch.long) pack_idx = torch.tensor([0, 2, 4, 6, 1, 3, 5, 7], device=device, dtype=torch.long) elem = tc_row[:, None] + offsets[None, :] row = elem // _PACK_FACTOR_4BIT pos = elem % _PACK_FACTOR_4BIT col1 = (warp_id * 16 + tc_col)[:, None].expand(-1, 4) col2 = col1 + 8 source_index = torch.cat( [row * _PACKED_TILE_N_SIZE + col1, row * _PACKED_TILE_N_SIZE + col2], dim=1, )[:, pack_idx] source_shift = torch.cat([pos, pos], dim=1)[:, pack_idx] * 4 result = ( torch.empty(packed_shape, device=device, dtype=torch.int32) if out is None else out ) result.zero_() for slot in range(8): gather_index = source_index[:, slot].view(1, 1, 128).expand( k_tiles, n_tiles, 128, ) shift = source_shift[:, slot].view(1, 1, 128) if gather_scratch is None: gathered = flat.gather(2, gather_index) nibble = (gathered >> shift) & 0xF result |= nibble << (slot * 4) else: torch.gather(flat, 2, gather_index, out=gather_scratch) torch.bitwise_right_shift(gather_scratch, shift, out=gather_scratch) torch.bitwise_and(gather_scratch, 0xF, out=gather_scratch) if slot: torch.bitwise_left_shift( gather_scratch, slot * 4, out=gather_scratch, ) torch.bitwise_or(result, gather_scratch, out=result) return result.reshape(k_tiles, n_tiles * 128).contiguous() def _repack_weight( weight: torch.Tensor, *, size_k: int, size_n: int, row_rotation: int | None = None, reuse_input_storage: bool = False, ) -> torch.Tensor: num_experts = int(weight.shape[0]) if tuple(weight.shape[1:]) != (size_n, size_k // 2): raise ValueError( f"expected packed weight shape {(num_experts, size_n, size_k // 2)}, " f"got {tuple(weight.shape)}" ) if size_k % _PACKED_TILE_SIZE != 0 or size_n % _PACKED_TILE_N_SIZE != 0: raise ValueError( f"W4A16 repack requires K,N multiples of 16,64; got {size_k},{size_n}" ) packed_shape = ( num_experts, size_k // _PACKED_TILE_SIZE, (size_n // _PACKED_TILE_N_SIZE) * 128, ) if reuse_input_storage: if not weight.is_contiguous(): raise ValueError("reuse_input_storage requires contiguous packed weights") packed = weight.view(torch.int32).reshape(packed_shape) else: packed = torch.empty(packed_shape, device=weight.device, dtype=torch.int32) k_tiles = size_k // _PACKED_TILE_SIZE n_tiles = size_n // _PACKED_TILE_N_SIZE qweight_scratch = torch.empty( (size_k // _PACK_FACTOR_4BIT, size_n), device=weight.device, dtype=torch.int32, ) flat_scratch = torch.empty( (k_tiles, n_tiles, 128), device=weight.device, dtype=torch.int32, ) gather_scratch = torch.empty_like(flat_scratch) for expert in range(num_experts): expert_weight = weight[expert].view(torch.int32) if row_rotation is not None: rotated_rows = int(size_n) - int(row_rotation) qweight_scratch[:, :rotated_rows].copy_(expert_weight[row_rotation:].T) qweight_scratch[:, rotated_rows:].copy_(expert_weight[:row_rotation].T) else: qweight_scratch.copy_(expert_weight.T) _repack_4bit_no_perm( qweight_scratch, size_k=size_k, size_n=size_n, out=packed[expert].view(k_tiles, n_tiles, 128), flat_scratch=flat_scratch, gather_scratch=gather_scratch, ) return packed def _permute_nvfp4_scales( scales: torch.Tensor, global_scales: torch.Tensor, *, size_k: int, size_n: int, a_dtype: torch.dtype, row_rotation: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: combined_scale_factor = _nvfp4_compute_scale_factor(scales, a_dtype) packed_scales: torch.Tensor | None = None for expert in range(scales.shape[0]): expert_source = scales[expert].to(a_dtype) if row_rotation is not None: expert_source = torch.cat( [expert_source[row_rotation:], expert_source[:row_rotation]], dim=0, ) expert_scales = _permute_packed_scales( expert_source.T, size_k=size_k, size_n=size_n, group_size=16, ) expert_packed = _process_nvfp4_packed_scales( expert_scales, scale_factor=combined_scale_factor, ) if packed_scales is None: packed_scales = torch.empty( (int(scales.shape[0]), *expert_packed.shape), dtype=expert_packed.dtype, device=expert_packed.device, ) packed_scales[expert].copy_(expert_packed) if packed_scales is None: packed_scales = torch.empty( (0, size_k // _PACKED_TILE_SIZE, size_n // 2), dtype=torch.float8_e4m3fn, device=scales.device, ) packed_global = _process_nvfp4_packed_global_scale( global_scales, a_dtype=a_dtype, ).to(torch.float32) packed_global = packed_global / combined_scale_factor return packed_scales, packed_global.contiguous() # --------------------------------------------------------------------------- # NF3 ("nf3_2p1") packing. # # The kernel's B operand is a stream of 32-code "units" (2 N-columns x 16 K), # each a 12-byte (lo0, lo1, hi) triple. The GEMM reads unit index # unit = cur_group_id * (tile_n // 2) + (tid % (tile_n // 2)) # from the pipe's B SMEM region (see kernel._load_b_scale_registers, # NF3-MAPPING-V1) which -- verified algebraically -- equals the stock staged # unit index; cur_group_id is the K16 row within the CTA-K tile, (tid % # (tile_n // 2)) is the N-pair. The flat-span staging copies the packer's # global bytes verbatim, so the packer must place, at global unit # I = nt*(K/16)*(tile_n/2) + R*(tile_n/2) + p (n_tile-major) # the codes for CTA N-tile nt, K16-row-global R, N-pair p. The 4 jj words in # the triple correspond to the stock 4-u32 group; per jj the 8 codes decode # (packed_dequant_nf3x8_to_bfloat2x4) to fragments # frag[0,0]=(cb[c0],cb[c1]) frag[0,1]=(cb[c2],cb[c3]) # frag[1,0]=(cb[c4],cb[c5]) frag[1,1]=(cb[c6],cb[c7]) # which must equal the stock fragment identity for this thread/jj: # frag[0,0]=(C1@K0, C1@K0+1) frag[0,1]=(C1@K0+8, C1@K0+9) # frag[1,0]=(C2@K0, C2@K0+1) frag[1,1]=(C2@K0+8, C2@K0+9) # where, from _repack_4bit_no_perm composed with the read mapping, # wru = p + (tile_n/2)*nt ; n_tile_64 = wru // 32 ; th_id = wru % 32 # tc_col = th_id // 4 ; tc_row = (th_id % 4) * 2 # C1 = n_tile_64*64 + jj*16 + tc_col ; C2 = C1 + 8 ; K0 = R*16 + tc_row # => code order [c0..c7] = # [ (C1,K0),(C1,K0+1),(C1,K0+8),(C1,K0+9), # (C2,K0),(C2,K0+1),(C2,K0+8),(C2,K0+9) ]. # This whole chain is asserted by _nf3_pack_selftest against an independent # torch simulation of the kernel read+dequant. # --------------------------------------------------------------------------- def _nf3_check_shapes(size_k: int, size_n: int, tile_n: int) -> None: if int(size_k) % 32 != 0: raise ValueError(f"NF3 requires K % 32 == 0, got {size_k}") if int(size_n) % 64 != 0: raise ValueError(f"NF3 requires N % 64 == 0, got {size_n}") if int(tile_n) % 16 != 0 or int(tile_n) < 64: raise ValueError(f"NF3 tile_n must be a multiple of 16 and >= 64, got {tile_n}") if int(size_n) % int(tile_n) != 0: raise ValueError(f"NF3 requires N ({size_n}) divisible by tile_n ({tile_n})") def _nf3_code_gather_index( size_k: int, size_n: int, tile_n: int, device: torch.device ) -> torch.Tensor: """[units, 4, 8] flat indices into a contiguous [N, K] code tensor. Entry [I, jj, v] is the (C * size_k + K) index of the v-th code (NF3 order) of triple word jj at global unit I. Built from the exact _repack / read-mapping math documented above so the packer and the kernel agree. """ _nf3_check_shapes(size_k, size_n, tile_n) npairs = int(tile_n) // 2 k16 = int(size_k) // 16 units = k16 * (int(size_n) // 2) ntile_units = k16 * npairs idx_i = torch.arange(units, device=device, dtype=torch.long) nt = idx_i // ntile_units within = idx_i % ntile_units R = within // npairs p = within % npairs wru = p + npairs * nt n_tile_64 = wru // 32 th_id = wru % 32 tc_col = th_id // 4 tc_row = (th_id % 4) * 2 jj = torch.arange(4, device=device, dtype=torch.long) col1 = jj[None, :] * 16 + tc_col[:, None] # [units, 4] col2 = col1 + 8 C1 = n_tile_64[:, None] * 64 + col1 # [units, 4] C2 = n_tile_64[:, None] * 64 + col2 K0 = (R * 16 + tc_row)[:, None] # [units, 1] k_off = torch.tensor([0, 1, 8, 9], device=device, dtype=torch.long) idx = torch.empty((units, 4, 8), device=device, dtype=torch.long) for t in range(4): idx[:, :, t] = C1 * int(size_k) + (K0 + k_off[t]) idx[:, :, 4 + t] = C2 * int(size_k) + (K0 + k_off[t]) return idx def _nf3_pack_codes( codes_nk: torch.Tensor, *, size_k: int, size_n: int, tile_n: int ) -> torch.Tensor: """Pack one expert's [N, K] codes (0..7) into int32 [units*3] NF3 planes.""" if tuple(codes_nk.shape) != (int(size_n), int(size_k)): raise ValueError( f"expected codes shape {(int(size_n), int(size_k))}, " f"got {tuple(codes_nk.shape)}" ) device = codes_nk.device flat = codes_nk.reshape(-1).to(torch.int64) if bool((flat < 0).any()) or bool((flat > 7).any()): raise ValueError("NF3 codes must be integers in [0, 7]") idx = _nf3_code_gather_index(size_k, size_n, tile_n, device) # [units,4,8] g = flat[idx.reshape(-1)].reshape(idx.shape) # [units,4,8] sh_lo = (torch.arange(4, device=device, dtype=torch.int64) * 2).view(1, 1, 4) sh_hi = torch.arange(4, device=device, dtype=torch.int64).view(1, 1, 4) la = ((g[:, :, 0:4] & 3) << sh_lo).sum(-1) # [units,4] lb = ((g[:, :, 4:8] & 3) << sh_lo).sum(-1) lo16 = la | (lb << 8) ha = (((g[:, :, 0:4] >> 2) & 1) << sh_hi).sum(-1) hb = (((g[:, :, 4:8] >> 2) & 1) << sh_hi).sum(-1) hi8 = ha | (hb << 4) lo0 = lo16[:, 0] | (lo16[:, 1] << 16) lo1 = lo16[:, 2] | (lo16[:, 3] << 16) hi = hi8[:, 0] | (hi8[:, 1] << 8) | (hi8[:, 2] << 16) | (hi8[:, 3] << 24) words = torch.stack([lo0, lo1, hi], dim=1).reshape(-1) # int64 in [0, 2**32) words = torch.where(words >= 2**31, words - 2**32, words) return words.to(torch.int32).contiguous() def _process_nf3_packed_scales(packed_scales: torch.Tensor) -> torch.Tensor: """[rows, N] float t_s -> [rows, N] uint8 e4m3-style byte = highbyte(f16(t_s*2**-4)<<1).""" packed_scales = packed_scales.to(torch.float16) packed_scales = packed_scales.view(-1, 4)[:, [0, 2, 1, 3]].view( packed_scales.size(0), -1, ) packed_scales = (packed_scales.float() * (2.0**-4)).to(torch.float16) packed_scales = packed_scales.view(torch.int16) << 1 packed_scales = packed_scales.view(torch.uint8) return packed_scales[:, 1::2].contiguous() def _nf3_pack_scales( t_s: torch.Tensor, *, size_k: int, size_n: int ) -> torch.Tensor: """Pack one expert's [N, K//32] scales into uint8 [K//32, N] (permuted).""" if tuple(t_s.shape) != (int(size_n), int(size_k) // 32): raise ValueError( f"expected scale shape {(int(size_n), int(size_k) // 32)}, " f"got {tuple(t_s.shape)}" ) t_s = t_s.to(torch.float32) # zero scales (whole group exactly zero in the checkpoint) MUST stay zero: # stored byte 0 decodes to exact +0.0. Flooring them to _NF3_SCALE_FLOOR # injects noise into all-zero groups (real early-layer experts have many) # -> compounding early-layer error -> prompt-dependent degeneration. zero_mask = t_s == 0 t_s = t_s.clamp(min=_NF3_SCALE_FLOOR) t_s[zero_mask] = 0.0 if bool((t_s >= 32.0).any()): raise ValueError("NF3 scales must be < 32 for the e4m3 K/32 encoding") permuted = _permute_packed_scales( t_s.T.contiguous(), size_k=size_k, size_n=size_n, group_size=32, ) # [K//32, N] float return _process_nf3_packed_scales(permuted) def _nf3_pack_code_experts( codes: torch.Tensor, *, size_k: int, size_n: int, tile_n: int ) -> torch.Tensor: packed = [ _nf3_pack_codes(codes[e], size_k=size_k, size_n=size_n, tile_n=tile_n) for e in range(int(codes.shape[0])) ] return torch.stack(packed, dim=0).contiguous() def _nf3_pack_scale_experts( t_s: torch.Tensor, *, size_k: int, size_n: int ) -> torch.Tensor: packed = [ _nf3_pack_scales(t_s[e], size_k=size_k, size_n=size_n) for e in range(int(t_s.shape[0])) ] return torch.stack(packed, dim=0).contiguous() def _prepare_w4a16_packed_weights( w13_fp4: torch.Tensor, w13_blockscale: torch.Tensor, w13_global_scale: torch.Tensor, w2_fp4: torch.Tensor, w2_blockscale: torch.Tensor, w2_global_scale: torch.Tensor, *, activation: str, params_dtype: torch.dtype = torch.bfloat16, source_format: str, w13_layout: str = "w13", reuse_input_storage: bool = False, ) -> W4A16PackedWeights: source_format = _normalize_source_format(source_format) w13_layout = _normalize_w13_layout(w13_layout) shape = validate_w4a16_packed_inputs( w13_fp4, w13_global_scale, w2_fp4, w2_global_scale, activation=activation, ) num_experts = shape.num_experts hidden_size = shape.hidden_size intermediate_size = shape.intermediate_size w13_rows = shape.w13_rows is_gated = shape.is_gated w13 = w13_fp4 w13_scale = unswizzle_expert_scales( w13_blockscale, rows=w13_rows, cols=hidden_size, ) w13_row_rotation = None if is_gated and w13_layout == "w13": # In-place: the half-swap is folded into the repack via row_rotation; # never materialize a second copy of the weights/scales. w13_row_rotation = intermediate_size w2_scale = unswizzle_expert_scales( w2_blockscale, rows=hidden_size, cols=intermediate_size, ) packed_w13 = _repack_weight( w13 if reuse_input_storage else w13.contiguous(), size_k=hidden_size, size_n=w13_rows, row_rotation=w13_row_rotation, reuse_input_storage=reuse_input_storage, ) packed_w2 = _repack_weight( w2_fp4 if reuse_input_storage else w2_fp4.contiguous(), size_k=intermediate_size, size_n=hidden_size, reuse_input_storage=reuse_input_storage, ) native_w13_global_scale = _source_global_scale( w13_global_scale, source_format=source_format, ) native_w2_global_scale = _source_global_scale( w2_global_scale, source_format=source_format, ) packed_w13_scale, packed_w13_global_scale = _permute_nvfp4_scales( w13_scale, native_w13_global_scale, size_k=hidden_size, size_n=w13_rows, a_dtype=params_dtype, row_rotation=w13_row_rotation, ) packed_w2_scale, packed_w2_global_scale = _permute_nvfp4_scales( w2_scale, native_w2_global_scale, size_k=intermediate_size, size_n=hidden_size, a_dtype=params_dtype, ) return W4A16PackedWeights( w13=packed_w13, w13_scale=packed_w13_scale, w13_global_scale=packed_w13_global_scale, w2=packed_w2, w2_scale=packed_w2_scale, w2_global_scale=packed_w2_global_scale, workspace=_make_workspace(w13_fp4.device, max_blocks_per_sm=4), hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, is_gated=is_gated, params_dtype=params_dtype, source_format=source_format, w13_layout=w13_layout, scale_format="e4m3_k16", ) def prepare_w4a16_modelopt_nvfp4_weights( w13_fp4: torch.Tensor, w13_blockscale: torch.Tensor, w13_global_scale: torch.Tensor, w2_fp4: torch.Tensor, w2_blockscale: torch.Tensor, w2_global_scale: torch.Tensor, *, activation: str, params_dtype: torch.dtype = torch.bfloat16, w13_layout: str = "w13", reuse_input_storage: bool = False, ) -> W4A16PackedWeights: """Prepare ModelOpt NVFP4 tensors into the W4A16 packed runtime layout. The per-block scales are the normal NVFP4 K/16 scale grid in b12x swizzled storage. The global scales are raw ModelOpt weight global scales; activation input scales are not folded into W4A16 weight preparation. For gated activations, ``w13_layout`` describes whether fused W13 rows arrive in checkpoint/logical W13 order or already swapped W31 order. """ return _prepare_w4a16_packed_weights( w13_fp4, w13_blockscale, w13_global_scale, w2_fp4, w2_blockscale, w2_global_scale, activation=activation, params_dtype=params_dtype, source_format="modelopt_nvfp4", w13_layout=w13_layout, reuse_input_storage=reuse_input_storage, ) def prepare_w4a16_modelopt_native_weights( w13_fp4: torch.Tensor, w13_blockscale: torch.Tensor, w13_global_scale: torch.Tensor, w2_fp4: torch.Tensor, w2_blockscale: torch.Tensor, w2_global_scale: torch.Tensor, *, activation: str, params_dtype: torch.dtype = torch.bfloat16, source_format: str = "modelopt_nvfp4", w13_layout: str = "w13", ) -> W4A16ModelOptWeights: """Prepare W4A16 metadata while keeping ModelOpt FP4 weights native. This is the memory-safe path for GLM serving that needs A4 prefill and A16 decode in the same process. It keeps the checkpoint FP4 tensors resident instead of materializing a second full W4A16 packed copy. """ source_format = _normalize_source_format(source_format) if source_format not in _MODEL_OPT_NVFP4_FORMATS: raise ValueError( "native W4A16 ModelOpt weights require source_format " "'modelopt_nvfp4'" ) w13_layout = _normalize_w13_layout(w13_layout) shape = validate_w4a16_packed_inputs( w13_fp4, w13_global_scale, w2_fp4, w2_global_scale, activation=activation, ) num_experts = shape.num_experts hidden_size = shape.hidden_size intermediate_size = shape.intermediate_size w13_rows = shape.w13_rows is_gated = shape.is_gated w13_scale = unswizzle_expert_scales( w13_blockscale, rows=w13_rows, cols=hidden_size, ) w2_scale = unswizzle_expert_scales( w2_blockscale, rows=hidden_size, cols=intermediate_size, ) native_w13_global_scale = _source_global_scale( w13_global_scale, source_format=source_format, ) native_w2_global_scale = _source_global_scale( w2_global_scale, source_format=source_format, ) # The W4A16 activation consumes FC1 output in gate/up logical order. # Checkpoint-native ModelOpt GLM tensors are up/gate, while vLLM/FI can # hand over gate/up tensors after its own W13 reorder. Keep that physical # order explicit so source_format never implies a layout transformation. w13_row_rotation = ( intermediate_size if is_gated and w13_layout == "w13" else None ) packed_w13_scale, packed_w13_global_scale = _permute_nvfp4_scales( w13_scale, native_w13_global_scale, size_k=hidden_size, size_n=w13_rows, a_dtype=params_dtype, row_rotation=w13_row_rotation, ) packed_w2_scale, packed_w2_global_scale = _permute_nvfp4_scales( w2_scale, native_w2_global_scale, size_k=intermediate_size, size_n=hidden_size, a_dtype=params_dtype, ) return W4A16ModelOptWeights( w13=w13_fp4, w13_scale=packed_w13_scale, w13_global_scale=packed_w13_global_scale, w2=w2_fp4, w2_scale=packed_w2_scale, w2_global_scale=packed_w2_global_scale, workspace=_make_workspace(w13_fp4.device, max_blocks_per_sm=4), hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, is_gated=is_gated, params_dtype=params_dtype, source_format=source_format, micro_w13_scale=packed_w13_scale, micro_w13_global_scale=_process_nvfp4_micro_global_scale_from_packed( packed_w13_global_scale, a_dtype=params_dtype, ), micro_w2_scale=packed_w2_scale, micro_w2_global_scale=_process_nvfp4_micro_global_scale_from_packed( packed_w2_global_scale, a_dtype=params_dtype, ), w13_layout=w13_layout, ) def prepare_w4a16_e8m0_native_weights( w13_fp4: torch.Tensor, w13_e8m0_scale: torch.Tensor, w13_global_scale: torch.Tensor, w2_fp4: torch.Tensor, w2_e8m0_scale: torch.Tensor, w2_global_scale: torch.Tensor, *, activation: str, params_dtype: torch.dtype = torch.bfloat16, w13_layout: str = "w13", ) -> W4A16ModelOptWeights: """Prepare native MXFP4 (E8M0 K/32) weights for the W4A16 path. Keeps the FP4 weights resident as a single copy (``weight_layout="modelopt"``) and carries two small scale forms so one object serves both kernels: ``w13_scale``/``w2_scale`` are the packed E8M0 grid the main W4A16 GEMM reads at med/large M, and ``micro_*`` are packed E8M0 grids in the native row order that the small-M micro decode kernel reads. ``run_w4a16_moe`` routes small M to micro and the rest to the main W4A16 kernel automatically. """ w13_layout = _normalize_w13_layout(w13_layout) shape = validate_w4a16_packed_inputs( w13_fp4, w13_global_scale, w2_fp4, w2_global_scale, activation=activation, ) num_experts = shape.num_experts hidden_size = shape.hidden_size intermediate_size = shape.intermediate_size w13_rows = shape.w13_rows is_gated = shape.is_gated allow_w2_k_tail = intermediate_size % 32 != 0 padded_w13_scale_n = ( _e8m0_logical_tail_scale_n(w13_rows) if allow_w2_k_tail else w13_rows ) w13_scale = _validate_e8m0_k32_scales( w13_e8m0_scale, rows=w13_rows, cols=hidden_size, name="w13_e8m0_scale", ) w2_scale = _validate_e8m0_k32_scales( w2_e8m0_scale, rows=hidden_size, cols=intermediate_size, name="w2_e8m0_scale", allow_k_tail=allow_w2_k_tail, ) # Main-GEMM (med/large M) packed E8M0 scales. The "w13" (up_gate) layout # needs the FC1 half-swap folded into the scale grid; the kernel applies the # matching source_n_rotation to the native weights. micro reads the un-rotated # grid and handles the layout itself (w13_gate_first). w13_row_rotation = intermediate_size if (is_gated and w13_layout == "w13") else None packed_w13_scale = _pack_e8m0_k32_scales( w13_scale, size_k=hidden_size, size_n=w13_rows, row_rotation=w13_row_rotation, padded_size_n=padded_w13_scale_n, ) micro_w13_scale = ( packed_w13_scale if w13_row_rotation is None else _pack_e8m0_k32_scales( w13_scale, size_k=hidden_size, size_n=w13_rows, padded_size_n=padded_w13_scale_n, ) ) packed_w2_scale = _pack_e8m0_k32_scales( w2_scale, size_k=intermediate_size, size_n=hidden_size, allow_k_tail=allow_w2_k_tail, ) # Storage-compatible single grid: micro reads the SAME packed _pack_e8m0_k32 # scales the main GEMM reads (no separate K/16 micro grid). w13_global = w13_global_scale.contiguous() w2_global = w2_global_scale.contiguous() return W4A16ModelOptWeights( w13=w13_fp4, w13_scale=packed_w13_scale, w13_global_scale=w13_global, w2=w2_fp4, w2_scale=packed_w2_scale, w2_global_scale=w2_global, workspace=_make_workspace(w13_fp4.device, max_blocks_per_sm=4), hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, is_gated=is_gated, params_dtype=params_dtype, source_format="fp4_e8m0_k32", scale_format="e8m0_k32", micro_w13_scale=micro_w13_scale, micro_w13_global_scale=w13_global, micro_w2_scale=packed_w2_scale, micro_w2_global_scale=w2_global, w13_layout=w13_layout, ) def prepare_w4a16_compressed_tensors_weights( w13_fp4: torch.Tensor, w13_blockscale: torch.Tensor, w13_global_scale: torch.Tensor, w2_fp4: torch.Tensor, w2_blockscale: torch.Tensor, w2_global_scale: torch.Tensor, *, activation: str, params_dtype: torch.dtype = torch.bfloat16, w13_layout: str = "w13", reuse_input_storage: bool = False, ) -> W4A16PackedWeights: """Prepare CompressedTensors NVFP4 tensors into the W4A16 packed runtime layout. The per-block scales are the normal NVFP4 K/16 scale grid in b12x swizzled storage. The CT global scales are stored inverted relative to the ModelOpt weight global scale convention, so they are inverted before packing. """ return _prepare_w4a16_packed_weights( w13_fp4, w13_blockscale, w13_global_scale, w2_fp4, w2_blockscale, w2_global_scale, activation=activation, params_dtype=params_dtype, source_format="compressed_tensors", w13_layout=w13_layout, reuse_input_storage=reuse_input_storage, ) def prepare_w4a16_fp4_e8m0_k32_weights( w13_fp4: torch.Tensor, w13_e8m0_scale: torch.Tensor, w13_global_scale: torch.Tensor, w2_fp4: torch.Tensor, w2_e8m0_scale: torch.Tensor, w2_global_scale: torch.Tensor, *, activation: str, params_dtype: torch.dtype = torch.bfloat16, w13_layout: str = "w13", reuse_input_storage: bool = False, ) -> W4A16PackedWeights: """Prepare FP4 weights with E8M0 K/32 scales for W4A16. The per-block source scales are [E, N, K/32] E8M0 bytes. They are only saturated to the BF16 kernel's supported byte range and rearranged for kernel access; they are not expanded to K/16 or folded into global scales. """ w13_layout = _normalize_w13_layout(w13_layout) shape = validate_w4a16_packed_inputs( w13_fp4, w13_global_scale, w2_fp4, w2_global_scale, activation=activation, ) num_experts = shape.num_experts hidden_size = shape.hidden_size intermediate_size = shape.intermediate_size w13_rows = shape.w13_rows is_gated = shape.is_gated w13 = w13_fp4 w13_scale = _validate_e8m0_k32_scales( w13_e8m0_scale, rows=w13_rows, cols=hidden_size, name="w13_e8m0_scale", ) w13_row_rotation = None if is_gated and w13_layout != "w31": w13_row_rotation = intermediate_size w2_scale = _validate_e8m0_k32_scales( w2_e8m0_scale, rows=hidden_size, cols=intermediate_size, name="w2_e8m0_scale", ) packed_w13 = _repack_weight( w13 if reuse_input_storage else w13.contiguous(), size_k=hidden_size, size_n=w13_rows, row_rotation=w13_row_rotation, reuse_input_storage=reuse_input_storage, ) packed_w2 = _repack_weight( w2_fp4 if reuse_input_storage else w2_fp4.contiguous(), size_k=intermediate_size, size_n=hidden_size, reuse_input_storage=reuse_input_storage, ) packed_w13_scale = _pack_e8m0_k32_scales( w13_scale, size_k=hidden_size, size_n=w13_rows, row_rotation=w13_row_rotation, reuse_input_storage=reuse_input_storage, ) packed_w2_scale = _pack_e8m0_k32_scales( w2_scale, size_k=intermediate_size, size_n=hidden_size, reuse_input_storage=reuse_input_storage, ) return W4A16PackedWeights( w13=packed_w13, w13_scale=packed_w13_scale, w13_global_scale=w13_global_scale.contiguous(), w2=packed_w2, w2_scale=packed_w2_scale, w2_global_scale=w2_global_scale.contiguous(), workspace=_make_workspace(w13_fp4.device, max_blocks_per_sm=4), hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, is_gated=is_gated, params_dtype=params_dtype, source_format="fp4_e8m0_k32", w13_layout=w13_layout, scale_format="e8m0_k32", ) def prepare_w4a16_packed_weights( *args, source_format: str = "modelopt_nvfp4", w13_layout: str = "w13", **kwargs, ) -> W4A16PackedWeights: source_format = _normalize_source_format(source_format) w13_layout = _normalize_w13_layout(w13_layout) if source_format == "modelopt_nvfp4": return prepare_w4a16_modelopt_nvfp4_weights( *args, w13_layout=w13_layout, **kwargs ) if source_format == "compressed_tensors": return prepare_w4a16_compressed_tensors_weights( *args, w13_layout=w13_layout, **kwargs ) if source_format == "fp4_e8m0_k32": return prepare_w4a16_fp4_e8m0_k32_weights( *args, w13_layout=w13_layout, **kwargs ) raise AssertionError(f"unhandled W4A16 source_format {source_format!r}") def make_w4a16_packed_buffers( prepared: W4A16PackedWeights | W4A16ModelOptWeights, *, m: int, topk: int, dtype: torch.dtype, device: torch.device, route_num_experts: int | None = None, ) -> W4A16PackedBuffers: return _make_w4a16_packed_buffers( prepared, m=m, topk=topk, dtype=dtype, device=device, route_num_experts=route_num_experts, ) def prepare_nf3_moe_weights( w13_codes: torch.Tensor, w13_scale: torch.Tensor, w2_codes: torch.Tensor, w2_scale: torch.Tensor, *, activation: str, fc1_tile_n: int, fc2_tile_n: int, params_dtype: torch.dtype = torch.bfloat16, ) -> PreparedNF3MoeWeights: """Pack NF3 codes/scales into the runtime "nf3_2p1" W4A16 layout. Inputs are per-expert integer codes (0..7) in kernel-native output order (gate/up already resolved for FC1) and per-group float scales ``t_s``: w13_codes [E, 2*I, hidden] w13_scale [E, 2*I, hidden//32] w2_codes [E, hidden, I] w2_scale [E, hidden, I//32] ``fc1_tile_n`` / ``fc2_tile_n`` MUST equal the CTA N-tiles the kernel will use (read them from the compiled W4A16FusedMoeCompileResult) -- the flat-span layout is tile_n specific. """ if params_dtype != torch.bfloat16: raise ValueError("nf3_2p1 W4A16 weights are bf16-only for v1") shape = validate_nf3_moe_inputs( w13_codes, w13_scale, w2_codes, w2_scale, activation=activation, ) device = w13_codes.device packed_w13 = _nf3_pack_code_experts( w13_codes, size_k=shape.hidden_size, size_n=shape.w13_rows, tile_n=fc1_tile_n ) packed_w2 = _nf3_pack_code_experts( w2_codes, size_k=shape.intermediate_size, size_n=shape.hidden_size, tile_n=fc2_tile_n, ) packed_w13_scale = _nf3_pack_scale_experts( w13_scale, size_k=shape.hidden_size, size_n=shape.w13_rows ) packed_w2_scale = _nf3_pack_scale_experts( w2_scale, size_k=shape.intermediate_size, size_n=shape.hidden_size ) global_scale = torch.full( (shape.num_experts,), _NF3_SCALE_GLOBAL, dtype=torch.float32, device=device ) return PreparedNF3MoeWeights( w13=packed_w13, w13_scale=packed_w13_scale, w13_global_scale=global_scale, w2=packed_w2, w2_scale=packed_w2_scale, w2_global_scale=global_scale.clone(), workspace=_make_workspace(device, max_blocks_per_sm=4), hidden_size=shape.hidden_size, intermediate_size=shape.intermediate_size, num_experts=shape.num_experts, is_gated=shape.is_gated, params_dtype=params_dtype, fc1_tile_n=int(fc1_tile_n), fc2_tile_n=int(fc2_tile_n), ) def _nf3_pack_selftest() -> None: """Prove the packer matches the kernel's read+dequant chain (torch/CPU). Packs random codes, then INDEPENDENTLY simulates the kernel's flat-span staging + register read (unit = cur_group_id*(tile_n//2) + tid%(tile_n//2)) + packed_dequant_nf3x8 fragment placement in pure torch, and asserts the reconstructed dequantized matrix equals codebook[codes] at every (N, K). """ codebook = _NF3_CODEBOOK def simulate(packed, size_k, size_n, tile_n, tile_k): cta_threads = tile_n * tile_k // 64 b_sh_stride = tile_n // 2 cta_k_blocks = tile_k // 16 b_sh_stage = b_sh_stride * cta_k_blocks b_sh_wr_iters = b_sh_stage // cta_threads tb_n_warps = tile_n // 64 b_chunks = b_sh_stage * 12 // 16 wr_iters_nf3 = (b_chunks + cta_threads - 1) // cta_threads ntile_stride = (size_k // 16) * (tile_n // 2) n_tiles = size_n // tile_n k_tiles = size_k // tile_k packed = packed.tolist() w = [[None] * size_n for _ in range(size_k)] def u32(x): return x & 0xFFFFFFFF for nt in range(n_tiles): for tile_idx in range(k_tiles): span_base_unit = ( nt * ntile_stride + tile_idx * cta_k_blocks * b_sh_stride ) smem = {} for i in range(wr_iters_nf3): for tid in range(cta_threads): c = i * cta_threads + tid if c >= b_chunks: continue for word in range(4): smem[c * 4 + word] = u32( packed[span_base_unit * 3 + c * 4 + word] ) for tid in range(cta_threads): warp_row = (tid // 32) // tb_n_warps for kk in range(b_sh_wr_iters): cur = b_sh_wr_iters * warp_row + kk unit = cur * b_sh_stride + (tid % b_sh_stride) lo0 = smem[unit * 3 + 0] lo1 = smem[unit * 3 + 1] hi = smem[unit * 3 + 2] R = tile_idx * cta_k_blocks + cur p = tid % b_sh_stride for jj in range(4): lo_w = lo0 if (jj // 2) == 0 else lo1 lo16 = (lo_w >> (16 * (jj % 2))) & 0xFFFF hi8 = (hi >> (8 * jj)) & 0xFF la = lo16 & 0xFF lb = (lo16 >> 8) & 0xFF ha = hi8 & 0xF hb = (hi8 >> 4) & 0xF def code(byte2, nib1, j): return ((byte2 >> (2 * j)) & 3) | ( ((nib1 >> j) & 1) << 2 ) cds = [code(la, ha, j) for j in range(4)] + [ code(lb, hb, j) for j in range(4) ] wru = p + (tile_n // 2) * nt n64 = wru // 32 th = wru % 32 tc_col = th // 4 tc_row = (th % 4) * 2 col1 = jj * 16 + tc_col c1 = n64 * 64 + col1 c2 = c1 + 8 k0 = R * 16 + tc_row w[k0 + 0][c1] = codebook[cds[0]] w[k0 + 1][c1] = codebook[cds[1]] w[k0 + 8][c1] = codebook[cds[2]] w[k0 + 9][c1] = codebook[cds[3]] w[k0 + 0][c2] = codebook[cds[4]] w[k0 + 1][c2] = codebook[cds[5]] w[k0 + 8][c2] = codebook[cds[6]] w[k0 + 9][c2] = codebook[cds[7]] return w torch.manual_seed(0) cases = [ (256, 128, 128, 128), (256, 128, 128, 64), (64, 256, 256, 64), (512, 256, 256, 64), ] for size_k, size_n, tile_n, tile_k in cases: codes = torch.randint(0, 8, (size_n, size_k), dtype=torch.int64) packed = _nf3_pack_codes( codes, size_k=size_k, size_n=size_n, tile_n=tile_n ) w = simulate(packed, size_k, size_n, tile_n, tile_k) for k in range(size_k): for n in range(size_n): expect = _NF3_CODEBOOK[int(codes[n, k])] got = w[k][n] assert got is not None, f"unfilled ({k},{n}) for {(size_k, size_n, tile_n, tile_k)}" assert abs(got - expect) < 1e-9, ( f"mismatch at ({k},{n}) for {(size_k, size_n, tile_n, tile_k)}: " f"{got} != {expect}" ) print("nf3 pack self-test PASSED") __all__ = [ "PreparedNF3MoeWeights", "W4A16PackedBuffers", "W4A16ModelOptWeights", "W4A16PackedWeights", "make_w4a16_packed_buffers", "prepare_nf3_moe_weights", "prepare_w4a16_compressed_tensors_weights", "prepare_w4a16_e8m0_native_weights", "prepare_w4a16_fp4_e8m0_k32_weights", "prepare_w4a16_modelopt_native_weights", "prepare_w4a16_modelopt_nvfp4_weights", "prepare_w4a16_packed_weights", ]