From 0a0aeec30ff0bea89358194436b6044b53f60d48 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 26 Sep 2026 09:16:12 +0000 Subject: [PATCH 1/4] Prepare EXL3 A16 experts and bounded W4A16 prefill Preserve source-coordinate expert extents, retained two-bit lookup tables and declared row capacities. Keep block-codec epilogues and pipeline tuning intact. Optional FP32 atomic prefill reduction does not guarantee bitwise summation order. Validation: Host eligibility checks pass for packed, ModelOpt, EXL3 and three block codecs. Syntax and whitespace checks pass. GPU validation of the canonical-base composition remains pending. Co-authored-by: Martin Vit Co-authored-by: Codex --- .lil/changes/b12x-415.json | 18 + b12x/_lib/intrinsics.py | 75 +++- b12x/moe/_shared/kernels/w4a16/host.py | 97 ++++- b12x/moe/_shared/kernels/w4a16/kernel.py | 363 ++++++++++++++---- b12x/moe/_shared/kernels/w4a16/route_pack.py | 125 +++++- b12x/moe/fused_moe/__init__.py | 4 + b12x/moe/fused_moe/_impl.py | 266 ++++++++----- b12x/moe/fused_moe/_preparation.py | 17 + b12x/moe/fused_moe/_tuning.py | 4 +- b12x/moe/fused_moe/api.py | 9 +- b12x/moe/fused_moe/planning.py | 68 +++- b12x/moe/fused_moe/source.py | 15 +- b12x/moe/fused_moe/weights.py | 27 ++ docs/trellis-decode-table.md | 39 ++ docs/w4a16-prefill-reduction.md | 45 +++ tests/moe/test_exl3_prepare.py | 136 +++++++ tests/moe/test_exl3_schema.py | 48 +++ tests/moe/test_fused_moe_variant_selection.py | 61 +++ tests/moe/test_nvfp4_auto.py | 2 +- tests/moe/test_tp_moe_scratch_bindings.py | 95 +++++ tests/moe/test_trellis_direct_lut_decode.py | 48 ++- tests/moe/test_w4a16_e2e.py | 170 ++++++++ tests/moe/test_w4a16_mixed_trellis.py | 3 +- tests/moe/test_w4a16_route_pack.py | 13 +- 24 files changed, 1531 insertions(+), 217 deletions(-) create mode 100644 .lil/changes/b12x-415.json create mode 100644 docs/trellis-decode-table.md create mode 100644 docs/w4a16-prefill-reduction.md diff --git a/.lil/changes/b12x-415.json b/.lil/changes/b12x-415.json new file mode 100644 index 000000000..36ac2b841 --- /dev/null +++ b/.lil/changes/b12x-415.json @@ -0,0 +1,18 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "b12x-415", + "category": "feature", + "summary": "Prepare EXL3 expert extents and bounded W4A16 prefill workspaces", + "models": [ + "Kimi-K3" + ], + "compatibility": "No user action required.", + "details": [ + "Uses an exact codeword lookup table for eligible two-bit EXL3 kernels.", + "Preserves IQ2_XS, IQ2_XXS and Q8_0 epilogues and pipeline tuning." + ], + "pull_requests": [ + 415 + ], + "requires": [] +} diff --git a/b12x/_lib/intrinsics.py b/b12x/_lib/intrinsics.py index 4a36457fa..f549fbb6f 100644 --- a/b12x/_lib/intrinsics.py +++ b/b12x/_lib/intrinsics.py @@ -1634,6 +1634,37 @@ def red_add_global_f32(addr: Int64, val: Float32, *, loc=None, ip=None): ) +@dsl_user_op +def red_add_global_v4_f32( + addr: Int64, + val0: Float32, + val1: Float32, + val2: Float32, + val3: Float32, + *, + loc=None, + ip=None, +): + """Reduce-add four FP32 elements at a 16-byte-aligned address.""" + llvm.inline_asm( + None, + [ + Int64(addr).ir_value(loc=loc, ip=ip), + Float32(val0).ir_value(loc=loc, ip=ip), + Float32(val1).ir_value(loc=loc, ip=ip), + Float32(val2).ir_value(loc=loc, ip=ip), + Float32(val3).ir_value(loc=loc, ip=ip), + ], + "red.relaxed.gpu.global.v4.f32.add [$0], {$1, $2, $3, $4};", + "l,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + @dsl_user_op def cvt_bf16x2_to_f16x2(packed: Uint32, *, loc=None, ip=None) -> Uint32: """Convert a u32 holding two bf16 (lo, hi) into an f16x2 u32 (lo, hi).""" @@ -6803,6 +6834,7 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( bits: int = 3, rate_indexed: bool = True, *, + in_shared: bool = False, loc=None, ip=None, ): @@ -6810,8 +6842,8 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( This is the lookup lower-bound primitive: one byte load per reconstructed weight. Every slot owns its own address/code registers so all eight - global gathers are in flight together; a shared scratch pair would chain - them behind full global-memory latency each. + gathers are in flight together. With ``in_shared``, the address names a + caller-staged shared-memory table instead of a global-memory table. """ bits = int(bits) @@ -6822,41 +6854,47 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( table_offset = ((bits - 2) << 16) if rate_indexed else 0 extract_lines: list[str] = [] load_lines: list[str] = [] - pack_lines: list[str] = [] for index in range(8): source = "$3" if index < 4 else "$2" shift = (3 - (index & 3)) * bits # Accumulate into declared scratch registers, never directly into # the "=r" outputs: without early-clobber the outputs may alias the # window inputs that later iterations still read. - target = "out0" if index < 4 else "out1" - byte_shift = 8 * (index & 3) if shift: extract_lines.append( f"bfe.u32 w{index}, {source}, {shift}, 16;" ) else: extract_lines.append(f"and.b32 w{index}, {source}, 0xffff;") - load_lines.append( - f""" + if in_shared: + load_lines.append( + f"add.u32 w{index}, w{index}, $4;" + f" ld.shared.u8 w{index}, [w{index}];" + ) + else: + load_lines.append( + f""" cvt.u64.u32 addr{index}, w{index}; add.u64 addr{index}, addr{index}, $4; ld.global.u8 w{index}, [addr{index}]; - """ - ) - if byte_shift: - pack_lines.append( - f"shl.b32 w{index}, w{index}, {byte_shift};" + """ ) - pack_lines.append(f"or.b32 {target}, {target}, w{index};") + # Each load returns one zero-extended byte. Byte permutations pack four + # independent results without a serial shift/or accumulator. + pack_lines = [ + "prmt.b32 w0, w0, w1, 0x4040;", + "prmt.b32 w2, w2, w3, 0x4040;", + "prmt.b32 out0, w0, w2, 0x5410;", + "prmt.b32 w4, w4, w5, 0x4040;", + "prmt.b32 w6, w6, w7, 0x4040;", + "prmt.b32 out1, w4, w6, 0x5410;", + ] asm = ( """ { .reg .b32 out0,out1; .reg .b32 w0,w1,w2,w3,w4,w5,w6,w7; .reg .b64 addr0,addr1,addr2,addr3,addr4,addr5,addr6,addr7; - mov.b32 out0, 0; - mov.b32 out1, 0; """ + "\n".join(extract_lines) + "\n".join(load_lines) @@ -6866,7 +6904,10 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( mov.b32 $1, out1; }""" ) - table_base = Int64(direct_lut_addr) + Int64(table_offset) + table_base = ( + Int32(direct_lut_addr) + Int32(table_offset) + if in_shared else Int64(direct_lut_addr) + Int64(table_offset) + ) result = llvm.inline_asm( llvm.StructType.get_literal([T.i32(), T.i32()]), [ @@ -6875,7 +6916,7 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( table_base.ir_value(loc=loc, ip=ip), ], asm, - "=r,=r,r,r,l", + "=r,=r,r,r,r" if in_shared else "=r,=r,r,r,l", has_side_effects=False, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, diff --git a/b12x/moe/_shared/kernels/w4a16/host.py b/b12x/moe/_shared/kernels/w4a16/host.py index b148506fe..7a12f9e2a 100644 --- a/b12x/moe/_shared/kernels/w4a16/host.py +++ b/b12x/moe/_shared/kernels/w4a16/host.py @@ -2,12 +2,13 @@ from __future__ import annotations -from b12x._lib.quant.block_codec import BLOCK_CODECS - +import os from dataclasses import dataclass import torch +from b12x._lib.env import env_flag +from b12x._lib.quant.block_codec import BLOCK_CODECS from b12x.moe._shared.kernels.activations import ( SUPPORTED_MOE_ACTIVATIONS, is_gated_moe_activation, @@ -20,6 +21,58 @@ _SUPPORTED_ACTIVATIONS = SUPPORTED_MOE_ACTIVATIONS +def trellis_decode_table() -> str: + """Read the plan-time trellis decode-table selection; replay never reads it.""" + value = os.environ.get("B12X_TRELLIS_DECODE_TABLE", "auto") + if value not in {"auto", "compact", "full"}: + raise ValueError("B12X_TRELLIS_DECODE_TABLE must be auto, compact, or full") + return value + + +def prefill_fused_sum_enabled() -> bool: + """Enable direct FP32 route reduction for large-M W4A16 launches. + + The FC2 epilogue uses relaxed FP32 global reductions. Disable this path + when bitwise-identical route accumulation order is required. + """ + return env_flag("W4A16_PREFILL_FUSED_SUM") + + +def prefill_fused_sum_eligible( + *, + dtype: torch.dtype | str, + m: int, + full_rotation: bool, + weight_layout: str, + collect_activation_amax: bool, + enabled: bool | None = None, +) -> bool: + """Return whether one W4A16 launch may use direct route reduction. + + Args: + dtype: Activation element dtype as a PyTorch dtype or kernel dtype name. + m: Planned token capacity, or exact rows for an unprepared launch. + full_rotation: Whether the launch uses full-rotation Trellis weights. + weight_layout: Prepared W4A16 weight layout. + collect_activation_amax: Whether the launch records activation maxima. + enabled: Frozen feature selection. ``None`` reads the process setting. + + Returns: + ``True`` when planning and execution may use the FP32 accumulator. + """ + if enabled is None: + enabled = prefill_fused_sum_enabled() + element_dtype = str(dtype).removeprefix("torch.") + return bool( + enabled + and element_dtype in {"bfloat16", "bf16"} + and int(m) > 8 + and not full_rotation + and weight_layout in {"packed", "modelopt"} + and not collect_activation_amax + ) + + @dataclass(frozen=True) class W4A16PackedShape: num_experts: int @@ -34,6 +87,7 @@ class W4A16PackedBuffers: intermediate_cache13: torch.Tensor intermediate_cache2: torch.Tensor output: torch.Tensor + prefill_sum_accum: torch.Tensor | None = None fc1_c_tmp: torch.Tensor | None = None fc2_c_tmp: torch.Tensor | None = None packed_route_indices: torch.Tensor | None = None @@ -57,6 +111,7 @@ class W4A16BufferPlan: intermediate_cache2_elements: int block_size_m: int rotation_a_elements: int = 0 + prefill_sum_accum_elements: int = 0 def validate_activation(activation: str) -> bool: @@ -280,8 +335,12 @@ def plan_w4a16_buffers( topk: int, route_num_experts: int | None = None, sms: int, + dtype: torch.dtype | None = None, full_rotation: bool = False, block_size_m: int | None = None, + weight_layout: str | None = None, + collect_activation_amax: bool = False, + prefill_fused_sum: bool | None = None, ) -> W4A16BufferPlan: routed_rows = int(m) * int(topk) route_num_experts = ( @@ -312,6 +371,19 @@ def plan_w4a16_buffers( if int(m) <= 8 and (bool(prepared.is_gated) or prepared.weight_layout in BLOCK_CODECS): gemm_route_slots = max(gemm_route_slots, routed_rows * block_size_m) scratch_sms = int(sms) + weight_layout = str( + weight_layout + if weight_layout is not None + else getattr(prepared, "weight_layout", "packed") + ) + use_prefill_fused_sum = prefill_fused_sum_eligible( + dtype=dtype if dtype is not None else "", + m=m, + full_rotation=full_rotation, + weight_layout=weight_layout, + collect_activation_amax=collect_activation_amax, + enabled=prefill_fused_sum, + ) return W4A16BufferPlan( routed_rows=routed_rows, fc1_cols=fc1_cols, @@ -331,10 +403,19 @@ def plan_w4a16_buffers( sms=scratch_sms, weight_layout=prepared.weight_layout, ), - intermediate_cache13_elements=routed_rows * max(fc1_cols, hidden_size), + intermediate_cache13_elements=( + routed_rows * fc1_cols + if use_prefill_fused_sum + else routed_rows * max(fc1_cols, hidden_size) + ), intermediate_cache2_elements=routed_rows * intermediate_size, block_size_m=block_size_m, rotation_a_elements=(routed_rows * hidden_size if full_rotation else 0), + prefill_sum_accum_elements=( + int(m) * hidden_size + if use_prefill_fused_sum + else 0 + ), ) @@ -361,6 +442,7 @@ def make_w4a16_packed_buffers( topk=topk, route_num_experts=route_num_experts, sms=sms, + dtype=dtype, full_rotation=full_rotation, block_size_m=block_size_m, ) @@ -410,6 +492,15 @@ def make_w4a16_packed_buffers( dtype=torch.float32 if full_rotation else dtype, device=device, ), + prefill_sum_accum=( + torch.empty( + (plan.prefill_sum_accum_elements,), + dtype=torch.float32, + device=device, + ) + if plan.prefill_sum_accum_elements + else None + ), fc1_c_tmp=fc1_c_tmp, fc2_c_tmp=fc2_c_tmp, packed_route_indices=torch.empty( diff --git a/b12x/moe/_shared/kernels/w4a16/kernel.py b/b12x/moe/_shared/kernels/w4a16/kernel.py index aa471b123..cc48cf3dc 100644 --- a/b12x/moe/_shared/kernels/w4a16/kernel.py +++ b/b12x/moe/_shared/kernels/w4a16/kernel.py @@ -72,10 +72,12 @@ packed_decode_lut_fp16_to_bfloat2x4, packed_decode_lut_fp16_to_half2x4, packed_decode_lut_e4m3_to_e4m3x8, + packed_decode_lut_e4m3_direct_to_e4m3x8, ld_global_nc_v4_u32, pack_f32x2_to_bfloat2, pack_f32x2_to_f16x2, red_add_global_bf16x2, + red_add_global_v4_f32, red_add_global_release_i32, red_max_global_f32_nonnegative, shared_ptr_to_u32, @@ -94,7 +96,10 @@ warp_reduce, ) from b12x._lib.quant.iq2_xs import iq2_xs_execution_lut -from b12x._lib.quant.lut_e4m3 import lut_e4m3_value_table +from b12x._lib.quant.lut_e4m3 import ( + lut_e4m3_direct_table, + lut_e4m3_value_table, +) from b12x.moe._shared.kernels.trellis_ring import ( trellis256_lane_geom_bits as _trellis_ring_lane_geom_bits, ) @@ -116,6 +121,8 @@ max_packed_route_slots, packed_gemm_scratch_elements, plan_w4a16_buffers, + prefill_fused_sum_eligible, + prefill_fused_sum_enabled, select_route_block_size_m, validate_activation, ) @@ -295,10 +302,12 @@ def _gather_native_scale_bytes(a: Uint32, b: Uint32, c: Uint32, d: Uint32, byte: def _trellis256_execution_lut( - device: torch.device | str, codebook: str + device: torch.device | str, codebook: str, *, direct_lut: bool = False ) -> torch.Tensor: if codebook == LUT_FP16: return lut_fp16_segment_table(device) + if direct_lut: + return lut_e4m3_direct_table(device) return lut_e4m3_value_table(device) # TC-decode runs on the packed W4A16 object and folds the top-k sum into the FC2 @@ -860,6 +869,7 @@ class W4A16FusedMoeCompileResult: use_expert_map: bool = False scale_format: str = "e4m3_k16" tc_decode_fused_sum: bool = False + prefill_fused_sum_fp32: bool = False collect_activation_amax: bool = False schedule_whole_tiles: bool = False intermediate_rotation: bool = False @@ -875,6 +885,7 @@ class W4A16FusedMoeCompileResult: shared_memory_bytes: int = -1 broadcast_suh: bool = False small_m_direct_launches: tuple[_W4A16SmallMDirectLaunch, ...] = () + trellis_direct_lut: bool = False @dataclass(frozen=True) @@ -1021,6 +1032,7 @@ def __init__( dual_a: bool = False, route_major_a: bool = False, fused_topk_sum: bool = False, + fused_sum_fp32: bool = False, fused_sum_topk: int = 1, schedule_whole_tiles: bool = False, dynamic_num_experts: bool = False, @@ -1210,6 +1222,7 @@ def __init__( static_pair_rates.get(trellis_pair_kind, (3, 3)) ) self.lut_e4m3_smem = False + self.trellis_direct_lut = False # Small-M stripe split-K: opt out of the one-tile-per-CTA fast path # so decode-heavy small-M phases spread each mn-tile's K range across # multiple CTAs (existing tail scheduling plus cross-CTA finalize). @@ -1289,6 +1302,7 @@ def __init__( if self.route_major_a and not self.dual_a: raise ValueError("route_major_a requires the exact dual-A FC1 path") self.fused_topk_sum = bool(fused_topk_sum) + self.fused_sum_fp32 = bool(fused_sum_fp32) self.fused_sum_topk = int(fused_sum_topk) # Whole-tile persistent scheduling: every mn-tile is computed by one # CTA over the full K (grid-strided waves, ragged last wave), skipping @@ -1315,10 +1329,10 @@ def __init__( raise ValueError( "schedule_whole_tiles requires direct_topk_routes or trellis_t256" ) - if self.fused_topk_sum and not self.direct_topk_routes: - raise ValueError("fused_topk_sum requires direct_topk_routes") if self.fused_topk_sum and self.fused_sum_topk < 1: raise ValueError("fused_sum_topk must be >= 1") + if self.fused_sum_fp32 and not self.fused_topk_sum: + raise ValueError("fused_sum_fp32 requires fused_topk_sum") self.cta_m_blocks = int(_covering_count(moe_block_size, 16)) self.uses_m_block_8 = moe_block_size == 8 self.stages = _w4a16_pipeline_stages( @@ -1336,6 +1350,7 @@ def __init__( else: self.sms = 120 max_shared_mem = _DEFAULT_MAX_SHARED_MEM + self.max_shared_mem = max_shared_mem self.blocks_per_sm = _determine_blocks_per_sm( problem_m=self.size_m, problem_n=self.covered_size_n, @@ -1485,7 +1500,9 @@ def __cache_key__(self) -> tuple[object, ...]: self.dual_a, self.route_major_a, self.fused_topk_sum, + self.fused_sum_fp32, self.fused_sum_topk, + self.size_m if self.fused_sum_fp32 else None, self.cta_m_blocks, self.uses_m_block_8, self.shared_words, @@ -1496,6 +1513,7 @@ def __cache_key__(self) -> tuple[object, ...]: self.schedule_whole_tiles, self.schedule_route_block_factor, self.lut_e4m3_smem, + self.trellis_direct_lut, self.small_m_splitk, ) @@ -4088,13 +4106,19 @@ def _scaled_dequant_b_fragment_trellis256_bits( win_a, win_b, trellis_lut_addr, int(bits) ) else: - e_lo, e_hi = packed_decode_lut_e4m3_to_e4m3x8( - win_a, - win_b, - trellis_lut_addr, - int(bits), - value_table_in_shared=self.lut_e4m3_smem, - ) + if cutlass.const_expr(self.trellis_direct_lut): + e_lo, e_hi = packed_decode_lut_e4m3_direct_to_e4m3x8( + win_a, win_b, trellis_lut_addr, int(bits), + in_shared=self.lut_e4m3_smem, + ) + else: + e_lo, e_hi = packed_decode_lut_e4m3_to_e4m3x8( + win_a, + win_b, + trellis_lut_addr, + int(bits), + value_table_in_shared=self.lut_e4m3_smem, + ) if cutlass.const_expr(self.is_fp16): o0, o1 = fp8x4_e4m3_to_half2x2(e_lo) o2, o3 = fp8x4_e4m3_to_half2x2(e_hi) @@ -5587,14 +5611,32 @@ def _drain_output_smem( # to token = route_index // top_k. bf16x2 add lands two # consecutive hidden lanes per word. token_idx = route_index // Int32(self.fused_sum_topk) - out_idx = Int64(token_idx) * Int64(c_gl_stride) + Int64( - c_gl_wr % c_gl_stride - ) - out_addr = get_ptr_as_int64(c_bf16_flat, out_idx * Int64(8)) - red_add_global_bf16x2(out_addr, q0) - red_add_global_bf16x2(out_addr + Int64(4), q1) - red_add_global_bf16x2(out_addr + Int64(8), q2) - red_add_global_bf16x2(out_addr + Int64(12), q3) + col_word = c_gl_wr % c_gl_stride + if cutlass.const_expr(self.fused_sum_fp32): + out_elem = ( + Int64(token_idx) * Int64(self.size_n) + + Int64(col_word) * Int64(8) + ) + out_addr = get_ptr_as_int64(c_bf16_flat, out_elem) + q00, q01 = self._elem2_to_f32x2(q0) + q10, q11 = self._elem2_to_f32x2(q1) + q20, q21 = self._elem2_to_f32x2(q2) + q30, q31 = self._elem2_to_f32x2(q3) + red_add_global_v4_f32(out_addr, q00, q01, q10, q11) + red_add_global_v4_f32( + out_addr + Int64(16), q20, q21, q30, q31 + ) + else: + out_idx = Int64(token_idx) * Int64(c_gl_stride) + Int64( + col_word + ) + out_addr = get_ptr_as_int64( + c_bf16_flat, out_idx * Int64(8) + ) + red_add_global_bf16x2(out_addr, q0) + red_add_global_bf16x2(out_addr + Int64(4), q1) + red_add_global_bf16x2(out_addr + Int64(8), q2) + red_add_global_bf16x2(out_addr + Int64(12), q3) else: st_global_v4_u32( get_ptr_as_int64(c_bf16_flat, true_idx * Int64(8)), @@ -5649,12 +5691,32 @@ def _drain_output_smem_tail( q3 = self._relu2_elem2(q3) if cutlass.const_expr(self.fused_topk_sum): token_idx = route_index // Int32(self.fused_sum_topk) - out_idx = Int64(token_idx) * Int64(c_gl_stride) + Int64(col_word) - out_addr = get_ptr_as_int64(c_bf16_flat, out_idx * Int64(8)) - red_add_global_bf16x2(out_addr, q0) - red_add_global_bf16x2(out_addr + Int64(4), q1) - red_add_global_bf16x2(out_addr + Int64(8), q2) - red_add_global_bf16x2(out_addr + Int64(12), q3) + if cutlass.const_expr(self.fused_sum_fp32): + out_elem = ( + Int64(token_idx) * Int64(self.size_n) + + Int64(col_word) * Int64(8) + ) + out_addr = get_ptr_as_int64(c_bf16_flat, out_elem) + q00, q01 = self._elem2_to_f32x2(q0) + q10, q11 = self._elem2_to_f32x2(q1) + q20, q21 = self._elem2_to_f32x2(q2) + q30, q31 = self._elem2_to_f32x2(q3) + red_add_global_v4_f32(out_addr, q00, q01, q10, q11) + red_add_global_v4_f32( + out_addr + Int64(16), q20, q21, q30, q31 + ) + else: + out_idx = ( + Int64(token_idx) * Int64(c_gl_stride) + + Int64(col_word) + ) + out_addr = get_ptr_as_int64( + c_bf16_flat, out_idx * Int64(8) + ) + red_add_global_bf16x2(out_addr, q0) + red_add_global_bf16x2(out_addr + Int64(4), q1) + red_add_global_bf16x2(out_addr + Int64(8), q2) + red_add_global_bf16x2(out_addr + Int64(12), q3) else: st_global_v4_u32( get_ptr_as_int64(c_bf16_flat, true_idx * Int64(8)), @@ -6231,6 +6293,7 @@ def __init__( direct_topk_routes: bool = False, use_expert_map: bool = False, tc_decode_fused_sum: bool = False, + prefill_fused_sum_fp32: bool = False, tc_zero_output: bool = True, collect_activation_amax: bool = False, schedule_whole_tiles: bool = False, @@ -6240,6 +6303,7 @@ def __init__( rotation_input_dtype: str = "fp16", broadcast_suh: bool = False, pipeline_stages: int | None = None, + trellis_decode_table: str = "auto", ): activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) @@ -6261,20 +6325,33 @@ def __init__( else: w13_layout = "packed" self.tc_decode_fused_sum = bool(tc_decode_fused_sum) + self.prefill_fused_sum_fp32 = bool(prefill_fused_sum_fp32) + if self.tc_decode_fused_sum and self.prefill_fused_sum_fp32: + raise ValueError( + "TC-decode and large-M FP32 route reduction are mutually exclusive" + ) # When two TC-decode launches share one pre-zeroed output, only the # first must zero it. Default True preserves single-launch behavior. self.tc_zero_output = bool(tc_zero_output) self.collect_activation_amax = bool(collect_activation_amax) if self.collect_activation_amax and bool(direct_topk_routes): raise ValueError("activation amax collection requires route-packed W4A16") - if self.collect_activation_amax and self.tc_decode_fused_sum: + if self.collect_activation_amax and ( + self.tc_decode_fused_sum or self.prefill_fused_sum_fp32 + ): raise ValueError( - "activation amax collection is incompatible with TC-decode" + "activation amax collection is incompatible with fused route reduction" ) if self.tc_decode_fused_sum and not bool(direct_topk_routes): raise ValueError("tc_decode_fused_sum requires direct_topk_routes") if self.tc_decode_fused_sum and element_dtype != "bf16": raise ValueError("tc_decode_fused_sum currently requires bf16 activations") + if self.prefill_fused_sum_fp32 and element_dtype != "bf16": + raise ValueError("prefill_fused_sum_fp32 requires bf16 activations") + if self.prefill_fused_sum_fp32 and int(size_m) <= _TC_DECODE_MAX_M: + raise ValueError( + "prefill_fused_sum_fp32 requires a token capacity above the decode range" + ) fc1_cols = int(intermediate_size) * (2 if is_gated else 1) routed_rows = int(size_m) * int(top_k) self.size_m = int(size_m) @@ -6428,8 +6505,8 @@ def __init__( raise ValueError("full_rotation requires fp16 GEMM operands") if self.rotation_input_dtype not in {"bf16", "fp16"}: raise ValueError("full_rotation input dtype must be 'bf16' or 'fp16'") - if self.tc_decode_fused_sum: - raise ValueError("full_rotation is incompatible with TC decode") + if self.tc_decode_fused_sum or self.prefill_fused_sum_fp32: + raise ValueError("full_rotation is incompatible with fused route reduction") if self.apply_router_weight_on_input: raise ValueError( "full_rotation applies router weights only in the fp32 top-k sum" @@ -6513,7 +6590,11 @@ def __init__( ), single_token_route_fast_path=size_m == 1 and not self.direct_topk_routes, direct_topk_routes=self.direct_topk_routes, - fused_topk_sum=(self.tc_decode_fused_sum and weight_layout not in BLOCK_CODECS), + fused_topk_sum=( + (self.tc_decode_fused_sum and weight_layout not in BLOCK_CODECS) + or self.prefill_fused_sum_fp32 + ), + fused_sum_fp32=self.prefill_fused_sum_fp32, fused_sum_topk=int(top_k), schedule_whole_tiles=self.schedule_whole_tiles, dynamic_num_experts=self.dynamic_num_experts, @@ -6534,10 +6615,38 @@ def __init__( self.shared_words += self.iq2_lut_bytes // 4 self.fc1.iq2_xs_smem_lut = True self.fc2.iq2_xs_smem_lut = True - self.lut_e4m3_smem = ( + if trellis_decode_table not in {"auto", "compact", "full"}: + raise ValueError("trellis decode table must be auto, compact, or full") + compact_smem = ( + self.trellis_codebook == LUT_E4M3 + and _lut_e4m3_smem_enabled() + ) + lut_offset = (self.shared_words * 4 + 15) // 16 * 16 + # The cooperative grid must retain its occupancy contract. A direct + # table is eligible only when it fits without reducing resident CTAs. + direct_eligible = ( self.weight_layout == "trellis_t256" and self.trellis_codebook == LUT_E4M3 - and _lut_e4m3_smem_enabled() + and self.intermediate_hadamard + and self.trellis_bits == 2 + and self.fc1_trellis_pair_kind is None + and self.fc2_trellis_pair_kind is None + and self.blocks_per_sm == 1 + and lut_offset + 65536 <= self.fc1.max_shared_mem + ) + if trellis_decode_table == "full" and not direct_eligible: + raise ValueError( + "full trellis decode table is incompatible with the planned kernel geometry" + ) + self.trellis_direct_lut = direct_eligible and ( + trellis_decode_table == "full" + or (trellis_decode_table == "auto" and compact_smem and size_m <= 16) + ) + self.fc1.trellis_direct_lut = self.trellis_direct_lut + self.fc2.trellis_direct_lut = self.trellis_direct_lut + self.lut_e4m3_smem = compact_smem or self.trellis_direct_lut + self.trellis_lut_smem_bytes = ( + 65536 if self.trellis_direct_lut else _LUT_E4M3_SMEM_REGION_BYTES ) self.lut_e4m3_smem_off = 0 if self.lut_e4m3_smem: @@ -6546,8 +6655,10 @@ def __init__( ) // 16 * 16 self.shared_words = ( self.lut_e4m3_smem_off - + _LUT_E4M3_SMEM_REGION_BYTES + + self.trellis_lut_smem_bytes ) // 4 + if self.shared_words * 4 > self.fc1.max_shared_mem: + raise ValueError(f"full trellis decode table requires {self.shared_words * 4} bytes") self.fc1.lut_e4m3_smem = True self.fc2.lut_e4m3_smem = True self.barrier_count_off = self.sms * 4 @@ -6591,6 +6702,7 @@ def __cache_key__(self) -> tuple[object, ...]: self.broadcast_suh, self.rotation_input_dtype, self.lut_e4m3_smem, + self.trellis_direct_lut, self.small_m_splitk, self.fc1.__cache_key__, self.fc2.__cache_key__, @@ -6982,8 +7094,8 @@ class Storage: fc1_trellis_lut_addr = get_ptr_as_int64(fc1_trellis_lut_flat, Int32(0)) fc2_trellis_lut_addr = get_ptr_as_int64(fc2_trellis_lut_flat, Int32(0)) - # The emit hooks receive the staged value table's shared byte offset - # through the LUT ABI slot. + # The LUT ABI slot carries the shared byte offset for either the + # compact value table or the complete 2-bit direct table. fc1_phase_lut_addr = fc1_trellis_lut_addr fc2_phase_lut_addr = fc2_trellis_lut_addr if cutlass.const_expr(self.weight_layout in IQ2_CODECS): @@ -7004,7 +7116,7 @@ class Storage: self._lut_smem_copy( fc1_trellis_lut_addr, smem_base + Int32(self.lut_e4m3_smem_off), - _LUT_E4M3_SMEM_REGION_BYTES, + self.trellis_lut_smem_bytes, tid, ) cute.arch.sync_threads() @@ -7197,34 +7309,35 @@ def _moe_body( active_m, ) self._grid_barrier(locks_i32_flat, tid, grid_x) - if cutlass.const_expr(self.tc_decode_fused_sum and self.weight_layout not in BLOCK_CODECS): - # The TC-decode FC2 epilogue atomically accumulates per-route - # partials directly into the per-token output, so the output must be - # pre-zeroed. Previously this was a SEPARATE host-side output.zero_() - # kernel launch on the latency-bound decode critical path (an extra - # launch + its grid-fill memset before the fused kernel even starts). - # Fold it into the fused kernel prologue here: every CTA zeroes a - # grid-strided slice of the output BEFORE FC1, and the EXISTING - # post-FC1 grid barrier (already required to order FC1 writes before - # the activation/FC2 read) makes all zero stores globally visible - # before the first FC2 atomic -- so no extra barrier is added. The - # tiny m*hidden bf16 memset (decode: <=4*4096 elems) is dwarfed by - # FC1's whole-K FP4-weight stream, but we delete one whole kernel - # launch from the per-decode chain. The TC-decode output is per-token - # (top_k routes atomically summed into the SAME token row), so the - # zero span is active_m*hidden_size -- NOT the per-route - # active_m*top_k*hidden_size of _zero_fc2_output. - # tc_zero_output=False skips the zero (a paired earlier launch has - # already zeroed the shared output); the grid barrier below is - # unconditional so ordering is preserved either way. + if cutlass.const_expr( + (self.tc_decode_fused_sum and self.weight_layout not in BLOCK_CODECS) + or self.prefill_fused_sum_fp32 + ): + # FC2 route reduction atomically accumulates into one row per token. + # Every CTA zeroes a grid-strided slice before FC1. The mandatory + # post-FC1 grid barrier orders these stores before every FC2 atomic. + # ``tc_zero_output=False`` is valid only when a paired launch has + # already zeroed the same output and participates in that barrier. if cutlass.const_expr(self.tc_zero_output): zidx = cta * Int32(self.cta_threads) + tid zstride = grid_x * Int32(self.cta_threads) - ztotal = active_m * Int32(self.hidden_size) - zzero = self._cast_elem(cutlass.Float32(0.0)) - while zidx < ztotal: - fc2_bf16_flat[zidx] = zzero - zidx += zstride + zzero = ( + cutlass.Float32(0.0) + if cutlass.const_expr(self.prefill_fused_sum_fp32) + else self._cast_elem(cutlass.Float32(0.0)) + ) + if cutlass.const_expr(self.prefill_fused_sum_fp32): + zidx_i64 = Int64(zidx) + zstride_i64 = Int64(zstride) + ztotal_i64 = Int64(active_m) * Int64(self.hidden_size) + while zidx_i64 < ztotal_i64: + fc2_bf16_flat[zidx_i64] = zzero + zidx_i64 += zstride_i64 + else: + ztotal = active_m * Int32(self.hidden_size) + while zidx < ztotal: + fc2_bf16_flat[zidx] = zzero + zidx += zstride if cutlass.const_expr(self.activation_is_gated): self.fc1._run_persistent_gemm( @@ -7330,7 +7443,7 @@ def _moe_body( grid_x, active_m, ) - if cutlass.const_expr(self.zero_fc2_output): + if cutlass.const_expr(self.zero_fc2_output and not self.prefill_fused_sum_fp32): self._zero_fc2_output(fc2_bf16_flat, tid, cta, grid_x, active_m) self._grid_barrier(locks_i32_flat, tid, grid_x) fc2_store = fc2_bf16_flat @@ -9692,6 +9805,7 @@ def compile_w4a16_fused_moe( fast_math: bool = True, sms: int, max_shared_mem: int, + direct_token_capacity: int | None = None, swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, @@ -9705,6 +9819,7 @@ def compile_w4a16_fused_moe( direct_topk_routes: bool = False, use_expert_map: bool = False, tc_decode_fused_sum: bool = False, + prefill_fused_sum_fp32: bool = False, collect_activation_amax: bool = False, force_tile_config: tuple[int, int, int, int] | None = None, pipeline_stages: int | None = None, @@ -9713,6 +9828,7 @@ def compile_w4a16_fused_moe( intermediate_hadamard: bool = False, rotation_input_dtype: str | None = None, broadcast_suh: bool = False, + trellis_decode_table: str = "auto", _require_cached: bool = False, ) -> W4A16FusedMoeCompileResult: scale_format = _normalize_scale_format(scale_format) @@ -9762,10 +9878,17 @@ def compile_w4a16_fused_moe( direct_topk_routes = bool(direct_topk_routes) use_expert_map = bool(use_expert_map) tc_decode_fused_sum = bool(tc_decode_fused_sum) + prefill_fused_sum_fp32 = bool(prefill_fused_sum_fp32) + if tc_decode_fused_sum and prefill_fused_sum_fp32: + raise ValueError( + "TC-decode and large-M FP32 route reduction are mutually exclusive" + ) if use_expert_map and not direct_topk_routes: raise ValueError("use_expert_map requires direct_topk_routes") collect_activation_amax = bool(collect_activation_amax) - if collect_activation_amax and (direct_topk_routes or tc_decode_fused_sum): + if collect_activation_amax and ( + direct_topk_routes or tc_decode_fused_sum or prefill_fused_sum_fp32 + ): raise ValueError( "W4A16 activation amax collection requires the route-packed fused path" ) @@ -9784,8 +9907,8 @@ def compile_w4a16_fused_moe( raise ValueError( "rotation_input_dtype must be 'bf16' or 'fp16' for full_rotation" ) - if tc_decode_fused_sum: - raise ValueError("full_rotation is incompatible with TC decode") + if tc_decode_fused_sum or prefill_fused_sum_fp32: + raise ValueError("full_rotation is incompatible with fused route reduction") if apply_router_weight_on_input: raise ValueError( "full_rotation requires apply_router_weight_on_input=False" @@ -10092,6 +10215,7 @@ def compile_w4a16_fused_moe( direct_topk_routes=direct_topk_routes, use_expert_map=use_expert_map, tc_decode_fused_sum=tc_decode_fused_sum, + prefill_fused_sum_fp32=prefill_fused_sum_fp32, collect_activation_amax=collect_activation_amax, intermediate_rotation=intermediate_rotation, full_rotation=full_rotation, @@ -10099,6 +10223,7 @@ def compile_w4a16_fused_moe( rotation_input_dtype=rotation_input_dtype, broadcast_suh=broadcast_suh, pipeline_stages=pipeline_stages, + trellis_decode_table=trellis_decode_table, ) cache_key = ( "w4a16_fused_moe", @@ -10114,9 +10239,14 @@ def compile_w4a16_fused_moe( "count before capturing" ) + # Packed routes use rounded storage capacity; native direct launches use + # the exact planned row count, which may be smaller than that bucket. + direct_m = size_m if direct_token_capacity is None else int(direct_token_capacity) + if not 0 < direct_m <= size_m: + raise ValueError("direct token capacity must be within packed capacity") small_m_direct_launches = [] if (not collect_activation_amax) and _small_m_direct_supported( - m=size_m, + m=direct_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, @@ -10133,7 +10263,7 @@ def compile_w4a16_fused_moe( ): for ids_dtype in (torch.int32, torch.int64): direct = _compile_w4a16_small_m_direct( - m=size_m, + m=direct_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, @@ -10202,8 +10332,9 @@ def compile_w4a16_fused_moe( assumed_align=16, ) fc2_fake = cute.runtime.make_fake_compact_tensor( - cutlass_dtype, - (compile_routed_rows * hidden_size,), + cutlass.Float32 if kernel.prefill_fused_sum_fp32 else cutlass_dtype, + ((compile_size_m if kernel.prefill_fused_sum_fp32 else compile_routed_rows) + * hidden_size,), assumed_align=16, ) pair_metadata_cutlass_dtype = ( @@ -10359,6 +10490,7 @@ def compile_w4a16_fused_moe( use_expert_map=kernel.use_expert_map, scale_format=scale_format, tc_decode_fused_sum=bool(tc_decode_fused_sum), + prefill_fused_sum_fp32=bool(prefill_fused_sum_fp32), collect_activation_amax=collect_activation_amax, schedule_whole_tiles=kernel.schedule_whole_tiles, intermediate_rotation=intermediate_rotation, @@ -10374,6 +10506,7 @@ def compile_w4a16_fused_moe( shared_memory_bytes=kernel.shared_words * 4, broadcast_suh=bool(broadcast_suh), small_m_direct_launches=tuple(small_m_direct_launches), + trellis_direct_lut=kernel.trellis_direct_lut, ) attach_programs(result, compiled, *(launch.compiled for launch in small_m_direct_launches)) _FUSED_CACHE[cache_key] = result @@ -10948,6 +11081,7 @@ def _w4a16_fused_moe_launch_flat( fc2_tile_n: int, direct_topk_routes: bool, tc_decode_fused_sum: bool, + prefill_fused_sum_fp32: bool, collect_activation_amax: bool, stream_int: int, expert_map: torch.Tensor | None = None, @@ -11055,6 +11189,7 @@ def _w4a16_fused_moe_launch_flat( direct_topk_routes=bool(direct_topk_routes), use_expert_map=use_expert_map, tc_decode_fused_sum=bool(tc_decode_fused_sum), + prefill_fused_sum_fp32=bool(prefill_fused_sum_fp32), collect_activation_amax=collect_activation_amax, # The custom-op boundary cannot carry the compiled launch object. Re-pin # its selected geometry so tile-specific packs resolve the @@ -11079,7 +11214,7 @@ def _w4a16_fused_moe_launch_flat( fc2_trellis_lut_addr = iq2_lut.data_ptr() elif weight_layout == "trellis_t256" and trellis_codebook != "mcg": trellis_rank_lut = _trellis256_execution_lut( - a_input.device, trellis_codebook + a_input.device, trellis_codebook, direct_lut=fused.trellis_direct_lut ) fc1_trellis_lut_addr = trellis_rank_lut.data_ptr() fc2_trellis_lut_addr = trellis_rank_lut.data_ptr() @@ -11393,6 +11528,7 @@ def _w4a16_fused_moe_launch_op( fc2_tile_n=fc2_tile_n, direct_topk_routes=direct_topk_routes, tc_decode_fused_sum=tc_decode_fused_sum, + prefill_fused_sum_fp32=False, collect_activation_amax=False, stream_int=stream_int, ) @@ -11550,6 +11686,7 @@ def _w4a16_fused_moe_calibrated_launch_op( fc2_tile_n=fc2_tile_n, direct_topk_routes=False, tc_decode_fused_sum=False, + prefill_fused_sum_fp32=False, collect_activation_amax=True, stream_int=stream_int, ) @@ -12377,6 +12514,7 @@ def run_w4a16_moe( intermediate_cache13: torch.Tensor, intermediate_cache2: torch.Tensor, output: torch.Tensor, + prefill_sum_accum: torch.Tensor | None = None, fc1_c_tmp: torch.Tensor | None = None, fc2_c_tmp: torch.Tensor | None = None, packed_route_indices: torch.Tensor | None = None, @@ -12828,6 +12966,9 @@ def run_w4a16_moe( # without materializing remapped IDs or masked router weights. # Prepared launches retain this contract in tc_decode_fused_sum. preplanned_tc_decode = bool(getattr(fused_launch, "tc_decode_fused_sum", False)) + preplanned_prefill_fused_sum = bool( + getattr(fused_launch, "prefill_fused_sum_fp32", False) + ) prefer_tc_decode = route_mode == "direct" or ( route_mode == "auto" and _w4a16_tc_decode_preferred( @@ -12895,6 +13036,38 @@ def run_w4a16_moe( "planned W4A16 direct routing is unsupported for this launch shape" ) + prefill_fused_sum_requested = ( + preplanned_prefill_fused_sum + if fused_launch is not None + else prefill_fused_sum_enabled() + ) + use_prefill_fused_sum = prefill_fused_sum_eligible( + dtype=element_dtype, + m=int(fused_launch.size_m) if fused_launch is not None else m, + full_rotation=full_rotation, + weight_layout=weight_layout, + collect_activation_amax=collect_activation_amax, + enabled=prefill_fused_sum_requested, + ) + use_fused_topk_sum = bool(use_tc_decode or use_prefill_fused_sum) + + if use_prefill_fused_sum: + required_accum_elements = ( + int(fused_launch.size_m) if fused_launch is not None else int(m) + ) * hidden_size + if ( + prefill_sum_accum is None + or prefill_sum_accum.dtype != torch.float32 + or prefill_sum_accum.device != a_input.device + or not prefill_sum_accum.is_contiguous() + or prefill_sum_accum.numel() < required_accum_elements + ): + raise ValueError( + "W4A16 prefill fused sum requires a contiguous FP32 accumulator " + f"with at least {required_accum_elements} elements on " + f"{a_input.device}" + ) + # Prepared TC decode owns the per-token reduction and omits the separate # top-k sum launch. A mismatched routing contract would corrupt output. if preplanned_tc_decode and not use_tc_decode: @@ -12902,6 +13075,11 @@ def run_w4a16_moe( "preplanned TC-decode W4A16 launch requires small-M packed bf16 " f"decode (m <= {_TC_DECODE_MAX_M}, cuda int32/int64 topk_ids)" ) + if preplanned_prefill_fused_sum and not use_prefill_fused_sum: + raise RuntimeError( + "preplanned W4A16 prefill fused-sum launch requires the enabled " + "large-M packed or modelopt BF16 route-reduction contract" + ) route_slots_for_scratch = int(m) * int(topk) * int(block_size_m) required_m_blocks = int(m) * int(topk) if use_direct_topk_routes else 0 @@ -12911,6 +13089,13 @@ def run_w4a16_moe( int(block_size_m), route_num_experts, ) + if route_pack_launches is not None: + # A retained small-M GEMM can use a route packer prepared for a + # larger capacity. Its kernels initialize every planned route slot; + # do not shrink those views to the selected GEMM's live-row bucket. + route_slots_capacity = max( + route_slots_capacity, route_pack_launches.max_packed_routes + ) route_blocks_capacity = (route_slots_capacity + int(block_size_m) - 1) // int( block_size_m ) @@ -12991,8 +13176,12 @@ def run_w4a16_moe( topk=topk, route_num_experts=route_num_experts, sms=sms, + dtype=(prepared_dtype if full_rotation else a_input.dtype), full_rotation=full_rotation, block_size_m=block_size_m, + weight_layout=weight_layout, + collect_activation_amax=collect_activation_amax, + prefill_fused_sum=use_prefill_fused_sum, ) intermediate_size = int(prepared.intermediate_size) fc1_cols = buffer_plan.fc1_cols @@ -13058,6 +13247,7 @@ def run_w4a16_moe( direct_topk_routes=use_direct_topk_routes, use_expert_map=mapped_direct and use_direct_topk_routes, tc_decode_fused_sum=use_tc_decode, + prefill_fused_sum_fp32=use_prefill_fused_sum, collect_activation_amax=collect_activation_amax, intermediate_rotation=intermediate_rotation_scales is not None, full_rotation=full_rotation, @@ -13102,6 +13292,8 @@ def run_w4a16_moe( fc2_trellis_pair_kind, bool(use_direct_topk_routes), mapped_direct and use_direct_topk_routes, + bool(use_tc_decode), + bool(use_prefill_fused_sum), bool(collect_activation_amax), block_size_m, bool(intermediate_rotation_scales is not None), @@ -13144,6 +13336,8 @@ def run_w4a16_moe( getattr(fused_launch, "fc2_trellis_pair_kind", None), bool(getattr(fused_launch, "direct_topk_routes", False)), bool(getattr(fused_launch, "use_expert_map", False)), + bool(getattr(fused_launch, "tc_decode_fused_sum", False)), + bool(getattr(fused_launch, "prefill_fused_sum_fp32", False)), bool(getattr(fused_launch, "collect_activation_amax", False)), int(fused_launch.moe_block_size), bool(getattr(fused_launch, "intermediate_rotation", False)), @@ -13163,12 +13357,22 @@ def run_w4a16_moe( fused = fused_launch capacity_m = int(fused.size_m) capacity_routed_rows = capacity_m * topk - if intermediate_cache13_flat.numel() < capacity_routed_rows * max( - fc1_cols, hidden_size - ): + required_cache13_elements = ( + capacity_routed_rows * fc1_cols + if use_prefill_fused_sum + else capacity_routed_rows * max(fc1_cols, hidden_size) + ) + if intermediate_cache13_flat.numel() < required_cache13_elements: raise ValueError( "intermediate_cache13 is smaller than the selected W4A16 launch capacity: " - f"capacity_rows={capacity_m}, topk={topk}" + f"capacity_rows={capacity_m}, topk={topk}, " + f"available_elements={intermediate_cache13_flat.numel()}, " + f"required_elements={required_cache13_elements}, " + f"fused_topk_sum={use_fused_topk_sum}, " + f"prefill_fused_sum={use_prefill_fused_sum}, " + f"collect_activation_amax={collect_activation_amax}, " + f"full_rotation={full_rotation}, weight_layout={weight_layout}, " + f"element_dtype={element_dtype}" ) if intermediate_cache2_flat.numel() < capacity_routed_rows * intermediate_size: raise ValueError( @@ -13182,7 +13386,10 @@ def run_w4a16_moe( ) fc1_out = intermediate_cache13_flat[: capacity_routed_rows * fc1_store_cols] activated = intermediate_cache2_flat[: capacity_routed_rows * intermediate_size] - if use_tc_decode: + if use_prefill_fused_sum: + assert prefill_sum_accum is not None + fc2_out = prefill_sum_accum[: capacity_m * hidden_size] + elif use_tc_decode: # The fused kernel owns per-token reduction and output initialization. fc2_out = output.view(-1) else: @@ -13319,6 +13526,7 @@ def run_w4a16_moe( or _intermediate_rotation or weight_layout in {"trellis_t256", "iq2_xs", "iq2_xxs", "q8_0"} or (mapped_direct and use_direct_topk_routes) + or use_prefill_fused_sum ): # Native t256 bypasses the registered torch op so its shape-derived # bitrate reaches compilation without widening the stable public op ABI. @@ -13425,6 +13633,7 @@ def run_w4a16_moe( fc2_tile_n=_lt_fc2tn, direct_topk_routes=bool(use_direct_topk_routes), tc_decode_fused_sum=bool(use_tc_decode), + prefill_fused_sum_fp32=bool(use_prefill_fused_sum), collect_activation_amax=False, stream_int=int(stream), expert_map=expert_map if use_direct_topk_routes else None, @@ -13451,6 +13660,10 @@ def run_w4a16_moe( int(stream), ) + if use_prefill_fused_sum: + assert prefill_sum_accum is not None + output.copy_(prefill_sum_accum[: m * hidden_size].view(m, hidden_size)) + return output if use_tc_decode: # FC2 already wrote the top-k-summed result into `output`. return output diff --git a/b12x/moe/_shared/kernels/w4a16/route_pack.py b/b12x/moe/_shared/kernels/w4a16/route_pack.py index 10ea8ab00..cd4f1502d 100644 --- a/b12x/moe/_shared/kernels/w4a16/route_pack.py +++ b/b12x/moe/_shared/kernels/w4a16/route_pack.py @@ -10,6 +10,7 @@ import torch from b12x._lib.compile_plan import launch_triton +from b12x._lib.env import env_flag from b12x.moe._shared.kernels.w4a16.host import route_pack_capacity @@ -22,6 +23,9 @@ _FAST_COUNT_BLOCK_T = 1024 +_STABLE_SORT_MIN_ROUTES = 4096 +_STABLE_SORT_BLOCK_T = 4096 +_STABLE_SORT_EXPERTS_PER_PROGRAM = 1 @triton.jit @@ -296,6 +300,73 @@ def _pack_topk_routes_sort_kernel( tl.store(packed_route_indices + ranks, offsets, mask=valid) +@triton.jit +def _pack_topk_routes_stable_kernel( + topk_ids, + expert_map, + packed_route_indices, + expert_offsets, + live_numel, + NUMEL_CAPACITY: tl.constexpr, + NUM_EXPERTS: tl.constexpr, + HAS_EXPERT_MAP: tl.constexpr, + BLOCK_T: tl.constexpr, + EXPERTS_PER_PROGRAM: tl.constexpr, +): + """Pack each expert's routes in ascending token-major route order.""" + expert_ids = ( + tl.program_id(0) * EXPERTS_PER_PROGRAM + + tl.arange(0, EXPERTS_PER_PROGRAM) + ) + expert_mask = expert_ids < NUM_EXPERTS + output_starts = tl.load( + expert_offsets + expert_ids, + mask=expert_mask, + other=0, + ) + output_counts = tl.zeros((EXPERTS_PER_PROGRAM,), dtype=tl.int32) + lanes = tl.arange(0, BLOCK_T) + + for start in tl.range(0, NUMEL_CAPACITY, BLOCK_T): + offsets = start + lanes + raw_ids = tl.load( + topk_ids + offsets, + mask=offsets < live_numel, + other=-1, + ).to(tl.int32) + valid = ( + (offsets < live_numel) + & (raw_ids >= 0) + & (raw_ids < NUM_EXPERTS) + ) + ids = raw_ids + if HAS_EXPERT_MAP: + safe_ids = tl.minimum(tl.maximum(raw_ids, 0), NUM_EXPERTS - 1) + ids = tl.load(expert_map + safe_ids, mask=valid, other=-1).to( + tl.int32 + ) + valid = valid & (ids >= 0) & (ids < NUM_EXPERTS) + + matches = ( + expert_mask[:, None] + & valid[None, :] + & (expert_ids[:, None] == ids[None, :]) + ) + match_i32 = matches.to(tl.int32) + local_ranks = tl.cumsum(match_i32, axis=1) - 1 + output_indices = ( + output_starts[:, None] + + output_counts[:, None] + + local_ranks + ) + tl.store( + packed_route_indices + output_indices, + offsets[None, :], + mask=matches, + ) + output_counts += tl.sum(match_i32, axis=1) + + @dataclass(frozen=True) class W4A16RoutePackLaunches: """Exact prepared Triton objects for one W4A16 packed-route geometry.""" @@ -306,6 +377,7 @@ class W4A16RoutePackLaunches: max_packed_routes: int max_route_blocks: int use_small_prefix: bool + stable_order: bool programs: MappingProxyType def program(self, name: str, topk_ids: torch.Tensor, has_expert_map: bool): @@ -328,10 +400,13 @@ def compile_w4a16_route_pack_launches( block_size: int, num_experts: int, ordinal: int, + stable_order: bool | None = None, + bucket_tokens: bool = True, ) -> W4A16RoutePackLaunches: """Compile every legal route-pack helper for one planned W4A16 capacity.""" numel_capacity, max_packed_routes, max_route_blocks = route_pack_capacity( - int(tokens) * int(topk), int(block_size), int(num_experts), topk=int(topk) + int(tokens) * int(topk), int(block_size), int(num_experts), topk=int(topk), + bucket_tokens=bucket_tokens, ) max_packed_routes = max(int(max_packed_routes), 1) max_route_blocks = max(int(max_route_blocks), 1) @@ -342,6 +417,9 @@ def compile_w4a16_route_pack_launches( block_route_init <= _SMALL_PREFIX_MAX_PACKED_ROUTES and block_m <= _SMALL_PREFIX_MAX_ROUTE_BLOCKS ) + if stable_order is None: + stable_order = env_flag("W4A16_STABLE_ROUTE_PACK") + stable_order = bool(stable_order and numel_capacity >= _STABLE_SORT_MIN_ROUTES) programs = {} # Avoid Triton's divisibility and value-one specialization of the live bound. live_numel = 3 @@ -369,12 +447,22 @@ def compile_w4a16_route_pack_launches( BLOCK_T=_FAST_COUNT_BLOCK_T, num_warps=4, grid=(triton.cdiv(int(tokens) * int(topk), _FAST_COUNT_BLOCK_T), 1, 1), ) - programs[("sort", *key)] = _pack_topk_routes_sort_kernel.warmup( - ids_dtype, expert_map_dtype, torch.int32, torch.int32, live_numel, - NUM_EXPERTS=int(num_experts), HAS_EXPERT_MAP=has_expert_map, - BLOCK_T=_SORT_BLOCK_T, num_warps=4, - grid=(triton.cdiv(int(tokens) * int(topk), _SORT_BLOCK_T), 1, 1), - ) + if stable_order: + programs[("sort", *key)] = _pack_topk_routes_stable_kernel.warmup( + ids_dtype, expert_map_dtype, torch.int32, torch.int32, live_numel, + NUMEL_CAPACITY=int(numel_capacity), NUM_EXPERTS=int(num_experts), + HAS_EXPERT_MAP=has_expert_map, BLOCK_T=_STABLE_SORT_BLOCK_T, + EXPERTS_PER_PROGRAM=_STABLE_SORT_EXPERTS_PER_PROGRAM, + num_warps=8, + grid=(triton.cdiv(num_experts, _STABLE_SORT_EXPERTS_PER_PROGRAM), 1, 1), + ) + else: + programs[("sort", *key)] = _pack_topk_routes_sort_kernel.warmup( + ids_dtype, expert_map_dtype, torch.int32, torch.int32, live_numel, + NUM_EXPERTS=int(num_experts), HAS_EXPERT_MAP=has_expert_map, + BLOCK_T=_SORT_BLOCK_T, num_warps=4, + grid=(triton.cdiv(int(tokens) * int(topk), _SORT_BLOCK_T), 1, 1), + ) if not use_small_prefix: programs[("prefix", torch.int32, False)] = _w4a16_route_prefix_from_counts_kernel.warmup( torch.int32, torch.int32, torch.int32, BLOCK_SIZE=int(block_size), @@ -400,6 +488,7 @@ def compile_w4a16_route_pack_launches( numel_capacity=int(numel_capacity), block_size=int(block_size), num_experts=int(num_experts), max_packed_routes=max_packed_routes, max_route_blocks=max_route_blocks, use_small_prefix=use_small_prefix, + stable_order=stable_order, programs=MappingProxyType(programs), ) @@ -676,7 +765,27 @@ def pack_topk_routes_by_expert( int(block_size), int(num_experts), max_packed_routes, max_route_blocks, _POST_PREFIX_BLOCK_T, block_e.bit_length(), ) - if launches is None: + stable_order = ( + launches.stable_order if launches is not None else + env_flag("W4A16_STABLE_ROUTE_PACK") and numel >= _STABLE_SORT_MIN_ROUTES + ) + if stable_order: + stable_grid = (triton.cdiv(num_experts, _STABLE_SORT_EXPERTS_PER_PROGRAM),) + if launches is None: + launch_triton(_pack_topk_routes_stable_kernel, stable_grid, + topk_ids, expert_map_tensor, packed_route_indices, expert_offsets, numel, + NUMEL_CAPACITY=numel_capacity, NUM_EXPERTS=int(num_experts), + HAS_EXPERT_MAP=expert_map is not None, BLOCK_T=_STABLE_SORT_BLOCK_T, + EXPERTS_PER_PROGRAM=_STABLE_SORT_EXPERTS_PER_PROGRAM, num_warps=8, + ) + else: + _launch_prepared( + launches.program("sort", topk_ids, expert_map is not None), stable_grid, + topk_ids, expert_map_tensor, packed_route_indices, expert_offsets, numel, + numel_capacity, int(num_experts), expert_map is not None, + _STABLE_SORT_BLOCK_T, _STABLE_SORT_EXPERTS_PER_PROGRAM, + ) + elif launches is None: launch_triton(_pack_topk_routes_sort_kernel, sort_grid, topk_ids, expert_map_tensor, packed_route_indices, expert_offsets, numel, NUM_EXPERTS=int(num_experts), HAS_EXPERT_MAP=expert_map is not None, diff --git a/b12x/moe/fused_moe/__init__.py b/b12x/moe/fused_moe/__init__.py index 8ddc08c26..12aa5d0e7 100644 --- a/b12x/moe/fused_moe/__init__.py +++ b/b12x/moe/fused_moe/__init__.py @@ -61,6 +61,8 @@ group="moe", api_style="planned", entry_points=( + "Exl3Source", + "Exl3Weights", "ActivationMode", "ActivationSpec", "Binding", @@ -127,6 +129,8 @@ if TYPE_CHECKING: # static analysis only; runtime resolution is lazy from .api import ( # noqa: F401 + Exl3Source, + Exl3Weights, ActivationMode, ActivationSpec, Binding, diff --git a/b12x/moe/fused_moe/_impl.py b/b12x/moe/fused_moe/_impl.py index 07056e154..9ab7d6eea 100644 --- a/b12x/moe/fused_moe/_impl.py +++ b/b12x/moe/fused_moe/_impl.py @@ -273,6 +273,7 @@ class TPW4A16Workspace: routed_rows_capacity: int intermediate_cache13: torch.Tensor intermediate_cache2: torch.Tensor + prefill_sum_accum: torch.Tensor | None fc1_c_tmp: torch.Tensor fc2_c_tmp: torch.Tensor packed_route_indices: torch.Tensor @@ -298,6 +299,7 @@ class TPW4A16Workspace: planned_swiglu_beta: float = 0.0 planned_scale_format: str = "e4m3_k16" planned_collect_activation_amax: bool = False + planned_prefill_fused_sum_fp32: bool = False planned_fused_moe_launches: dict[object, object] = field(default_factory=dict) planned_topk_sum_launches: dict[object, object] = field(default_factory=dict) # Mapped direct-route launches, keyed by exact live token count. These @@ -767,6 +769,7 @@ class _TPCoreWorkspacePlan: trellis_codebook: str | None = None intermediate_hadamard: bool = False route_block_size_m: int | None = None + prefill_fused_sum_fp32: bool = False tensor_specs: Tuple[_TensorAllocSpec, ...] = () @@ -835,6 +838,9 @@ class TPMoEScratchCaps: deterministic_output: bool | None = None w4a16_block_size_m: int | None = None w4a16_fast_math: bool = True + w4a16_prefill_fused_sum: bool | None = None + w4a16_stable_route_pack: bool | None = None + trellis_decode_table: str = "auto" frozen: bool = True def __post_init__(self) -> None: @@ -888,6 +894,19 @@ def __post_init__(self) -> None: raise ValueError("w4a16_block_size_m must be one of 8, 16, 32, 48, 64") object.__setattr__(self, "w4a16_block_size_m", block_size_m) object.__setattr__(self, "w4a16_fast_math", bool(self.w4a16_fast_math)) + if self.trellis_decode_table not in {"auto", "compact", "full"}: + raise ValueError("trellis decode table must be auto, compact, or full") + if self.w4a16_prefill_fused_sum is None: + from b12x.moe._shared.kernels.w4a16.host import prefill_fused_sum_enabled + + object.__setattr__( + self, "w4a16_prefill_fused_sum", prefill_fused_sum_enabled() + ) + if self.w4a16_stable_route_pack is None: + object.__setattr__( + self, "w4a16_stable_route_pack", + _env_flag("B12X_W4A16_STABLE_ROUTE_PACK", default=False), + ) object.__setattr__(self, "frozen", bool(self.frozen)) @property @@ -951,6 +970,7 @@ class TPMoEScratchPlan: _mixed_trellis_launches: tuple[tuple[torch.dtype, bool, bool, object], ...] = field( default=(), repr=False ) + _prewarmed_route_pack_launches: object | None = field(default=None, repr=False) @property def full_rotation(self) -> bool: @@ -1059,6 +1079,7 @@ def bind( activation_amax=activation_amax, ) elif self.caps.quant_mode == "w4a16" and self._core_workspace_plan.full_rotation: + route_pack_launches = self._prewarmed_route_pack_launches tokens = int(a.shape[0]) prepared = experts.representation_for("w4a16") broadcast_suh = prepared.gate_suh.numel() == self.caps.k @@ -1193,6 +1214,7 @@ class TPMoEFP4Binding: materialized_intermediate: torch.Tensor | None = None intermediate_cache13: torch.Tensor | None = None intermediate_cache2: torch.Tensor | None = None + prefill_sum_accum: torch.Tensor | None = None fc1_c_tmp: torch.Tensor | None = None fc2_c_tmp: torch.Tensor | None = None packed_route_indices: torch.Tensor | None = None @@ -3065,6 +3087,7 @@ def _build_tp_moe_fp4_binding_from_views( routed_rows_capacity=plan.routed_rows, intermediate_cache13=tensors["intermediate_cache13"], intermediate_cache2=tensors["intermediate_cache2"], + prefill_sum_accum=tensors.get("prefill_sum_accum"), fc1_c_tmp=tensors["fc1_c_tmp"], fc2_c_tmp=tensors["fc2_c_tmp"], packed_route_indices=tensors["packed_route_indices"], @@ -3219,6 +3242,7 @@ def _plan_core_workspace( w4a16_scale_format: str | None = None, route_num_experts: int | None = None, w4a16_block_size_m: int | None = None, + w4a16_prefill_fused_sum: bool | None = None, trellis_bits: int = 3, trellis_tile_config: tuple[int, int, int, int] | None = None, trellis_pair_kinds: frozenset[str] | None = None, @@ -3226,6 +3250,7 @@ def _plan_core_workspace( intermediate_hadamard: bool = False, projection_mixed_trellis: bool = False, apply_router_weight_on_input: bool = False, + collect_activation_amax: bool = False, deterministic_output: bool = False, swiglu_limit: float | None = None, swiglu_alpha: float | None = None, @@ -3245,6 +3270,7 @@ def _plan_core_workspace( from b12x.moe._shared.kernels.w4a16.host import ( max_packed_route_slots, packed_gemm_scratch_elements, + prefill_fused_sum_eligible, route_block_sizes_for_capacity, select_route_block_size_m, ) @@ -3582,10 +3608,26 @@ def _plan_core_workspace( // _dtype_nbytes(dtype), ) cache_dtype = torch.float16 if full_rotation else dtype + use_prefill_fused_sum = prefill_fused_sum_eligible( + dtype=dtype, + m=token_capacity, + full_rotation=full_rotation, + weight_layout=weight_layout, + collect_activation_amax=collect_activation_amax, + enabled=w4a16_prefill_fused_sum, + ) + intermediate_cache13_elements = ( + max( + routed_capacity * fc1_cols, + token_capacity * int(k), + ) + if use_prefill_fused_sum + else routed_capacity * max(fc1_cols, int(k)) + ) tensor_specs = [ _TensorAllocSpec( "intermediate_cache13", - (routed_capacity * max(fc1_cols, int(k)),), + (intermediate_cache13_elements,), cache_dtype, ), _TensorAllocSpec( @@ -3603,6 +3645,14 @@ def _plan_core_workspace( _TensorAllocSpec("expert_offsets", (route_E + 1,), torch.int32), _TensorAllocSpec("expert_counts", (route_E,), torch.int32), ] + if use_prefill_fused_sum: + tensor_specs.append( + _TensorAllocSpec( + "prefill_sum_accum", + (token_capacity * int(k),), + torch.float32, + ) + ) if full_rotation: max_tokens = max(routed_capacity // max(int(num_topk), 1), 1) tensor_specs.extend( @@ -3653,6 +3703,7 @@ def _plan_core_workspace( trellis_codebook=trellis_codebook, intermediate_hadamard=bool(intermediate_hadamard), route_block_size_m=w4a16_block_size_m, + prefill_fused_sum_fp32=use_prefill_fused_sum, tensor_specs=tuple(tensor_specs), ) @@ -4118,6 +4169,7 @@ def _materialize_workspace_from_core_arena( routed_rows_capacity=plan.routed_rows, intermediate_cache13=tensors["intermediate_cache13"], intermediate_cache2=tensors["intermediate_cache2"], + prefill_sum_accum=tensors.get("prefill_sum_accum"), fc1_c_tmp=tensors["fc1_c_tmp"], fc2_c_tmp=tensors["fc2_c_tmp"], packed_route_indices=tensors["packed_route_indices"], @@ -4140,6 +4192,7 @@ def _materialize_workspace_from_core_arena( trellis_codebook=plan.trellis_codebook, intermediate_hadamard=plan.intermediate_hadamard, route_block_size_m=plan.route_block_size_m, + planned_prefill_fused_sum_fp32=plan.prefill_fused_sum_fp32, volatile_launch_state=bool(volatile_launch_state), ) if a1_gscale is None or a2_gscale is None: @@ -7488,13 +7541,17 @@ def _validate_frozen_w4a16_launch( f"scale_format={scale_format!r}, " f"collect_activation_amax={bool(collect_activation_amax)}" ) + fused_reduces_routes = bool( + getattr(fused, "tc_decode_fused_sum", False) + or getattr(fused, "prefill_fused_sum_fp32", False) + ) has_topk_sum = planned_capacity in workspace.planned_topk_sum_launches if workspace.full_rotation: has_topk_sum = any( isinstance(key, tuple) and key[0] == planned_capacity for key in workspace.planned_topk_sum_launches ) - if not has_topk_sum: + if not fused_reduces_routes and not has_topk_sum: raise RuntimeError( "frozen W4A16 MoE workspace is missing its preplanned top-k sum launch " f"for capacity={planned_capacity}" @@ -7588,6 +7645,11 @@ def _w4a16_preplanned_launches( fused = workspace.planned_fused_moe_launches.get( (weight_layout, scale_format, planned_capacity, collect_activation_amax) ) + if fused is not None and bool( + getattr(fused, "tc_decode_fused_sum", False) + or getattr(fused, "prefill_fused_sum_fp32", False) + ): + return fused, None topk_sum_key: object = planned_capacity if workspace.full_rotation: topk_sum_key = ( @@ -7817,6 +7879,7 @@ def plan_tp_moe_arena_layout( collect_activation_amax: bool = False, deterministic_output: bool | None = None, w4a16_block_size_m: int | None = None, + w4a16_prefill_fused_sum: bool | None = None, decode_config: MoeDecodeConfig, ) -> TPMoEArenaLayout: """Compute the byte layout needed by one lane-owned MoE pool.""" @@ -7915,6 +7978,7 @@ def plan_tp_moe_arena_layout( w4a16_scale_format=w4a16_scale_format, route_num_experts=route_num_experts, w4a16_block_size_m=w4a16_block_size_m, + w4a16_prefill_fused_sum=w4a16_prefill_fused_sum, trellis_bits=weight_plan.trellis_bits or 3, trellis_tile_config=weight_plan.trellis_tile_config, trellis_pair_kinds=weight_plan.trellis_pair_kinds, @@ -7924,6 +7988,7 @@ def plan_tp_moe_arena_layout( intermediate_hadamard=weight_plan.intermediate_hadamard, projection_mixed_trellis=_uses_projection_mixed_trellis(weight_plan), apply_router_weight_on_input=apply_router_weight_on_input, + collect_activation_amax=collect_activation_amax, deterministic_output=plan.deterministic_output, swiglu_limit=plan.swiglu_limit, swiglu_alpha=plan.swiglu_alpha, @@ -7959,10 +8024,8 @@ def _plan_full_rotation_w4a16_launches( ]: """Compile the fixed Trellis launches before serving memory profiling. - The caller-owned scratch path cannot use the mutable workspace prewarm - dictionaries. Keeping the launch objects on the immutable scratch plan - preserves the old Trellis API contract and prevents first-use compilation - from being charged as peak activation memory by vLLM. + The immutable scratch plan owns compiled programs independently of module + cache eviction. Runtime binding performs no kernel resolution or allocation. """ if not core_plan.full_rotation or core_plan.projection_mixed_trellis: return (), () @@ -7975,15 +8038,11 @@ def _plan_full_rotation_w4a16_launches( if torch.cuda.is_current_stream_capturing(): raise RuntimeError("Trellis launch planning cannot run during capture") - from b12x.moe._shared.kernels.w4a16.host import ( - route_pack_capacity, - route_pack_warmup_token_counts, - ) + from b12x.moe._shared.kernels.w4a16.host import route_pack_capacity from b12x.moe._shared.kernels.w4a16.kernel import ( _DEFAULT_MAX_SHARED_MEM, compile_w4a16_fused_moe, compile_w4a16_topk_sum, - pack_topk_routes_by_expert, ) capacity_tokens = max(int(capacity_tokens), 1) @@ -8009,7 +8068,7 @@ def _plan_full_rotation_w4a16_launches( f"got weight_layout={weight_layout!r}, scale_format={scale_format!r}" ) w13_layout = "trellis_t256_proj" - _, capacity_route_slots, capacity_m_blocks = route_pack_capacity( + _, _, capacity_m_blocks = route_pack_capacity( capacity_tokens * core_plan.num_topk, block_size_m, core_plan.route_E, @@ -8064,6 +8123,7 @@ def compile_fused(token_count: int, broadcast_suh: bool) -> object: w13_layout=w13_layout, trellis_bits=core_plan.trellis_bits, trellis_codebook=core_plan.trellis_codebook or LUT_E4M3, + trellis_decode_table=caps.trellis_decode_table, force_tile_config=core_plan.trellis_tile_config, intermediate_rotation=True, full_rotation=True, @@ -8099,83 +8159,6 @@ def compile_fused(token_count: int, broadcast_suh: bool) -> object: for mapped in (False, True) for broadcast_svh in (False, True) ) - packed_route_indices = torch.empty( - capacity_route_slots, - dtype=torch.int32, - device=core_plan.device, - ) - block_expert_ids = torch.empty( - capacity_m_blocks, - dtype=torch.int32, - device=core_plan.device, - ) - packed_route_count = torch.empty( - 1, - dtype=torch.int32, - device=core_plan.device, - ) - expert_offsets = torch.empty( - core_plan.route_E + 1, - dtype=torch.int32, - device=core_plan.device, - ) - expert_counts = torch.empty( - core_plan.route_E, - dtype=torch.int32, - device=core_plan.device, - ) - pending_route_pack_keys: list[tuple[object, ...]] = [] - for route_ids_dtype in (torch.int32, torch.int64): - dummy_topk_ids = torch.zeros( - capacity_tokens, - core_plan.num_topk, - dtype=route_ids_dtype, - device=core_plan.device, - ) - for mapped in (False, True): - expert_map = None - if mapped: - expert_map = torch.arange( - core_plan.route_E, - dtype=torch.int32, - device=core_plan.device, - ) - for token_count in route_pack_warmup_token_counts(capacity_tokens): - route_pack_key = route_pack_prewarm_key( - core_plan.device.type, - int(torch.cuda.current_device()), - route_ids_dtype, - token_count, - core_plan.num_topk, - capacity_route_slots, - capacity_m_blocks, - int(block_size_m), - int(core_plan.route_E), - mapped, - ) - if route_pack_key in _W4A16_ROUTE_PACK_PREWARMED: - continue - pack_topk_routes_by_expert( - dummy_topk_ids[:token_count], - block_size_m, - core_plan.route_E, - expert_map=expert_map, - 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, - ) - pending_route_pack_keys.append(route_pack_key) - torch.cuda.current_stream(core_plan.device).synchronize() - _W4A16_ROUTE_PACK_PREWARMED.update(pending_route_pack_keys) - if pending_route_pack_keys: - logger.info( - "Prewarmed %d full-rotation W4A16 route-pack variant(s) " - "for token capacity %d.", - len(pending_route_pack_keys), - capacity_tokens, - ) return fused_launches, topk_sum_launches @@ -8547,11 +8530,51 @@ def _resolve_trellis_route_block_size(caps: TPMoEScratchCaps) -> int | None: bucket_w4a16_tokens=False, ) route_experts = caps.route_num_experts or caps.weight_E - return select_route_block_size_m( + preferred = select_route_block_size_m( max(token_counts), caps.num_topk, route_experts, ) + tile = caps.weight_plan.trellis_tile_config + if tile is None or torch.device(caps.device).type != "cuda": + return preferred + # Native Trellis packing fixes both projection tiles. The route block + # must fit those tiles' shared memory, not a different automatically + # selected GEMM tile with the same logical weight dimensions. + from b12x.moe._shared.kernels.w4a16.host import ( + route_block_sizes_for_capacity, + ) + from b12x.moe._shared.kernels.w4a16.kernel import _candidate_tile_fits + + props = torch.cuda.get_device_properties(caps.device) + budget = int(props.shared_memory_per_block_optin) - 512 + candidates = route_block_sizes_for_capacity( + max_tokens=max(token_counts), + topk=caps.num_topk, + num_experts=route_experts, + ) + dimensions = ( + (_activation_w1_rows(caps.weight_plan.activation, caps.n), caps.k), + (caps.k, caps.n), + ) + for block_rows in reversed(candidates): + if all( + _candidate_tile_fits( + problem_n=n, + problem_k=k, + cta_m_blocks=(block_rows + 15) // 16, + tile_k=tk, + tile_n=tn, + cta_threads=tk * tn // 64, + max_shared_mem=budget, + scale_format="e4m3_k32", + weight_layout="trellis_t256", + weight_bits=max(4, caps.weight_plan.trellis_bits or 3), + ) + for (n, k), (tk, tn) in zip(dimensions, (tile[:2], tile[2:])) + ): + return block_rows + raise ValueError("Native Trellis projection tiles exceed shared-memory capacity") def _plan_tp_moe_arena_layout_from_caps( @@ -8579,6 +8602,7 @@ def _plan_tp_moe_arena_layout_from_caps( collect_activation_amax=caps.collect_activation_amax, deterministic_output=deterministic_output, w4a16_block_size_m=_resolve_trellis_route_block_size(caps), + w4a16_prefill_fused_sum=caps.w4a16_prefill_fused_sum, decode_config=caps.decode_config, ) @@ -8638,6 +8662,7 @@ def plan_tp_moe_scratch( w4a16_scale_format=caps.w4a16_scale_format, route_num_experts=caps.route_num_experts, w4a16_block_size_m=resolved_block_size_m, + w4a16_prefill_fused_sum=caps.w4a16_prefill_fused_sum, trellis_bits=caps.weight_plan.trellis_bits or 3, trellis_tile_config=caps.weight_plan.trellis_tile_config, trellis_pair_kinds=caps.weight_plan.trellis_pair_kinds, @@ -8647,6 +8672,7 @@ def plan_tp_moe_scratch( intermediate_hadamard=caps.weight_plan.intermediate_hadamard, projection_mixed_trellis=_uses_projection_mixed_trellis(caps.weight_plan), apply_router_weight_on_input=caps.apply_router_weight_on_input, + collect_activation_amax=caps.collect_activation_amax, deterministic_output=launch_plan.deterministic_output, swiglu_limit=launch_plan.swiglu_limit, swiglu_alpha=launch_plan.swiglu_alpha, @@ -8667,6 +8693,30 @@ def plan_tp_moe_scratch( fused_launches = () topk_sum_launches = () mixed_trellis_launches = () + route_pack_launches = None + if ( + prewarm_launches + and core_workspace_plan.full_rotation + and not core_workspace_plan.projection_mixed_trellis + and core_workspace_plan.device.type == "cuda" + ): + from b12x.moe._shared.kernels.w4a16.route_pack import ( + compile_w4a16_route_pack_launches, + ) + + route_pack_launches = compile_w4a16_route_pack_launches( + tokens=capacity_tokens, + topk=core_workspace_plan.num_topk, + block_size=resolved_block_size_m, + num_experts=core_workspace_plan.route_E, + ordinal=( + core_workspace_plan.device.index + if core_workspace_plan.device.index is not None + else torch.cuda.current_device() + ), + stable_order=caps.w4a16_stable_route_pack, + bucket_tokens=False, + ) return TPMoEScratchPlan( caps=caps, layout=layout, @@ -8679,13 +8729,11 @@ def plan_tp_moe_scratch( device=torch.device(caps.device), ), ), - # Keep strong references to the launches primed in the module caches, - # but leave the runtime binding unresolved. ``run_w4a16_moe`` uses a - # None launch as a dispatch signal and then gets these exact objects - # from its compile cache without doing first-use JIT work. + # Retain every route/GEMM/reduction program required after cache eviction. _prewarmed_fused_launches=fused_launches, _prewarmed_topk_sum_launches=topk_sum_launches, _mixed_trellis_launches=mixed_trellis_launches, + _prewarmed_route_pack_launches=route_pack_launches, ) @@ -8709,6 +8757,7 @@ def _prewarm_w4a16_planned_launches( weight_layout: str = "packed", w13_layout: str = "w13", collect_activation_amax: bool = False, + prefill_fused_sum: bool = False, ) -> None: """Resolve every W4A16 kernel shape owned by a frozen arena. @@ -8729,6 +8778,7 @@ def _prewarm_w4a16_planned_launches( collect_activation_amax = bool(collect_activation_amax) from b12x.moe._shared.kernels.w4a16.host import ( + prefill_fused_sum_eligible, route_pack_capacity, select_route_block_size_m, ) @@ -8794,6 +8844,14 @@ def _prewarm_w4a16_planned_launches( token_count, collect_activation_amax, ) + build_prefill_fused_sum = prefill_fused_sum_eligible( + dtype=element_dtype, + m=token_count, + full_rotation=full_rotation, + weight_layout=weight_layout, + collect_activation_amax=collect_activation_amax, + enabled=prefill_fused_sum, + ) # Rotation-table row count belongs to the prepared artifact, which # is not bound until after the workspace is frozen. Resolve both # full-rotation specializations now so either artifact contract is @@ -8820,6 +8878,7 @@ def _prewarm_w4a16_planned_launches( weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, + prefill_fused_sum_fp32=build_prefill_fused_sum, collect_activation_amax=collect_activation_amax, trellis_bits=workspace.trellis_bits, trellis_codebook=workspace.trellis_codebook or LUT_E4M3, @@ -8856,7 +8915,7 @@ def _prewarm_w4a16_planned_launches( topk_sum_launches[(token_count, ids_dtype, mapped)] = ( resolved_topk_sum ) - else: + elif not build_prefill_fused_sum: topk_sum_launches[token_count] = compile_w4a16_topk_sum( m=token_count, topk=workspace.num_topk, @@ -9202,6 +9261,7 @@ def materialize_tp_moe_arena_workspaces( w4a16_scale_format=w4a16_scale_format, route_num_experts=caps.route_num_experts, w4a16_block_size_m=resolved_block_size_m, + w4a16_prefill_fused_sum=caps.w4a16_prefill_fused_sum, trellis_bits=weight_plan.trellis_bits or 3, trellis_tile_config=weight_plan.trellis_tile_config, trellis_pair_kinds=weight_plan.trellis_pair_kinds, @@ -9211,6 +9271,7 @@ def materialize_tp_moe_arena_workspaces( intermediate_hadamard=weight_plan.intermediate_hadamard, projection_mixed_trellis=_uses_projection_mixed_trellis(weight_plan), apply_router_weight_on_input=apply_router_weight_on_input, + collect_activation_amax=collect_activation_amax, deterministic_output=plan.deterministic_output, swiglu_limit=plan.swiglu_limit, swiglu_alpha=plan.swiglu_alpha, @@ -9256,6 +9317,14 @@ def materialize_tp_moe_arena_workspaces( != w4a16_scale_format or bool(getattr(existing, "planned_collect_activation_amax", False)) != collect_activation_amax + or bool( + getattr( + existing, + "planned_prefill_fused_sum_fp32", + False, + ) + ) + != core_plan.prefill_fused_sum_fp32 ): pass else: @@ -9312,6 +9381,9 @@ def materialize_tp_moe_arena_workspaces( materialized.planned_swiglu_alpha = plan.swiglu_alpha materialized.planned_swiglu_beta = plan.swiglu_beta materialized.planned_collect_activation_amax = collect_activation_amax + materialized.planned_prefill_fused_sum_fp32 = ( + core_plan.prefill_fused_sum_fp32 + ) t_prewarm0 = time.perf_counter() if _B12X_TIMING else 0.0 _prewarm_w4a16_planned_launches( materialized, @@ -9324,6 +9396,7 @@ def materialize_tp_moe_arena_workspaces( weight_layout=w4a16_weight_layout, w13_layout=w13_layout, collect_activation_amax=collect_activation_amax, + prefill_fused_sum=core_plan.prefill_fused_sum_fp32, ) if _B12X_TIMING: prewarm_ms += (time.perf_counter() - t_prewarm0) * 1000.0 @@ -9617,6 +9690,7 @@ def build_tp_moe_fp4_binding( routed_rows_capacity=workspace.routed_rows_capacity, intermediate_cache13=workspace.intermediate_cache13, intermediate_cache2=workspace.intermediate_cache2, + prefill_sum_accum=workspace.prefill_sum_accum, fc1_c_tmp=workspace.fc1_c_tmp, fc2_c_tmp=workspace.fc2_c_tmp, packed_route_indices=workspace.packed_route_indices, @@ -12950,6 +13024,7 @@ def b12x_moe_fp4(*, binding: TPMoEFP4Binding) -> torch.Tensor: ) intermediate_cache13 = _require_binding_field(binding, "intermediate_cache13") intermediate_cache2 = _require_binding_field(binding, "intermediate_cache2") + prefill_sum_accum = binding.prefill_sum_accum fc1_c_tmp = _require_binding_field(binding, "fc1_c_tmp") fc2_c_tmp = _require_binding_field(binding, "fc2_c_tmp") packed_route_indices = _require_binding_field(binding, "packed_route_indices") @@ -12987,6 +13062,7 @@ def b12x_moe_fp4(*, binding: TPMoEFP4Binding) -> torch.Tensor: fast_math=fast_math, intermediate_cache13=intermediate_cache13, intermediate_cache2=intermediate_cache2, + prefill_sum_accum=prefill_sum_accum, output=scatter_output, fc1_c_tmp=fc1_c_tmp, fc2_c_tmp=fc2_c_tmp, diff --git a/b12x/moe/fused_moe/_preparation.py b/b12x/moe/fused_moe/_preparation.py index 9c9734c60..a9baafecd 100644 --- a/b12x/moe/fused_moe/_preparation.py +++ b/b12x/moe/fused_moe/_preparation.py @@ -93,10 +93,16 @@ def _quant_mode(experts: PreparedExperts, config: MoeDecodeConfig) -> str: def _control_snapshot() -> FrozenMapping: """Capture host controls once while declaring the immutable query.""" from . import _impl + from b12x.moe._shared.kernels.w4a16.host import ( + prefill_fused_sum_enabled, trellis_decode_table, + ) tile = _impl._dynamic_tile_mn_override() raw_materialized = _impl.os.environ.get(_impl._DYNAMIC_NVFP4_MATERIALIZED_ENV) return FrozenMapping({ + "trellis_decode_table": trellis_decode_table(), + "w4a16_prefill_fused_sum": prefill_fused_sum_enabled(), + "w4a16_stable_route_pack": _impl._env_flag("B12X_W4A16_STABLE_ROUTE_PACK", default=False), "dynamic_nvfp4_materialized": ( None if raw_materialized is None else raw_materialized not in ("", "0", "false", "False") ), @@ -246,6 +252,11 @@ def _lower_caps( if config.w4a16_block_size_m is not None else query.w4a16_block_size_m), w4a16_fast_math=query.fast_math, + w4a16_prefill_fused_sum=bool( + query.controls.get("w4a16_prefill_fused_sum", False) + ), + w4a16_stable_route_pack=bool(query.controls.get("w4a16_stable_route_pack", False)), + trellis_decode_table=str(query.controls.get("trellis_decode_table", "auto")), swiglu_beta=_decode_scalar(query.swiglu_beta), ) @@ -370,6 +381,7 @@ def _w4a16_primary_launches(scratch, caps) -> _W4A16PrimaryLaunches: props = torch.cuda.get_device_properties(core.device) compiler_args = dict( size_m=tokens, hidden_size=core.k, intermediate_size=core.n, + direct_token_capacity=int(caps.max_tokens), num_experts=core.weight_E, top_k=core.num_topk, activation=core.activation, apply_router_weight_on_input=caps.apply_router_weight_on_input, @@ -387,9 +399,11 @@ def _w4a16_primary_launches(scratch, caps) -> _W4A16PrimaryLaunches: ) packed = compile_w4a16_fused_moe( **compiler_args, zero_fc2_output=False, max_m_blocks=packed_blocks, + prefill_fused_sum_fp32=core.prefill_fused_sum_fp32, ) packed_mapped = compile_w4a16_fused_moe( **compiler_args, zero_fc2_output=True, max_m_blocks=packed_blocks, + prefill_fused_sum_fp32=core.prefill_fused_sum_fp32, ) direct = direct_mapped = None if weight_layout in BLOCK_CODECS: @@ -430,6 +444,7 @@ def _w4a16_primary_launches(scratch, caps) -> _W4A16PrimaryLaunches: route_pack = compile_w4a16_route_pack_launches( tokens=tokens, topk=core.num_topk, block_size=int(route_block), num_experts=core.route_E, ordinal=core.device.index, + stable_order=caps.w4a16_stable_route_pack, ) # Direct routing requires exact M; packed routing accepts live M up to capacity. return _W4A16PrimaryLaunches( @@ -580,6 +595,8 @@ def _program_carriers( launches = [item[-1] for item in scratch._prewarmed_fused_launches] launches.extend(item[-1] for item in scratch._prewarmed_topk_sum_launches) launches.extend(item[-1] for item in scratch._mixed_trellis_launches) + if scratch._prewarmed_route_pack_launches is not None: + launches.extend(scratch._prewarmed_route_pack_launches.carriers()) plan = scratch.launch_plan if plan.implementation == "w4a16" and not scratch.full_rotation: launches.extend( diff --git a/b12x/moe/fused_moe/_tuning.py b/b12x/moe/fused_moe/_tuning.py index b1a9f8e38..bbace1042 100644 --- a/b12x/moe/fused_moe/_tuning.py +++ b/b12x/moe/fused_moe/_tuning.py @@ -233,6 +233,8 @@ def _validate_query(query: MoeDecodeQuery, _device: DeviceIdentity | None) -> No raise TypeError("query must be MoeDecodeQuery") if not isinstance(query.controls, FrozenMapping): raise TypeError("MoE controls must be frozen declaration metadata") + if query.controls.get("trellis_decode_table", "auto") not in {"auto", "compact", "full"}: + raise ValueError("trellis decode table must be auto, compact, or full") if query.io_dtype not in {"bfloat16", "float16"}: raise TypeError("MoE I/O dtype must be bfloat16 or float16") if query.route_logits_dtype not in {None, "float16", "bfloat16", "float32"}: @@ -634,7 +636,7 @@ def _validate_fc2_query(query: MoeFC2Query, _device) -> None: TUNING = TuningContract( component_id="moe.decode", - query_schema_version=11, + query_schema_version=13, config_schema_version=8, query_fields=frozenset(MoeDecodeQuery.__dataclass_fields__), config_fields=frozenset(MoeDecodeConfig.__dataclass_fields__), diff --git a/b12x/moe/fused_moe/api.py b/b12x/moe/fused_moe/api.py index 2f9d4f43a..6c7a68f75 100644 --- a/b12x/moe/fused_moe/api.py +++ b/b12x/moe/fused_moe/api.py @@ -26,6 +26,7 @@ RouteTopKInvocation, plan_fc2 as _plan_fc2, plan_route_topk as _plan_route_topk, + variant_for, ) from .execution import ExecutionCapacity, RoutingSpec, plan_execution as _plan_execution from .planning import ( @@ -37,8 +38,9 @@ plan_weights as _plan_weights, prepare_weights as _prepare_weights, ) -from .source import PackedSource, PackedSourceFormat, W13Layout, WeightSource +from .source import Exl3Source, PackedSource, PackedSourceFormat, W13Layout, WeightSource from .weights import ( + Exl3Weights, PackedWeights, IQ2XSWeights, BlockQuantWeights, @@ -70,7 +72,7 @@ def plan_weights( def prepare_weights( - *, plan: WeightPlan, weights: PackedWeights | TrellisWeights | IQ2XSWeights + *, plan: WeightPlan, weights: PackedWeights | TrellisWeights | IQ2XSWeights | Exl3Weights ) -> PreparedExperts: """Prepare the canonical weight representation owned by this layer.""" return _prepare_weights(plan=plan, weights=weights) @@ -190,6 +192,8 @@ def is_supported(device=None) -> bool: __all__ = [ + "Exl3Source", + "Exl3Weights", "ActivationMode", "ActivationSpec", "Binding", @@ -233,4 +237,3 @@ def is_supported(device=None) -> bool: "run_sparse", "run", ] - diff --git a/b12x/moe/fused_moe/planning.py b/b12x/moe/fused_moe/planning.py index 068320eb6..a03284945 100644 --- a/b12x/moe/fused_moe/planning.py +++ b/b12x/moe/fused_moe/planning.py @@ -17,8 +17,9 @@ prepare_b12x_iq2_xs_weights, ) from .config import TrellisConfig -from .source import PackedSource, WeightSource +from .source import Exl3Source, PackedSource, WeightSource from .weights import ( + Exl3Weights, PackedWeights, IQ2XSWeights, PreparedExperts, @@ -176,7 +177,7 @@ def _prepared_format( f"required packing {packing.value!r} is not available; " f"planner produced {sorted(value.value for value in available)}" ) - if isinstance(source, TrellisConfig): + if isinstance(source, (TrellisConfig, Exl3Source)): weights = WeightEncoding.TRELLIS scales = ScaleEncoding.TRELLIS_SCALES else: @@ -251,6 +252,54 @@ def plan_weights( w13_layout=source.w13_layout.value, w4a16_layout=requested_layout, ) + elif isinstance(source, Exl3Source): + manifest = source.manifest + if activation.mode is not ActivationMode.A16: + raise ValueError("EXL3 canonical preparation requires A16 activations") + if constraints.required_packing not in { + None, + WeightPacking.TRELLIS_NATIVE, + }: + raise ValueError("EXL3 weights require trellis_native packing") + if manifest.rates.structure != "uniform": + raise NotImplementedError( + "EXL3 canonical preparation requires uniform rates" + ) + if (geometry.num_experts, geometry.hidden_size) != ( + manifest.geometry.num_experts, + manifest.geometry.hidden_size, + ): + raise ValueError( + "EXL3 expert count and hidden width must match the manifest" + ) + if geometry.intermediate_size > manifest.geometry.intermediate_size: + raise ValueError("EXL3 local intermediate width exceeds the source width") + recipe = "w4a16" + raw_plan = plan_b12x_fp4_moe_weights( + quant_modes=recipe, + source_format="exl3", + activation=activation.nonlinearity, + params_dtype=activation.io_dtype, + num_experts=geometry.num_experts, + hidden_size=geometry.hidden_size, + intermediate_size=geometry.intermediate_size, + w13_layout="w31", + w4a16_layout="trellis_native", + trellis_bits=manifest.rates.bits, + trellis_tile_config=( + (128, 128, 128, 128) + if manifest.hadamard.intermediate_hadamard and manifest.rates.bits == 2 + else (64, 256, 64, 256) + ), + intermediate_hadamard=manifest.hadamard.intermediate_hadamard, + trellis_codebook=manifest.codebook, + trellis_rate_granularity=manifest.rates.structure, + intermediate_hadamard_blocks=( + (manifest.hadamard.pre_block, manifest.hadamard.post_block) + if manifest.hadamard.intermediate_hadamard + else None + ), + ) elif isinstance(source, TrellisConfig): if activation.mode is not ActivationMode.A16: raise ValueError("Trellis fused MoE currently requires A16 activations") @@ -288,7 +337,7 @@ def plan_weights( ), ) else: - raise TypeError("source must be a PackedSource or TrellisConfig") + raise TypeError("source must be a PackedSource, TrellisConfig or Exl3Source") return WeightPlan( source=source, @@ -307,7 +356,7 @@ def plan_weights( def prepare_weights( *, plan: WeightPlan, - weights: PackedWeights | TrellisWeights | IQ2XSWeights, + weights: PackedWeights | TrellisWeights | IQ2XSWeights | Exl3Weights, ) -> PreparedExperts: """Materialize the in-memory representation selected by ``plan_weights``.""" @@ -317,6 +366,17 @@ def prepare_weights( if not isinstance(weights, IQ2XSWeights) or weights.codec != plan.source.format.value: raise TypeError(f"{plan.source.format.value} preparation requires matching BlockQuantWeights") prepared = prepare_b12x_iq2_xs_weights(plan=plan._impl, weights=weights) + elif isinstance(plan.source, Exl3Source): + if not isinstance(weights, Exl3Weights): + raise TypeError("EXL3 preparation requires Exl3Weights") + if weights.layer.manifest != plan.source.manifest: + raise ValueError("EXL3 extent manifest differs from its weight plan") + prepared = prepare_b12x_fp4_moe_weights( + plan=plan._impl, + params_dtype=plan.activation.io_dtype, + exl3_layer=weights.layer, + exl3_device=weights.device, + ) elif isinstance(plan.source, TrellisConfig): if not isinstance(weights, TrellisWeights): raise TypeError("Trellis preparation requires TrellisWeights") diff --git a/b12x/moe/fused_moe/source.py b/b12x/moe/fused_moe/source.py index c9d4c5ee7..0c32132ed 100644 --- a/b12x/moe/fused_moe/source.py +++ b/b12x/moe/fused_moe/source.py @@ -6,6 +6,7 @@ from enum import Enum from typing import TypeAlias +from .._shared.exl3_schema import Exl3Manifest from .config import TrellisConfig @@ -40,10 +41,22 @@ def __post_init__(self) -> None: object.__setattr__(self, "w13_layout", W13Layout(self.w13_layout)) -WeightSource: TypeAlias = PackedSource | TrellisConfig +@dataclass(frozen=True, kw_only=True) +class Exl3Source: + """Trellis slot encoding with checkpoint-global transform coordinates.""" + + manifest: Exl3Manifest + + def __post_init__(self) -> None: + if not isinstance(self.manifest, Exl3Manifest): + raise TypeError("Exl3Source.manifest must be a Exl3Manifest") + + +WeightSource: TypeAlias = PackedSource | TrellisConfig | Exl3Source __all__ = [ + "Exl3Source", "PackedSource", "PackedSourceFormat", "W13Layout", diff --git a/b12x/moe/fused_moe/weights.py b/b12x/moe/fused_moe/weights.py index 3f5771544..be6951e7d 100644 --- a/b12x/moe/fused_moe/weights.py +++ b/b12x/moe/fused_moe/weights.py @@ -11,6 +11,7 @@ from b12x._lib.quant.block_codec import block_codec if TYPE_CHECKING: + from .._shared.kernels.w4a16.exl3 import Exl3Layer from ._impl import B12XFP4ExpertWeights from .planning import WeightPlan @@ -121,6 +122,31 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True, kw_only=True) +class Exl3Weights: + """One validated CPU slot extent and its destination CUDA device. + + Preparation stages bounded expert batches. The complete source payload + need not reside on the GPU beside its prepared representation. + """ + + layer: "Exl3Layer" + device: torch.device + + def __post_init__(self) -> None: + from .._shared.kernels.w4a16.exl3 import Exl3Layer + + if not isinstance(self.layer, Exl3Layer): + raise TypeError("Exl3Weights.layer must be a Exl3Layer") + device = torch.device(self.device) + if device.type != "cuda": + raise ValueError("Exl3Weights.device must identify a CUDA device") + object.__setattr__(self, "device", device) + self.layer.manifest.validate_extent( + self.layer.first_slot, self.layer.slot_count + ) + + @dataclass(frozen=True) class BlockQuantWeights: """Raw uint8[E,N,K/block_weights,block_bytes] checkpoint blocks. @@ -220,6 +246,7 @@ def device(self) -> torch.device: __all__ = [ + "Exl3Weights", "PackedWeights", "IQ2XSWeights", "BlockQuantWeights", diff --git a/docs/trellis-decode-table.md b/docs/trellis-decode-table.md new file mode 100644 index 000000000..6273678f3 --- /dev/null +++ b/docs/trellis-decode-table.md @@ -0,0 +1,39 @@ +# Trellis decode table + +Fused W4A16 MoE can decode uniform 2-bit `lut_e4m3` trellis weights with an +intermediate Hadamard through a 64 KiB table in each thread block's shared +memory. The table stores the complete codeword-to-E4M3 mapping. It replaces +the per-window permutation arithmetic; it does not change checkpoint weights, +activation precision, matrix multiplication, or reduction order. + +`B12X_TRELLIS_DECODE_TABLE` controls canonical fused-MoE preparation: + +| Value | Behavior | +| --- | --- | +| `auto` (default) | Use the full table for planned capacities of at most 16 tokens when the shared-memory and cooperative-residency contracts permit it. Other capacities and weight formats keep compact decoding. | +| `compact` | Compute the permutation per window and read the 4 KiB value table. | +| `full` | Require the full table; fail preparation for an incompatible kernel geometry. This is intended for component qualification. | + +For an explicit reference run, set the variable before declaring plans: + +```bash +export B12X_TRELLIS_DECODE_TABLE=compact +``` + +Canonical preparation saves the setting in its immutable query and propagates +it to compiler workers. Changing the environment after declaration cannot +change that plan. Binding and replay retain compiled programs and use +caller-owned scratch. The table is process-shared immutable execution data, +not a decoded copy of model weights. + +The full table is admitted only for uniform 2-bit `lut_e4m3` weights with an +intermediate Hadamard, one resident cooperative CTA per SM, and enough shared +memory for both the original GEMM staging and the 64 KiB table. It never +reduces the planned resident-CTA count. Large-capacity prefill keeps its +existing staging footprint. + +`tests/moe/test_trellis_direct_lut_decode.py` checks all 65,536 codewords and +random overlapping windows against the compact decoder, including shared +tables for each supported bit width. Only uniform 2-bit weights use the full +table in fused MoE; testing the lookup primitive at 3 and 4 bits does not +enable those execution paths. diff --git a/docs/w4a16-prefill-reduction.md b/docs/w4a16-prefill-reduction.md new file mode 100644 index 000000000..eef785df8 --- /dev/null +++ b/docs/w4a16-prefill-reduction.md @@ -0,0 +1,45 @@ +# Bounded W4A16 prefill route reduction + +Status: implemented and component-qualified. Full-model serving qualification +is separate from the kernel checks described here. + +Set `B12X_W4A16_PREFILL_FUSED_SUM=1` before declaring a fused-MoE plan to use +one FP32 output accumulator per token instead of a BF16 output per routed +expert. This applies to BF16 activations, packed or checkpoint-native MXFP4 / +NVFP4 weights, and planned token capacities above eight. Trellis rotation, +FP16 activations, activation-max calibration and decode-capacity plans retain +their existing reduction paths. The default remains disabled. + +The FC2 epilogue retains the BF16-rounded, router-weighted contribution of +each expert, converts that contribution to FP32, and reduces it into the +caller-owned accumulator. The fused kernel zeroes that accumulator before +FC1; its mandatory grid barrier orders initialization before FC2 reduction. +The output is cast to BF16 once after FC2. Relaxed FP32 atomic ordering is +not a bitwise-deterministic summation contract. + +Set `B12X_W4A16_STABLE_ROUTE_PACK=1` before plan declaration to retain +ascending token-major route order within each expert for capacities of at +least 4096 routed rows. This stabilizes expert packing, not floating-point +atomic accumulation order. Smaller capacity plans retain atomic packing. + +Both options are captured in the immutable MoE query. Changing the process +environment after declaration does not alter scratch geometry or dispatch. +Preparation compiles and retains capacity-specialized programs. Binding +maps caller-owned storage and accepts live token counts within that capacity; +it does not allocate an arena or select another kernel specialization. + +Validation covers prepared eager execution, CUDA Graph replay, changing +inputs, mapped/disabled experts, finite nonzero outputs, allocation-free +replay, and fixed-capacity dispatch at several live row counts. Run: + +```bash +.venv/bin/python -m pytest \ + tests/moe/test_w4a16_e2e.py::test_w4a16_prefill_reduction_prepared_capacity_and_graph \ + tests/moe/test_tp_moe_scratch_bindings.py::test_w4a16_prefill_reduction_freezes_caller_scratch_contract \ + tests/moe/test_w4a16_route_pack.py +``` + +For 4096 tokens, 896 experts, top-16 routing, hidden width 7168 and local +intermediate width 192, the planned scratch size decreases from +1,071,913,568 to 308,550,240 bytes. This is a scratch-memory reduction, not +a claim about end-to-end model throughput. diff --git a/tests/moe/test_exl3_prepare.py b/tests/moe/test_exl3_prepare.py index 4d500aef3..2b3f136f3 100644 --- a/tests/moe/test_exl3_prepare.py +++ b/tests/moe/test_exl3_prepare.py @@ -309,3 +309,139 @@ def test_exl3_uniform_padding_must_be_zero(tmp_path) -> None: activation="situ", device=_device(), ) + + +@requires_cuda +@pytest.mark.parametrize("first_slot, slots", [(0, 8), (8, 4), (48, 12), (60, 8)]) +@pytest.mark.parametrize("activation", ["silu", "situ"]) +@pytest.mark.parametrize("capacity", [1, 8, 129]) +def test_canonical_exl3_preparation_preserves_source_extent( + tmp_path, first_slot, slots, activation, capacity +): + """Uneven TP extents preserve rotations in preparation, execution and graphs.""" + from b12x._lib.runtime_control import kernel_resolution_guard + from b12x.moe import fused_moe + from b12x.moe._shared.kernels.w4a16.exl3 import read_exl3_layer + from b12x.moe._shared.kernels.w4a16.exl3_synth import write_exl3_checkpoint + from b12x.preparation import PreparationSession, PreparedCall + from b12x.preparation.types import require_prepared + from tests.moe.test_w4a16_mixed_trellis import _serial_tier + + config = Exl3SynthConfig( + codebook="lut_e4m3", + num_experts=2, + hidden_size=512, + intermediate_size=3072, + moe_layer_indices=(1,), + bits=2, + intermediate_hadamard=True, + pre_block=512, + post_block=128, + per_expert_input_rotations=True, + extent_alignment_slots=4, + extent_barriers=(48,), + seed=53, + ) + manifest = write_exl3_checkpoint(tmp_path, config) + layer = read_exl3_layer( + tmp_path, manifest, 1, first_slot=first_slot, slot_count=slots + ) + layer.sign_pattern.copy_(torch.tensor([0, 6], dtype=torch.uint8)) + expected = prepare_exl3_moe_weights(layer, activation=activation, device=_device()) + plan = fused_moe.plan_weights( + source=fused_moe.Exl3Source(manifest=manifest), + activation=fused_moe.ActivationSpec( + mode="a16", nonlinearity=activation, io_dtype=torch.bfloat16 + ), + geometry=fused_moe.MoEGeometry( + num_experts=2, hidden_size=512, intermediate_size=slots * 32 + ), + ) + prepared = fused_moe.prepare_weights( + plan=plan, weights=fused_moe.Exl3Weights(layer=layer, device=_device()) + ) + actual = prepared._impl.representation.value + for name in ( + "w13", + "w2", + "intermediate_rotations", + "gate_suh", + "up_suh", + "down_svh", + ): + torch.testing.assert_close( + getattr(actual, name), getattr(expected, name), rtol=0, atol=0 + ) + assert actual.intermediate_hadamard + + topk = 2 + torch.manual_seed(617) + source = (torch.randn(capacity, 512, device=_device()) * 0.01).bfloat16() + route_ids = torch.tensor([0, 1], device=_device(), dtype=torch.int32) + route_ids = route_ids.repeat(capacity, 1) + route_weights = torch.softmax(torch.randn(capacity, topk, device=_device()), -1) + expert_map = torch.arange(2, device=_device(), dtype=torch.int32) + output = torch.empty_like(source) + execution = fused_moe.plan_execution( + experts=prepared, + capacity=fused_moe.ExecutionCapacity(max_tokens=capacity, top_k=topk), + routing=fused_moe.RoutingSpec(), + ) + + def allocate(state): + return tuple(torch.empty(s.shape, dtype=s.dtype, device=s.device) + for s in state.scratch.scratch_specs()) + + def bind(state, scratch, rows): + return state.bind( + scratch=scratch, a=source[:rows], experts=prepared, + topk_weights=route_weights[:rows], topk_ids=route_ids[:rows], + output=output[:rows], + ) + + def primer(state): + scratch = allocate(state) + binding = bind(state, scratch, capacity) + return PreparedCall(run=lambda: state.run(binding), owners=scratch) + + def oracle(rows): + return _serial_tier( + source[:rows], expected, route_weights[:rows], route_ids[:rows], + expert_map, block_size_m=8, activation=activation, + ).to(output.dtype) + + with PreparationSession(device=source.device, autotune=False) as session: + request = execution.request(name="exl3-source-extent", prepare_call=primer) + session.prepare((request,)) + state = require_prepared(request.plan, "moe.decode") + scratch = allocate(state) + pointers = tuple(t.data_ptr() for t in (*scratch, output)) + for rows in sorted({1, min(8, capacity), capacity, min(17, capacity)}): + reference = oracle(rows) + with kernel_resolution_guard("EXL3 source extent"): + binding = bind(state, scratch, rows) + result = state.run(binding) + assert result.data_ptr() == output.data_ptr() + assert torch.isfinite(result).all() and torch.count_nonzero(result) + torch.testing.assert_close(result, reference, rtol=2e-3, atol=2e-3) + graph = torch.cuda.CUDAGraph() + try: + with session.capture(), torch.cuda.graph(graph): + # Bind initializes synchronization scalars in caller scratch; + # capture those writes together with the consuming kernels. + captured = state.run(bind(state, scratch, min(17, capacity))) + for factor in (-0.5, 2.0): + source.mul_(factor) + route_weights.copy_(route_weights.flip(-1)) + reference = oracle(min(17, capacity)) + for tensor in scratch: + tensor.fill_(255) + before = torch.cuda.memory_allocated() + graph.replay() + torch.cuda.synchronize() + assert torch.cuda.memory_allocated() == before + assert pointers == tuple(t.data_ptr() for t in (*scratch, output)) + assert torch.isfinite(captured).all() and torch.count_nonzero(captured) + torch.testing.assert_close(captured, reference, rtol=2e-3, atol=2e-3) + finally: + graph.reset() diff --git a/tests/moe/test_exl3_schema.py b/tests/moe/test_exl3_schema.py index 600de8e6a..48b355630 100644 --- a/tests/moe/test_exl3_schema.py +++ b/tests/moe/test_exl3_schema.py @@ -83,6 +83,54 @@ def test_uniform_manifest_parses() -> None: assert manifest.layout.extent_barriers == (8,) +@pytest.mark.parametrize("bits", [2, 3]) +def test_canonical_weight_plan_preserves_exl3_transform_and_rate(bits: int) -> None: + from b12x.moe import fused_moe + + manifest = Exl3Manifest.from_dict(_uniform_manifest_dict(bits=bits, intermediate_hadamard=True)) + plan = fused_moe.plan_weights( + source=fused_moe.Exl3Source(manifest=manifest), + activation=fused_moe.ActivationSpec( + mode="a16", nonlinearity="silu", io_dtype=torch.bfloat16 + ), + geometry=fused_moe.MoEGeometry( + num_experts=8, hidden_size=512, intermediate_size=256 + ), + ) + assert plan.prepared_format.weights is fused_moe.WeightEncoding.TRELLIS + assert plan.prepared_format.packing is fused_moe.WeightPacking.TRELLIS_NATIVE + assert plan._impl.trellis_bits == bits + assert plan._impl.intermediate_hadamard_blocks == (512, 128) + assert plan.geometry.intermediate_size == 256 + assert plan.source.manifest.geometry.intermediate_size == 512 + + +@pytest.mark.parametrize( + ("mode", "experts", "intermediate", "message"), + [ + ("a8", 8, 256, "A16 activations"), + ("a16", 9, 256, "must match the manifest"), + ("a16", 8, 1024, "exceeds the source width"), + ], +) +def test_canonical_exl3_plan_rejects_incompatible_contract( + mode: str, experts: int, intermediate: int, message: str +) -> None: + from b12x.moe import fused_moe + + manifest = Exl3Manifest.from_dict(_uniform_manifest_dict(intermediate_hadamard=True)) + with pytest.raises(ValueError, match=message): + fused_moe.plan_weights( + source=fused_moe.Exl3Source(manifest=manifest), + activation=fused_moe.ActivationSpec( + mode=mode, nonlinearity="silu", io_dtype=torch.bfloat16 + ), + geometry=fused_moe.MoEGeometry( + num_experts=experts, hidden_size=512, intermediate_size=intermediate + ), + ) + + def test_per_expert_manifest_parses() -> None: manifest = Exl3Manifest.from_dict(_per_expert_manifest_dict()) assert manifest.rates.pair_kinds == frozenset({"P33", "P43"}) diff --git a/tests/moe/test_fused_moe_variant_selection.py b/tests/moe/test_fused_moe_variant_selection.py index 9f17c2e6a..fd6c084f6 100644 --- a/tests/moe/test_fused_moe_variant_selection.py +++ b/tests/moe/test_fused_moe_variant_selection.py @@ -11,9 +11,56 @@ _FusedMoeCapacityState, _W4A16PrimaryLaunches, variant_for, + _control_snapshot, ) +@pytest.mark.parametrize("mode", ["auto", "compact", "full"]) +def test_trellis_decode_table_is_frozen_at_declaration(monkeypatch, mode): + monkeypatch.setenv("B12X_TRELLIS_DECODE_TABLE", mode) + controls = _control_snapshot() + monkeypatch.setenv("B12X_TRELLIS_DECODE_TABLE", "invalid") + assert controls["trellis_decode_table"] == mode + with pytest.raises(ValueError, match="B12X_TRELLIS_DECODE_TABLE"): + _control_snapshot() + + +@pytest.mark.parametrize("rows,mode,expected,shared_limit", [ + (1, "auto", True, 101376), (8, "auto", True, 101376), + (16, "auto", True, 101376), (17, "auto", False, 101376), + (8, "compact", False, 101376), (8, "full", True, 101376), + (8, "auto", False, 65536), +]) +def test_trellis_direct_table_preserves_cooperative_residency( + monkeypatch, rows, mode, expected, shared_limit, +): + from b12x.moe._shared.kernels.w4a16.kernel import W4A16FusedMoeKernel + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace( + multi_processor_count=188, shared_memory_per_block_optin=shared_limit, + )) + monkeypatch.setenv("B12X_LUT_E4M3_SMEM", "1") + kernel = W4A16FusedMoeKernel( + size_m=rows, hidden_size=3584, intermediate_size=384, + num_experts=512, top_k=16, activation="silu", + apply_router_weight_on_input=False, zero_fc2_output=False, + fc1_tile_n=128, fc1_tile_k=128, fc2_tile_n=128, fc2_tile_k=128, + moe_block_size=8, max_m_blocks=1024, element_dtype="fp16", + weight_layout="trellis_t256", scale_format="e4m3_k32", + w13_layout="trellis_t256_proj", trellis_bits=2, + full_rotation=True, intermediate_rotation=True, intermediate_hadamard=True, + trellis_decode_table=mode, + ) + assert kernel.trellis_direct_lut == expected + assert kernel.shared_words * 4 <= kernel.fc1.max_shared_mem + if expected: + assert kernel.blocks_per_sm == 1 + assert kernel.lut_e4m3_smem + assert kernel.fc1.trellis_direct_lut and kernel.fc2.trellis_direct_lut + + class _Variant: def __init__(self, count): self.count = count @@ -139,6 +186,20 @@ def resolve_grid(namespace, **_kwargs): _impl._launch_dynamic_flat(**arguments) +@pytest.mark.parametrize("tokens, expected", [(4, 4), (11, 128)]) +def test_public_route_dispatch_selects_retained_capacity(monkeypatch, tokens, expected): + from b12x.moe.fused_moe import api + + plan = object() + binding = SimpleNamespace(plan=plan, hidden_states=torch.empty(tokens, 16)) + variants = { + count: SimpleNamespace(route=lambda bound, count=count: (count, bound)) + for count in (4, 128) + } + monkeypatch.setattr(api, "require_prepared", lambda *_args: SimpleNamespace(variants=variants)) + assert api.route(plan, binding=binding) == (expected, binding) + + def _launches(*, direct, route_pack, route_mode="auto"): return _W4A16PrimaryLaunches( tokens=8, route_mode=route_mode, packed="packed", packed_mapped="packed_mapped", diff --git a/tests/moe/test_nvfp4_auto.py b/tests/moe/test_nvfp4_auto.py index 483a5d67f..40e789651 100644 --- a/tests/moe/test_nvfp4_auto.py +++ b/tests/moe/test_nvfp4_auto.py @@ -145,7 +145,7 @@ def test_auto_native_decode_retains_launches_and_replays_shared_storage( ): assert tensor.data_ptr() == original.data_ptr() originals = [tensor.clone() for tensor in raw] - counts = (1, 2, 4, 8) + counts = (1, 2, 4, 7, 8) x = torch.randn(8, hidden_size, dtype=torch.bfloat16, device="cuda") * 0.25 ids = torch.stack([torch.randperm(num_experts, device="cuda")[:topk] for _ in range(8)]) probabilities = torch.softmax(torch.randn(8, topk, device="cuda"), dim=-1) diff --git a/tests/moe/test_tp_moe_scratch_bindings.py b/tests/moe/test_tp_moe_scratch_bindings.py index a346c9aed..d0e9c9e05 100644 --- a/tests/moe/test_tp_moe_scratch_bindings.py +++ b/tests/moe/test_tp_moe_scratch_bindings.py @@ -20,6 +20,7 @@ ) from b12x.moe._shared.execution import PreparedWeightLayout from b12x.moe._shared.kernels.w4a8.weights import repack_w4a8_weights +from b12x.moe.fused_moe._tuning import MoeDecodeConfig def _weight_plan( quant_mode: str = "nvfp4", *, @@ -64,6 +65,70 @@ def _clear_moe_force_env(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv(name, raising=False) +@pytest.mark.parametrize("weight_layout", ["iq2_xs", "iq2_xxs", "q8_0", "trellis_t256"]) +def test_prefill_route_reduction_does_not_change_other_codec_epilogues(weight_layout): + from b12x.moe._shared.kernels.w4a16.host import prefill_fused_sum_eligible + + assert not prefill_fused_sum_eligible( + dtype=torch.bfloat16, + m=4096, + full_rotation=False, + weight_layout=weight_layout, + collect_activation_amax=False, + enabled=True, + ) + + +def test_w4a16_prefill_reduction_freezes_caller_scratch_contract(monkeypatch): + from b12x.moe.fused_moe._tuning import MoeDecodeConfig + + monkeypatch.setattr(tp_moe_impl, "get_num_sm", lambda _device: 188) + weight_plan = _weight_plan( + "w4a16", + source_format="fp4_e8m0_k32", + experts=896, + k=7168, + n=192, + activation="situ", + ) + kwargs = dict( + max_tokens=4096, + core_token_counts=(4096,), + num_topk=16, + route_num_experts=896, + device="cpu", + weight_plan=weight_plan, + quant_mode="w4a16", + decode_config=MoeDecodeConfig( + backend="w4a16", + route_planner="internal", + max_active_clusters=None, + w4a16_route_mode="packed", + ), + ) + monkeypatch.setenv("B12X_W4A16_PREFILL_FUSED_SUM", "0") + materialized_caps = TPMoEScratchCaps(**kwargs) + monkeypatch.setenv("B12X_W4A16_PREFILL_FUSED_SUM", "1") + fused_caps = TPMoEScratchCaps(**kwargs) + calibrated_caps = TPMoEScratchCaps(**kwargs, collect_activation_amax=True) + monkeypatch.setenv("B12X_W4A16_PREFILL_FUSED_SUM", "0") + materialized, fused, calibrated = [ + plan_tp_moe_scratch(caps, prewarm_launches=False) + for caps in (materialized_caps, fused_caps, calibrated_caps) + ] + specs = {spec.name: spec for spec in fused._core_workspace_plan.tensor_specs} + assert specs["intermediate_cache13"].shape == (max(4096 * 16 * 384, 4096 * 7168),) + assert specs["prefill_sum_accum"].shape == (4096 * 7168,) + assert specs["prefill_sum_accum"].dtype == torch.float32 + assert fused._core_workspace_plan.prefill_fused_sum_fp32 + assert not materialized._core_workspace_plan.prefill_fused_sum_fp32 + assert not calibrated._core_workspace_plan.prefill_fused_sum_fp32 + assert ( + fused.layout.core_workspace_nbytes < materialized.layout.core_workspace_nbytes + ) + assert calibrated.layout.core_workspace_nbytes == materialized.layout.core_workspace_nbytes + + def test_dynamic_deterministic_output_is_opt_in( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -633,6 +698,9 @@ def test_trellis_scratch_plan_resolves_default_route_block( device="cpu", weight_plan=weight_plan, quant_mode="w4a16", + decode_config=MoeDecodeConfig( + backend="w4a16", route_planner="internal", max_active_clusters=None, + ), ) plan = plan_tp_moe_scratch(caps) @@ -642,6 +710,33 @@ def test_trellis_scratch_plan_resolves_default_route_block( assert 4096 not in plan.layout.core_token_counts +@pytest.mark.parametrize( + "shared_memory, expected_rows", [(101376, 48), (65536, 16), (131072, 64)] +) +def test_trellis_route_block_respects_fixed_tile_shared_memory( + monkeypatch: pytest.MonkeyPatch, shared_memory: int, expected_rows: int +) -> None: + """Packed projection tiles constrain route size before launches are compiled.""" + monkeypatch.setattr( + torch.cuda, "get_device_properties", + lambda _device: SimpleNamespace(shared_memory_per_block_optin=shared_memory), + ) + plan = plan_b12x_fp4_moe_weights( + quant_modes="w4a16", source_format="exl3", trellis_codebook="lut_e4m3", + activation="silu", params_dtype=torch.bfloat16, + num_experts=2, hidden_size=512, intermediate_size=256, + trellis_bits=2, trellis_tile_config=(128, 128, 128, 128), + ) + caps = TPMoEScratchCaps( + max_tokens=129, num_topk=2, device="cuda:0", weight_plan=plan, + quant_mode="w4a16", + decode_config=MoeDecodeConfig( + backend="w4a16", route_planner="internal", max_active_clusters=None, + ), + ) + assert tp_moe_impl._resolve_trellis_route_block_size(caps) == expected_rows + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") def test_trellis_launch_planner_compiles_fixed_launch_matrix() -> None: """The real planner must cover every fixed decode and route-pack variant.""" diff --git a/tests/moe/test_trellis_direct_lut_decode.py b/tests/moe/test_trellis_direct_lut_decode.py index 77f3129d8..443e4a091 100644 --- a/tests/moe/test_trellis_direct_lut_decode.py +++ b/tests/moe/test_trellis_direct_lut_decode.py @@ -10,6 +10,7 @@ from __future__ import annotations import cuda.bindings.driver as cuda +import cutlass import cutlass.cute as cute import pytest import torch @@ -20,6 +21,7 @@ from b12x._lib.intrinsics import ( packed_decode_lut_e4m3_to_e4m3x8, packed_decode_lut_e4m3_direct_to_e4m3x8, + shared_ptr_to_u32, ) from b12x._lib.quant.lut_e4m3 import ( lut_e4m3_direct_table_cpu, @@ -32,8 +34,9 @@ class _DecodeProbe: - def __init__(self, bits: int): + def __init__(self, bits: int, shared: bool = False): self.bits = int(bits) + self.shared = shared @cute.jit def __call__( @@ -45,7 +48,7 @@ def __call__( stream: cuda.CUstream, ): self.kernel(wins, value_table, direct, out).launch( - grid=(1, 1, 1), block=[32, 1, 1], stream=stream + grid=(cute.size(wins) // 64, 1, 1), block=[32, 1, 1], stream=stream ) @cute.kernel @@ -57,16 +60,26 @@ def kernel( out: cute.Tensor, ): tidx, _, _ = cute.arch.thread_idx() - lane = Int32(tidx) + bidx, _, _ = cute.arch.block_idx() + lane = Int32(tidx) + Int32(bidx) * 32 value_table_addr = value_table.iterator.toint() dir_addr = direct.iterator.toint() + if cutlass.const_expr(self.shared): + table = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Uint8, cute.make_layout(65536), byte_alignment=16 + ) + for offset in range(Int32(tidx), 65536, 32): + table[offset] = direct[(self.bits - 2) * 65536 + offset] + cute.arch.sync_threads() + dir_addr = shared_ptr_to_u32(table.iterator) wa = Uint32(wins[2 * lane]) wb = Uint32(wins[2 * lane + 1]) lo_t, hi_t = packed_decode_lut_e4m3_to_e4m3x8( wa, wb, value_table_addr, self.bits, value_table_in_shared=False ) lo_d, hi_d = packed_decode_lut_e4m3_direct_to_e4m3x8( - wa, wb, dir_addr, self.bits, rate_indexed=True + wa, wb, dir_addr, self.bits, rate_indexed=not self.shared, + in_shared=self.shared, ) out[4 * lane] = Int32(lo_t) out[4 * lane + 1] = Int32(hi_t) @@ -76,7 +89,8 @@ def kernel( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("bits", [2, 3, 4]) -def test_direct_table_decode_bit_equals_value_table(bits: int) -> None: +@pytest.mark.parametrize("shared", [False, True]) +def test_direct_table_decode_bit_equals_value_table(bits: int, shared: bool) -> None: device = torch.device("cuda") torch.manual_seed(20260812 + bits) wins = torch.zeros(64, dtype=torch.int32, device=device) @@ -93,7 +107,7 @@ def args(): current_cuda_stream(), ) - compiled = b12x_compile(_DecodeProbe(bits), *args()) + compiled = b12x_compile(_DecodeProbe(bits, shared), *args()) mismatched = 0 for _ in range(512): wins.copy_( @@ -107,3 +121,25 @@ def args(): mismatched += int((o[:, 0] != o[:, 2]).sum()) mismatched += int((o[:, 1] != o[:, 3]).sum()) assert mismatched == 0, mismatched + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("bits", [2, 3, 4]) +def test_shared_direct_lut_covers_all_codewords(bits: int) -> None: + # Both 16-bit windows visit every codeword; random high bits also exercise + # the shifted windows that cross the low-16-bit boundary. + device = torch.device("cuda") + torch.manual_seed(731 + bits) + states = torch.arange(65536, device=device, dtype=torch.int32) + upper = torch.randint(0, 65536, (65536, 2), device=device, dtype=torch.int32) + wins = ((upper << 16) | states[:, None]).flatten() + value_table = lut_e4m3_value_table_cpu().to(device) + direct = lut_e4m3_direct_table_cpu().to(device) + out = torch.empty(65536 * 4, device=device, dtype=torch.int32) + args = tuple(from_dlpack(tensor, assumed_align=16) + for tensor in (wins, value_table, direct, out)) + (current_cuda_stream(),) + compiled = b12x_compile(_DecodeProbe(bits, shared=True), *args) + compiled(*args) + torch.cuda.synchronize() + result = out.view(-1, 4) + torch.testing.assert_close(result[:, :2], result[:, 2:], rtol=0, atol=0) diff --git a/tests/moe/test_w4a16_e2e.py b/tests/moe/test_w4a16_e2e.py index 230f0cc59..54158ac48 100644 --- a/tests/moe/test_w4a16_e2e.py +++ b/tests/moe/test_w4a16_e2e.py @@ -64,6 +64,176 @@ def test_w4a16_small_m_host_barrier_reset_kill_switch( assert not _small_m_direct_host_barrier_reset_enabled() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("source_format", ["modelopt_nvfp4", "fp4_e8m0_k32"]) +@pytest.mark.parametrize("fused_sum", [False, True]) +@pytest.mark.parametrize("mapped", [False, True]) +@pytest.mark.parametrize("capacity", [14, 129]) +def test_w4a16_prefill_reduction_prepared_capacity_and_graph( + source_format, fused_sum, mapped, capacity, monkeypatch +): + """One declared capacity serves changed inputs and live lengths without JIT.""" + from b12x.moe import fused_moe + from b12x.preparation import PreparationSession, PreparedCall + from b12x.preparation.types import require_prepared + + monkeypatch.setenv("B12X_W4A16_PREFILL_FUSED_SUM", str(int(fused_sum))) + torch.manual_seed(81281) + experts, hidden, intermediate, topk = 8, 256, 192, 4 + x = (torch.randn(capacity, hidden, device="cuda") * 0.125).to(torch.bfloat16) + ids = torch.randint(experts, (capacity, topk), device="cuda", dtype=torch.int32) + expert_map = None + if mapped: + expert_map = torch.arange(experts - 1, -1, -1, dtype=torch.int32, device="cuda") + expert_map[-1] = -1 + weights = torch.softmax(torch.randn(capacity, topk, device="cuda"), dim=-1) + if source_format == "modelopt_nvfp4": + raw = _make_weights( + experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + activation="silu", + ) + else: + raw = ( + torch.randint( + 256, + (experts, intermediate * 2, hidden // 2), + dtype=torch.uint8, + device="cuda", + ), + _pattern_e8m0((experts, intermediate * 2, hidden // 32)), + torch.ones(experts, device="cuda"), + torch.randint( + 256, + (experts, hidden, intermediate // 2), + dtype=torch.uint8, + device="cuda", + ), + _pattern_e8m0((experts, hidden, intermediate // 32), offset=1), + torch.ones(experts, device="cuda"), + ) + weight_plan = fused_moe.plan_weights( + source=fused_moe.PackedSource( + format=fused_moe.PackedSourceFormat(source_format) + ), + activation=fused_moe.ActivationSpec( + mode=fused_moe.ActivationMode.A16, + nonlinearity="silu", + io_dtype=x.dtype, + ), + geometry=fused_moe.MoEGeometry( + num_experts=experts, + hidden_size=hidden, + intermediate_size=intermediate, + ), + ) + prepared = fused_moe.prepare_weights( + plan=weight_plan, + weights=fused_moe.PackedWeights( + w13=raw[0].clone(), + w13_block_scales=raw[1].clone(), + w13_global_scales=raw[2].clone(), + w2=raw[3].clone(), + w2_block_scales=raw[4].clone(), + w2_global_scales=raw[5].clone(), + ), + ) + plan = fused_moe.plan_execution( + experts=prepared, + capacity=fused_moe.ExecutionCapacity(max_tokens=capacity, top_k=topk), + routing=fused_moe.RoutingSpec(), + ) + # Numerical controls must come from the declaration, not from bind-time env. + monkeypatch.setenv("B12X_W4A16_PREFILL_FUSED_SUM", "0") + output = torch.empty_like(x) + + def allocate(state): + return tuple( + torch.empty(s.shape, dtype=s.dtype, device=s.device) + for s in state.scratch.scratch_specs() + ) + + def primer(state): + scratch = allocate(state) + binding = state.bind( + scratch=scratch, + a=x, + experts=prepared, + topk_weights=weights, + topk_ids=ids, + output=output, + route_expert_map=expert_map, + ) + return PreparedCall(run=lambda: state.run(binding), owners=scratch) + + def expected(rows): + oracle_ids, oracle_weights = ids[:rows], weights[:rows] + if expert_map is not None: + oracle_ids = expert_map[oracle_ids.long()] + oracle_weights = oracle_weights.masked_fill(oracle_ids < 0, 0.0) + oracle_ids = oracle_ids.clamp_min(0) + if source_format == "fp4_e8m0_k32": + return moe_reference_w4a16_fp4_e8m0_k32( + x[:rows], + *raw, + oracle_ids, + oracle_weights, + experts, + hidden, + intermediate, + activation="silu", + swiglu_limit=None, + w13_layout="w13", + ) + return _reference_w4a16( + x[:rows], *raw, oracle_ids, oracle_weights, activation="silu" + ) + + with PreparationSession(device=x.device, autotune=False) as session: + request = plan.request(name="prefill-reduction-capacity", prepare_call=primer) + session.prepare((request,)) + state = require_prepared(request.plan, "moe.decode") + assert state.scratch._core_workspace_plan.prefill_fused_sum_fp32 == fused_sum + scratch = allocate(state) + pointers = tuple(t.data_ptr() for t in scratch) + launchers = set() + for rows in (capacity, min(17, capacity), 1, capacity - 1): + oracle = expected(rows) + with kernel_resolution_guard(): + binding = fused_moe.bind( + request.plan, + scratch=scratch, + a=x[:rows], + experts=prepared, + topk_weights=weights[:rows], + topk_ids=ids[:rows], + output=output[:rows], + route_expert_map=expert_map, + ) + actual = binding.run() + assert torch.isfinite(actual).all() and torch.count_nonzero(actual) + _assert_matches_oracle(actual, oracle, activation="silu") + launchers.add(id(binding.fused_launch)) + assert len(launchers) == 1 + graph = torch.cuda.CUDAGraph() + with session.capture(), torch.cuda.graph(graph): + captured = binding.run() + for factor in (-0.5, 2.0): + x.mul_(factor) + oracle = expected(capacity - 1) + for tensor in scratch: + tensor.fill_(255) + before = torch.cuda.memory_allocated() + graph.replay() + torch.cuda.synchronize() + assert torch.cuda.memory_allocated() == before + assert pointers == tuple(t.data_ptr() for t in scratch) + assert torch.isfinite(captured).all() and torch.count_nonzero(captured) + _assert_matches_oracle(captured, oracle, activation="silu") + del graph + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("host_barrier_reset", [False, True]) def test_w4a16_small_m_direct_barrier_modes_eager_and_graph( diff --git a/tests/moe/test_w4a16_mixed_trellis.py b/tests/moe/test_w4a16_mixed_trellis.py index 2bf34c3a7..4d8cc0915 100644 --- a/tests/moe/test_w4a16_mixed_trellis.py +++ b/tests/moe/test_w4a16_mixed_trellis.py @@ -359,6 +359,7 @@ def _serial_tier( topk_ids: torch.Tensor, expert_map: torch.Tensor, block_size_m: int = 8, + activation: str = "silu", ) -> torch.Tensor: m, topk = int(topk_ids.shape[0]), int(topk_ids.shape[1]) buffers = make_w4a16_packed_buffers( @@ -378,7 +379,7 @@ def _serial_tier( prepared, topk_weights, topk_ids, - activation="silu", + activation=activation, intermediate_cache13=buffers.intermediate_cache13, intermediate_cache2=buffers.intermediate_cache2, output=buffers.output, diff --git a/tests/moe/test_w4a16_route_pack.py b/tests/moe/test_w4a16_route_pack.py index 2a398ab2e..b136dab0c 100644 --- a/tests/moe/test_w4a16_route_pack.py +++ b/tests/moe/test_w4a16_route_pack.py @@ -154,14 +154,19 @@ def test_route_pack_reuses_provided_fixed_capacity_for_prefill_tail() -> None: @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("capacity", (128, 8192)) +@pytest.mark.parametrize("capacity, bucket_tokens", ((128, True), (129, False), (8192, True))) @pytest.mark.parametrize("ids_dtype", (torch.int32, torch.int64)) @pytest.mark.parametrize("mapped", (False, True)) -def test_prepared_route_pack_uses_live_bounds_with_fixed_geometry(capacity, ids_dtype, mapped): +@pytest.mark.parametrize("stable", (False, True)) +def test_prepared_route_pack_uses_live_bounds_with_fixed_geometry(capacity, bucket_tokens, ids_dtype, mapped, stable): launches = route_pack_module.compile_w4a16_route_pack_launches( tokens=capacity, topk=2, block_size=8, num_experts=32, ordinal=torch.cuda.current_device(), + stable_order=stable, + bucket_tokens=bucket_tokens, ) + if not bucket_tokens: + assert launches.numel_capacity == capacity * 2 buffers = dict( packed_route_indices=torch.empty(launches.max_packed_routes, dtype=torch.int32, device="cuda"), block_expert_ids=torch.empty(launches.max_route_blocks, dtype=torch.int32, device="cuda"), @@ -194,6 +199,10 @@ def check(rows, output): payload = host_routes[host_routes < rows * 2] expected_payload = torch.nonzero(valid).flatten() torch.testing.assert_close(payload.sort().values, expected_payload, rtol=0, atol=0) + if launches.stable_order: + for expert in range(32): + expert_routes = payload[expected_ids[payload] == expert] + assert torch.equal(expert_routes, expert_routes.sort().values) for block, expert in enumerate(expected_blocks.tolist()): block_routes = host_routes[block * 8:(block + 1) * 8] block_payload = block_routes[block_routes < rows * 2] From 7fca2aa55d5eeb7116a4c09633632cc5f02d850d Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 26 Sep 2026 17:12:23 +0000 Subject: [PATCH 2/4] Share trellis preparation across native and EXL3 containers Normalize uniform EXL3 slots into TrellisSource and TrellisWeights. Preserve codewords, FP16 scale boundaries, global Hadamard coordinates and bounded CUDA staging. Route the compatibility loader through common preparation while retaining asymmetric-pair support. Native config callers keep their schema and projection-tiered path. CPU layout and adapter validation: 98 tests passed. GPU preparation, graph replay, performance and full-model fidelity remain unqualified in this environment. Assisted-by: OpenAI Codex --- b12x/moe/_shared/kernels/w4a16/exl3.py | 416 +++---------------------- b12x/moe/checkpoints/__init__.py | 1 + b12x/moe/checkpoints/exl3.py | 345 ++++++++++++++++++++ b12x/moe/fused_moe/__init__.py | 10 +- b12x/moe/fused_moe/_impl.py | 18 +- b12x/moe/fused_moe/api.py | 15 +- b12x/moe/fused_moe/planning.py | 150 ++++----- b12x/moe/fused_moe/source.py | 74 ++++- b12x/moe/fused_moe/trellis.py | 177 ++++++----- b12x/moe/fused_moe/trellis_layout.py | 168 ++++++++++ b12x/moe/fused_moe/weights.py | 35 +-- docs/exl3-checkpoint-format.md | 6 +- docs/trellis-preparation.md | 105 +++++++ tests/moe/test_exl3_prepare.py | 70 ++++- tests/moe/test_exl3_schema.py | 25 +- tests/moe/test_trellis_adapter.py | 344 ++++++++++++++++++++ 16 files changed, 1343 insertions(+), 616 deletions(-) create mode 100644 b12x/moe/checkpoints/__init__.py create mode 100644 b12x/moe/checkpoints/exl3.py create mode 100644 b12x/moe/fused_moe/trellis_layout.py create mode 100644 docs/trellis-preparation.md create mode 100644 tests/moe/test_trellis_adapter.py diff --git a/b12x/moe/_shared/kernels/w4a16/exl3.py b/b12x/moe/_shared/kernels/w4a16/exl3.py index 26e9216ab..028980b86 100644 --- a/b12x/moe/_shared/kernels/w4a16/exl3.py +++ b/b12x/moe/_shared/kernels/w4a16/exl3.py @@ -1,240 +1,34 @@ -"""EXL3 container reading and W4A16 MoE weight preparation. +"""Low-level EXL3 compatibility and asymmetric-pair preparation. -One metadata-driven load path serves every declared configuration of the -`exl3-v1` container: the manifest and rate tables locate every byte, -the extent rules come from the manifest's layout section, and preparation -reuses the shared trellis machinery (`prepare_trellis256_moe_weights` -for uniform rate structures, the pair finalizer for per-expert ones). -Nothing in this module depends on a specific model geometry, profile -name, or rate-placement convention. +Uniform extents forward through the public checkpoint adapter and common +trellis preparation. Pair extents retain their distinct record addressing; +they are not coerced into the canonical uniform-rate API. """ from __future__ import annotations -import hashlib -import json -import pathlib -from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import torch from b12x.moe._shared.exl3_schema import ( SLOTS_PER_PAIR, - EXL3_MANIFEST_FILENAME, - EXL3_SCHEMA, - Exl3Manifest, RATE_CODE_PAIR_KINDS, - RATE_STRUCTURE_PER_EXPERT_PAIR, RATE_STRUCTURE_UNIFORM, matrix_slot_bytes, rate_code, rate_code_bits, ) -from b12x.moe._shared.kernels.w4a16.prepare import ( - PreparedW4A16MoeWeights, - _finalize_prepared_trellis_weights, - _intermediate_hadamard_signs, - _restore_plane_words, - prepare_trellis256_moe_weights, -) - -_EXPERT_CHUNK = 64 - - -@dataclass(frozen=True) -class Exl3Layer: - """One layer's extent-sliced content, validated against the manifest.""" - - manifest: Exl3Manifest - layer_index: int - first_slot: int - slot_count: int - codes: torch.Tensor - rotations: torch.Tensor - gate_suh: torch.Tensor - up_suh: torch.Tensor - down_svh: torch.Tensor - rates_fc1: torch.Tensor | None - rates_fc2: torch.Tensor | None - sign_pattern: torch.Tensor | None - - @property - def local_intermediate_size(self) -> int: - return self.slot_count * self.manifest.geometry.slot_channels - - -def read_exl3_manifest(root: str | pathlib.Path) -> Exl3Manifest: - root = pathlib.Path(root) - data = json.loads((root / EXL3_MANIFEST_FILENAME).read_text()) - return Exl3Manifest.from_dict(data) - - -def read_exl3_layer( - root: str | pathlib.Path, - manifest: Exl3Manifest, - layer_index: int, - *, - first_slot: int, - slot_count: int, - verify_sha: bool = False, -) -> Exl3Layer: - """Load one rank extent of one layer as CPU tensors. - - The extent is validated against the manifest's layout declarations and - the safetensors metadata is cross-checked against the manifest before - any tensor is interpreted. - """ - - from safetensors import safe_open - - root = pathlib.Path(root) - manifest.validate_extent(first_slot, slot_count) - if layer_index not in manifest.layers: - raise ValueError(f"EXL3 manifest does not declare layer {layer_index}") - ref = manifest.layers[layer_index] - path = root / ref.file - if verify_sha: - digest = hashlib.sha256(path.read_bytes()).hexdigest() - if digest != ref.sha256: - raise ValueError( - f"EXL3 layer {layer_index} sha256 mismatch: manifest " - f"{ref.sha256}, file {digest}" - ) - - geometry = manifest.geometry - per_expert = manifest.rates.structure == RATE_STRUCTURE_PER_EXPERT_PAIR - with safe_open(str(path), framework="pt") as handle: - metadata = handle.metadata() or {} - expected = { - "schema": EXL3_SCHEMA, - "codebook": manifest.codebook, - "layer": str(int(layer_index)), - "num_experts": str(geometry.num_experts), - "hidden_size": str(geometry.hidden_size), - "intermediate_size": str(geometry.intermediate_size), - "slot_channels": str(geometry.slot_channels), - } - for key, value in expected.items(): - if metadata.get(key) != value: - raise ValueError( - f"EXL3 layer {layer_index} metadata {key!r} is " - f"{metadata.get(key)!r}; the manifest declares {value!r}" - ) - names = set(handle.keys()) - required = {"codes", "rotations", "gate_suh", "up_suh", "down_svh"} - if per_expert: - required |= {"rates_fc1", "rates_fc2"} - if manifest.hadamard.intermediate_hadamard: - required |= {"sign_pattern"} - if names != required: - raise ValueError( - f"EXL3 layer {layer_index} tensors {sorted(names)} do not " - f"match the declared set {sorted(required)}" - ) - - codes_slice = handle.get_slice("codes") - codes_shape = codes_slice.get_shape() - if ( - len(codes_shape) != 2 - or codes_shape[0] != geometry.num_slots - or codes_shape[1] % manifest.layout.row_alignment - ): - raise ValueError( - f"EXL3 layer {layer_index} codes shape {codes_shape} violates " - "the declared geometry or row alignment" - ) - codes = codes_slice[first_slot : first_slot + slot_count] - - rotations_slice = handle.get_slice("rotations") - if tuple(rotations_slice.get_shape()) != ( - geometry.num_slots, - geometry.num_experts, - 3, - geometry.slot_channels, - ): - raise ValueError( - f"EXL3 layer {layer_index} rotations shape is not " - "[num_slots, num_experts, 3, slot_channels]" - ) - rotations = rotations_slice[first_slot : first_slot + slot_count] - - h_shapes = ( - (geometry.num_experts, geometry.hidden_size) - if manifest.hadamard.per_expert_input_rotations - else (geometry.hidden_size,) - ) - sides = {} - for name in ("gate_suh", "up_suh", "down_svh"): - tensor = handle.get_tensor(name) - if tuple(tensor.shape) != h_shapes or tensor.dtype != torch.float16: - raise ValueError( - f"EXL3 layer {layer_index} {name} must be fp16 {h_shapes}" - ) - sides[name] = tensor - rates_fc1 = rates_fc2 = None - if per_expert: - pairs = geometry.num_slots // SLOTS_PER_PAIR - first_pair = first_slot // SLOTS_PER_PAIR - pair_count = slot_count // SLOTS_PER_PAIR - declared = manifest.rates.pair_kinds or frozenset() - observed: set[str] = set() - tables = {} - for name in ("rates_fc1", "rates_fc2"): - table_slice = handle.get_slice(name) - if tuple(table_slice.get_shape()) != ( - pairs, - geometry.num_experts, - ): - raise ValueError( - f"EXL3 layer {layer_index} {name} must be " - "[num_slots/8, num_experts]" - ) - table = table_slice[first_pair : first_pair + pair_count] - for code in table.unique().tolist(): - kind = RATE_CODE_PAIR_KINDS.get(int(code)) - if kind is None: - raise ValueError( - f"EXL3 layer {layer_index} {name} contains " - f"unknown rate code {int(code):#x}" - ) - observed.add(kind) - tables[name] = table - if not observed <= set(declared): - raise ValueError( - f"EXL3 layer {layer_index} rate tables use kinds " - f"{sorted(observed)} outside the declared " - f"{sorted(declared)}" - ) - rates_fc1, rates_fc2 = tables["rates_fc1"], tables["rates_fc2"] +if TYPE_CHECKING: + from b12x.moe._shared.kernels.w4a16.prepare import PreparedW4A16MoeWeights - sign_pattern = None - if manifest.hadamard.intermediate_hadamard: - sign_pattern = handle.get_tensor("sign_pattern") - if ( - tuple(sign_pattern.shape) != (geometry.num_experts,) - or sign_pattern.dtype != torch.uint8 - or bool(torch.any(sign_pattern > 7)) - ): - raise ValueError( - f"EXL3 layer {layer_index} sign_pattern must be " - "uint8[num_experts] in 0..7" - ) - - return Exl3Layer( - manifest=manifest, - layer_index=layer_index, - first_slot=first_slot, - slot_count=slot_count, - codes=codes, - rotations=rotations, - gate_suh=sides["gate_suh"], - up_suh=sides["up_suh"], - down_svh=sides["down_svh"], - rates_fc1=rates_fc1, - rates_fc2=rates_fc2, - sign_pattern=sign_pattern, - ) +from b12x.moe.checkpoints.exl3 import ( + Exl3Layer, + read_exl3_layer as read_exl3_layer, + read_exl3_manifest as read_exl3_manifest, + trellis_from_exl3, +) def _extent_rotation_tables( @@ -265,106 +59,6 @@ def _side(tensor: torch.Tensor) -> torch.Tensor: ) -def _intermediate_hadamard_rotation_rows( - layer: Exl3Layer, intermediate: torch.Tensor, device: torch.device -) -> torch.Tensor: - """Append the intermediate-Hadamard sign rows: [values 3I | pre 2I | post I].""" - - manifest = layer.manifest - if (manifest.hadamard.pre_block, manifest.hadamard.post_block) != ( - 512, - 128, - ): - raise ValueError( - "intermediate-Hadamard EXL3 preparation currently implements pre/post Hadamard " - "blocks (512, 128); the manifest declares " - f"({manifest.hadamard.pre_block}, {manifest.hadamard.post_block})" - ) - assert layer.sign_pattern is not None - experts = manifest.geometry.num_experts - global_i = manifest.geometry.intermediate_size - local = layer.local_intermediate_size - pre_begin = 2 * layer.first_slot * manifest.geometry.slot_channels - post_begin = layer.first_slot * manifest.geometry.slot_channels - signs = torch.empty((experts, 3 * local), dtype=torch.float16) - sign_patterns = layer.sign_pattern - for sign_pattern in sorted(set(int(value) for value in sign_patterns.tolist())): - rows = torch.nonzero(sign_patterns == sign_pattern, as_tuple=False).flatten() - pre = _intermediate_hadamard_signs(2 * global_i, sign_pattern=sign_pattern, axis=1)[ - pre_begin : pre_begin + 2 * local - ] - post = _intermediate_hadamard_signs(global_i, sign_pattern=sign_pattern, axis=2)[ - post_begin : post_begin + local - ] - signs.index_copy_( - 0, - rows, - torch.cat((pre, post)).to(torch.float16).expand(rows.numel(), -1), - ) - return torch.cat( - (intermediate, signs.to(device=device)), dim=1 - ).contiguous() - - -def _uniform_native_tensors( - layer: Exl3Layer, device: torch.device -) -> tuple[torch.Tensor, torch.Tensor]: - """Assemble native trellis_t256 tensors from a uniform-rate extent. - - Returns projection-major FC1 ``[2, E, H/16, I_local/16, 16*bits]`` and - FC2 ``[E, I_local/16, H/16, 16*bits]`` int16 tensors. - """ - - manifest = layer.manifest - geometry = manifest.geometry - bits = manifest.rates.bits - assert bits is not None - experts = geometry.num_experts - hidden_tiles = geometry.hidden_size // 16 - slots = layer.slot_count - section = matrix_slot_bytes(geometry.hidden_size, bits, bits) - payload = experts * 3 * section - if layer.codes.shape[1] < payload: - raise ValueError("EXL3 codes rows are shorter than their expert bundles") - if bool(torch.any(layer.codes[:, payload:] != 0)): - raise ValueError("EXL3 codes row padding must be zero") - - w13 = torch.empty( - (2, experts, hidden_tiles, 2 * slots, 16 * bits), - dtype=torch.int16, - device=device, - ) - w2 = torch.empty( - (experts, 2 * slots, hidden_tiles, 16 * bits), - dtype=torch.int16, - device=device, - ) - bundles = layer.codes[:, :payload].reshape(slots, experts, 3 * section) - for first in range(0, experts, _EXPERT_CHUNK): - count = min(_EXPERT_CHUNK, experts - first) - chunk = ( - bundles[:, first : first + count] - .contiguous() - .to(device=device) - .view(torch.int16) - .reshape(slots, count, 3, 2, hidden_tiles, 16 * bits) - ) - for matrix in range(2): - # FC1 planes are the slot's two consecutive N16 columns. - w13[matrix, first : first + count].copy_( - chunk[:, :, matrix].permute(1, 3, 0, 2, 4).reshape( - count, hidden_tiles, 2 * slots, 16 * bits - ) - ) - # FC2 planes are the slot's two consecutive K16 rows. - w2[first : first + count].copy_( - chunk[:, :, 2].permute(1, 0, 2, 3, 4).reshape( - count, 2 * slots, hidden_tiles, 16 * bits - ) - ) - return w13, w2 - - def prepare_exl3_moe_weights( layer: Exl3Layer, *, @@ -375,79 +69,34 @@ def prepare_exl3_moe_weights( dummy_scale: torch.Tensor | None = None, workspace: torch.Tensor | None = None, ) -> PreparedW4A16MoeWeights: - """Prepare one EXL3 rank extent for the fused W4A16 serving path.""" + """Prepare an EXL3 extent; uniform payloads use common trellis preparation.""" manifest = layer.manifest device = torch.device(device) - gate_suh, up_suh, down_svh, intermediate = _extent_rotation_tables( - layer, device - ) - rotations = intermediate - if manifest.hadamard.intermediate_hadamard: - rotations = _intermediate_hadamard_rotation_rows(layer, intermediate, device) - # The intermediate-Hadamard residual transform interleaves gate/up into one - # length-2I axis whose two stored halves each carry one input-side - # table; both physical FC1 slots of a rank use the half its extent - # lies in. - pre_half_slots = manifest.geometry.num_slots // 2 - source_suh = gate_suh if layer.first_slot < pre_half_slots else up_suh - gate_suh = source_suh - up_suh = source_suh - if manifest.rates.structure == RATE_STRUCTURE_UNIFORM: - assert manifest.rates.bits is not None - w13, w2 = _uniform_native_tensors(layer, device) - if tile_config is None: - tile_config = ( - (128, 128, 128, 128) - if manifest.hadamard.intermediate_hadamard and manifest.rates.bits == 2 - else (64, 256, 64, 256) - ) - prepared = prepare_trellis256_moe_weights( - w13=w13, - w2=w2, + from b12x.moe.fused_moe.trellis import prepare_trellis_weights + + source, weights = trellis_from_exl3(layer) + return prepare_trellis_weights( + source, weights, + activation=activation, params_dtype=params_dtype, + num_experts=manifest.geometry.num_experts, hidden_size=manifest.geometry.hidden_size, intermediate_size=layer.local_intermediate_size, - num_experts=manifest.geometry.num_experts, - activation=activation, - fc1_tile_n=tile_config[1], - fc2_tile_n=tile_config[3], device=device, - params_dtype=params_dtype, - w13_layout="trellis_t256_proj", - trellis_bits=manifest.rates.bits, - codebook=manifest.codebook, - gate_suh=gate_suh, - up_suh=up_suh, - intermediate_rotations=intermediate, - down_svh=down_svh, - tile_config=tile_config, - dummy_scale=dummy_scale, - workspace=workspace, - ) - if not manifest.hadamard.intermediate_hadamard: - return prepared - assert prepared.trellis is not None - return replace( - prepared, - trellis=replace( - prepared.trellis, - intermediate_hadamard=True, - intermediate_rotations=rotations, + tile_config=tile_config or ( + (128, 128, 128, 128) + if manifest.hadamard.intermediate_hadamard and manifest.rates.bits == 2 + else (64, 256, 64, 256) ), + dummy_scale=dummy_scale, workspace=workspace, ) - + gate, up, down, rotations = _extent_rotation_tables(layer, device) return _prepare_exl3_pair_extent( - layer, - device=device, - gate_suh=gate_suh, - up_suh=up_suh, - down_svh=down_svh, - rotations=rotations, - params_dtype=params_dtype, + layer, device=device, gate_suh=gate, up_suh=up, down_svh=down, + rotations=rotations, params_dtype=params_dtype, tile_config=tile_config or (64, 256, 64, 256), - dummy_scale=dummy_scale, - workspace=workspace, + dummy_scale=dummy_scale, workspace=workspace, ) @@ -471,6 +120,11 @@ def _prepare_exl3_pair_extent( is exactly one pair of slots. """ + from b12x.moe._shared.kernels.w4a16.prepare import ( + _finalize_prepared_trellis_weights, + _restore_plane_words, + ) + manifest = layer.manifest geometry = manifest.geometry if layer.slot_count != SLOTS_PER_PAIR: diff --git a/b12x/moe/checkpoints/__init__.py b/b12x/moe/checkpoints/__init__.py new file mode 100644 index 000000000..600101df9 --- /dev/null +++ b/b12x/moe/checkpoints/__init__.py @@ -0,0 +1 @@ +"""Checkpoint adapters producing container-independent fused-MoE inputs.""" diff --git a/b12x/moe/checkpoints/exl3.py b/b12x/moe/checkpoints/exl3.py new file mode 100644 index 000000000..db38fcd5a --- /dev/null +++ b/b12x/moe/checkpoints/exl3.py @@ -0,0 +1,345 @@ +"""Normalize uniform ``exl3-v1`` slot extents without requantization. + +Only this adapter interprets manifest row padding and interleaved FC1 side +tables. Preparation and execution consume ordinary TrellisSource/Weights. +Per-expert-pair containers remain available through the low-level reader but +are not accepted by this uniform-rate adapter. +""" + +from __future__ import annotations + +import hashlib +import json +import pathlib +from dataclasses import dataclass + +import torch + +from b12x.moe._shared.exl3_schema import ( + SLOTS_PER_PAIR, + EXL3_MANIFEST_FILENAME, + EXL3_SCHEMA, + Exl3Manifest, + RATE_CODE_PAIR_KINDS, + RATE_STRUCTURE_PER_EXPERT_PAIR, +) +from b12x.moe.fused_moe.config import TrellisConfig +from b12x.moe.fused_moe.source import TrellisExtent, TrellisSource +from b12x.moe.fused_moe.weights import ScaleFactors, TrellisWeights + + +@dataclass(frozen=True) +class Exl3Layer: + """One layer's extent-sliced content, validated against the manifest.""" + + manifest: Exl3Manifest + layer_index: int + first_slot: int + slot_count: int + codes: torch.Tensor + rotations: torch.Tensor + gate_suh: torch.Tensor + up_suh: torch.Tensor + down_svh: torch.Tensor + rates_fc1: torch.Tensor | None + rates_fc2: torch.Tensor | None + sign_pattern: torch.Tensor | None + + @property + def local_intermediate_size(self) -> int: + return self.slot_count * self.manifest.geometry.slot_channels + + +def read_exl3_manifest(root: str | pathlib.Path) -> Exl3Manifest: + root = pathlib.Path(root) + data = json.loads((root / EXL3_MANIFEST_FILENAME).read_text()) + return Exl3Manifest.from_dict(data) + + +def read_exl3_layer( + root: str | pathlib.Path, + manifest: Exl3Manifest, + layer_index: int, + *, + first_slot: int, + slot_count: int, + verify_sha: bool = False, +) -> Exl3Layer: + """Load one rank extent of one layer as CPU tensors. + + The extent is validated against the manifest's layout declarations and + the safetensors metadata is cross-checked against the manifest before + any tensor is interpreted. + """ + + from safetensors import safe_open + + root = pathlib.Path(root) + manifest.validate_extent(first_slot, slot_count) + if layer_index not in manifest.layers: + raise ValueError(f"EXL3 manifest does not declare layer {layer_index}") + ref = manifest.layers[layer_index] + path = root / ref.file + if verify_sha: + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != ref.sha256: + raise ValueError( + f"EXL3 layer {layer_index} sha256 mismatch: manifest " + f"{ref.sha256}, file {digest}" + ) + + geometry = manifest.geometry + per_expert = manifest.rates.structure == RATE_STRUCTURE_PER_EXPERT_PAIR + with safe_open(str(path), framework="pt") as handle: + metadata = handle.metadata() or {} + expected = { + "schema": EXL3_SCHEMA, + "codebook": manifest.codebook, + "layer": str(int(layer_index)), + "num_experts": str(geometry.num_experts), + "hidden_size": str(geometry.hidden_size), + "intermediate_size": str(geometry.intermediate_size), + "slot_channels": str(geometry.slot_channels), + } + for key, value in expected.items(): + if metadata.get(key) != value: + raise ValueError( + f"EXL3 layer {layer_index} metadata {key!r} is " + f"{metadata.get(key)!r}; the manifest declares {value!r}" + ) + names = set(handle.keys()) + required = {"codes", "rotations", "gate_suh", "up_suh", "down_svh"} + if per_expert: + required |= {"rates_fc1", "rates_fc2"} + if manifest.hadamard.intermediate_hadamard: + required |= {"sign_pattern"} + if names != required: + raise ValueError( + f"EXL3 layer {layer_index} tensors {sorted(names)} do not " + f"match the declared set {sorted(required)}" + ) + + codes_slice = handle.get_slice("codes") + codes_shape = codes_slice.get_shape() + if ( + len(codes_shape) != 2 + or codes_shape[0] != geometry.num_slots + or codes_shape[1] % manifest.layout.row_alignment + ): + raise ValueError( + f"EXL3 layer {layer_index} codes shape {codes_shape} violates " + "the declared geometry or row alignment" + ) + codes = codes_slice[first_slot : first_slot + slot_count] + + rotations_slice = handle.get_slice("rotations") + if tuple(rotations_slice.get_shape()) != ( + geometry.num_slots, + geometry.num_experts, + 3, + geometry.slot_channels, + ): + raise ValueError( + f"EXL3 layer {layer_index} rotations shape is not " + "[num_slots, num_experts, 3, slot_channels]" + ) + rotations = rotations_slice[first_slot : first_slot + slot_count] + + h_shapes = ( + (geometry.num_experts, geometry.hidden_size) + if manifest.hadamard.per_expert_input_rotations + else (geometry.hidden_size,) + ) + sides = {} + for name in ("gate_suh", "up_suh", "down_svh"): + tensor = handle.get_tensor(name) + if tuple(tensor.shape) != h_shapes or tensor.dtype != torch.float16: + raise ValueError( + f"EXL3 layer {layer_index} {name} must be fp16 {h_shapes}" + ) + sides[name] = tensor + + rates_fc1 = rates_fc2 = None + if per_expert: + pairs = geometry.num_slots // SLOTS_PER_PAIR + first_pair = first_slot // SLOTS_PER_PAIR + pair_count = slot_count // SLOTS_PER_PAIR + declared = manifest.rates.pair_kinds or frozenset() + observed: set[str] = set() + tables = {} + for name in ("rates_fc1", "rates_fc2"): + table_slice = handle.get_slice(name) + if tuple(table_slice.get_shape()) != ( + pairs, + geometry.num_experts, + ): + raise ValueError( + f"EXL3 layer {layer_index} {name} must be " + "[num_slots/8, num_experts]" + ) + table = table_slice[first_pair : first_pair + pair_count] + for code in table.unique().tolist(): + kind = RATE_CODE_PAIR_KINDS.get(int(code)) + if kind is None: + raise ValueError( + f"EXL3 layer {layer_index} {name} contains " + f"unknown rate code {int(code):#x}" + ) + observed.add(kind) + tables[name] = table + if not observed <= set(declared): + raise ValueError( + f"EXL3 layer {layer_index} rate tables use kinds " + f"{sorted(observed)} outside the declared " + f"{sorted(declared)}" + ) + rates_fc1, rates_fc2 = tables["rates_fc1"], tables["rates_fc2"] + + sign_pattern = None + if manifest.hadamard.intermediate_hadamard: + sign_pattern = handle.get_tensor("sign_pattern") + if ( + tuple(sign_pattern.shape) != (geometry.num_experts,) + or sign_pattern.dtype != torch.uint8 + or bool(torch.any(sign_pattern > 7)) + ): + raise ValueError( + f"EXL3 layer {layer_index} sign_pattern must be " + "uint8[num_experts] in 0..7" + ) + + return Exl3Layer( + manifest=manifest, + layer_index=layer_index, + first_slot=first_slot, + slot_count=slot_count, + codes=codes, + rotations=rotations, + gate_suh=sides["gate_suh"], + up_suh=sides["up_suh"], + down_svh=sides["down_svh"], + rates_fc1=rates_fc1, + rates_fc2=rates_fc2, + sign_pattern=sign_pattern, + ) + + +def trellis_from_exl3(layer: Exl3Layer) -> tuple[TrellisSource, TrellisWeights]: + """Validate a CPU extent and expose codeword views plus FP16 scale tables. + + The payload is not moved to CUDA here. ``prepare_weights(device=...)`` + controls destination ownership and bounded staging. No codeword or scale + is numerically converted. + """ + + if not isinstance(layer, Exl3Layer): + raise TypeError("layer must be an Exl3Layer") + manifest = layer.manifest + manifest.validate_extent(layer.first_slot, layer.slot_count) + if manifest.rates.structure != "uniform": + raise NotImplementedError("the trellis adapter requires uniform EXL3 rates") + bits = manifest.rates.bits + assert bits is not None + geometry = manifest.geometry + experts, hidden = geometry.num_experts, geometry.hidden_size + local = layer.local_intermediate_size + payload = experts * 3 * (hidden // 16) * 64 * bits + codes = layer.codes + if ( + codes.device.type != "cpu" + or codes.dtype != torch.uint8 + or codes.ndim != 2 + or codes.shape[0] != layer.slot_count + or codes.shape[1] < payload + or codes.stride(1) != 1 + ): + raise ValueError( + "EXL3 codes must be CPU uint8 slot rows containing every expert" + ) + if bool(torch.any(codes[:, payload:] != 0)): + raise ValueError("EXL3 codes row padding must be zero") + if layer.rotations.shape != (layer.slot_count, experts, 3, 32): + raise ValueError("EXL3 rotations must have shape [slots,experts,3,32]") + for name in ("rotations", "gate_suh", "up_suh", "down_svh"): + tensor = getattr(layer, name) + if tensor.dtype != torch.float16 or tensor.device.type != "cpu": + raise ValueError(f"EXL3 {name} must be an FP16 CPU tensor") + per_expert = manifest.hadamard.per_expert_input_rotations + side_shape = (experts, hidden) if per_expert else (hidden,) + if any( + getattr(layer, name).shape != side_shape + for name in ("gate_suh", "up_suh", "down_svh") + ): + raise ValueError(f"EXL3 hidden-axis tables must have shape {side_shape}") + gate, up = layer.gate_suh, layer.up_suh + expert_transform: dict[str, object] = {"kind": "none"} + if manifest.hadamard.intermediate_hadamard: + half = geometry.num_slots // 2 + if layer.first_slot < half < layer.first_slot + layer.slot_count: + raise ValueError( + "EXL3 transformed extent cannot cross the FC1 side-table boundary" + ) + if ( + layer.sign_pattern is None + or layer.sign_pattern.dtype != torch.uint8 + or layer.sign_pattern.device.type != "cpu" + or layer.sign_pattern.shape != (experts,) + or bool(torch.any(layer.sign_pattern > 7)) + ): + raise ValueError("EXL3 sign_pattern must be CPU uint8 [experts] in 0..7") + # The stored length-2I transform interleaves gate/up within each half. + # Both physical FC1 projections therefore use the same source side table. + gate = up = gate if layer.first_slot < half else up + expert_transform = { + "kind": "intermediate_hadamard", + "pre_block_size": manifest.hadamard.pre_block, + "post_block_size": manifest.hadamard.post_block, + "sign_pattern_granularity": "per_expert", + } + elif layer.sign_pattern is not None: + raise ValueError("EXL3 sign_pattern is invalid without an expert transform") + side_granularity = "per_expert" if per_expert else "per_layer" + config = TrellisConfig.from_dict( + { + "version": 2, + "codebook": manifest.codebook, + "rate": {"granularity": "uniform"}, + "scale": { + "input_scales": {"vectors": side_granularity, "gains": "none"}, + "intermediate_scales": {"vectors": "per_expert", "gains": "none"}, + "output_scales": {"vectors": side_granularity, "gains": "none"}, + }, + "transform": { + "projection": {"kind": "scaled_hadamard", "block_size": 128}, + "expert": expert_transform, + }, + } + ) + return ( + TrellisSource( + config=config, + uniform_bits=bits, + extent=TrellisExtent( + global_intermediate_size=geometry.intermediate_size, + first_slot=layer.first_slot, + slot_count=layer.slot_count, + ), + ), + TrellisWeights( + codes=codes[:, :payload], + rate=torch.tensor([bits * 17], dtype=torch.uint8), + input_scales=ScaleFactors( + gate.contiguous() if gate is up else torch.stack((gate, up), dim=-2) + ), + intermediate_scales=ScaleFactors( + layer.rotations.permute(1, 2, 0, 3) + .reshape(experts, 3, local) + .contiguous() + ), + output_scales=ScaleFactors(layer.down_svh.contiguous()), + expert_sign_patterns=layer.sign_pattern, + ), + ) + + +__all__ = ["read_exl3_manifest", "read_exl3_layer", "trellis_from_exl3"] diff --git a/b12x/moe/fused_moe/__init__.py b/b12x/moe/fused_moe/__init__.py index 12aa5d0e7..975605ae1 100644 --- a/b12x/moe/fused_moe/__init__.py +++ b/b12x/moe/fused_moe/__init__.py @@ -61,8 +61,9 @@ group="moe", api_style="planned", entry_points=( - "Exl3Source", - "Exl3Weights", + "TrellisExtent", + "TrellisSource", + "TrellisStaging", "ActivationMode", "ActivationSpec", "Binding", @@ -129,8 +130,9 @@ if TYPE_CHECKING: # static analysis only; runtime resolution is lazy from .api import ( # noqa: F401 - Exl3Source, - Exl3Weights, + TrellisExtent, + TrellisSource, + TrellisStaging, ActivationMode, ActivationSpec, Binding, diff --git a/b12x/moe/fused_moe/_impl.py b/b12x/moe/fused_moe/_impl.py index 9ab7d6eea..514860d99 100644 --- a/b12x/moe/fused_moe/_impl.py +++ b/b12x/moe/fused_moe/_impl.py @@ -6795,17 +6795,20 @@ def prepare_b12x_iq2_xs_weights(*, plan, weights) -> B12XFP4ExpertWeights: def prepare_b12x_trellis_v2_weights( *, plan: MoEWeightPreparationPlan, - config: object, + source: object, weights: object, + device: torch.device | str | None = None, + staging: object | None = None, + rotation_dtype: torch.dtype | None = None, ) -> B12XFP4ExpertWeights: """Prepare canonical v2 trellis tensors into the normal expert owner.""" - from b12x.moe.fused_moe.config import TrellisConfig + from b12x.moe.fused_moe.source import TrellisSource from b12x.moe.fused_moe.trellis import prepare_trellis_weights from b12x.moe.fused_moe.weights import TrellisWeights - if not isinstance(config, TrellisConfig): - raise TypeError("config must be a TrellisConfig") + if not isinstance(source, TrellisSource): + raise TypeError("source must be a TrellisSource") if not isinstance(weights, TrellisWeights): raise TypeError("trellis preparation requires TrellisWeights") params_dtype = { @@ -6815,13 +6818,16 @@ def prepare_b12x_trellis_v2_weights( if params_dtype is None: raise TypeError(f"unsupported trellis activation dtype {plan.io_dtype!r}") value = prepare_trellis_weights( - config, + source, weights, activation=plan.activation, - params_dtype=params_dtype, + params_dtype=rotation_dtype or params_dtype, num_experts=plan.num_experts, hidden_size=plan.hidden_size, intermediate_size=plan.intermediate_size, + device=device, + staging=staging, + tile_config=plan.trellis_tile_config, ) representation = _PreparedWeightRepresentation( quant_mode="w4a16", diff --git a/b12x/moe/fused_moe/api.py b/b12x/moe/fused_moe/api.py index 6c7a68f75..dda73f60b 100644 --- a/b12x/moe/fused_moe/api.py +++ b/b12x/moe/fused_moe/api.py @@ -38,9 +38,9 @@ plan_weights as _plan_weights, prepare_weights as _prepare_weights, ) -from .source import Exl3Source, PackedSource, PackedSourceFormat, W13Layout, WeightSource +from .source import PackedSource, PackedSourceFormat, TrellisExtent, TrellisSource, W13Layout, WeightSource +from .trellis_layout import TrellisStaging from .weights import ( - Exl3Weights, PackedWeights, IQ2XSWeights, BlockQuantWeights, @@ -72,10 +72,12 @@ def plan_weights( def prepare_weights( - *, plan: WeightPlan, weights: PackedWeights | TrellisWeights | IQ2XSWeights | Exl3Weights + *, plan: WeightPlan, weights: PackedWeights | TrellisWeights | IQ2XSWeights, + device: torch.device | str | None = None, + staging: TrellisStaging | None = None, ) -> PreparedExperts: """Prepare the canonical weight representation owned by this layer.""" - return _prepare_weights(plan=plan, weights=weights) + return _prepare_weights(plan=plan, weights=weights, device=device, staging=staging) def plan_execution( @@ -192,8 +194,9 @@ def is_supported(device=None) -> bool: __all__ = [ - "Exl3Source", - "Exl3Weights", + "TrellisExtent", + "TrellisSource", + "TrellisStaging", "ActivationMode", "ActivationSpec", "Binding", diff --git a/b12x/moe/fused_moe/planning.py b/b12x/moe/fused_moe/planning.py index a03284945..06a13a451 100644 --- a/b12x/moe/fused_moe/planning.py +++ b/b12x/moe/fused_moe/planning.py @@ -17,9 +17,9 @@ prepare_b12x_iq2_xs_weights, ) from .config import TrellisConfig -from .source import Exl3Source, PackedSource, WeightSource +from .source import PackedSource, TrellisSource, WeightSource +from .trellis_layout import TrellisStaging from .weights import ( - Exl3Weights, PackedWeights, IQ2XSWeights, PreparedExperts, @@ -42,7 +42,12 @@ class ActivationMode(str, Enum): @dataclass(frozen=True, kw_only=True) class ActivationSpec: - """Activation precision, nonlinearity, and public I/O dtype.""" + """Activation precision, nonlinearity, and public I/O dtype. + + ``rotation_dtype`` selects the internal full-rotation trellis arithmetic, + independently of BF16/FP16 public inputs. None retains the I/O dtype. + Only the FP16 full-rotation implementation is supported explicitly. + """ mode: ActivationMode nonlinearity: str @@ -50,12 +55,15 @@ class ActivationSpec: swiglu_limit: float | None = None swiglu_alpha: float | None = None swiglu_beta: float | None = None + rotation_dtype: torch.dtype | None = None def __post_init__(self) -> None: object.__setattr__(self, "mode", ActivationMode(self.mode)) object.__setattr__(self, "nonlinearity", str(self.nonlinearity).lower()) if self.io_dtype not in {torch.bfloat16, torch.float16}: raise TypeError("io_dtype must be torch.bfloat16 or torch.float16") + if self.rotation_dtype not in (None, torch.float16): + raise TypeError("rotation_dtype must be torch.float16 or None") @dataclass(frozen=True, kw_only=True) @@ -122,11 +130,10 @@ def _packed_recipe(source: PackedSource, mode: ActivationMode) -> str: ) from exc -def _validate_trellis_runtime(source: TrellisConfig) -> None: - if source.codebook.value == "lut_fp16": +def _validate_trellis_runtime(source: TrellisConfig, *, explicit_uniform: bool = False) -> None: + if source.codebook.value == "lut_fp16" and not explicit_uniform: raise NotImplementedError( - "lut_fp16 is defined by the checkpoint schema but is not " - "implemented by the routed fused MoE runtime" + "lut_fp16 preparation requires an explicit uniform_bits declaration" ) if source.rate.group_size is not None: raise NotImplementedError( @@ -177,7 +184,7 @@ def _prepared_format( f"required packing {packing.value!r} is not available; " f"planner produced {sorted(value.value for value in available)}" ) - if isinstance(source, (TrellisConfig, Exl3Source)): + if isinstance(source, TrellisSource): weights = WeightEncoding.TRELLIS scales = ScaleEncoding.TRELLIS_SCALES else: @@ -212,6 +219,10 @@ def plan_weights( constraints = constraints or WeightPlanConstraints() if not isinstance(constraints, WeightPlanConstraints): raise TypeError("constraints must be WeightPlanConstraints") + if isinstance(source, TrellisConfig): + source = TrellisSource(config=source) + if activation.rotation_dtype is not None and not isinstance(source, TrellisSource): + raise ValueError("rotation_dtype is only valid for trellis weights") if isinstance(source, PackedSource): automatic = activation.mode is ActivationMode.AUTO @@ -252,65 +263,37 @@ def plan_weights( w13_layout=source.w13_layout.value, w4a16_layout=requested_layout, ) - elif isinstance(source, Exl3Source): - manifest = source.manifest + elif isinstance(source, TrellisSource): + config = source.config if activation.mode is not ActivationMode.A16: - raise ValueError("EXL3 canonical preparation requires A16 activations") - if constraints.required_packing not in { - None, - WeightPacking.TRELLIS_NATIVE, - }: - raise ValueError("EXL3 weights require trellis_native packing") - if manifest.rates.structure != "uniform": - raise NotImplementedError( - "EXL3 canonical preparation requires uniform rates" - ) - if (geometry.num_experts, geometry.hidden_size) != ( - manifest.geometry.num_experts, - manifest.geometry.hidden_size, + raise ValueError("Trellis fused MoE requires A16 activations") + if constraints.required_packing not in {None, WeightPacking.TRELLIS_NATIVE}: + raise ValueError("Trellis weights require trellis_native packing") + _validate_trellis_runtime(config, explicit_uniform=source.uniform_bits is not None) + expert = config.transform.expert + if source.extent is not None: + if source.extent.intermediate_size != geometry.intermediate_size: + raise ValueError("trellis extent width differs from the weight geometry") + if expert.kind == "intermediate_hadamard" and ( + 32 * source.extent.first_slot % expert.post_block_size + or geometry.intermediate_size % expert.post_block_size + ): + raise ValueError("trellis extent must contain complete post-Hadamard blocks") + if activation.rotation_dtype is not None and ( + config.codebook.value == "mcg" and source.uniform_bits is None ): - raise ValueError( - "EXL3 expert count and hidden width must match the manifest" - ) - if geometry.intermediate_size > manifest.geometry.intermediate_size: - raise ValueError("EXL3 local intermediate width exceeds the source width") + raise ValueError("rotation_dtype requires a uniform full-rotation path") recipe = "w4a16" - raw_plan = plan_b12x_fp4_moe_weights( - quant_modes=recipe, - source_format="exl3", - activation=activation.nonlinearity, - params_dtype=activation.io_dtype, - num_experts=geometry.num_experts, - hidden_size=geometry.hidden_size, - intermediate_size=geometry.intermediate_size, - w13_layout="w31", - w4a16_layout="trellis_native", - trellis_bits=manifest.rates.bits, - trellis_tile_config=( - (128, 128, 128, 128) - if manifest.hadamard.intermediate_hadamard and manifest.rates.bits == 2 - else (64, 256, 64, 256) - ), - intermediate_hadamard=manifest.hadamard.intermediate_hadamard, - trellis_codebook=manifest.codebook, - trellis_rate_granularity=manifest.rates.structure, - intermediate_hadamard_blocks=( - (manifest.hadamard.pre_block, manifest.hadamard.post_block) - if manifest.hadamard.intermediate_hadamard - else None - ), + bits = source.uniform_bits or 3 + # Match the native K2 intermediate-Hadamard and projection-tiered + # geometries; container identity does not select an execution policy. + tile_config = ( + (128, 128, 128, 128) + if expert.kind == "intermediate_hadamard" and bits == 2 + else (128, 256, 64, 256) + if config.codebook.value == "mcg" and source.uniform_bits is None + else (64, 256, 64, 256) ) - elif isinstance(source, TrellisConfig): - if activation.mode is not ActivationMode.A16: - raise ValueError("Trellis fused MoE currently requires A16 activations") - if constraints.required_packing not in { - None, - WeightPacking.TRELLIS_NATIVE, - }: - raise ValueError("Trellis weights require trellis_native packing") - _validate_trellis_runtime(source) - expert = source.transform.expert - recipe = "w4a16" raw_plan = plan_b12x_fp4_moe_weights( quant_modes=recipe, source_format="b12x_trellis", @@ -321,23 +304,18 @@ def plan_weights( intermediate_size=geometry.intermediate_size, w13_layout="w31", w4a16_layout="trellis_native", - trellis_bits=3, - trellis_tile_config=( - (128, 256, 64, 256) - if source.codebook.value == "mcg" - else (64, 256, 64, 256) - ), + trellis_bits=bits, + trellis_tile_config=tile_config, intermediate_hadamard=expert.kind == "intermediate_hadamard", - trellis_codebook=source.codebook.value, - trellis_rate_granularity=source.rate.granularity.value, + trellis_codebook=config.codebook.value, + trellis_rate_granularity=config.rate.granularity.value, intermediate_hadamard_blocks=( - None - if expert.kind == "none" + None if expert.kind == "none" else (expert.pre_block_size, expert.post_block_size) ), ) else: - raise TypeError("source must be a PackedSource, TrellisConfig or Exl3Source") + raise TypeError("source must be a PackedSource, TrellisConfig or TrellisSource") return WeightPlan( source=source, @@ -356,34 +334,30 @@ def plan_weights( def prepare_weights( *, plan: WeightPlan, - weights: PackedWeights | TrellisWeights | IQ2XSWeights | Exl3Weights, + weights: PackedWeights | TrellisWeights | IQ2XSWeights, + device: torch.device | str | None = None, + staging: TrellisStaging | None = None, ) -> PreparedExperts: """Materialize the in-memory representation selected by ``plan_weights``.""" if not isinstance(plan, WeightPlan): raise TypeError("plan must be a WeightPlan") + if (device is not None or staging is not None) and not isinstance(plan.source, TrellisSource): + raise ValueError("device and staging are only supported for trellis preparation") if isinstance(plan.source, PackedSource) and plan.source.format.value in BLOCK_CODECS: if not isinstance(weights, IQ2XSWeights) or weights.codec != plan.source.format.value: raise TypeError(f"{plan.source.format.value} preparation requires matching BlockQuantWeights") prepared = prepare_b12x_iq2_xs_weights(plan=plan._impl, weights=weights) - elif isinstance(plan.source, Exl3Source): - if not isinstance(weights, Exl3Weights): - raise TypeError("EXL3 preparation requires Exl3Weights") - if weights.layer.manifest != plan.source.manifest: - raise ValueError("EXL3 extent manifest differs from its weight plan") - prepared = prepare_b12x_fp4_moe_weights( - plan=plan._impl, - params_dtype=plan.activation.io_dtype, - exl3_layer=weights.layer, - exl3_device=weights.device, - ) - elif isinstance(plan.source, TrellisConfig): + elif isinstance(plan.source, TrellisSource): if not isinstance(weights, TrellisWeights): raise TypeError("Trellis preparation requires TrellisWeights") prepared = prepare_b12x_trellis_v2_weights( plan=plan._impl, - config=plan.source, + source=plan.source, weights=weights, + device=device, + staging=staging, + rotation_dtype=plan.activation.rotation_dtype, ) else: if not isinstance(weights, PackedWeights): diff --git a/b12x/moe/fused_moe/source.py b/b12x/moe/fused_moe/source.py index 0c32132ed..bc85fd2f9 100644 --- a/b12x/moe/fused_moe/source.py +++ b/b12x/moe/fused_moe/source.py @@ -6,8 +6,8 @@ from enum import Enum from typing import TypeAlias -from .._shared.exl3_schema import Exl3Manifest -from .config import TrellisConfig +from .._shared.trellis_codebooks import validate_codebook_bits +from .config import RateGranularity, TrellisConfig class PackedSourceFormat(str, Enum): @@ -42,21 +42,77 @@ def __post_init__(self) -> None: @dataclass(frozen=True, kw_only=True) -class Exl3Source: - """Trellis slot encoding with checkpoint-global transform coordinates.""" +class TrellisExtent: + """Contiguous 32-channel slots in the checkpoint-global intermediate axis. - manifest: Exl3Manifest + Sign sequences are generated at ``global_intermediate_size`` and sliced, + never regenerated using a tensor-parallel rank's shorter local length. + """ + + global_intermediate_size: int + first_slot: int + slot_count: int def __post_init__(self) -> None: - if not isinstance(self.manifest, Exl3Manifest): - raise TypeError("Exl3Source.manifest must be a Exl3Manifest") + for name in ("global_intermediate_size", "first_slot", "slot_count"): + if type(getattr(self, name)) is not int: + raise TypeError(f"TrellisExtent.{name} must be an integer") + if self.global_intermediate_size <= 0 or self.global_intermediate_size % 32: + raise ValueError( + "global_intermediate_size must be a positive multiple of 32" + ) + if self.first_slot < 0 or self.slot_count <= 0: + raise ValueError( + "a trellis extent requires first_slot >= 0 and slot_count > 0" + ) + if 32 * (self.first_slot + self.slot_count) > self.global_intermediate_size: + raise ValueError("trellis extent exceeds the global intermediate axis") + + @property + def intermediate_size(self) -> int: + return 32 * self.slot_count + + +@dataclass(frozen=True, kw_only=True) +class TrellisSource: + """Container-independent trellis encoding and optional global coordinates. + + ``uniform_bits`` declares the scalar rate before weights are transferred; + preparation verifies it against the rate tensor. Omitting it retains the + native config's K3 LUT or K3/K4/K5 projection-tiered MCG contract. + """ + config: TrellisConfig + uniform_bits: int | None = None + extent: TrellisExtent | None = None -WeightSource: TypeAlias = PackedSource | TrellisConfig | Exl3Source + def __post_init__(self) -> None: + if not isinstance(self.config, TrellisConfig): + raise TypeError("TrellisSource.config must be a TrellisConfig") + if self.uniform_bits is not None: + if type(self.uniform_bits) is not int or self.uniform_bits not in ( + 2, + 3, + 4, + 5, + 6, + ): + raise ValueError("uniform_bits must be one of 2, 3, 4, 5, 6") + if self.config.rate.granularity is not RateGranularity.UNIFORM: + raise ValueError("uniform_bits requires uniform rate granularity") + validate_codebook_bits(self.config.codebook.value, self.uniform_bits) + if self.config.codebook.value == "mcg" and self.uniform_bits < 3: + raise ValueError("MCG trellis execution requires at least K3") + if self.extent is not None and not isinstance(self.extent, TrellisExtent): + raise TypeError("TrellisSource.extent must be a TrellisExtent or None") + + +WeightSource: TypeAlias = PackedSource | TrellisConfig | TrellisSource __all__ = [ - "Exl3Source", + "TrellisExtent", + "TrellisSource", "PackedSource", "PackedSourceFormat", "W13Layout", diff --git a/b12x/moe/fused_moe/trellis.py b/b12x/moe/fused_moe/trellis.py index c6a80212c..713acc0be 100644 --- a/b12x/moe/fused_moe/trellis.py +++ b/b12x/moe/fused_moe/trellis.py @@ -3,19 +3,13 @@ from __future__ import annotations from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import torch -from b12x.moe._shared.kernels.w4a16.mixed_trellis import ( - MixedTrellisRotations, - build_projection_tiered_maps, -) -from b12x.moe._shared.kernels.w4a16.prepare import ( - PreparedW4A16MoeWeights, - _intermediate_hadamard_signs, - _finalize_prepared_trellis_weights, - prepare_trellis256_moe_weights, -) +if TYPE_CHECKING: + from b12x.moe._shared.kernels.w4a16.mixed_trellis import MixedTrellisRotations + from b12x.moe._shared.kernels.w4a16.prepare import PreparedW4A16MoeWeights from .config import ( RateGranularity, @@ -25,6 +19,8 @@ TrellisScaleFactorsConfig, ) from .weights import ScaleFactors, TrellisWeights +from .source import TrellisExtent, TrellisSource +from .trellis_layout import TrellisStaging, append_intermediate_signs, assemble_uniform_slots _SLOT_CHANNELS = 32 @@ -137,8 +133,7 @@ def _effective_input_scales( vectors = _selected_scale_tensor( factors.vectors, name="input_scales.vectors", device=device ) - vector_experts = num_experts if _expert_axis(declaration.vectors) else 1 - if vector_experts == 1: + if not _expert_axis(declaration.vectors): if tuple(vectors.shape) == (hidden_size,): vectors = vectors.reshape(1, 1, hidden_size) elif tuple(vectors.shape) == (2, hidden_size): @@ -308,7 +303,9 @@ def _local_rate_matrix( return rate -def _symmetric_bits(config: TrellisConfig, rates: torch.Tensor) -> torch.Tensor: +def _symmetric_bits( + config: TrellisConfig, rates: torch.Tensor, *, uniform_bits: int | None = None +) -> torch.Tensor: host = rates.detach().cpu().to(torch.int64) low = host & 0x0F high = host >> 4 @@ -316,7 +313,9 @@ def _symmetric_bits(config: TrellisConfig, rates: torch.Tensor) -> torch.Tensor: raise ValueError( "fused MoE trellis rates must encode symmetric K values (0xKK)" ) - if config.codebook is TrellisCodebook.MCG: + if uniform_bits is not None: + allowed = (uniform_bits,) + elif config.codebook is TrellisCodebook.MCG: allowed = (3, 4, 5) elif config.codebook is TrellisCodebook.LUT_E4M3: allowed = (3,) @@ -402,22 +401,13 @@ def _intermediate_hadamard_rows( intermediate_size: int, device: torch.device, ) -> torch.Tensor: - sign_patterns_host = sign_patterns.detach().cpu() - signs = torch.empty( - (int(values.shape[0]), 3 * intermediate_size), dtype=torch.float16 + return append_intermediate_signs( + values.to(device=device), sign_patterns, + extent=TrellisExtent( + global_intermediate_size=intermediate_size, + first_slot=0, slot_count=intermediate_size // 32, + ), ) - for sign_pattern in sorted(set(int(value) for value in sign_patterns_host.tolist())): - if not 0 <= sign_pattern < 8: - raise ValueError("expert_sign_patterns values must be in 0..7") - rows = torch.nonzero(sign_patterns_host == sign_pattern, as_tuple=False).flatten() - pre = _intermediate_hadamard_signs(2 * intermediate_size, sign_pattern=sign_pattern, axis=1) - post = _intermediate_hadamard_signs(intermediate_size, sign_pattern=sign_pattern, axis=2) - signs.index_copy_( - 0, - rows, - torch.cat((pre, post)).to(torch.float16).expand(rows.numel(), -1), - ) - return torch.cat((values, signs.to(device=device)), dim=1).contiguous() def _uniform_prepared( @@ -434,49 +424,31 @@ def _uniform_prepared( up_suh: torch.Tensor, intermediate: torch.Tensor, down_svh: torch.Tensor, + extent: TrellisExtent | None = None, + staging: TrellisStaging = TrellisStaging(), + tile_config: tuple[int, int, int, int] = (64, 256, 64, 256), + dummy_scale: torch.Tensor | None = None, + workspace: torch.Tensor | None = None, ) -> PreparedW4A16MoeWeights: unique = sorted(set(bits.reshape(-1).tolist())) if len(unique) != 1: raise ValueError("uniform trellis preparation requires one K value") bit = int(unique[0]) - offsets = _bundle_offsets(bits, hidden_size) - experts = list(range(num_experts)) - gate = _projection_native( - weights.codes, - experts=experts, - projection=0, - bits=bit, - offsets=offsets, - hidden_size=hidden_size, - fc1=True, - ) - up = _projection_native( - weights.codes, - experts=experts, - projection=1, - bits=bit, - offsets=offsets, - hidden_size=hidden_size, - fc1=True, - ) - down = _projection_native( - weights.codes, - experts=experts, - projection=2, - bits=bit, - offsets=offsets, - hidden_size=hidden_size, - fc1=False, + from b12x.moe._shared.kernels.w4a16.prepare import prepare_trellis256_moe_weights + + w13, down = assemble_uniform_slots( + weights.codes, num_experts=num_experts, hidden_size=hidden_size, + bits=bit, device=gate_suh.device, staging=staging, ) prepared = prepare_trellis256_moe_weights( - w13=torch.stack((gate, up)), + w13=w13, w2=down, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, activation=activation, - fc1_tile_n=256, - fc2_tile_n=256, + fc1_tile_n=tile_config[1], + fc2_tile_n=tile_config[3], params_dtype=params_dtype, w13_layout="trellis_t256_proj", trellis_bits=bit, @@ -485,7 +457,9 @@ def _uniform_prepared( up_suh=up_suh, intermediate_rotations=intermediate.reshape(num_experts, -1), down_svh=down_svh, - tile_config=(64, 256, 64, 256), + tile_config=tile_config, + dummy_scale=dummy_scale, + workspace=workspace, ) if config.transform.expert.kind == "none": if weights.expert_sign_patterns is not None: @@ -500,11 +474,11 @@ def _uniform_prepared( sign_patterns, name="expert_sign_patterns", dtype=torch.uint8, - device=weights.codes.device, + device=gate_suh.device, ) if tuple(sign_patterns.shape) != (num_experts,): raise ValueError(f"expert_sign_patterns must be uint8[{num_experts}]") - if torch.count_nonzero(sign_patterns).item(): + if extent is None and torch.count_nonzero(sign_patterns).item(): raise NotImplementedError( "rank-local intermediate_hadamard preparation currently requires " "all-zero expert_sign_patterns" @@ -515,11 +489,13 @@ def _uniform_prepared( trellis=replace( prepared.trellis, intermediate_hadamard=True, - intermediate_rotations=_intermediate_hadamard_rows( + intermediate_rotations=append_intermediate_signs( intermediate.reshape(num_experts, -1), sign_patterns, - intermediate_size=intermediate_size, - device=weights.codes.device, + extent=extent or TrellisExtent( + global_intermediate_size=intermediate_size, + first_slot=0, slot_count=intermediate_size // 32, + ), ), ), ) @@ -539,6 +515,13 @@ def _projection_prepared( intermediate: torch.Tensor, down_svh: torch.Tensor, ) -> PreparedProjectionTrellisWeights: + from b12x.moe._shared.kernels.w4a16.mixed_trellis import ( + MixedTrellisRotations, build_projection_tiered_maps, + ) + from b12x.moe._shared.kernels.w4a16.prepare import ( + _finalize_prepared_trellis_weights, + ) + if config.codebook is not TrellisCodebook.MCG: raise ValueError("projection-tiered trellis execution requires codebook 'mcg'") if config.transform.expert.kind != "none": @@ -707,7 +690,7 @@ def _projection_prepared( def prepare_trellis_weights( - config: TrellisConfig, + source: TrellisSource | TrellisConfig, weights: TrellisWeights, *, activation: str, @@ -715,16 +698,39 @@ def prepare_trellis_weights( num_experts: int, hidden_size: int, intermediate_size: int, + device: torch.device | str | None = None, + staging: TrellisStaging | None = None, + tile_config: tuple[int, int, int, int] | None = None, + dummy_scale: torch.Tensor | None = None, + workspace: torch.Tensor | None = None, ) -> PreparedW4A16MoeWeights | PreparedProjectionTrellisWeights: """Prepare one rank-local canonical Trellis MoE layer.""" + if isinstance(source, TrellisConfig): + source = TrellisSource(config=source) + config = source.config codes = weights.codes if codes.dtype != torch.uint8: raise TypeError(f"codes must be torch.uint8, got {codes.dtype}") - if codes.device.type != "cuda": - raise ValueError("trellis preparation requires CUDA-resident codes") - if codes.ndim != 2 or not codes.is_contiguous(): - raise ValueError("codes must be contiguous uint8 [I_local/32,row_stride]") + device = torch.device(device) if device is not None else codes.device + if device.type != "cuda": + raise ValueError("CPU trellis payloads require an explicit CUDA destination") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + if codes.device.type != "cpu" and codes.device != device: + raise ValueError("trellis codes must reside on CPU or the destination CUDA device") + if codes.ndim != 2 or codes.stride(1) != 1: + raise ValueError("codes must have contiguous uint8 logical rows") + uniform = source.uniform_bits is not None or config.codebook is not TrellisCodebook.MCG + if not uniform and (codes.device != device or not codes.is_contiguous()): + raise ValueError("projection-tiered MCG preparation requires contiguous CUDA codes") + if not uniform and staging is not None: + raise ValueError("bounded host staging is supported only for uniform rates") + staging = staging or TrellisStaging() + if not isinstance(staging, TrellisStaging): + raise TypeError("staging must be a TrellisStaging") + if source.extent is not None and source.extent.intermediate_size != intermediate_size: + raise ValueError("trellis extent width differs from the local intermediate size") if int(codes.shape[0]) * _SLOT_CHANNELS != intermediate_size: raise ValueError( "codes first dimension must equal intermediate_size/32: " @@ -737,11 +743,10 @@ def prepare_trellis_weights( "fused MoE trellis preparation implements scaled_hadamard(128)" ) - device = codes.device rates = _local_rate_matrix( - config, weights.rate, num_experts=num_experts, device=device + config, weights.rate, num_experts=num_experts, device=weights.rate.device ) - bits = _symmetric_bits(config, rates) + bits = _symmetric_bits(config, rates, uniform_bits=source.uniform_bits) offsets = _bundle_offsets(bits, hidden_size) required_row_bytes = max( offset @@ -755,6 +760,23 @@ def prepare_trellis_weights( f"canonical projection payload ({required_row_bytes} bytes)" ) + def move_tensor(tensor: torch.Tensor | None) -> torch.Tensor | None: + if tensor is None: + return None + if tensor.device.type != "cpu" and tensor.device != device: + raise ValueError("trellis metadata must reside on CPU or the destination device") + return tensor.to(device=device) + + def move_scales(factors: ScaleFactors) -> ScaleFactors: + return ScaleFactors(move_tensor(factors.vectors), move_tensor(factors.gains)) + + weights = replace( + weights, + input_scales=move_scales(weights.input_scales), + intermediate_scales=move_scales(weights.intermediate_scales), + output_scales=move_scales(weights.output_scales), + expert_sign_patterns=move_tensor(weights.expert_sign_patterns), + ) scale = config.scale gate_suh, up_suh = _effective_input_scales( weights.input_scales, @@ -778,7 +800,7 @@ def prepare_trellis_weights( device=device, ) - if config.codebook is not TrellisCodebook.MCG: + if uniform: return _uniform_prepared( config, weights, @@ -792,6 +814,11 @@ def prepare_trellis_weights( up_suh=up_suh, intermediate=intermediate, down_svh=down_svh, + extent=source.extent, + staging=staging, + tile_config=tile_config or (64, 256, 64, 256), + dummy_scale=dummy_scale, + workspace=workspace, ) return _projection_prepared( config, diff --git a/b12x/moe/fused_moe/trellis_layout.py b/b12x/moe/fused_moe/trellis_layout.py new file mode 100644 index 000000000..373f54980 --- /dev/null +++ b/b12x/moe/fused_moe/trellis_layout.py @@ -0,0 +1,168 @@ +"""Lossless trellis slot assembly and checkpoint-global sign coordinates. + +These load-time tensor operations have no kernel/compiler dependency. CPU +execution is useful for format validation; serving prepares onto CUDA before +any execution plan or graph is built. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from .source import TrellisExtent + + +@dataclass(frozen=True, kw_only=True) +class TrellisStaging: + """Limits on a temporary uniform-codeword batch, excluding final weights. + + Transfers contain at most ``max_experts`` experts and ``max_bytes`` bytes. + A single expert larger than the byte limit is rejected before allocation. + Permuted copy operands can add up to one projection (one third of a batch) + to this bound. Scale/sign tables and CUDA allocator overhead are separate. + """ + + max_experts: int = 64 + max_bytes: int = 64 * 1024 * 1024 + + def __post_init__(self) -> None: + for name in ("max_experts", "max_bytes"): + if type(getattr(self, name)) is not int or getattr(self, name) <= 0: + raise ValueError(f"TrellisStaging.{name} must be a positive integer") + + def batch_size(self, *, slots: int, hidden_size: int, bits: int) -> int: + expert_bytes = slots * 3 * (hidden_size // 16) * 64 * bits + if expert_bytes <= 0: + raise ValueError("trellis staging geometry must be positive") + count = min(self.max_experts, self.max_bytes // expert_bytes) + if count == 0: + raise ValueError( + f"one trellis expert requires {expert_bytes} staging bytes; " + f"max_bytes={self.max_bytes} is insufficient" + ) + return count + + +def assemble_uniform_slots( + codes: torch.Tensor, + *, + num_experts: int, + hidden_size: int, + bits: int, + device: torch.device | str, + staging: TrellisStaging = TrellisStaging(), +) -> tuple[torch.Tensor, torch.Tensor]: + """Assemble FC1/FC2 int16 carriers without decoding or rounding weights. + + Input rows are slot-major expert bundles, each containing gate, up, down. + The uint8 payload may have a larger physical row stride, but no logical + padding columns. CPU payloads are copied in bounded expert batches. + """ + + device = torch.device(device) + if codes.dtype != torch.uint8 or codes.ndim != 2 or codes.shape[0] == 0: + raise ValueError("codes must be nonempty uint8 [slots,payload_bytes]") + if hidden_size <= 0 or hidden_size % 32 or num_experts <= 0: + raise ValueError("trellis assembly requires E > 0 and H divisible by 32") + if bits not in (2, 3, 4, 5, 6): + raise ValueError("trellis assembly supports symmetric rates K2 through K6") + if codes.device.type not in ("cpu", "cuda"): + raise ValueError("trellis payload must reside on CPU or CUDA") + if codes.device.type == "cuda" and codes.device != device: + raise ValueError("CUDA trellis payload must reside on its destination device") + slots = codes.shape[0] + hidden_tiles = hidden_size // 16 + section = hidden_tiles * 64 * bits + payload = num_experts * 3 * section + if codes.shape[1] != payload or codes.stride(1) != 1: + raise ValueError(f"codes must have {payload} contiguous bytes per logical row") + batch_size = staging.batch_size(slots=slots, hidden_size=hidden_size, bits=bits) + w13 = torch.empty( + (2, num_experts, hidden_tiles, 2 * slots, 16 * bits), + dtype=torch.int16, + device=device, + ) + w2 = torch.empty( + (num_experts, 2 * slots, hidden_tiles, 16 * bits), + dtype=torch.int16, + device=device, + ) + bundles = codes.unflatten(1, (num_experts, 3 * section)) + for first in range(0, num_experts, batch_size): + count = min(batch_size, num_experts - first) + chunk = ( + bundles[:, first : first + count] + .contiguous() + .to(device=device) + .view(torch.int16) + .reshape(slots, count, 3, 2, hidden_tiles, 16 * bits) + ) + for matrix in range(2): + w13[matrix, first : first + count].copy_( + chunk[:, :, matrix] + .permute(1, 3, 0, 2, 4) + .reshape(count, hidden_tiles, 2 * slots, 16 * bits) + ) + w2[first : first + count].copy_( + chunk[:, :, 2] + .permute(1, 0, 2, 3, 4) + .reshape(count, 2 * slots, hidden_tiles, 16 * bits) + ) + # Do not retain the previous staging allocation during the next H2D copy. + del chunk + return w13, w2 + + +def intermediate_hadamard_signs( + length: int, *, sign_pattern: int, axis: int +) -> torch.Tensor: + """Deterministic CPU sign sequence used by the intermediate transform.""" + + if length <= 0 or not 0 <= sign_pattern < 8: + raise ValueError("sign generation requires length > 0 and pattern in 0..7") + if sign_pattern == 0: + return torch.ones(length, dtype=torch.float32) + generator = torch.Generator(device="cpu") + generator.manual_seed( + (0x6A09E667F3BCC909 * sign_pattern + 0xBB67AE8584CAA73B * axis) + & ((1 << 63) - 1) + ) + return torch.randint(0, 2, (length,), generator=generator).mul_(2).sub_(1).float() + + +def append_intermediate_signs( + values: torch.Tensor, + sign_patterns: torch.Tensor, + *, + extent: TrellisExtent, +) -> torch.Tensor: + """Append [pre 2I | post I] signs sliced in checkpoint-global coordinates.""" + + local = extent.intermediate_size + if ( + values.ndim != 2 + or values.shape[1] != 3 * local + or values.dtype != torch.float16 + ): + raise ValueError("intermediate values must be fp16 [experts,3*I_local]") + if sign_patterns.dtype != torch.uint8 or sign_patterns.shape != (values.shape[0],): + raise ValueError("expert_sign_patterns must be uint8 [experts]") + patterns = sign_patterns.detach().cpu() + if bool(torch.any(patterns > 7)): + raise ValueError("expert_sign_patterns values must be in 0..7") + signs = torch.empty((values.shape[0], 3 * local), dtype=torch.float16) + begin = 32 * extent.first_slot + for pattern in sorted(set(patterns.tolist())): + rows = torch.nonzero(patterns == pattern, as_tuple=False).flatten() + pre = intermediate_hadamard_signs( + 2 * extent.global_intermediate_size, sign_pattern=pattern, axis=1 + )[2 * begin : 2 * (begin + local)] + post = intermediate_hadamard_signs( + extent.global_intermediate_size, sign_pattern=pattern, axis=2 + )[begin : begin + local] + signs.index_copy_( + 0, rows, torch.cat((pre, post)).half().expand(rows.numel(), -1) + ) + return torch.cat((values, signs.to(device=values.device)), dim=1).contiguous() diff --git a/b12x/moe/fused_moe/weights.py b/b12x/moe/fused_moe/weights.py index be6951e7d..571585693 100644 --- a/b12x/moe/fused_moe/weights.py +++ b/b12x/moe/fused_moe/weights.py @@ -11,7 +11,6 @@ from b12x._lib.quant.block_codec import block_codec if TYPE_CHECKING: - from .._shared.kernels.w4a16.exl3 import Exl3Layer from ._impl import B12XFP4ExpertWeights from .planning import WeightPlan @@ -88,11 +87,15 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class TrellisWeights: - """Layer-local views of the canonical ``b12x_trellis`` tensors. + """Container-independent layer-local trellis tensors. - ``codes`` is the rank-local ``[I_local/32, row_stride]`` uint8 payload. + ``codes`` is the rank-local ``[I_local/32, payload_bytes]`` uint8 payload. + Its physical row stride may include padding outside the logical columns. ``rate`` is a view selected from the single model-level uint8 rate tensor; it is never copied merely to give each layer its own rate parameter. + CPU payloads may be prepared onto an explicitly selected CUDA device. + Row padding is removed and checked by the checkpoint adapter, not by the + execution planner. All scale vectors and gains retain FP16 storage. """ codes: torch.Tensor @@ -122,31 +125,6 @@ def __post_init__(self) -> None: ) -@dataclass(frozen=True, kw_only=True) -class Exl3Weights: - """One validated CPU slot extent and its destination CUDA device. - - Preparation stages bounded expert batches. The complete source payload - need not reside on the GPU beside its prepared representation. - """ - - layer: "Exl3Layer" - device: torch.device - - def __post_init__(self) -> None: - from .._shared.kernels.w4a16.exl3 import Exl3Layer - - if not isinstance(self.layer, Exl3Layer): - raise TypeError("Exl3Weights.layer must be a Exl3Layer") - device = torch.device(self.device) - if device.type != "cuda": - raise ValueError("Exl3Weights.device must identify a CUDA device") - object.__setattr__(self, "device", device) - self.layer.manifest.validate_extent( - self.layer.first_slot, self.layer.slot_count - ) - - @dataclass(frozen=True) class BlockQuantWeights: """Raw uint8[E,N,K/block_weights,block_bytes] checkpoint blocks. @@ -246,7 +224,6 @@ def device(self) -> torch.device: __all__ = [ - "Exl3Weights", "PackedWeights", "IQ2XSWeights", "BlockQuantWeights", diff --git a/docs/exl3-checkpoint-format.md b/docs/exl3-checkpoint-format.md index 4f376c9f3..89902f0e8 100644 --- a/docs/exl3-checkpoint-format.md +++ b/docs/exl3-checkpoint-format.md @@ -141,8 +141,10 @@ rotation blocks; an intermediate-Hadamard pre-activation transform whose halves split contributes a barrier. The reader enforces the declarations without knowing the reasons. -A rank shorter than the serving kernel's minimum local width is padded at -preparation time, exactly as any other legal extent. +The serving kernel must also support the rank-local shape. The common trellis +preparer rejects a width that does not fit its tile/transform contract; it +does not silently pad expert payloads. See [trellis preparation](trellis-preparation.md) +for the container-independent loading API and its staging controls. ## Support status diff --git a/docs/trellis-preparation.md b/docs/trellis-preparation.md new file mode 100644 index 000000000..5ac08e575 --- /dev/null +++ b/docs/trellis-preparation.md @@ -0,0 +1,105 @@ +# Trellis preparation across checkpoint containers + +Use `TrellisSource` for encoding and source coordinates, `TrellisWeights` for +tensor data, and `prepare_weights` to create GPU-resident `PreparedExperts`. +The execution API does not dispatch on a checkpoint filename or manifest type. + +For an EXL3 checkpoint, the adapter interprets its manifest and returns these +common objects without requantizing the payload. This example prepares eight +32-channel slots from layer 1 of a Kimi-K3 EXL3 checkpoint on GPU 0: + +```python +import torch +from b12x.moe import fused_moe +from b12x.moe.checkpoints.exl3 import ( + read_exl3_manifest, read_exl3_layer, trellis_from_exl3, +) + +root = "/data/kimi-k3-exl3-2bit" +manifest = read_exl3_manifest(root) +layer = read_exl3_layer(root, manifest, 1, first_slot=0, slot_count=8) +source, weights = trellis_from_exl3(layer) +plan = fused_moe.plan_weights( + source=source, + activation=fused_moe.ActivationSpec( + mode="a16", nonlinearity="situ", io_dtype=torch.bfloat16, + rotation_dtype=torch.float16, + ), + geometry=fused_moe.MoEGeometry( + num_experts=manifest.geometry.num_experts, + hidden_size=manifest.geometry.hidden_size, + intermediate_size=layer.local_intermediate_size, + ), +) +experts = fused_moe.prepare_weights( + plan=plan, weights=weights, device="cuda:0", + staging=fused_moe.TrellisStaging(max_experts=64, max_bytes=64 * 1024 * 1024), +) +``` + +The caller assigns a legal extent to each tensor-parallel rank before this +step. B12X validates source bounds, tensor shapes, rates, and supported +transforms. This API implementation requires GPU qualification before it +replaces an existing serving deployment; CPU layout tests alone do not verify +CUDA execution or its performance. + +## Encoding, data, and arithmetic + +- `TrellisSource.config` reuses the native version-2 `TrellisConfig` schema. + Existing native callers may still pass `TrellisConfig` directly; the planner + wraps it in a source descriptor. No checkpoint JSON migration is required. +- `TrellisSource.uniform_bits` declares a uniform K value before preparation. + The actual uint8 rate tensor must agree. Omitting this field retains the + native K3 LUT or projection-tiered MCG preparation contract. The EXL3 adapter + sets it from the manifest, including K2 for the converted Kimi-K3 model. +- `TrellisExtent` records global intermediate width and the rank's first/count + of 32-channel slots. Nonzero Hadamard sign patterns are generated globally + and sliced, not regenerated from local width. Native callers without global + extent metadata retain their zero-sign-pattern restriction. +- `TrellisWeights` contains uint8 codes, rate bytes, FP16 scale vectors and + optional gains, and optional per-expert sign-pattern IDs. There is no + `Exl3Weights` or `Exl3Source` branch in canonical planning. +- `ActivationSpec.rotation_dtype=torch.float16` makes the FP16 full-rotation + arithmetic explicit while retaining BF16 public inputs/outputs. This is the + arithmetic contract used by the EXL3 serving integration, not an additional + quantization step. Its default is `None`, which retains the I/O dtype for + existing callers. It is rejected for non-trellis or projection-tiered MCG + preparation. + +The adapter checks zero row padding and removes it through a view. It maps +the two physical FC1 projections to the appropriate hidden-axis scale table +for the interleaved intermediate transform. Codewords and FP16 scale values +are copied or permuted, never decoded and re-encoded. + +## Staging and lifetime + +`prepare_weights(device=...)` accepts CPU uniform-codeword inputs and moves +them to the specified CUDA device. CUDA inputs must already belong to that +device. CPU weights do not remain in the inference path. + +`TrellisStaging` defaults to at most 64 experts and 64 MiB of codeword transfer +data per batch. The smaller limit wins. For a 32 MiB transfer budget, pass +`TrellisStaging(max_bytes=32 * 1024 * 1024)`. A single expert that exceeds the +budget is rejected before final weight allocation; increase the budget or +use a smaller legal source extent. This setting does not limit total model +memory: final weights, scale/sign tables, allocator overhead, and at most one +projection-sized copy operand are separate allocations. + +Preparation completes before CUDA Graph capture. Prepared experts own final +GPU storage; bind/run/replay do not read checkpoint files, transfer source +weights, or choose another kernel. Projection-tiered native MCG retains its +CUDA-resident input contract; bounded CPU staging applies to uniform rates. + +The low-level EXL3 entry point forwards uniform inputs into common preparation. +Its asymmetric-pair implementation remains separate because its rate and +record contracts differ. The common adapter rejects pair-rate containers +rather than flattening their rates. + +## Validation + +`tests/moe/test_trellis_adapter.py` checks byte-exact word assembly, FP16 tables, +source-global signs, staging bounds, corrupt inputs, and safetensors reads. +`tests/moe/test_trellis_config.py` checks the native schema and projection layout. +`tests/moe/test_exl3_prepare.py` compares CUDA preparation against independent +word/rotation assembly and exercises execution and graph replay. Run all three +plus `tests/moe/test_exl3_schema.py` in a B12X CUDA environment before deployment. diff --git a/tests/moe/test_exl3_prepare.py b/tests/moe/test_exl3_prepare.py index 2b3f136f3..a9ff00e2d 100644 --- a/tests/moe/test_exl3_prepare.py +++ b/tests/moe/test_exl3_prepare.py @@ -8,19 +8,18 @@ from __future__ import annotations import json +from dataclasses import replace import pytest import torch from b12x.moe._shared.exl3_schema import ( EXL3_MANIFEST_FILENAME, - matrix_slot_bytes, rate_code, ) from b12x.moe._shared.kernels.w4a16.exl3 import ( prepare_exl3_moe_weights, read_exl3_layer, - read_exl3_manifest, ) from b12x.moe._shared.kernels.w4a16.exl3_synth import ( Exl3SynthConfig, @@ -311,6 +310,59 @@ def test_exl3_uniform_padding_must_be_zero(tmp_path) -> None: ) +def _uniform_reference(layer, *, activation, device): + """Independent word/rotation assembly, bound to the shared compute kernel. + + Neither the checkpoint adapter nor common trellis preparation participates + in this oracle. Global sign coordinates and side-table selection are + reconstructed from the source extent explicitly. + """ + from tests.moe.test_trellis_adapter import _independent_planes + + geometry = layer.manifest.geometry + experts, local = geometry.num_experts, layer.local_intermediate_size + w13, w2 = _independent_planes(layer) + rotations = torch.stack([ + torch.cat([layer.rotations[s, e] for s in range(layer.slot_count)], dim=1) + for e in range(experts) + ]).reshape(experts, 3 * local).to(device) + gate, up = layer.gate_suh.to(device), layer.up_suh.to(device) + hadamard = layer.manifest.hadamard.intermediate_hadamard + if hadamard: + gate = up = gate if layer.first_slot < geometry.num_slots // 2 else up + tiles = (128, 128, 128, 128) if hadamard and layer.manifest.rates.bits == 2 else (64, 256, 64, 256) + result = prepare_trellis256_moe_weights( + w13.to(device), w2.to(device), hidden_size=geometry.hidden_size, + intermediate_size=local, num_experts=experts, activation=activation, + fc1_tile_n=tiles[1], fc2_tile_n=tiles[3], params_dtype=torch.float16, + w13_layout="trellis_t256_proj", trellis_bits=layer.manifest.rates.bits, + codebook=layer.manifest.codebook, gate_suh=gate, up_suh=up, + intermediate_rotations=rotations, down_svh=layer.down_svh.to(device), + tile_config=tiles, + ) + if not hadamard: + return result + signs = [] + for pattern in layer.sign_pattern.tolist(): + parts = [] + for axis, multiplier in ((1, 2), (2, 1)): + length = multiplier * geometry.intermediate_size + begin = multiplier * 32 * layer.first_slot + if pattern == 0: + full = torch.ones(length) + else: + gen = torch.Generator().manual_seed( + (0x6A09E667F3BCC909 * pattern + 0xBB67AE8584CAA73B * axis) % (2**63) + ) + full = torch.randint(2, (length,), generator=gen) * 2 - 1 + parts.append(full[begin:begin + multiplier * local]) + signs.append(torch.cat(parts)) + return replace(result, trellis=replace( + result.trellis, intermediate_hadamard=True, + intermediate_rotations=torch.cat((rotations, torch.stack(signs).half().to(device)), dim=1), + )) + + @requires_cuda @pytest.mark.parametrize("first_slot, slots", [(0, 8), (8, 4), (48, 12), (60, 8)]) @pytest.mark.parametrize("activation", ["silu", "situ"]) @@ -321,6 +373,7 @@ def test_canonical_exl3_preparation_preserves_source_extent( """Uneven TP extents preserve rotations in preparation, execution and graphs.""" from b12x._lib.runtime_control import kernel_resolution_guard from b12x.moe import fused_moe + from b12x.moe.checkpoints.exl3 import trellis_from_exl3 from b12x.moe._shared.kernels.w4a16.exl3 import read_exl3_layer from b12x.moe._shared.kernels.w4a16.exl3_synth import write_exl3_checkpoint from b12x.preparation import PreparationSession, PreparedCall @@ -347,20 +400,24 @@ def test_canonical_exl3_preparation_preserves_source_extent( tmp_path, manifest, 1, first_slot=first_slot, slot_count=slots ) layer.sign_pattern.copy_(torch.tensor([0, 6], dtype=torch.uint8)) - expected = prepare_exl3_moe_weights(layer, activation=activation, device=_device()) + expected = _uniform_reference(layer, activation=activation, device=_device()) + source_metadata, source_weights = trellis_from_exl3(layer) plan = fused_moe.plan_weights( - source=fused_moe.Exl3Source(manifest=manifest), + source=source_metadata, activation=fused_moe.ActivationSpec( - mode="a16", nonlinearity=activation, io_dtype=torch.bfloat16 + mode="a16", nonlinearity=activation, io_dtype=torch.bfloat16, + rotation_dtype=torch.float16, ), geometry=fused_moe.MoEGeometry( num_experts=2, hidden_size=512, intermediate_size=slots * 32 ), ) prepared = fused_moe.prepare_weights( - plan=plan, weights=fused_moe.Exl3Weights(layer=layer, device=_device()) + plan=plan, weights=source_weights, device=_device(), ) actual = prepared._impl.representation.value + assert actual.params_dtype == expected.params_dtype == torch.float16 + assert actual.tile_config == expected.tile_config for name in ( "w13", "w2", @@ -373,6 +430,7 @@ def test_canonical_exl3_preparation_preserves_source_extent( getattr(actual, name), getattr(expected, name), rtol=0, atol=0 ) assert actual.intermediate_hadamard + assert actual.gate_suh.data_ptr() == actual.up_suh.data_ptr() topk = 2 torch.manual_seed(617) diff --git a/tests/moe/test_exl3_schema.py b/tests/moe/test_exl3_schema.py index 48b355630..43e8816d0 100644 --- a/tests/moe/test_exl3_schema.py +++ b/tests/moe/test_exl3_schema.py @@ -86,42 +86,47 @@ def test_uniform_manifest_parses() -> None: @pytest.mark.parametrize("bits", [2, 3]) def test_canonical_weight_plan_preserves_exl3_transform_and_rate(bits: int) -> None: from b12x.moe import fused_moe + from tests.moe.test_trellis_adapter import _layer + from b12x.moe.checkpoints.exl3 import trellis_from_exl3 - manifest = Exl3Manifest.from_dict(_uniform_manifest_dict(bits=bits, intermediate_hadamard=True)) + source, _ = trellis_from_exl3(_layer(first=8, slots=4, bits=bits)) plan = fused_moe.plan_weights( - source=fused_moe.Exl3Source(manifest=manifest), + source=source, activation=fused_moe.ActivationSpec( - mode="a16", nonlinearity="silu", io_dtype=torch.bfloat16 + mode="a16", nonlinearity="silu", io_dtype=torch.bfloat16, + rotation_dtype=torch.float16, ), geometry=fused_moe.MoEGeometry( - num_experts=8, hidden_size=512, intermediate_size=256 + num_experts=3, hidden_size=512, intermediate_size=128 ), ) assert plan.prepared_format.weights is fused_moe.WeightEncoding.TRELLIS assert plan.prepared_format.packing is fused_moe.WeightPacking.TRELLIS_NATIVE assert plan._impl.trellis_bits == bits assert plan._impl.intermediate_hadamard_blocks == (512, 128) - assert plan.geometry.intermediate_size == 256 - assert plan.source.manifest.geometry.intermediate_size == 512 + assert plan.geometry.intermediate_size == 128 + assert plan.source.extent.global_intermediate_size == 3072 + assert plan._impl.source_format == "b12x_trellis" @pytest.mark.parametrize( ("mode", "experts", "intermediate", "message"), [ ("a8", 8, 256, "A16 activations"), - ("a16", 9, 256, "must match the manifest"), - ("a16", 8, 1024, "exceeds the source width"), + ("a16", 8, 1024, "extent width differs"), ], ) def test_canonical_exl3_plan_rejects_incompatible_contract( mode: str, experts: int, intermediate: int, message: str ) -> None: from b12x.moe import fused_moe + from tests.moe.test_trellis_adapter import _layer + from b12x.moe.checkpoints.exl3 import trellis_from_exl3 - manifest = Exl3Manifest.from_dict(_uniform_manifest_dict(intermediate_hadamard=True)) + source, _ = trellis_from_exl3(_layer(first=0, slots=8, experts=8)) with pytest.raises(ValueError, match=message): fused_moe.plan_weights( - source=fused_moe.Exl3Source(manifest=manifest), + source=source, activation=fused_moe.ActivationSpec( mode=mode, nonlinearity="silu", io_dtype=torch.bfloat16 ), diff --git a/tests/moe/test_trellis_adapter.py b/tests/moe/test_trellis_adapter.py new file mode 100644 index 000000000..84c0eb30d --- /dev/null +++ b/tests/moe/test_trellis_adapter.py @@ -0,0 +1,344 @@ +"""CPU qualification of lossless slot assembly and EXL3 normalization. + +These tests validate physical layouts, not CUDA kernel execution. The CUDA +preparation/execution/graph comparisons are in test_exl3_prepare.py. +""" + +from dataclasses import replace + +import pytest +import torch + +from b12x.moe._shared.exl3_schema import Exl3Manifest +from b12x.moe.checkpoints.exl3 import Exl3Layer, read_exl3_layer, trellis_from_exl3 +from b12x.moe.fused_moe.source import TrellisExtent, TrellisSource +from b12x.moe.fused_moe.trellis import ( + _effective_input_scales, + _effective_intermediate_scales, + _effective_output_scales, + _local_rate_matrix, + _symmetric_bits, +) +from b12x.moe.fused_moe.trellis_layout import ( + TrellisStaging, + append_intermediate_signs, + assemble_uniform_slots, +) + + +def _layer(*, first=48, slots=12, experts=3, bits=2, per_expert=True, hadamard=True): + hidden, global_i = 512, 3072 + had = {"intermediate_hadamard": hadamard, "per_expert_input_rotations": per_expert} + if hadamard: + had.update(pre_block=512, post_block=128) + manifest = Exl3Manifest.from_dict( + { + "kind": "exl3-manifest", + "schema": "exl3-v1", + "codebook": "lut_e4m3", + "geometry": { + "num_experts": experts, + "hidden_size": hidden, + "intermediate_size": global_i, + "slot_channels": 32, + "num_slots": global_i // 32, + "moe_layer_indices": [1], + }, + "rates": {"structure": "uniform", "bits": bits}, + "hadamard": had, + "layout": { + "row_alignment": 65536, + "extent_alignment_slots": 4, + "extent_barriers": [48], + }, + "layers": { + "1": {"file": "exl3-layer-00001.safetensors", "sha256": "0" * 64} + }, + } + ) + payload = experts * 3 * (hidden // 16) * 64 * bits + row_bytes = (payload + 65535) // 65536 * 65536 + codes = torch.zeros((slots, row_bytes), dtype=torch.uint8) + codes[:, :payload] = torch.randint(0, 256, (slots, payload), dtype=torch.uint8) + side_shape = (experts, hidden) if per_expert else (hidden,) + return Exl3Layer( + manifest=manifest, + layer_index=1, + first_slot=first, + slot_count=slots, + codes=codes, + rotations=torch.randn(slots, experts, 3, 32).half(), + gate_suh=torch.randn(side_shape).half(), + up_suh=torch.randn(side_shape).half(), + down_svh=torch.randn(side_shape).half(), + rates_fc1=None, + rates_fc2=None, + sign_pattern=(torch.arange(experts) % 8).byte() if hadamard else None, + ) + + +def _independent_planes(layer): + """Explicit slot/expert/matrix addressing, independent of batched assembly.""" + experts, hidden = ( + layer.manifest.geometry.num_experts, + layer.manifest.geometry.hidden_size, + ) + bits, slots = layer.manifest.rates.bits, layer.slot_count + tiles = hidden // 16 + w13 = torch.empty(2, experts, tiles, 2 * slots, 16 * bits, dtype=torch.int16) + w2 = torch.empty(experts, 2 * slots, tiles, 16 * bits, dtype=torch.int16) + section = tiles * 64 * bits + for slot in range(slots): + for expert in range(experts): + for matrix in range(3): + start = (expert * 3 + matrix) * section + words = ( + layer.codes[slot, start : start + section] + .contiguous() + .view(torch.int16) + ) + for half in range(2): + plane = words[ + half * tiles * 16 * bits : (half + 1) * tiles * 16 * bits + ] + if matrix < 2: + w13[matrix, expert, :, 2 * slot + half] = plane.view( + tiles, 16 * bits + ) + else: + w2[expert, 2 * slot + half] = plane.view(tiles, 16 * bits) + return w13, w2 + + +@pytest.mark.parametrize("first,slots", [(0, 8), (8, 4), (48, 12), (60, 8)]) +@pytest.mark.parametrize("bits", [2, 3, 4]) +@pytest.mark.parametrize("batch", [1, 2, 64]) +def test_uniform_assembly_is_bit_exact(first, slots, bits, batch): + layer = _layer(first=first, slots=slots, bits=bits) + source, weights = trellis_from_exl3(layer) + assert weights.codes.data_ptr() == layer.codes.data_ptr() + before = layer.codes.clone() + actual = assemble_uniform_slots( + weights.codes, + num_experts=3, + hidden_size=512, + bits=bits, + device="cpu", + staging=TrellisStaging(max_experts=batch), + ) + for result, expected in zip(actual, _independent_planes(layer)): + assert torch.equal(result, expected) + assert torch.equal(before, layer.codes) + assert source.extent.first_slot == first + assert source.uniform_bits == bits + assert source.config.to_dict()["version"] == 2 + + +@pytest.mark.parametrize("first", [0, 8, 48, 60]) +@pytest.mark.parametrize("per_expert", [False, True]) +@pytest.mark.parametrize("experts", [1, 3]) +@pytest.mark.parametrize("hadamard", [False, True]) +def test_side_and_intermediate_tables_are_bit_exact( + first, per_expert, experts, hadamard +): + layer = _layer( + first=first, slots=8, experts=experts, per_expert=per_expert, hadamard=hadamard + ) + source, weights = trellis_from_exl3(layer) + scales = source.config.scale + gate, up = _effective_input_scales( + weights.input_scales, + scales.input_scales, + num_experts=experts, + hidden_size=512, + device=torch.device("cpu"), + ) + expected_gate, expected_up = layer.gate_suh, layer.up_suh + if hadamard: + expected_gate = expected_up = layer.gate_suh if first < 48 else layer.up_suh + assert torch.equal(gate, expected_gate.reshape(-1, 512)) + assert torch.equal(up, expected_up.reshape(-1, 512)) + if hadamard: + assert gate.data_ptr() == up.data_ptr() + assert weights.input_scales.vectors.numel() == expected_gate.numel() + intermediate = _effective_intermediate_scales( + weights.intermediate_scales, + scales.intermediate_scales, + num_experts=experts, + intermediate_size=256, + device=torch.device("cpu"), + ) + expected = torch.stack( + [ + torch.cat([layer.rotations[s, e] for s in range(8)], dim=1) + for e in range(experts) + ] + ) + assert torch.equal(intermediate, expected) + down = _effective_output_scales( + weights.output_scales, + scales.output_scales, + num_experts=experts, + hidden_size=512, + device=torch.device("cpu"), + ) + assert torch.equal(down, layer.down_svh.reshape(-1, 512)) + + +@pytest.mark.parametrize("first,slots", [(0, 8), (8, 4), (48, 12), (60, 8)]) +def test_signs_are_sliced_from_global_sequence(first, slots): + layer = _layer(first=first, slots=slots, experts=8) + source, weights = trellis_from_exl3(layer) + local = slots * 32 + values = weights.intermediate_scales.vectors.reshape(8, 3 * local) + actual = append_intermediate_signs( + values, weights.expert_sign_patterns, extent=source.extent + ) + assert torch.equal(actual[:, : 3 * local], values) + for pattern in range(8): + signs = [] + for axis, width, offset, size in ( + (1, 6144, first * 64, local * 2), + (2, 3072, first * 32, local), + ): + if pattern == 0: + full = torch.ones(width) + else: + gen = torch.Generator().manual_seed( + (0x6A09E667F3BCC909 * pattern + 0xBB67AE8584CAA73B * axis) % (2**63) + ) + full = torch.randint(2, (width,), generator=gen) * 2 - 1 + signs.append(full[offset : offset + size]) + assert torch.equal(actual[pattern, 3 * local :], torch.cat(signs).half()) + + +def test_staging_limit_selects_complete_experts_and_rejects_insufficient_budget(): + layer = _layer() + _, weights = trellis_from_exl3(layer) + one = layer.slot_count * 3 * 32 * 64 * 2 + staging = TrellisStaging(max_bytes=one * 2 + 7) + assert staging.batch_size(slots=12, hidden_size=512, bits=2) == 2 + actual = assemble_uniform_slots( + weights.codes, + num_experts=3, + hidden_size=512, + bits=2, + device="cpu", + staging=staging, + ) + for result, expected in zip(actual, _independent_planes(layer)): + assert torch.equal(result, expected) + with pytest.raises(ValueError, match="insufficient"): + assemble_uniform_slots( + weights.codes, + num_experts=3, + hidden_size=512, + bits=2, + device="cpu", + staging=TrellisStaging(max_bytes=one - 1), + ) + + +@pytest.mark.parametrize( + "change,message", + [ + ("padding", "padding"), + ("rotation_dtype", "FP16"), + ("sign_pattern", "sign_pattern"), + ("short_codes", "containing every expert"), + ], +) +def test_adapter_rejects_corrupt_payload(change, message): + layer = _layer() + if change == "padding": + layer.codes[0, -1] = 1 + elif change == "rotation_dtype": + layer = replace(layer, rotations=layer.rotations.bfloat16()) + elif change == "sign_pattern": + layer.sign_pattern[0] = 8 + else: + layer = replace(layer, codes=layer.codes[:, :100]) + with pytest.raises(ValueError, match=message): + trellis_from_exl3(layer) + + +def test_preparation_verifies_declared_uniform_rate(): + source, weights = trellis_from_exl3(_layer()) + rates = _local_rate_matrix( + source.config, weights.rate, num_experts=3, device=torch.device("cpu") + ) + assert torch.equal( + _symmetric_bits(source.config, rates, uniform_bits=2), torch.full((3, 3), 2) + ) + with pytest.raises(ValueError, match="observed"): + _symmetric_bits(source.config, rates, uniform_bits=3) + with pytest.raises(ValueError, match="symmetric"): + _symmetric_bits( + source.config, torch.tensor([[0x23]], dtype=torch.uint8), uniform_bits=2 + ) + + +@pytest.mark.parametrize("first,count", [(-1, 4), (96, 4), (4, 0)]) +def test_extent_rejects_invalid_bounds(first, count): + with pytest.raises(ValueError): + TrellisExtent(global_intermediate_size=3072, first_slot=first, slot_count=count) + + +def test_uniform_rate_rejects_undefined_codebook_range(): + source, _ = trellis_from_exl3(_layer()) + with pytest.raises(ValueError, match="K2/K3/K4"): + TrellisSource(config=source.config, uniform_bits=5) + + +@pytest.mark.parametrize("first,slots", [(0, 8), (8, 4), (48, 12), (60, 8)]) +def test_safetensors_reader_and_adapter_keep_extent_bytes(tmp_path, first, slots): + from safetensors.torch import save_file + + full = _layer(first=0, slots=96) + tensors = { + name: getattr(full, name) + for name in ( + "codes", + "rotations", + "gate_suh", + "up_suh", + "down_svh", + "sign_pattern", + ) + } + path = tmp_path / full.manifest.layers[1].file + save_file( + tensors, + path, + metadata={ + "schema": "exl3-v1", + "codebook": "lut_e4m3", + "layer": "1", + "num_experts": "3", + "hidden_size": "512", + "intermediate_size": "3072", + "slot_channels": "32", + }, + ) + loaded = read_exl3_layer( + tmp_path, full.manifest, 1, first_slot=first, slot_count=slots + ) + assert torch.equal(loaded.codes, full.codes[first : first + slots]) + assert torch.equal(loaded.rotations, full.rotations[first : first + slots]) + _, weights = trellis_from_exl3(loaded) + for actual, expected in zip( + assemble_uniform_slots( + weights.codes, num_experts=3, hidden_size=512, bits=2, device="cpu" + ), + _independent_planes(loaded), + ): + assert torch.equal(actual, expected) + with pytest.raises(ValueError, match="sha256 mismatch"): + read_exl3_layer( + tmp_path, + full.manifest, + 1, + first_slot=first, + slot_count=slots, + verify_sha=True, + ) From 437cead4d902dc549d5c9744630090fc142d9847 Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Sun, 27 Sep 2026 08:13:11 -0700 Subject: [PATCH 3/4] Keep common trellis preparation on compact decoding and atomic route packing Remove the full-codeword shared LUT and stable per-expert route ordering from the preparation interface change. Retain exact route capacities, prepared launch ownership and opt-in FP32 prefill reduction. Bump the MoE query schema for the reduced control set and scope the release fragment to PR 433. Validation: selected host suites report 181 passed and 28 failures; the same 28 missing-decode_config failures reproduce at 7fca2aa. Ten LUT-specific cases move with the LUT change. Changed Python sources parse and git diff --check passes. GPU execution and serving qualification were not run. Assisted-by: OpenAI Codex --- .lil/changes/b12x-415.json | 18 --- .lil/changes/b12x-433.json | 19 +++ b12x/_lib/intrinsics.py | 44 +++---- b12x/moe/_shared/kernels/w4a16/host.py | 8 -- b12x/moe/_shared/kernels/w4a16/kernel.py | 83 +++--------- b12x/moe/_shared/kernels/w4a16/route_pack.py | 121 +----------------- b12x/moe/fused_moe/_impl.py | 11 -- b12x/moe/fused_moe/_preparation.py | 7 +- b12x/moe/fused_moe/_tuning.py | 4 +- docs/trellis-decode-table.md | 39 ------ docs/w4a16-prefill-reduction.md | 7 +- tests/moe/test_fused_moe_variant_selection.py | 47 ------- tests/moe/test_trellis_direct_lut_decode.py | 48 +------ tests/moe/test_w4a16_route_pack.py | 8 +- 14 files changed, 69 insertions(+), 395 deletions(-) delete mode 100644 .lil/changes/b12x-415.json create mode 100644 .lil/changes/b12x-433.json delete mode 100644 docs/trellis-decode-table.md diff --git a/.lil/changes/b12x-415.json b/.lil/changes/b12x-415.json deleted file mode 100644 index 36ac2b841..000000000 --- a/.lil/changes/b12x-415.json +++ /dev/null @@ -1,18 +0,0 @@ -{ - "schema": "local-inference-release-change/v1", - "id": "b12x-415", - "category": "feature", - "summary": "Prepare EXL3 expert extents and bounded W4A16 prefill workspaces", - "models": [ - "Kimi-K3" - ], - "compatibility": "No user action required.", - "details": [ - "Uses an exact codeword lookup table for eligible two-bit EXL3 kernels.", - "Preserves IQ2_XS, IQ2_XXS and Q8_0 epilogues and pipeline tuning." - ], - "pull_requests": [ - 415 - ], - "requires": [] -} diff --git a/.lil/changes/b12x-433.json b/.lil/changes/b12x-433.json new file mode 100644 index 000000000..4ef5116c7 --- /dev/null +++ b/.lil/changes/b12x-433.json @@ -0,0 +1,19 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "b12x-433", + "category": "feature", + "summary": "Share native and EXL3 trellis preparation with bounded W4A16 workspaces", + "models": [ + "Kimi-K3" + ], + "compatibility": "Native TrellisConfig callers remain accepted. Canonical EXL3 callers use trellis_from_exl3 and pass device to prepare_weights. GPU serving qualification is pending.", + "details": [ + "Preserves codewords, FP16 scales and source-global transform coordinates through common trellis preparation.", + "Retains compact W4A16 trellis decoding and unordered atomic route packing.", + "Preserves IQ2_XS, IQ2_XXS and Q8_0 epilogues and pipeline tuning." + ], + "pull_requests": [ + 433 + ], + "requires": [] +} diff --git a/b12x/_lib/intrinsics.py b/b12x/_lib/intrinsics.py index f549fbb6f..3f2ecd551 100644 --- a/b12x/_lib/intrinsics.py +++ b/b12x/_lib/intrinsics.py @@ -6834,7 +6834,6 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( bits: int = 3, rate_indexed: bool = True, *, - in_shared: bool = False, loc=None, ip=None, ): @@ -6842,8 +6841,8 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( This is the lookup lower-bound primitive: one byte load per reconstructed weight. Every slot owns its own address/code registers so all eight - gathers are in flight together. With ``in_shared``, the address names a - caller-staged shared-memory table instead of a global-memory table. + global gathers are in flight together; a shared scratch pair would chain + them behind full global-memory latency each. """ bits = int(bits) @@ -6854,47 +6853,41 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( table_offset = ((bits - 2) << 16) if rate_indexed else 0 extract_lines: list[str] = [] load_lines: list[str] = [] + pack_lines: list[str] = [] for index in range(8): source = "$3" if index < 4 else "$2" shift = (3 - (index & 3)) * bits # Accumulate into declared scratch registers, never directly into # the "=r" outputs: without early-clobber the outputs may alias the # window inputs that later iterations still read. + target = "out0" if index < 4 else "out1" + byte_shift = 8 * (index & 3) if shift: extract_lines.append( f"bfe.u32 w{index}, {source}, {shift}, 16;" ) else: extract_lines.append(f"and.b32 w{index}, {source}, 0xffff;") - if in_shared: - load_lines.append( - f"add.u32 w{index}, w{index}, $4;" - f" ld.shared.u8 w{index}, [w{index}];" - ) - else: - load_lines.append( - f""" + load_lines.append( + f""" cvt.u64.u32 addr{index}, w{index}; add.u64 addr{index}, addr{index}, $4; ld.global.u8 w{index}, [addr{index}]; - """ + """ + ) + if byte_shift: + pack_lines.append( + f"shl.b32 w{index}, w{index}, {byte_shift};" ) - # Each load returns one zero-extended byte. Byte permutations pack four - # independent results without a serial shift/or accumulator. - pack_lines = [ - "prmt.b32 w0, w0, w1, 0x4040;", - "prmt.b32 w2, w2, w3, 0x4040;", - "prmt.b32 out0, w0, w2, 0x5410;", - "prmt.b32 w4, w4, w5, 0x4040;", - "prmt.b32 w6, w6, w7, 0x4040;", - "prmt.b32 out1, w4, w6, 0x5410;", - ] + pack_lines.append(f"or.b32 {target}, {target}, w{index};") asm = ( """ { .reg .b32 out0,out1; .reg .b32 w0,w1,w2,w3,w4,w5,w6,w7; .reg .b64 addr0,addr1,addr2,addr3,addr4,addr5,addr6,addr7; + mov.b32 out0, 0; + mov.b32 out1, 0; """ + "\n".join(extract_lines) + "\n".join(load_lines) @@ -6904,10 +6897,7 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( mov.b32 $1, out1; }""" ) - table_base = ( - Int32(direct_lut_addr) + Int32(table_offset) - if in_shared else Int64(direct_lut_addr) + Int64(table_offset) - ) + table_base = Int64(direct_lut_addr) + Int64(table_offset) result = llvm.inline_asm( llvm.StructType.get_literal([T.i32(), T.i32()]), [ @@ -6916,7 +6906,7 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( table_base.ir_value(loc=loc, ip=ip), ], asm, - "=r,=r,r,r,r" if in_shared else "=r,=r,r,r,l", + "=r,=r,r,r,l", has_side_effects=False, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, diff --git a/b12x/moe/_shared/kernels/w4a16/host.py b/b12x/moe/_shared/kernels/w4a16/host.py index 7a12f9e2a..7492b79c2 100644 --- a/b12x/moe/_shared/kernels/w4a16/host.py +++ b/b12x/moe/_shared/kernels/w4a16/host.py @@ -21,14 +21,6 @@ _SUPPORTED_ACTIVATIONS = SUPPORTED_MOE_ACTIVATIONS -def trellis_decode_table() -> str: - """Read the plan-time trellis decode-table selection; replay never reads it.""" - value = os.environ.get("B12X_TRELLIS_DECODE_TABLE", "auto") - if value not in {"auto", "compact", "full"}: - raise ValueError("B12X_TRELLIS_DECODE_TABLE must be auto, compact, or full") - return value - - def prefill_fused_sum_enabled() -> bool: """Enable direct FP32 route reduction for large-M W4A16 launches. diff --git a/b12x/moe/_shared/kernels/w4a16/kernel.py b/b12x/moe/_shared/kernels/w4a16/kernel.py index cc48cf3dc..e8ff2b369 100644 --- a/b12x/moe/_shared/kernels/w4a16/kernel.py +++ b/b12x/moe/_shared/kernels/w4a16/kernel.py @@ -72,7 +72,6 @@ packed_decode_lut_fp16_to_bfloat2x4, packed_decode_lut_fp16_to_half2x4, packed_decode_lut_e4m3_to_e4m3x8, - packed_decode_lut_e4m3_direct_to_e4m3x8, ld_global_nc_v4_u32, pack_f32x2_to_bfloat2, pack_f32x2_to_f16x2, @@ -96,10 +95,7 @@ warp_reduce, ) from b12x._lib.quant.iq2_xs import iq2_xs_execution_lut -from b12x._lib.quant.lut_e4m3 import ( - lut_e4m3_direct_table, - lut_e4m3_value_table, -) +from b12x._lib.quant.lut_e4m3 import lut_e4m3_value_table from b12x.moe._shared.kernels.trellis_ring import ( trellis256_lane_geom_bits as _trellis_ring_lane_geom_bits, ) @@ -302,12 +298,10 @@ def _gather_native_scale_bytes(a: Uint32, b: Uint32, c: Uint32, d: Uint32, byte: def _trellis256_execution_lut( - device: torch.device | str, codebook: str, *, direct_lut: bool = False + device: torch.device | str, codebook: str ) -> torch.Tensor: if codebook == LUT_FP16: return lut_fp16_segment_table(device) - if direct_lut: - return lut_e4m3_direct_table(device) return lut_e4m3_value_table(device) # TC-decode runs on the packed W4A16 object and folds the top-k sum into the FC2 @@ -885,7 +879,6 @@ class W4A16FusedMoeCompileResult: shared_memory_bytes: int = -1 broadcast_suh: bool = False small_m_direct_launches: tuple[_W4A16SmallMDirectLaunch, ...] = () - trellis_direct_lut: bool = False @dataclass(frozen=True) @@ -1222,7 +1215,6 @@ def __init__( static_pair_rates.get(trellis_pair_kind, (3, 3)) ) self.lut_e4m3_smem = False - self.trellis_direct_lut = False # Small-M stripe split-K: opt out of the one-tile-per-CTA fast path # so decode-heavy small-M phases spread each mn-tile's K range across # multiple CTAs (existing tail scheduling plus cross-CTA finalize). @@ -1350,7 +1342,6 @@ def __init__( else: self.sms = 120 max_shared_mem = _DEFAULT_MAX_SHARED_MEM - self.max_shared_mem = max_shared_mem self.blocks_per_sm = _determine_blocks_per_sm( problem_m=self.size_m, problem_n=self.covered_size_n, @@ -1513,7 +1504,6 @@ def __cache_key__(self) -> tuple[object, ...]: self.schedule_whole_tiles, self.schedule_route_block_factor, self.lut_e4m3_smem, - self.trellis_direct_lut, self.small_m_splitk, ) @@ -4106,19 +4096,13 @@ def _scaled_dequant_b_fragment_trellis256_bits( win_a, win_b, trellis_lut_addr, int(bits) ) else: - if cutlass.const_expr(self.trellis_direct_lut): - e_lo, e_hi = packed_decode_lut_e4m3_direct_to_e4m3x8( - win_a, win_b, trellis_lut_addr, int(bits), - in_shared=self.lut_e4m3_smem, - ) - else: - e_lo, e_hi = packed_decode_lut_e4m3_to_e4m3x8( - win_a, - win_b, - trellis_lut_addr, - int(bits), - value_table_in_shared=self.lut_e4m3_smem, - ) + e_lo, e_hi = packed_decode_lut_e4m3_to_e4m3x8( + win_a, + win_b, + trellis_lut_addr, + int(bits), + value_table_in_shared=self.lut_e4m3_smem, + ) if cutlass.const_expr(self.is_fp16): o0, o1 = fp8x4_e4m3_to_half2x2(e_lo) o2, o3 = fp8x4_e4m3_to_half2x2(e_hi) @@ -6303,7 +6287,6 @@ def __init__( rotation_input_dtype: str = "fp16", broadcast_suh: bool = False, pipeline_stages: int | None = None, - trellis_decode_table: str = "auto", ): activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) @@ -6615,38 +6598,10 @@ def __init__( self.shared_words += self.iq2_lut_bytes // 4 self.fc1.iq2_xs_smem_lut = True self.fc2.iq2_xs_smem_lut = True - if trellis_decode_table not in {"auto", "compact", "full"}: - raise ValueError("trellis decode table must be auto, compact, or full") - compact_smem = ( - self.trellis_codebook == LUT_E4M3 - and _lut_e4m3_smem_enabled() - ) - lut_offset = (self.shared_words * 4 + 15) // 16 * 16 - # The cooperative grid must retain its occupancy contract. A direct - # table is eligible only when it fits without reducing resident CTAs. - direct_eligible = ( + self.lut_e4m3_smem = ( self.weight_layout == "trellis_t256" and self.trellis_codebook == LUT_E4M3 - and self.intermediate_hadamard - and self.trellis_bits == 2 - and self.fc1_trellis_pair_kind is None - and self.fc2_trellis_pair_kind is None - and self.blocks_per_sm == 1 - and lut_offset + 65536 <= self.fc1.max_shared_mem - ) - if trellis_decode_table == "full" and not direct_eligible: - raise ValueError( - "full trellis decode table is incompatible with the planned kernel geometry" - ) - self.trellis_direct_lut = direct_eligible and ( - trellis_decode_table == "full" - or (trellis_decode_table == "auto" and compact_smem and size_m <= 16) - ) - self.fc1.trellis_direct_lut = self.trellis_direct_lut - self.fc2.trellis_direct_lut = self.trellis_direct_lut - self.lut_e4m3_smem = compact_smem or self.trellis_direct_lut - self.trellis_lut_smem_bytes = ( - 65536 if self.trellis_direct_lut else _LUT_E4M3_SMEM_REGION_BYTES + and _lut_e4m3_smem_enabled() ) self.lut_e4m3_smem_off = 0 if self.lut_e4m3_smem: @@ -6655,10 +6610,8 @@ def __init__( ) // 16 * 16 self.shared_words = ( self.lut_e4m3_smem_off - + self.trellis_lut_smem_bytes + + _LUT_E4M3_SMEM_REGION_BYTES ) // 4 - if self.shared_words * 4 > self.fc1.max_shared_mem: - raise ValueError(f"full trellis decode table requires {self.shared_words * 4} bytes") self.fc1.lut_e4m3_smem = True self.fc2.lut_e4m3_smem = True self.barrier_count_off = self.sms * 4 @@ -6702,7 +6655,6 @@ def __cache_key__(self) -> tuple[object, ...]: self.broadcast_suh, self.rotation_input_dtype, self.lut_e4m3_smem, - self.trellis_direct_lut, self.small_m_splitk, self.fc1.__cache_key__, self.fc2.__cache_key__, @@ -7094,8 +7046,8 @@ class Storage: fc1_trellis_lut_addr = get_ptr_as_int64(fc1_trellis_lut_flat, Int32(0)) fc2_trellis_lut_addr = get_ptr_as_int64(fc2_trellis_lut_flat, Int32(0)) - # The LUT ABI slot carries the shared byte offset for either the - # compact value table or the complete 2-bit direct table. + # The emit hooks receive the staged value table's shared byte offset + # through the LUT ABI slot. fc1_phase_lut_addr = fc1_trellis_lut_addr fc2_phase_lut_addr = fc2_trellis_lut_addr if cutlass.const_expr(self.weight_layout in IQ2_CODECS): @@ -7116,7 +7068,7 @@ class Storage: self._lut_smem_copy( fc1_trellis_lut_addr, smem_base + Int32(self.lut_e4m3_smem_off), - self.trellis_lut_smem_bytes, + _LUT_E4M3_SMEM_REGION_BYTES, tid, ) cute.arch.sync_threads() @@ -9828,7 +9780,6 @@ def compile_w4a16_fused_moe( intermediate_hadamard: bool = False, rotation_input_dtype: str | None = None, broadcast_suh: bool = False, - trellis_decode_table: str = "auto", _require_cached: bool = False, ) -> W4A16FusedMoeCompileResult: scale_format = _normalize_scale_format(scale_format) @@ -10223,7 +10174,6 @@ def compile_w4a16_fused_moe( rotation_input_dtype=rotation_input_dtype, broadcast_suh=broadcast_suh, pipeline_stages=pipeline_stages, - trellis_decode_table=trellis_decode_table, ) cache_key = ( "w4a16_fused_moe", @@ -10506,7 +10456,6 @@ def compile_w4a16_fused_moe( shared_memory_bytes=kernel.shared_words * 4, broadcast_suh=bool(broadcast_suh), small_m_direct_launches=tuple(small_m_direct_launches), - trellis_direct_lut=kernel.trellis_direct_lut, ) attach_programs(result, compiled, *(launch.compiled for launch in small_m_direct_launches)) _FUSED_CACHE[cache_key] = result @@ -11214,7 +11163,7 @@ def _w4a16_fused_moe_launch_flat( fc2_trellis_lut_addr = iq2_lut.data_ptr() elif weight_layout == "trellis_t256" and trellis_codebook != "mcg": trellis_rank_lut = _trellis256_execution_lut( - a_input.device, trellis_codebook, direct_lut=fused.trellis_direct_lut + a_input.device, trellis_codebook ) fc1_trellis_lut_addr = trellis_rank_lut.data_ptr() fc2_trellis_lut_addr = trellis_rank_lut.data_ptr() diff --git a/b12x/moe/_shared/kernels/w4a16/route_pack.py b/b12x/moe/_shared/kernels/w4a16/route_pack.py index cd4f1502d..6de60ce89 100644 --- a/b12x/moe/_shared/kernels/w4a16/route_pack.py +++ b/b12x/moe/_shared/kernels/w4a16/route_pack.py @@ -10,7 +10,6 @@ import torch from b12x._lib.compile_plan import launch_triton -from b12x._lib.env import env_flag from b12x.moe._shared.kernels.w4a16.host import route_pack_capacity @@ -23,9 +22,6 @@ _FAST_COUNT_BLOCK_T = 1024 -_STABLE_SORT_MIN_ROUTES = 4096 -_STABLE_SORT_BLOCK_T = 4096 -_STABLE_SORT_EXPERTS_PER_PROGRAM = 1 @triton.jit @@ -300,73 +296,6 @@ def _pack_topk_routes_sort_kernel( tl.store(packed_route_indices + ranks, offsets, mask=valid) -@triton.jit -def _pack_topk_routes_stable_kernel( - topk_ids, - expert_map, - packed_route_indices, - expert_offsets, - live_numel, - NUMEL_CAPACITY: tl.constexpr, - NUM_EXPERTS: tl.constexpr, - HAS_EXPERT_MAP: tl.constexpr, - BLOCK_T: tl.constexpr, - EXPERTS_PER_PROGRAM: tl.constexpr, -): - """Pack each expert's routes in ascending token-major route order.""" - expert_ids = ( - tl.program_id(0) * EXPERTS_PER_PROGRAM - + tl.arange(0, EXPERTS_PER_PROGRAM) - ) - expert_mask = expert_ids < NUM_EXPERTS - output_starts = tl.load( - expert_offsets + expert_ids, - mask=expert_mask, - other=0, - ) - output_counts = tl.zeros((EXPERTS_PER_PROGRAM,), dtype=tl.int32) - lanes = tl.arange(0, BLOCK_T) - - for start in tl.range(0, NUMEL_CAPACITY, BLOCK_T): - offsets = start + lanes - raw_ids = tl.load( - topk_ids + offsets, - mask=offsets < live_numel, - other=-1, - ).to(tl.int32) - valid = ( - (offsets < live_numel) - & (raw_ids >= 0) - & (raw_ids < NUM_EXPERTS) - ) - ids = raw_ids - if HAS_EXPERT_MAP: - safe_ids = tl.minimum(tl.maximum(raw_ids, 0), NUM_EXPERTS - 1) - ids = tl.load(expert_map + safe_ids, mask=valid, other=-1).to( - tl.int32 - ) - valid = valid & (ids >= 0) & (ids < NUM_EXPERTS) - - matches = ( - expert_mask[:, None] - & valid[None, :] - & (expert_ids[:, None] == ids[None, :]) - ) - match_i32 = matches.to(tl.int32) - local_ranks = tl.cumsum(match_i32, axis=1) - 1 - output_indices = ( - output_starts[:, None] - + output_counts[:, None] - + local_ranks - ) - tl.store( - packed_route_indices + output_indices, - offsets[None, :], - mask=matches, - ) - output_counts += tl.sum(match_i32, axis=1) - - @dataclass(frozen=True) class W4A16RoutePackLaunches: """Exact prepared Triton objects for one W4A16 packed-route geometry.""" @@ -377,7 +306,6 @@ class W4A16RoutePackLaunches: max_packed_routes: int max_route_blocks: int use_small_prefix: bool - stable_order: bool programs: MappingProxyType def program(self, name: str, topk_ids: torch.Tensor, has_expert_map: bool): @@ -400,7 +328,6 @@ def compile_w4a16_route_pack_launches( block_size: int, num_experts: int, ordinal: int, - stable_order: bool | None = None, bucket_tokens: bool = True, ) -> W4A16RoutePackLaunches: """Compile every legal route-pack helper for one planned W4A16 capacity.""" @@ -417,9 +344,6 @@ def compile_w4a16_route_pack_launches( block_route_init <= _SMALL_PREFIX_MAX_PACKED_ROUTES and block_m <= _SMALL_PREFIX_MAX_ROUTE_BLOCKS ) - if stable_order is None: - stable_order = env_flag("W4A16_STABLE_ROUTE_PACK") - stable_order = bool(stable_order and numel_capacity >= _STABLE_SORT_MIN_ROUTES) programs = {} # Avoid Triton's divisibility and value-one specialization of the live bound. live_numel = 3 @@ -447,22 +371,12 @@ def compile_w4a16_route_pack_launches( BLOCK_T=_FAST_COUNT_BLOCK_T, num_warps=4, grid=(triton.cdiv(int(tokens) * int(topk), _FAST_COUNT_BLOCK_T), 1, 1), ) - if stable_order: - programs[("sort", *key)] = _pack_topk_routes_stable_kernel.warmup( - ids_dtype, expert_map_dtype, torch.int32, torch.int32, live_numel, - NUMEL_CAPACITY=int(numel_capacity), NUM_EXPERTS=int(num_experts), - HAS_EXPERT_MAP=has_expert_map, BLOCK_T=_STABLE_SORT_BLOCK_T, - EXPERTS_PER_PROGRAM=_STABLE_SORT_EXPERTS_PER_PROGRAM, - num_warps=8, - grid=(triton.cdiv(num_experts, _STABLE_SORT_EXPERTS_PER_PROGRAM), 1, 1), - ) - else: - programs[("sort", *key)] = _pack_topk_routes_sort_kernel.warmup( - ids_dtype, expert_map_dtype, torch.int32, torch.int32, live_numel, - NUM_EXPERTS=int(num_experts), HAS_EXPERT_MAP=has_expert_map, - BLOCK_T=_SORT_BLOCK_T, num_warps=4, - grid=(triton.cdiv(int(tokens) * int(topk), _SORT_BLOCK_T), 1, 1), - ) + programs[("sort", *key)] = _pack_topk_routes_sort_kernel.warmup( + ids_dtype, expert_map_dtype, torch.int32, torch.int32, live_numel, + NUM_EXPERTS=int(num_experts), HAS_EXPERT_MAP=has_expert_map, + BLOCK_T=_SORT_BLOCK_T, num_warps=4, + grid=(triton.cdiv(int(tokens) * int(topk), _SORT_BLOCK_T), 1, 1), + ) if not use_small_prefix: programs[("prefix", torch.int32, False)] = _w4a16_route_prefix_from_counts_kernel.warmup( torch.int32, torch.int32, torch.int32, BLOCK_SIZE=int(block_size), @@ -488,7 +402,6 @@ def compile_w4a16_route_pack_launches( numel_capacity=int(numel_capacity), block_size=int(block_size), num_experts=int(num_experts), max_packed_routes=max_packed_routes, max_route_blocks=max_route_blocks, use_small_prefix=use_small_prefix, - stable_order=stable_order, programs=MappingProxyType(programs), ) @@ -765,27 +678,7 @@ def pack_topk_routes_by_expert( int(block_size), int(num_experts), max_packed_routes, max_route_blocks, _POST_PREFIX_BLOCK_T, block_e.bit_length(), ) - stable_order = ( - launches.stable_order if launches is not None else - env_flag("W4A16_STABLE_ROUTE_PACK") and numel >= _STABLE_SORT_MIN_ROUTES - ) - if stable_order: - stable_grid = (triton.cdiv(num_experts, _STABLE_SORT_EXPERTS_PER_PROGRAM),) - if launches is None: - launch_triton(_pack_topk_routes_stable_kernel, stable_grid, - topk_ids, expert_map_tensor, packed_route_indices, expert_offsets, numel, - NUMEL_CAPACITY=numel_capacity, NUM_EXPERTS=int(num_experts), - HAS_EXPERT_MAP=expert_map is not None, BLOCK_T=_STABLE_SORT_BLOCK_T, - EXPERTS_PER_PROGRAM=_STABLE_SORT_EXPERTS_PER_PROGRAM, num_warps=8, - ) - else: - _launch_prepared( - launches.program("sort", topk_ids, expert_map is not None), stable_grid, - topk_ids, expert_map_tensor, packed_route_indices, expert_offsets, numel, - numel_capacity, int(num_experts), expert_map is not None, - _STABLE_SORT_BLOCK_T, _STABLE_SORT_EXPERTS_PER_PROGRAM, - ) - elif launches is None: + if launches is None: launch_triton(_pack_topk_routes_sort_kernel, sort_grid, topk_ids, expert_map_tensor, packed_route_indices, expert_offsets, numel, NUM_EXPERTS=int(num_experts), HAS_EXPERT_MAP=expert_map is not None, diff --git a/b12x/moe/fused_moe/_impl.py b/b12x/moe/fused_moe/_impl.py index 514860d99..620fa7e10 100644 --- a/b12x/moe/fused_moe/_impl.py +++ b/b12x/moe/fused_moe/_impl.py @@ -839,8 +839,6 @@ class TPMoEScratchCaps: w4a16_block_size_m: int | None = None w4a16_fast_math: bool = True w4a16_prefill_fused_sum: bool | None = None - w4a16_stable_route_pack: bool | None = None - trellis_decode_table: str = "auto" frozen: bool = True def __post_init__(self) -> None: @@ -894,19 +892,12 @@ def __post_init__(self) -> None: raise ValueError("w4a16_block_size_m must be one of 8, 16, 32, 48, 64") object.__setattr__(self, "w4a16_block_size_m", block_size_m) object.__setattr__(self, "w4a16_fast_math", bool(self.w4a16_fast_math)) - if self.trellis_decode_table not in {"auto", "compact", "full"}: - raise ValueError("trellis decode table must be auto, compact, or full") if self.w4a16_prefill_fused_sum is None: from b12x.moe._shared.kernels.w4a16.host import prefill_fused_sum_enabled object.__setattr__( self, "w4a16_prefill_fused_sum", prefill_fused_sum_enabled() ) - if self.w4a16_stable_route_pack is None: - object.__setattr__( - self, "w4a16_stable_route_pack", - _env_flag("B12X_W4A16_STABLE_ROUTE_PACK", default=False), - ) object.__setattr__(self, "frozen", bool(self.frozen)) @property @@ -8129,7 +8120,6 @@ def compile_fused(token_count: int, broadcast_suh: bool) -> object: w13_layout=w13_layout, trellis_bits=core_plan.trellis_bits, trellis_codebook=core_plan.trellis_codebook or LUT_E4M3, - trellis_decode_table=caps.trellis_decode_table, force_tile_config=core_plan.trellis_tile_config, intermediate_rotation=True, full_rotation=True, @@ -8720,7 +8710,6 @@ def plan_tp_moe_scratch( if core_workspace_plan.device.index is not None else torch.cuda.current_device() ), - stable_order=caps.w4a16_stable_route_pack, bucket_tokens=False, ) return TPMoEScratchPlan( diff --git a/b12x/moe/fused_moe/_preparation.py b/b12x/moe/fused_moe/_preparation.py index a9baafecd..d9b9cf772 100644 --- a/b12x/moe/fused_moe/_preparation.py +++ b/b12x/moe/fused_moe/_preparation.py @@ -94,15 +94,13 @@ def _control_snapshot() -> FrozenMapping: """Capture host controls once while declaring the immutable query.""" from . import _impl from b12x.moe._shared.kernels.w4a16.host import ( - prefill_fused_sum_enabled, trellis_decode_table, + prefill_fused_sum_enabled, ) tile = _impl._dynamic_tile_mn_override() raw_materialized = _impl.os.environ.get(_impl._DYNAMIC_NVFP4_MATERIALIZED_ENV) return FrozenMapping({ - "trellis_decode_table": trellis_decode_table(), "w4a16_prefill_fused_sum": prefill_fused_sum_enabled(), - "w4a16_stable_route_pack": _impl._env_flag("B12X_W4A16_STABLE_ROUTE_PACK", default=False), "dynamic_nvfp4_materialized": ( None if raw_materialized is None else raw_materialized not in ("", "0", "false", "False") ), @@ -255,8 +253,6 @@ def _lower_caps( w4a16_prefill_fused_sum=bool( query.controls.get("w4a16_prefill_fused_sum", False) ), - w4a16_stable_route_pack=bool(query.controls.get("w4a16_stable_route_pack", False)), - trellis_decode_table=str(query.controls.get("trellis_decode_table", "auto")), swiglu_beta=_decode_scalar(query.swiglu_beta), ) @@ -444,7 +440,6 @@ def _w4a16_primary_launches(scratch, caps) -> _W4A16PrimaryLaunches: route_pack = compile_w4a16_route_pack_launches( tokens=tokens, topk=core.num_topk, block_size=int(route_block), num_experts=core.route_E, ordinal=core.device.index, - stable_order=caps.w4a16_stable_route_pack, ) # Direct routing requires exact M; packed routing accepts live M up to capacity. return _W4A16PrimaryLaunches( diff --git a/b12x/moe/fused_moe/_tuning.py b/b12x/moe/fused_moe/_tuning.py index bbace1042..2ae941cbc 100644 --- a/b12x/moe/fused_moe/_tuning.py +++ b/b12x/moe/fused_moe/_tuning.py @@ -233,8 +233,6 @@ def _validate_query(query: MoeDecodeQuery, _device: DeviceIdentity | None) -> No raise TypeError("query must be MoeDecodeQuery") if not isinstance(query.controls, FrozenMapping): raise TypeError("MoE controls must be frozen declaration metadata") - if query.controls.get("trellis_decode_table", "auto") not in {"auto", "compact", "full"}: - raise ValueError("trellis decode table must be auto, compact, or full") if query.io_dtype not in {"bfloat16", "float16"}: raise TypeError("MoE I/O dtype must be bfloat16 or float16") if query.route_logits_dtype not in {None, "float16", "bfloat16", "float32"}: @@ -636,7 +634,7 @@ def _validate_fc2_query(query: MoeFC2Query, _device) -> None: TUNING = TuningContract( component_id="moe.decode", - query_schema_version=13, + query_schema_version=14, config_schema_version=8, query_fields=frozenset(MoeDecodeQuery.__dataclass_fields__), config_fields=frozenset(MoeDecodeConfig.__dataclass_fields__), diff --git a/docs/trellis-decode-table.md b/docs/trellis-decode-table.md deleted file mode 100644 index 6273678f3..000000000 --- a/docs/trellis-decode-table.md +++ /dev/null @@ -1,39 +0,0 @@ -# Trellis decode table - -Fused W4A16 MoE can decode uniform 2-bit `lut_e4m3` trellis weights with an -intermediate Hadamard through a 64 KiB table in each thread block's shared -memory. The table stores the complete codeword-to-E4M3 mapping. It replaces -the per-window permutation arithmetic; it does not change checkpoint weights, -activation precision, matrix multiplication, or reduction order. - -`B12X_TRELLIS_DECODE_TABLE` controls canonical fused-MoE preparation: - -| Value | Behavior | -| --- | --- | -| `auto` (default) | Use the full table for planned capacities of at most 16 tokens when the shared-memory and cooperative-residency contracts permit it. Other capacities and weight formats keep compact decoding. | -| `compact` | Compute the permutation per window and read the 4 KiB value table. | -| `full` | Require the full table; fail preparation for an incompatible kernel geometry. This is intended for component qualification. | - -For an explicit reference run, set the variable before declaring plans: - -```bash -export B12X_TRELLIS_DECODE_TABLE=compact -``` - -Canonical preparation saves the setting in its immutable query and propagates -it to compiler workers. Changing the environment after declaration cannot -change that plan. Binding and replay retain compiled programs and use -caller-owned scratch. The table is process-shared immutable execution data, -not a decoded copy of model weights. - -The full table is admitted only for uniform 2-bit `lut_e4m3` weights with an -intermediate Hadamard, one resident cooperative CTA per SM, and enough shared -memory for both the original GEMM staging and the 64 KiB table. It never -reduces the planned resident-CTA count. Large-capacity prefill keeps its -existing staging footprint. - -`tests/moe/test_trellis_direct_lut_decode.py` checks all 65,536 codewords and -random overlapping windows against the compact decoder, including shared -tables for each supported bit width. Only uniform 2-bit weights use the full -table in fused MoE; testing the lookup primitive at 3 and 4 bits does not -enable those execution paths. diff --git a/docs/w4a16-prefill-reduction.md b/docs/w4a16-prefill-reduction.md index eef785df8..4bc2bd9d1 100644 --- a/docs/w4a16-prefill-reduction.md +++ b/docs/w4a16-prefill-reduction.md @@ -17,12 +17,7 @@ FC1; its mandatory grid barrier orders initialization before FC2 reduction. The output is cast to BF16 once after FC2. Relaxed FP32 atomic ordering is not a bitwise-deterministic summation contract. -Set `B12X_W4A16_STABLE_ROUTE_PACK=1` before plan declaration to retain -ascending token-major route order within each expert for capacities of at -least 4096 routed rows. This stabilizes expert packing, not floating-point -atomic accumulation order. Smaller capacity plans retain atomic packing. - -Both options are captured in the immutable MoE query. Changing the process +The option is captured in the immutable MoE query. Changing the process environment after declaration does not alter scratch geometry or dispatch. Preparation compiles and retains capacity-specialized programs. Binding maps caller-owned storage and accepts live token counts within that capacity; diff --git a/tests/moe/test_fused_moe_variant_selection.py b/tests/moe/test_fused_moe_variant_selection.py index fd6c084f6..ac2d9ea0e 100644 --- a/tests/moe/test_fused_moe_variant_selection.py +++ b/tests/moe/test_fused_moe_variant_selection.py @@ -11,56 +11,9 @@ _FusedMoeCapacityState, _W4A16PrimaryLaunches, variant_for, - _control_snapshot, ) -@pytest.mark.parametrize("mode", ["auto", "compact", "full"]) -def test_trellis_decode_table_is_frozen_at_declaration(monkeypatch, mode): - monkeypatch.setenv("B12X_TRELLIS_DECODE_TABLE", mode) - controls = _control_snapshot() - monkeypatch.setenv("B12X_TRELLIS_DECODE_TABLE", "invalid") - assert controls["trellis_decode_table"] == mode - with pytest.raises(ValueError, match="B12X_TRELLIS_DECODE_TABLE"): - _control_snapshot() - - -@pytest.mark.parametrize("rows,mode,expected,shared_limit", [ - (1, "auto", True, 101376), (8, "auto", True, 101376), - (16, "auto", True, 101376), (17, "auto", False, 101376), - (8, "compact", False, 101376), (8, "full", True, 101376), - (8, "auto", False, 65536), -]) -def test_trellis_direct_table_preserves_cooperative_residency( - monkeypatch, rows, mode, expected, shared_limit, -): - from b12x.moe._shared.kernels.w4a16.kernel import W4A16FusedMoeKernel - - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) - monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace( - multi_processor_count=188, shared_memory_per_block_optin=shared_limit, - )) - monkeypatch.setenv("B12X_LUT_E4M3_SMEM", "1") - kernel = W4A16FusedMoeKernel( - size_m=rows, hidden_size=3584, intermediate_size=384, - num_experts=512, top_k=16, activation="silu", - apply_router_weight_on_input=False, zero_fc2_output=False, - fc1_tile_n=128, fc1_tile_k=128, fc2_tile_n=128, fc2_tile_k=128, - moe_block_size=8, max_m_blocks=1024, element_dtype="fp16", - weight_layout="trellis_t256", scale_format="e4m3_k32", - w13_layout="trellis_t256_proj", trellis_bits=2, - full_rotation=True, intermediate_rotation=True, intermediate_hadamard=True, - trellis_decode_table=mode, - ) - assert kernel.trellis_direct_lut == expected - assert kernel.shared_words * 4 <= kernel.fc1.max_shared_mem - if expected: - assert kernel.blocks_per_sm == 1 - assert kernel.lut_e4m3_smem - assert kernel.fc1.trellis_direct_lut and kernel.fc2.trellis_direct_lut - - class _Variant: def __init__(self, count): self.count = count diff --git a/tests/moe/test_trellis_direct_lut_decode.py b/tests/moe/test_trellis_direct_lut_decode.py index 443e4a091..77f3129d8 100644 --- a/tests/moe/test_trellis_direct_lut_decode.py +++ b/tests/moe/test_trellis_direct_lut_decode.py @@ -10,7 +10,6 @@ from __future__ import annotations import cuda.bindings.driver as cuda -import cutlass import cutlass.cute as cute import pytest import torch @@ -21,7 +20,6 @@ from b12x._lib.intrinsics import ( packed_decode_lut_e4m3_to_e4m3x8, packed_decode_lut_e4m3_direct_to_e4m3x8, - shared_ptr_to_u32, ) from b12x._lib.quant.lut_e4m3 import ( lut_e4m3_direct_table_cpu, @@ -34,9 +32,8 @@ class _DecodeProbe: - def __init__(self, bits: int, shared: bool = False): + def __init__(self, bits: int): self.bits = int(bits) - self.shared = shared @cute.jit def __call__( @@ -48,7 +45,7 @@ def __call__( stream: cuda.CUstream, ): self.kernel(wins, value_table, direct, out).launch( - grid=(cute.size(wins) // 64, 1, 1), block=[32, 1, 1], stream=stream + grid=(1, 1, 1), block=[32, 1, 1], stream=stream ) @cute.kernel @@ -60,26 +57,16 @@ def kernel( out: cute.Tensor, ): tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - lane = Int32(tidx) + Int32(bidx) * 32 + lane = Int32(tidx) value_table_addr = value_table.iterator.toint() dir_addr = direct.iterator.toint() - if cutlass.const_expr(self.shared): - table = cutlass.utils.SmemAllocator().allocate_tensor( - cutlass.Uint8, cute.make_layout(65536), byte_alignment=16 - ) - for offset in range(Int32(tidx), 65536, 32): - table[offset] = direct[(self.bits - 2) * 65536 + offset] - cute.arch.sync_threads() - dir_addr = shared_ptr_to_u32(table.iterator) wa = Uint32(wins[2 * lane]) wb = Uint32(wins[2 * lane + 1]) lo_t, hi_t = packed_decode_lut_e4m3_to_e4m3x8( wa, wb, value_table_addr, self.bits, value_table_in_shared=False ) lo_d, hi_d = packed_decode_lut_e4m3_direct_to_e4m3x8( - wa, wb, dir_addr, self.bits, rate_indexed=not self.shared, - in_shared=self.shared, + wa, wb, dir_addr, self.bits, rate_indexed=True ) out[4 * lane] = Int32(lo_t) out[4 * lane + 1] = Int32(hi_t) @@ -89,8 +76,7 @@ def kernel( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("bits", [2, 3, 4]) -@pytest.mark.parametrize("shared", [False, True]) -def test_direct_table_decode_bit_equals_value_table(bits: int, shared: bool) -> None: +def test_direct_table_decode_bit_equals_value_table(bits: int) -> None: device = torch.device("cuda") torch.manual_seed(20260812 + bits) wins = torch.zeros(64, dtype=torch.int32, device=device) @@ -107,7 +93,7 @@ def args(): current_cuda_stream(), ) - compiled = b12x_compile(_DecodeProbe(bits, shared), *args()) + compiled = b12x_compile(_DecodeProbe(bits), *args()) mismatched = 0 for _ in range(512): wins.copy_( @@ -121,25 +107,3 @@ def args(): mismatched += int((o[:, 0] != o[:, 2]).sum()) mismatched += int((o[:, 1] != o[:, 3]).sum()) assert mismatched == 0, mismatched - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -@pytest.mark.parametrize("bits", [2, 3, 4]) -def test_shared_direct_lut_covers_all_codewords(bits: int) -> None: - # Both 16-bit windows visit every codeword; random high bits also exercise - # the shifted windows that cross the low-16-bit boundary. - device = torch.device("cuda") - torch.manual_seed(731 + bits) - states = torch.arange(65536, device=device, dtype=torch.int32) - upper = torch.randint(0, 65536, (65536, 2), device=device, dtype=torch.int32) - wins = ((upper << 16) | states[:, None]).flatten() - value_table = lut_e4m3_value_table_cpu().to(device) - direct = lut_e4m3_direct_table_cpu().to(device) - out = torch.empty(65536 * 4, device=device, dtype=torch.int32) - args = tuple(from_dlpack(tensor, assumed_align=16) - for tensor in (wins, value_table, direct, out)) + (current_cuda_stream(),) - compiled = b12x_compile(_DecodeProbe(bits, shared=True), *args) - compiled(*args) - torch.cuda.synchronize() - result = out.view(-1, 4) - torch.testing.assert_close(result[:, :2], result[:, 2:], rtol=0, atol=0) diff --git a/tests/moe/test_w4a16_route_pack.py b/tests/moe/test_w4a16_route_pack.py index b136dab0c..7151578bc 100644 --- a/tests/moe/test_w4a16_route_pack.py +++ b/tests/moe/test_w4a16_route_pack.py @@ -157,12 +157,10 @@ def test_route_pack_reuses_provided_fixed_capacity_for_prefill_tail() -> None: @pytest.mark.parametrize("capacity, bucket_tokens", ((128, True), (129, False), (8192, True))) @pytest.mark.parametrize("ids_dtype", (torch.int32, torch.int64)) @pytest.mark.parametrize("mapped", (False, True)) -@pytest.mark.parametrize("stable", (False, True)) -def test_prepared_route_pack_uses_live_bounds_with_fixed_geometry(capacity, bucket_tokens, ids_dtype, mapped, stable): +def test_prepared_route_pack_uses_live_bounds_with_fixed_geometry(capacity, bucket_tokens, ids_dtype, mapped): launches = route_pack_module.compile_w4a16_route_pack_launches( tokens=capacity, topk=2, block_size=8, num_experts=32, ordinal=torch.cuda.current_device(), - stable_order=stable, bucket_tokens=bucket_tokens, ) if not bucket_tokens: @@ -199,10 +197,6 @@ def check(rows, output): payload = host_routes[host_routes < rows * 2] expected_payload = torch.nonzero(valid).flatten() torch.testing.assert_close(payload.sort().values, expected_payload, rtol=0, atol=0) - if launches.stable_order: - for expert in range(32): - expert_routes = payload[expected_ids[payload] == expert] - assert torch.equal(expert_routes, expert_routes.sort().values) for block, expert in enumerate(expected_blocks.tolist()): block_routes = host_routes[block * 8:(block + 1) * 8] block_payload = block_routes[block_routes < rows * 2] From f2ecd274430e32014200b3e76e5944cdabba548e Mon Sep 17 00:00:00 2001 From: Luke Alonso Date: Sun, 27 Sep 2026 08:15:37 -0700 Subject: [PATCH 4/4] Isolate the prepared W4A16 full-codeword LUT experiment Restore the 64 KiB shared codeword table, immutable selection control, residency checks and decoding tests as an independently reviewable change. Preserve the original auto/compact/full behavior; qualification remains pending. Document the instruction-scheduling difference that a matched compact-decoder benchmark must control. Validation: adapter, config, schema and variant-selection suites: 170 passed. Python source matches the extracted implementation except for the query schema version; syntax and whitespace checks pass. CUDA execution, graph replay and performance were not run. Source-commit: 7fca2aa55d5eeb7116a4c09633632cc5f02d850d Assisted-by: OpenAI Codex --- b12x/_lib/intrinsics.py | 44 ++++++---- b12x/moe/_shared/kernels/w4a16/host.py | 8 ++ b12x/moe/_shared/kernels/w4a16/kernel.py | 83 +++++++++++++++---- b12x/moe/fused_moe/_impl.py | 4 + b12x/moe/fused_moe/_preparation.py | 4 +- b12x/moe/fused_moe/_tuning.py | 4 +- docs/trellis-decode-table.md | 49 +++++++++++ tests/moe/test_fused_moe_variant_selection.py | 47 +++++++++++ tests/moe/test_trellis_direct_lut_decode.py | 48 +++++++++-- 9 files changed, 250 insertions(+), 41 deletions(-) create mode 100644 docs/trellis-decode-table.md diff --git a/b12x/_lib/intrinsics.py b/b12x/_lib/intrinsics.py index 3f2ecd551..f549fbb6f 100644 --- a/b12x/_lib/intrinsics.py +++ b/b12x/_lib/intrinsics.py @@ -6834,6 +6834,7 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( bits: int = 3, rate_indexed: bool = True, *, + in_shared: bool = False, loc=None, ip=None, ): @@ -6841,8 +6842,8 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( This is the lookup lower-bound primitive: one byte load per reconstructed weight. Every slot owns its own address/code registers so all eight - global gathers are in flight together; a shared scratch pair would chain - them behind full global-memory latency each. + gathers are in flight together. With ``in_shared``, the address names a + caller-staged shared-memory table instead of a global-memory table. """ bits = int(bits) @@ -6853,41 +6854,47 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( table_offset = ((bits - 2) << 16) if rate_indexed else 0 extract_lines: list[str] = [] load_lines: list[str] = [] - pack_lines: list[str] = [] for index in range(8): source = "$3" if index < 4 else "$2" shift = (3 - (index & 3)) * bits # Accumulate into declared scratch registers, never directly into # the "=r" outputs: without early-clobber the outputs may alias the # window inputs that later iterations still read. - target = "out0" if index < 4 else "out1" - byte_shift = 8 * (index & 3) if shift: extract_lines.append( f"bfe.u32 w{index}, {source}, {shift}, 16;" ) else: extract_lines.append(f"and.b32 w{index}, {source}, 0xffff;") - load_lines.append( - f""" + if in_shared: + load_lines.append( + f"add.u32 w{index}, w{index}, $4;" + f" ld.shared.u8 w{index}, [w{index}];" + ) + else: + load_lines.append( + f""" cvt.u64.u32 addr{index}, w{index}; add.u64 addr{index}, addr{index}, $4; ld.global.u8 w{index}, [addr{index}]; - """ - ) - if byte_shift: - pack_lines.append( - f"shl.b32 w{index}, w{index}, {byte_shift};" + """ ) - pack_lines.append(f"or.b32 {target}, {target}, w{index};") + # Each load returns one zero-extended byte. Byte permutations pack four + # independent results without a serial shift/or accumulator. + pack_lines = [ + "prmt.b32 w0, w0, w1, 0x4040;", + "prmt.b32 w2, w2, w3, 0x4040;", + "prmt.b32 out0, w0, w2, 0x5410;", + "prmt.b32 w4, w4, w5, 0x4040;", + "prmt.b32 w6, w6, w7, 0x4040;", + "prmt.b32 out1, w4, w6, 0x5410;", + ] asm = ( """ { .reg .b32 out0,out1; .reg .b32 w0,w1,w2,w3,w4,w5,w6,w7; .reg .b64 addr0,addr1,addr2,addr3,addr4,addr5,addr6,addr7; - mov.b32 out0, 0; - mov.b32 out1, 0; """ + "\n".join(extract_lines) + "\n".join(load_lines) @@ -6897,7 +6904,10 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( mov.b32 $1, out1; }""" ) - table_base = Int64(direct_lut_addr) + Int64(table_offset) + table_base = ( + Int32(direct_lut_addr) + Int32(table_offset) + if in_shared else Int64(direct_lut_addr) + Int64(table_offset) + ) result = llvm.inline_asm( llvm.StructType.get_literal([T.i32(), T.i32()]), [ @@ -6906,7 +6916,7 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8( table_base.ir_value(loc=loc, ip=ip), ], asm, - "=r,=r,r,r,l", + "=r,=r,r,r,r" if in_shared else "=r,=r,r,r,l", has_side_effects=False, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, diff --git a/b12x/moe/_shared/kernels/w4a16/host.py b/b12x/moe/_shared/kernels/w4a16/host.py index 7492b79c2..7a12f9e2a 100644 --- a/b12x/moe/_shared/kernels/w4a16/host.py +++ b/b12x/moe/_shared/kernels/w4a16/host.py @@ -21,6 +21,14 @@ _SUPPORTED_ACTIVATIONS = SUPPORTED_MOE_ACTIVATIONS +def trellis_decode_table() -> str: + """Read the plan-time trellis decode-table selection; replay never reads it.""" + value = os.environ.get("B12X_TRELLIS_DECODE_TABLE", "auto") + if value not in {"auto", "compact", "full"}: + raise ValueError("B12X_TRELLIS_DECODE_TABLE must be auto, compact, or full") + return value + + def prefill_fused_sum_enabled() -> bool: """Enable direct FP32 route reduction for large-M W4A16 launches. diff --git a/b12x/moe/_shared/kernels/w4a16/kernel.py b/b12x/moe/_shared/kernels/w4a16/kernel.py index e8ff2b369..cc48cf3dc 100644 --- a/b12x/moe/_shared/kernels/w4a16/kernel.py +++ b/b12x/moe/_shared/kernels/w4a16/kernel.py @@ -72,6 +72,7 @@ packed_decode_lut_fp16_to_bfloat2x4, packed_decode_lut_fp16_to_half2x4, packed_decode_lut_e4m3_to_e4m3x8, + packed_decode_lut_e4m3_direct_to_e4m3x8, ld_global_nc_v4_u32, pack_f32x2_to_bfloat2, pack_f32x2_to_f16x2, @@ -95,7 +96,10 @@ warp_reduce, ) from b12x._lib.quant.iq2_xs import iq2_xs_execution_lut -from b12x._lib.quant.lut_e4m3 import lut_e4m3_value_table +from b12x._lib.quant.lut_e4m3 import ( + lut_e4m3_direct_table, + lut_e4m3_value_table, +) from b12x.moe._shared.kernels.trellis_ring import ( trellis256_lane_geom_bits as _trellis_ring_lane_geom_bits, ) @@ -298,10 +302,12 @@ def _gather_native_scale_bytes(a: Uint32, b: Uint32, c: Uint32, d: Uint32, byte: def _trellis256_execution_lut( - device: torch.device | str, codebook: str + device: torch.device | str, codebook: str, *, direct_lut: bool = False ) -> torch.Tensor: if codebook == LUT_FP16: return lut_fp16_segment_table(device) + if direct_lut: + return lut_e4m3_direct_table(device) return lut_e4m3_value_table(device) # TC-decode runs on the packed W4A16 object and folds the top-k sum into the FC2 @@ -879,6 +885,7 @@ class W4A16FusedMoeCompileResult: shared_memory_bytes: int = -1 broadcast_suh: bool = False small_m_direct_launches: tuple[_W4A16SmallMDirectLaunch, ...] = () + trellis_direct_lut: bool = False @dataclass(frozen=True) @@ -1215,6 +1222,7 @@ def __init__( static_pair_rates.get(trellis_pair_kind, (3, 3)) ) self.lut_e4m3_smem = False + self.trellis_direct_lut = False # Small-M stripe split-K: opt out of the one-tile-per-CTA fast path # so decode-heavy small-M phases spread each mn-tile's K range across # multiple CTAs (existing tail scheduling plus cross-CTA finalize). @@ -1342,6 +1350,7 @@ def __init__( else: self.sms = 120 max_shared_mem = _DEFAULT_MAX_SHARED_MEM + self.max_shared_mem = max_shared_mem self.blocks_per_sm = _determine_blocks_per_sm( problem_m=self.size_m, problem_n=self.covered_size_n, @@ -1504,6 +1513,7 @@ def __cache_key__(self) -> tuple[object, ...]: self.schedule_whole_tiles, self.schedule_route_block_factor, self.lut_e4m3_smem, + self.trellis_direct_lut, self.small_m_splitk, ) @@ -4096,13 +4106,19 @@ def _scaled_dequant_b_fragment_trellis256_bits( win_a, win_b, trellis_lut_addr, int(bits) ) else: - e_lo, e_hi = packed_decode_lut_e4m3_to_e4m3x8( - win_a, - win_b, - trellis_lut_addr, - int(bits), - value_table_in_shared=self.lut_e4m3_smem, - ) + if cutlass.const_expr(self.trellis_direct_lut): + e_lo, e_hi = packed_decode_lut_e4m3_direct_to_e4m3x8( + win_a, win_b, trellis_lut_addr, int(bits), + in_shared=self.lut_e4m3_smem, + ) + else: + e_lo, e_hi = packed_decode_lut_e4m3_to_e4m3x8( + win_a, + win_b, + trellis_lut_addr, + int(bits), + value_table_in_shared=self.lut_e4m3_smem, + ) if cutlass.const_expr(self.is_fp16): o0, o1 = fp8x4_e4m3_to_half2x2(e_lo) o2, o3 = fp8x4_e4m3_to_half2x2(e_hi) @@ -6287,6 +6303,7 @@ def __init__( rotation_input_dtype: str = "fp16", broadcast_suh: bool = False, pipeline_stages: int | None = None, + trellis_decode_table: str = "auto", ): activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) @@ -6598,10 +6615,38 @@ def __init__( self.shared_words += self.iq2_lut_bytes // 4 self.fc1.iq2_xs_smem_lut = True self.fc2.iq2_xs_smem_lut = True - self.lut_e4m3_smem = ( + if trellis_decode_table not in {"auto", "compact", "full"}: + raise ValueError("trellis decode table must be auto, compact, or full") + compact_smem = ( + self.trellis_codebook == LUT_E4M3 + and _lut_e4m3_smem_enabled() + ) + lut_offset = (self.shared_words * 4 + 15) // 16 * 16 + # The cooperative grid must retain its occupancy contract. A direct + # table is eligible only when it fits without reducing resident CTAs. + direct_eligible = ( self.weight_layout == "trellis_t256" and self.trellis_codebook == LUT_E4M3 - and _lut_e4m3_smem_enabled() + and self.intermediate_hadamard + and self.trellis_bits == 2 + and self.fc1_trellis_pair_kind is None + and self.fc2_trellis_pair_kind is None + and self.blocks_per_sm == 1 + and lut_offset + 65536 <= self.fc1.max_shared_mem + ) + if trellis_decode_table == "full" and not direct_eligible: + raise ValueError( + "full trellis decode table is incompatible with the planned kernel geometry" + ) + self.trellis_direct_lut = direct_eligible and ( + trellis_decode_table == "full" + or (trellis_decode_table == "auto" and compact_smem and size_m <= 16) + ) + self.fc1.trellis_direct_lut = self.trellis_direct_lut + self.fc2.trellis_direct_lut = self.trellis_direct_lut + self.lut_e4m3_smem = compact_smem or self.trellis_direct_lut + self.trellis_lut_smem_bytes = ( + 65536 if self.trellis_direct_lut else _LUT_E4M3_SMEM_REGION_BYTES ) self.lut_e4m3_smem_off = 0 if self.lut_e4m3_smem: @@ -6610,8 +6655,10 @@ def __init__( ) // 16 * 16 self.shared_words = ( self.lut_e4m3_smem_off - + _LUT_E4M3_SMEM_REGION_BYTES + + self.trellis_lut_smem_bytes ) // 4 + if self.shared_words * 4 > self.fc1.max_shared_mem: + raise ValueError(f"full trellis decode table requires {self.shared_words * 4} bytes") self.fc1.lut_e4m3_smem = True self.fc2.lut_e4m3_smem = True self.barrier_count_off = self.sms * 4 @@ -6655,6 +6702,7 @@ def __cache_key__(self) -> tuple[object, ...]: self.broadcast_suh, self.rotation_input_dtype, self.lut_e4m3_smem, + self.trellis_direct_lut, self.small_m_splitk, self.fc1.__cache_key__, self.fc2.__cache_key__, @@ -7046,8 +7094,8 @@ class Storage: fc1_trellis_lut_addr = get_ptr_as_int64(fc1_trellis_lut_flat, Int32(0)) fc2_trellis_lut_addr = get_ptr_as_int64(fc2_trellis_lut_flat, Int32(0)) - # The emit hooks receive the staged value table's shared byte offset - # through the LUT ABI slot. + # The LUT ABI slot carries the shared byte offset for either the + # compact value table or the complete 2-bit direct table. fc1_phase_lut_addr = fc1_trellis_lut_addr fc2_phase_lut_addr = fc2_trellis_lut_addr if cutlass.const_expr(self.weight_layout in IQ2_CODECS): @@ -7068,7 +7116,7 @@ class Storage: self._lut_smem_copy( fc1_trellis_lut_addr, smem_base + Int32(self.lut_e4m3_smem_off), - _LUT_E4M3_SMEM_REGION_BYTES, + self.trellis_lut_smem_bytes, tid, ) cute.arch.sync_threads() @@ -9780,6 +9828,7 @@ def compile_w4a16_fused_moe( intermediate_hadamard: bool = False, rotation_input_dtype: str | None = None, broadcast_suh: bool = False, + trellis_decode_table: str = "auto", _require_cached: bool = False, ) -> W4A16FusedMoeCompileResult: scale_format = _normalize_scale_format(scale_format) @@ -10174,6 +10223,7 @@ def compile_w4a16_fused_moe( rotation_input_dtype=rotation_input_dtype, broadcast_suh=broadcast_suh, pipeline_stages=pipeline_stages, + trellis_decode_table=trellis_decode_table, ) cache_key = ( "w4a16_fused_moe", @@ -10456,6 +10506,7 @@ def compile_w4a16_fused_moe( shared_memory_bytes=kernel.shared_words * 4, broadcast_suh=bool(broadcast_suh), small_m_direct_launches=tuple(small_m_direct_launches), + trellis_direct_lut=kernel.trellis_direct_lut, ) attach_programs(result, compiled, *(launch.compiled for launch in small_m_direct_launches)) _FUSED_CACHE[cache_key] = result @@ -11163,7 +11214,7 @@ def _w4a16_fused_moe_launch_flat( fc2_trellis_lut_addr = iq2_lut.data_ptr() elif weight_layout == "trellis_t256" and trellis_codebook != "mcg": trellis_rank_lut = _trellis256_execution_lut( - a_input.device, trellis_codebook + a_input.device, trellis_codebook, direct_lut=fused.trellis_direct_lut ) fc1_trellis_lut_addr = trellis_rank_lut.data_ptr() fc2_trellis_lut_addr = trellis_rank_lut.data_ptr() diff --git a/b12x/moe/fused_moe/_impl.py b/b12x/moe/fused_moe/_impl.py index 620fa7e10..ecc28ffc1 100644 --- a/b12x/moe/fused_moe/_impl.py +++ b/b12x/moe/fused_moe/_impl.py @@ -839,6 +839,7 @@ class TPMoEScratchCaps: w4a16_block_size_m: int | None = None w4a16_fast_math: bool = True w4a16_prefill_fused_sum: bool | None = None + trellis_decode_table: str = "auto" frozen: bool = True def __post_init__(self) -> None: @@ -892,6 +893,8 @@ def __post_init__(self) -> None: raise ValueError("w4a16_block_size_m must be one of 8, 16, 32, 48, 64") object.__setattr__(self, "w4a16_block_size_m", block_size_m) object.__setattr__(self, "w4a16_fast_math", bool(self.w4a16_fast_math)) + if self.trellis_decode_table not in {"auto", "compact", "full"}: + raise ValueError("trellis decode table must be auto, compact, or full") if self.w4a16_prefill_fused_sum is None: from b12x.moe._shared.kernels.w4a16.host import prefill_fused_sum_enabled @@ -8120,6 +8123,7 @@ def compile_fused(token_count: int, broadcast_suh: bool) -> object: w13_layout=w13_layout, trellis_bits=core_plan.trellis_bits, trellis_codebook=core_plan.trellis_codebook or LUT_E4M3, + trellis_decode_table=caps.trellis_decode_table, force_tile_config=core_plan.trellis_tile_config, intermediate_rotation=True, full_rotation=True, diff --git a/b12x/moe/fused_moe/_preparation.py b/b12x/moe/fused_moe/_preparation.py index d9b9cf772..a167b8d61 100644 --- a/b12x/moe/fused_moe/_preparation.py +++ b/b12x/moe/fused_moe/_preparation.py @@ -94,12 +94,13 @@ def _control_snapshot() -> FrozenMapping: """Capture host controls once while declaring the immutable query.""" from . import _impl from b12x.moe._shared.kernels.w4a16.host import ( - prefill_fused_sum_enabled, + prefill_fused_sum_enabled, trellis_decode_table, ) tile = _impl._dynamic_tile_mn_override() raw_materialized = _impl.os.environ.get(_impl._DYNAMIC_NVFP4_MATERIALIZED_ENV) return FrozenMapping({ + "trellis_decode_table": trellis_decode_table(), "w4a16_prefill_fused_sum": prefill_fused_sum_enabled(), "dynamic_nvfp4_materialized": ( None if raw_materialized is None else raw_materialized not in ("", "0", "false", "False") @@ -253,6 +254,7 @@ def _lower_caps( w4a16_prefill_fused_sum=bool( query.controls.get("w4a16_prefill_fused_sum", False) ), + trellis_decode_table=str(query.controls.get("trellis_decode_table", "auto")), swiglu_beta=_decode_scalar(query.swiglu_beta), ) diff --git a/b12x/moe/fused_moe/_tuning.py b/b12x/moe/fused_moe/_tuning.py index 2ae941cbc..f17cc30b1 100644 --- a/b12x/moe/fused_moe/_tuning.py +++ b/b12x/moe/fused_moe/_tuning.py @@ -233,6 +233,8 @@ def _validate_query(query: MoeDecodeQuery, _device: DeviceIdentity | None) -> No raise TypeError("query must be MoeDecodeQuery") if not isinstance(query.controls, FrozenMapping): raise TypeError("MoE controls must be frozen declaration metadata") + if query.controls.get("trellis_decode_table", "auto") not in {"auto", "compact", "full"}: + raise ValueError("trellis decode table must be auto, compact, or full") if query.io_dtype not in {"bfloat16", "float16"}: raise TypeError("MoE I/O dtype must be bfloat16 or float16") if query.route_logits_dtype not in {None, "float16", "bfloat16", "float32"}: @@ -634,7 +636,7 @@ def _validate_fc2_query(query: MoeFC2Query, _device) -> None: TUNING = TuningContract( component_id="moe.decode", - query_schema_version=14, + query_schema_version=15, config_schema_version=8, query_fields=frozenset(MoeDecodeQuery.__dataclass_fields__), config_fields=frozenset(MoeDecodeConfig.__dataclass_fields__), diff --git a/docs/trellis-decode-table.md b/docs/trellis-decode-table.md new file mode 100644 index 000000000..39ddd7b5e --- /dev/null +++ b/docs/trellis-decode-table.md @@ -0,0 +1,49 @@ +# Trellis decode table + +Status: research-only. CUDA execution, graph replay and performance require +qualification on the prepared execution path before serving adoption. + +Fused W4A16 MoE can decode uniform 2-bit `lut_e4m3` trellis weights with an +intermediate Hadamard through a 64 KiB table in each thread block's shared +memory. The table stores the complete codeword-to-E4M3 mapping. It replaces +the per-window permutation arithmetic; it does not change checkpoint weights, +activation precision, matrix multiplication, or reduction order. + +`B12X_TRELLIS_DECODE_TABLE` controls canonical fused-MoE preparation: + +| Value | Behavior | +| --- | --- | +| `auto` (default) | Use the full table for planned capacities of at most 16 tokens when the shared-memory and cooperative-residency contracts permit it. Other capacities and weight formats keep compact decoding. | +| `compact` | Compute the permutation per window and read the 4 KiB value table. | +| `full` | Require the full table; fail preparation for an incompatible kernel geometry. This is intended for component qualification. | + +For an explicit reference run, set the variable before declaring plans: + +```bash +export B12X_TRELLIS_DECODE_TABLE=compact +``` + +Canonical preparation saves the setting in its immutable query and propagates +it to compiler workers. Changing the environment after declaration cannot +change that plan. Binding and replay retain compiled programs and use +caller-owned scratch. The table is process-shared immutable execution data, +not a decoded copy of model weights. + +The full table is admitted only for uniform 2-bit `lut_e4m3` weights with an +intermediate Hadamard, one resident cooperative CTA per SM, and enough shared +memory for both the original GEMM staging and the 64 KiB table. It never +reduces the planned resident-CTA count. Large-capacity prefill keeps its +existing staging footprint. + +`tests/moe/test_trellis_direct_lut_decode.py` checks all 65,536 codewords and +random overlapping windows against the compact decoder, including shared +tables for each supported bit width. Only uniform 2-bit weights use the full +table in fused MoE; testing the lookup primitive at 3 and 4 bits does not +enable those execution paths. + +A performance comparison must account for instruction scheduling as well as +table size. The direct intrinsic exposes eight independent byte loads and +uses byte permutations to pack its results; the compact intrinsic processes +pairs of windows. Compare an equivalently scheduled compact decoder before +attributing a speedup solely to the larger table. Record per-launch table +staging cost, shared-memory usage and end-to-end kernel latency. diff --git a/tests/moe/test_fused_moe_variant_selection.py b/tests/moe/test_fused_moe_variant_selection.py index ac2d9ea0e..fd6c084f6 100644 --- a/tests/moe/test_fused_moe_variant_selection.py +++ b/tests/moe/test_fused_moe_variant_selection.py @@ -11,9 +11,56 @@ _FusedMoeCapacityState, _W4A16PrimaryLaunches, variant_for, + _control_snapshot, ) +@pytest.mark.parametrize("mode", ["auto", "compact", "full"]) +def test_trellis_decode_table_is_frozen_at_declaration(monkeypatch, mode): + monkeypatch.setenv("B12X_TRELLIS_DECODE_TABLE", mode) + controls = _control_snapshot() + monkeypatch.setenv("B12X_TRELLIS_DECODE_TABLE", "invalid") + assert controls["trellis_decode_table"] == mode + with pytest.raises(ValueError, match="B12X_TRELLIS_DECODE_TABLE"): + _control_snapshot() + + +@pytest.mark.parametrize("rows,mode,expected,shared_limit", [ + (1, "auto", True, 101376), (8, "auto", True, 101376), + (16, "auto", True, 101376), (17, "auto", False, 101376), + (8, "compact", False, 101376), (8, "full", True, 101376), + (8, "auto", False, 65536), +]) +def test_trellis_direct_table_preserves_cooperative_residency( + monkeypatch, rows, mode, expected, shared_limit, +): + from b12x.moe._shared.kernels.w4a16.kernel import W4A16FusedMoeKernel + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace( + multi_processor_count=188, shared_memory_per_block_optin=shared_limit, + )) + monkeypatch.setenv("B12X_LUT_E4M3_SMEM", "1") + kernel = W4A16FusedMoeKernel( + size_m=rows, hidden_size=3584, intermediate_size=384, + num_experts=512, top_k=16, activation="silu", + apply_router_weight_on_input=False, zero_fc2_output=False, + fc1_tile_n=128, fc1_tile_k=128, fc2_tile_n=128, fc2_tile_k=128, + moe_block_size=8, max_m_blocks=1024, element_dtype="fp16", + weight_layout="trellis_t256", scale_format="e4m3_k32", + w13_layout="trellis_t256_proj", trellis_bits=2, + full_rotation=True, intermediate_rotation=True, intermediate_hadamard=True, + trellis_decode_table=mode, + ) + assert kernel.trellis_direct_lut == expected + assert kernel.shared_words * 4 <= kernel.fc1.max_shared_mem + if expected: + assert kernel.blocks_per_sm == 1 + assert kernel.lut_e4m3_smem + assert kernel.fc1.trellis_direct_lut and kernel.fc2.trellis_direct_lut + + class _Variant: def __init__(self, count): self.count = count diff --git a/tests/moe/test_trellis_direct_lut_decode.py b/tests/moe/test_trellis_direct_lut_decode.py index 77f3129d8..443e4a091 100644 --- a/tests/moe/test_trellis_direct_lut_decode.py +++ b/tests/moe/test_trellis_direct_lut_decode.py @@ -10,6 +10,7 @@ from __future__ import annotations import cuda.bindings.driver as cuda +import cutlass import cutlass.cute as cute import pytest import torch @@ -20,6 +21,7 @@ from b12x._lib.intrinsics import ( packed_decode_lut_e4m3_to_e4m3x8, packed_decode_lut_e4m3_direct_to_e4m3x8, + shared_ptr_to_u32, ) from b12x._lib.quant.lut_e4m3 import ( lut_e4m3_direct_table_cpu, @@ -32,8 +34,9 @@ class _DecodeProbe: - def __init__(self, bits: int): + def __init__(self, bits: int, shared: bool = False): self.bits = int(bits) + self.shared = shared @cute.jit def __call__( @@ -45,7 +48,7 @@ def __call__( stream: cuda.CUstream, ): self.kernel(wins, value_table, direct, out).launch( - grid=(1, 1, 1), block=[32, 1, 1], stream=stream + grid=(cute.size(wins) // 64, 1, 1), block=[32, 1, 1], stream=stream ) @cute.kernel @@ -57,16 +60,26 @@ def kernel( out: cute.Tensor, ): tidx, _, _ = cute.arch.thread_idx() - lane = Int32(tidx) + bidx, _, _ = cute.arch.block_idx() + lane = Int32(tidx) + Int32(bidx) * 32 value_table_addr = value_table.iterator.toint() dir_addr = direct.iterator.toint() + if cutlass.const_expr(self.shared): + table = cutlass.utils.SmemAllocator().allocate_tensor( + cutlass.Uint8, cute.make_layout(65536), byte_alignment=16 + ) + for offset in range(Int32(tidx), 65536, 32): + table[offset] = direct[(self.bits - 2) * 65536 + offset] + cute.arch.sync_threads() + dir_addr = shared_ptr_to_u32(table.iterator) wa = Uint32(wins[2 * lane]) wb = Uint32(wins[2 * lane + 1]) lo_t, hi_t = packed_decode_lut_e4m3_to_e4m3x8( wa, wb, value_table_addr, self.bits, value_table_in_shared=False ) lo_d, hi_d = packed_decode_lut_e4m3_direct_to_e4m3x8( - wa, wb, dir_addr, self.bits, rate_indexed=True + wa, wb, dir_addr, self.bits, rate_indexed=not self.shared, + in_shared=self.shared, ) out[4 * lane] = Int32(lo_t) out[4 * lane + 1] = Int32(hi_t) @@ -76,7 +89,8 @@ def kernel( @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @pytest.mark.parametrize("bits", [2, 3, 4]) -def test_direct_table_decode_bit_equals_value_table(bits: int) -> None: +@pytest.mark.parametrize("shared", [False, True]) +def test_direct_table_decode_bit_equals_value_table(bits: int, shared: bool) -> None: device = torch.device("cuda") torch.manual_seed(20260812 + bits) wins = torch.zeros(64, dtype=torch.int32, device=device) @@ -93,7 +107,7 @@ def args(): current_cuda_stream(), ) - compiled = b12x_compile(_DecodeProbe(bits), *args()) + compiled = b12x_compile(_DecodeProbe(bits, shared), *args()) mismatched = 0 for _ in range(512): wins.copy_( @@ -107,3 +121,25 @@ def args(): mismatched += int((o[:, 0] != o[:, 2]).sum()) mismatched += int((o[:, 1] != o[:, 3]).sum()) assert mismatched == 0, mismatched + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("bits", [2, 3, 4]) +def test_shared_direct_lut_covers_all_codewords(bits: int) -> None: + # Both 16-bit windows visit every codeword; random high bits also exercise + # the shifted windows that cross the low-16-bit boundary. + device = torch.device("cuda") + torch.manual_seed(731 + bits) + states = torch.arange(65536, device=device, dtype=torch.int32) + upper = torch.randint(0, 65536, (65536, 2), device=device, dtype=torch.int32) + wins = ((upper << 16) | states[:, None]).flatten() + value_table = lut_e4m3_value_table_cpu().to(device) + direct = lut_e4m3_direct_table_cpu().to(device) + out = torch.empty(65536 * 4, device=device, dtype=torch.int32) + args = tuple(from_dlpack(tensor, assumed_align=16) + for tensor in (wins, value_table, direct, out)) + (current_cuda_stream(),) + compiled = b12x_compile(_DecodeProbe(bits, shared=True), *args) + compiled(*args) + torch.cuda.synchronize() + result = out.view(-1, 4) + torch.testing.assert_close(result[:, :2], result[:, 2:], rtol=0, atol=0)