perf(moe): full-codeword lookup for prepared K2 Trellis decode

b12x/_lib/intrinsics.py

@@ -6834,15 +6834,16 @@ def packed_decode_lut_e4m3_direct_to_e4m3x8(
     bits: int = 3,
     rate_indexed: bool = True,
     *,
+    in_shared: bool = False,
     loc=None,
     ip=None,
 ):
     """Decode eight 16-bit windows through a rate-specific direct E4M3 table.
 
     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,

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.
 

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()

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,

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),
     )
 

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__),