diff --git a/b12x/_lib/intrinsics.py b/b12x/_lib/intrinsics.py index d4ca937e8..ff473e72d 100644 --- a/b12x/_lib/intrinsics.py +++ b/b12x/_lib/intrinsics.py @@ -6262,6 +6262,7 @@ def packed_decode_sqg_xor_cheb_t12_to_e4m3x8( t12_lut_addr, bits: int = 3, t12_in_shared: bool = False, + stream_layout: bool = False, *, loc=None, ip=None, @@ -6276,10 +6277,17 @@ def packed_decode_sqg_xor_cheb_t12_to_e4m3x8( bits = int(bits) t12_in_shared = bool(t12_in_shared) - if bits not in (2, 3, 4): + stream_layout = bool(stream_layout) + if bits not in (2, 3, 4, 5, 6): raise ValueError( - f"unsupported SQG-XOR-Cheb-T12 bitrate {bits}; expected 2, 3, or 4" + f"unsupported SQG-XOR-Cheb-T12 bitrate {bits}; expected 2, 3, 4, 5, or 6" ) + if bits == 6 and not stream_layout: + raise ValueError( + "K6 SQG-XOR-Cheb-T12 requires the contiguous 64-bit stream layout" + ) + if stream_layout and bits != 6: + raise ValueError("the contiguous SQG stream layout is valid only for K6") width = 16 - bits phase_table_mask = (1 << (width - 4)) - 1 decode_blocks: list[str] = [] @@ -6291,21 +6299,41 @@ def packed_decode_sqg_xor_cheb_t12_to_e4m3x8( load_lines: list[str] = [] pack_lines: list[str] = [] for slot, index in enumerate(indices): - source = "$3" if index < 4 else "$2" - shift = (3 - (index & 3)) * bits target = "out0" if index < 4 else "out1" byte_shift = 8 * (index & 3) - extract_lines.append( - f""" + if stream_layout: + shift = (7 - index) * bits + if shift == 0: + extract = f"mov.b32 w{slot}, $2;" + elif shift < 32: + extract = f"shf.r.wrap.b32 w{slot}, $2, $3, {shift};" + else: + extract = f"shr.u32 w{slot}, $3, {shift - 32};" + extract_lines.append( + f""" + {extract} + and.b32 w{slot}, w{slot}, 0xffff; + """ + ) + else: + source = "$3" if index < 4 else "$2" + shift = (3 - (index & 3)) * bits + extract_lines.append( + f""" bfe.u32 w{slot}, {source}, {shift}, 16; + """ + ) + history_mix = "" + if width > 11: + history_mix = f""" + shr.u32 t{slot}, p{slot}, 11; + bfi.b32 p{slot}, p{slot}, p{slot}, 11, {width - 11}; + xor.b32 p{slot}, p{slot}, t{slot}; """ - ) product_lines.append( f""" shr.u32 p{slot}, w{slot}, {bits}; - shr.u32 t{slot}, p{slot}, 11; - bfi.b32 p{slot}, p{slot}, p{slot}, 11, {width - 11}; - xor.b32 p{slot}, p{slot}, t{slot}; + {history_mix} mad.lo.u32 p{slot}, p{slot}, 0x3fa7d929, 0xc928fd8e; """ ) @@ -6347,18 +6375,12 @@ def packed_decode_sqg_xor_cheb_t12_to_e4m3x8( ) decode_blocks.append( "\n".join( - extract_lines - + product_lines - + rank_lines - + load_lines - + pack_lines + extract_lines + product_lines + rank_lines + load_lines + pack_lines ) ) address_reg = ( - ".reg .b32 addr0,addr1;" - if t12_in_shared - else ".reg .b64 addr0,addr1;" + ".reg .b32 addr0,addr1;" if t12_in_shared else ".reg .b64 addr0,addr1;" ) asm = ( """ @@ -6447,9 +6469,7 @@ def packed_decode_sqg_fp16_d3l_to_half2x4( else: source = "$5" if index < 4 else "$4" bit_shift = (3 - (index & 3)) * bits - extract_lines.append( - f"bfe.u32 w{slot}, {source}, {bit_shift}, 16;" - ) + extract_lines.append(f"bfe.u32 w{slot}, {source}, {bit_shift}, 16;") graph_lines.append( f""" shr.u32 p{slot}, w{slot}, {bits}; @@ -6589,9 +6609,9 @@ def packed_decode_trellis_sqg_cheb_normal_e4m3_rank_lut_to_e4m3x8( if k2_q8h4 and bits != 2: raise ValueError("the virtual-octile graph is valid only for K2") if not global_lut: - # The packed form uses a 64-bit global pointer. Keep the scalar shared - # implementation available for the later staging experiment rather - # than mixing generic and shared address spaces in one PTX template. + # The packed form uses a 64-bit global pointer. Keep the scalar shared + # implementation separate so generic and shared address spaces do not + # enter the same PTX template. mask = Uint32(0xFFFF) source_a = Uint32(win_a) source_b = Uint32(win_b) @@ -6724,7 +6744,8 @@ def packed_decode_trellis_sqg_cheb_normal_e4m3_rank_lut_to_e4m3x8( """ ) - asm = """ + asm = ( + """ { .reg .b16 entry16; .reg .b32 w,h,b,phase,syndrome,syn,rev,stratum,rank; @@ -6733,7 +6754,10 @@ def packed_decode_trellis_sqg_cheb_normal_e4m3_rank_lut_to_e4m3x8( .reg .pred pneg,pnz,ptest; mov.b32 $0, 0; mov.b32 $1, 0; - """ + "\n".join(decode_blocks) + "\n}" + """ + + "\n".join(decode_blocks) + + "\n}" + ) result = llvm.inline_asm( llvm.StructType.get_literal([T.i32(), T.i32()]), [ @@ -6985,9 +7009,7 @@ def packed_decode_trellis_sqg_state_smem_to_e4m3x8( stratum_mask = (branch_mask << width) & 0xFFFF stratum_mult = (7 << width) & 0xFFFF state_global = graph_bits != 3 - state_blob_off = ( - SQG_STATE_BLOB_K4_OFF if graph_bits == 4 else SQG_STATE_BLOB_K2_OFF - ) + state_blob_off = SQG_STATE_BLOB_K4_OFF if graph_bits == 4 else SQG_STATE_BLOB_K2_OFF # PRMT byte tables giving the bit-reversed branch for selector values 0-7. # K4 falls back to BREV because PRMT indexes at most eight byte slots. if graph_bits == 2: @@ -7174,9 +7196,7 @@ def packed_decode_trellis_sqg_direct_lut_to_e4m3x8( 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;" - ) + 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( @@ -7187,9 +7207,7 @@ def packed_decode_trellis_sqg_direct_lut_to_e4m3x8( """ ) if byte_shift: - pack_lines.append( - f"shl.b32 w{index}, w{index}, {byte_shift};" - ) + pack_lines.append(f"shl.b32 w{index}, w{index}, {byte_shift};") pack_lines.append(f"or.b32 {target}, {target}, w{index};") asm = ( """ diff --git a/b12x/_lib/quant/sqg_e4m3.py b/b12x/_lib/quant/sqg_e4m3.py index ddee76acd..8b61d05c8 100644 --- a/b12x/_lib/quant/sqg_e4m3.py +++ b/b12x/_lib/quant/sqg_e4m3.py @@ -30,10 +30,9 @@ SQG_E4M3_RANK_LUT_ENTRIES = 1024 SQG_E4M3_DIRECT_LUT_ENTRIES = 3 * (1 << 16) +SQG_XOR_CHEB_T12_DIRECT_LUT_ENTRIES = 5 * (1 << 16) SQG_E4M3_STATE_ENTRIES = (1 << 14) + (1 << 13) + (1 << 12) -SQG_E4M3_STATE_LUT_ENTRIES = ( - SQG_E4M3_STATE_ENTRIES + SQG_E4M3_RANK_LUT_ENTRIES -) +SQG_E4M3_STATE_LUT_ENTRIES = SQG_E4M3_STATE_ENTRIES + SQG_E4M3_RANK_LUT_ENTRIES # Positive-half transition ranks for the E4M3-aware, full-tail normal # SQG-Cheb staircase. Its last 32-rank bucket contains three transitions. @@ -231,9 +230,9 @@ def _sqg_xor_cheb_t12_direct_lut_device( def sqg_xor_cheb_t12_direct_lut(device: torch.device | str) -> torch.Tensor: - """Return the process-lifetime rate-indexed 192 KiB direct state table. + """Return the process-lifetime rate-indexed direct state table. - Rows are the K2/K3/K4 slices in rate order: byte(state, bits) = + Rows are the K2-K6 slices in rate order: byte(state, bits) = table[((bits - 2) << 16) | state]. Each byte precomposes the frozen XOR-Cheb rank map with the modal T12 staircase, so lookups are bit-identical to the in-kernel T12 decode. @@ -251,7 +250,7 @@ def _sqg_xor_cheb_t12_rank_for_codewords( ) -> torch.Tensor: """Apply the frozen SQG-XOR graph to L16 codewords.""" - if bits not in (2, 3, 4): + if bits not in (2, 3, 4, 5, 6): raise ValueError(f"unsupported SQG-XOR-Cheb-T12 rate K{bits}") width = 16 - bits history_mask = (1 << width) - 1 @@ -275,14 +274,14 @@ def _sqg_xor_cheb_t12_rank_for_codewords( @functools.cache def sqg_xor_cheb_t12_direct_lut_cpu() -> torch.Tensor: - """Build independent K2/K3/K4 codeword tables for SQG-XOR-Cheb-T12.""" + """Build independent K2-K6 codeword tables for SQG-XOR-Cheb-T12.""" codewords = torch.arange(1 << 16, dtype=torch.int64) t12 = sqg_xor_cheb_t12_lut_cpu() return torch.cat( [ t12[_sqg_xor_cheb_t12_rank_for_codewords(codewords, bits) >> 4] - for bits in (2, 3, 4) + for bits in (2, 3, 4, 5, 6) ] ).contiguous() @@ -364,9 +363,7 @@ def sqg_cheb_normal_e4m3_state_lut_cpu() -> torch.Tensor: ) for bits in (2, 3, 4) ] - return torch.cat( - (*states, sqg_cheb_normal_e4m3_rank_lut_cpu()) - ).contiguous() + return torch.cat((*states, sqg_cheb_normal_e4m3_rank_lut_cpu())).contiguous() @functools.cache @@ -523,9 +520,9 @@ def _sqg_cheb_normal_k2_q8h4_w2_e4m3_direct_lut_device( device_index: int | None, ) -> torch.Tensor: device = torch.device(device_type, device_index) - return sqg_cheb_normal_k2_q8h4_w2_e4m3_direct_lut_cpu().to( - device=device - ).contiguous() + return ( + sqg_cheb_normal_k2_q8h4_w2_e4m3_direct_lut_cpu().to(device=device).contiguous() + ) def sqg_cheb_normal_k2_q8h4_w2_e4m3_direct_lut( @@ -537,13 +534,12 @@ def sqg_cheb_normal_k2_q8h4_w2_e4m3_direct_lut( index = resolved.index if resolved.type == "cuda" and index is None: index = torch.cuda.current_device() - return _sqg_cheb_normal_k2_q8h4_w2_e4m3_direct_lut_device( - resolved.type, index - ) + return _sqg_cheb_normal_k2_q8h4_w2_e4m3_direct_lut_device(resolved.type, index) __all__ = [ "SQG_E4M3_DIRECT_LUT_ENTRIES", + "SQG_XOR_CHEB_T12_DIRECT_LUT_ENTRIES", "SQG_E4M3_EXEC_LUT_ENTRIES", "SQG_E4M3_EXEC_LUT_K3_DIRECT_OFF", "SQG_E4M3_RANK_LUT_ENTRIES", diff --git a/b12x/gemm/trellis_linear/__init__.py b/b12x/gemm/trellis_linear/__init__.py index 4091faed7..36a281f29 100644 --- a/b12x/gemm/trellis_linear/__init__.py +++ b/b12x/gemm/trellis_linear/__init__.py @@ -23,6 +23,8 @@ "prepare_weight", "prepare_pair_weight", "run", + "run_sqg_k6_w6a16", + "sqg_k6_w6a16_scratch_elements", "is_supported", "clear_caches", ), @@ -55,6 +57,8 @@ prepare_weight, prepare_pair_weight, run, + run_sqg_k6_w6a16, + sqg_k6_w6a16_scratch_elements, ) install_lazy_api(globals(), META) diff --git a/b12x/gemm/trellis_linear/_small_m.py b/b12x/gemm/trellis_linear/_small_m.py index 4b3810396..be8214376 100644 --- a/b12x/gemm/trellis_linear/_small_m.py +++ b/b12x/gemm/trellis_linear/_small_m.py @@ -1,4 +1,4 @@ -"""JIT binding for the B12X-owned K6/MCG small-M CUDA kernel.""" +"""JIT binding for B12X K6 small-M CUDA kernels.""" from __future__ import annotations @@ -10,6 +10,8 @@ import torch from torch.utils.cpp_extension import load +from b12x._lib.quant.sqg_e4m3 import sqg_xor_cheb_t12_lut + _SOURCE_DIR = Path(__file__).resolve().parent / "csrc" _SOURCE = _SOURCE_DIR / "trellis_k6_small.cu" @@ -24,6 +26,8 @@ # 4096/2048-wide shared MLP dimensions. The budgets match the E2E-optimal # ExLlama autotuner result; using all 188 SMs serializes the graph branches. (6144, 1024): 64, + # Checkpoint-native SQG stores gate and up as separate TP4 projections. + (6144, 512): 48, (512, 6144): 96, } @@ -70,6 +74,46 @@ def _extension(): ) +def _validate_sm120(x: torch.Tensor) -> None: + capability = torch.cuda.get_device_capability(x.device) + if capability != (12, 0): + raise NotImplementedError( + "Trellis K6 small-M kernel is built for sm_120 only; " + f"device reports sm_{capability[0]}{capability[1]}" + ) + + +def _resolve_num_sms( + x: torch.Tensor, + output: torch.Tensor, + num_sms: int, +) -> int: + if num_sms > 0: + return int(num_sms) + device_index = x.device.index + if device_index is None: + device_index = torch.cuda.current_device() + return _default_num_sms( + int(x.shape[1]), + int(output.shape[1]), + _available_sms(int(device_index)), + ) + + +@lru_cache(maxsize=None) +def _configure_sqg_k6_lut(device_index: int) -> torch.Tensor: + """Install and retain the exact K6 SQG T12 table for one CUDA device.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "SQG K6 small-M decode table must be initialized before CUDA graph capture" + ) + device = torch.device("cuda", device_index) + labels = sqg_xor_cheb_t12_lut(device) + decoded = labels.view(torch.float8_e4m3fn).to(torch.float16).contiguous() + _extension().configure_k6_sqg_lut(decoded) + return decoded + + def run_k6_mcg( x: torch.Tensor, trellis: torch.Tensor, @@ -82,21 +126,8 @@ def run_k6_mcg( num_sms: int = 0, ) -> None: """Launch the capture-safe K6/MCG kernel on Torch's current stream.""" - capability = torch.cuda.get_device_capability(x.device) - if capability != (12, 0): - raise NotImplementedError( - "Trellis K6 small-M kernel is built for sm_120 only; " - f"device reports sm_{capability[0]}{capability[1]}" - ) - if num_sms <= 0: - device_index = x.device.index - if device_index is None: - device_index = torch.cuda.current_device() - num_sms = _default_num_sms( - int(x.shape[1]), - int(output.shape[1]), - _available_sms(int(device_index)), - ) + _validate_sm120(x) + num_sms = _resolve_num_sms(x, output, num_sms) _extension().launch_k6_mcg( x, trellis, @@ -109,4 +140,34 @@ def run_k6_mcg( ) -__all__ = ["run_k6_mcg"] +def run_k6_sqg( + x: torch.Tensor, + trellis: torch.Tensor, + output: torch.Tensor, + suh: torch.Tensor, + rotated_input: torch.Tensor, + svh: torch.Tensor, + locks: torch.Tensor, + *, + num_sms: int = 0, +) -> None: + """Launch the capture-safe K6/SQG kernel on Torch's current stream.""" + _validate_sm120(x) + device_index = x.device.index + if device_index is None: + device_index = torch.cuda.current_device() + _configure_sqg_k6_lut(int(device_index)) + num_sms = _resolve_num_sms(x, output, num_sms) + _extension().launch_k6_sqg( + x, + trellis, + output, + suh, + rotated_input, + svh, + locks, + int(num_sms), + ) + + +__all__ = ["run_k6_mcg", "run_k6_sqg"] diff --git a/b12x/gemm/trellis_linear/api.py b/b12x/gemm/trellis_linear/api.py index 4714453cd..4a556a3eb 100644 --- a/b12x/gemm/trellis_linear/api.py +++ b/b12x/gemm/trellis_linear/api.py @@ -11,6 +11,9 @@ clear_w4a16_kernel_cache, run_trellis256_dense, ) +from ...moe._shared.kernels.w4a16.host import ( + dense_trellis_gemm_scratch_elements_upper_bound, +) from ...moe._shared.kernels.w4a16.prepare import ( PreparedTrellis256DenseWeight, prepare_trellis256_dense_weight, @@ -31,6 +34,8 @@ def prepare_weight( codebook: Optional[str | int] = None, params_dtype: torch.dtype = torch.float16, dummy_scale: Optional[torch.Tensor] = None, + global_scale: Optional[torch.Tensor] = None, + workspace: Optional[torch.Tensor] = None, ) -> PreparedWeight: """Validate one native EXL3 dense weight and retain zero-copy views.""" return prepare_trellis256_dense_weight( @@ -42,6 +47,8 @@ def prepare_weight( codebook=codebook, params_dtype=params_dtype, dummy_scale=dummy_scale, + global_scale=global_scale, + workspace=workspace, ) @@ -107,6 +114,72 @@ def run( ) +def run_sqg_k6_w6a16( + x: torch.Tensor, + weight: PreparedWeight, + *, + output: Optional[torch.Tensor] = None, + gemm_output: Optional[torch.Tensor] = None, + c_tmp: Optional[torch.Tensor] = None, + input_f16: Optional[torch.Tensor] = None, + rotated_f16: Optional[torch.Tensor] = None, + rotated_compute: Optional[torch.Tensor] = None, + gemm_output_f16: Optional[torch.Tensor] = None, + output_f16: Optional[torch.Tensor] = None, + hadamard_128=None, + _moe_block_size: int | None = None, + _force_tile_config: tuple[int, int] | None = None, +) -> torch.Tensor: + """Execute a checkpoint-native K6 SQG matrix through W6A16 GEMM. + + The endpoint accepts only a uniform K6 ``sqg_xor_cheb_t12`` payload. Its + explicit contract prevents an MCG or paired payload from reaching the K6 + contiguous-stream decoder, where a format mismatch would change numerics. + """ + if str(getattr(weight, "trellis_codebook", "")).lower() != ("sqg_xor_cheb_t12"): + raise ValueError("SQG K6 W6A16 requires codebook='sqg_xor_cheb_t12'") + if int(getattr(weight, "trellis_bits", 0)) != 6: + raise ValueError("SQG K6 W6A16 requires a uniform K6 payload") + if getattr(weight, "trellis_pair_kind", None) is not None: + raise ValueError("SQG K6 W6A16 does not accept paired payloads") + if getattr(weight, "mcg", None) is not None: + raise ValueError("SQG K6 W6A16 does not accept an MCG marker") + if getattr(weight, "mul1_e4m3", None) is not None: + raise ValueError("SQG K6 W6A16 does not accept mul1 metadata") + return run_trellis256_dense( + x, + weight, + output=output, + gemm_output=gemm_output, + c_tmp=c_tmp, + input_f16=input_f16, + rotated_f16=rotated_f16, + rotated_compute=rotated_compute, + gemm_output_f16=gemm_output_f16, + output_f16=output_f16, + hadamard_128=hadamard_128, + _moe_block_size=_moe_block_size, + _force_tile_config=_force_tile_config, + ) + + +def sqg_k6_w6a16_scratch_elements( + rows: int, + out_features: int, + *, + device: torch.device | str, +) -> int: + """Return K6 GEMM scratch capacity for caller-owned graph storage.""" + + resolved = torch.device(device) + properties = torch.cuda.get_device_properties(resolved) + return dense_trellis_gemm_scratch_elements_upper_bound( + rows=int(rows), + size_n=int(out_features), + sms=int(properties.multi_processor_count), + ) + + def is_supported(device=None) -> bool: """True when the SM120/SM121 Trellis kernel stack is available.""" return default_is_supported(device, requires=META.requires) diff --git a/b12x/gemm/trellis_linear/csrc/trellis_k6_small.cu b/b12x/gemm/trellis_linear/csrc/trellis_k6_small.cu index bb9223521..c3fb09295 100644 --- a/b12x/gemm/trellis_linear/csrc/trellis_k6_small.cu +++ b/b12x/gemm/trellis_linear/csrc/trellis_k6_small.cu @@ -9,7 +9,9 @@ #include #include #include +#include #include +#include #include #include "vendor/util.h" @@ -17,12 +19,15 @@ namespace cg = cooperative_groups; +__device__ const half* b12x_sqg_k6_t12_lut = nullptr; + #include "vendor/quant/exl3_gemm_kernel.cuh" namespace { constexpr int kBits = 6; constexpr int kCodebookMcg = 1; +constexpr int kCodebookSqg = 3; constexpr int kTileM = 16; constexpr int kTileK = 32; constexpr int kTileN = 128; @@ -44,11 +49,66 @@ constexpr int kDynamicSmem = static_assert(kDynamicSmem <= EXL3_SMEM_MAX_BYTES, "Trellis K6 kernel exceeds the vendored shared-memory limit"); +std::mutex sqg_lut_mutex; +std::set sqg_lut_devices; + void check_cuda(cudaError_t status, const char* operation) { TORCH_CHECK(status == cudaSuccess, operation, ": ", cudaGetErrorString(status)); } -void launch_k6_mcg( +void configure_k6_sqg_lut(const torch::Tensor& lut) { + TORCH_CHECK(lut.is_cuda() && lut.scalar_type() == at::kHalf && + lut.is_contiguous() && lut.numel() == (1 << 12), + "SQG K6 T12 LUT must be a contiguous CUDA FP16 tensor with " + "4,096 entries"); + const int device = lut.get_device(); + const c10::cuda::CUDAGuard device_guard(lut.device()); + const half* lut_ptr = + reinterpret_cast(lut.data_ptr()); + check_cuda(cudaMemcpyToSymbol(b12x_sqg_k6_t12_lut, &lut_ptr, + sizeof(lut_ptr), 0, cudaMemcpyHostToDevice), + "cudaMemcpyToSymbol(SQG K6 T12 LUT)"); + std::lock_guard guard(sqg_lut_mutex); + sqg_lut_devices.insert(device); +} + +__global__ void decode_k6_sqg_codewords_kernel(const uint16_t* codewords, + half* output, int count) { + const int index = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (index < count) { + output[index] = decode_3inst(codewords[index]); + } +} + +torch::Tensor decode_k6_sqg_codewords(const torch::Tensor& codewords) { + TORCH_CHECK(codewords.is_cuda() && codewords.scalar_type() == at::kShort && + codewords.is_contiguous(), + "SQG K6 codewords must be a contiguous CUDA int16 tensor"); + const int device = codewords.get_device(); + const c10::cuda::CUDAGuard device_guard(codewords.device()); + { + std::lock_guard guard(sqg_lut_mutex); + TORCH_CHECK(sqg_lut_devices.count(device) != 0, + "SQG K6 T12 LUT is not configured on CUDA device ", device); + } + auto output = torch::empty(codewords.sizes(), + codewords.options().dtype(at::kHalf)); + const int64_t count64 = codewords.numel(); + TORCH_CHECK(count64 <= std::numeric_limits::max(), + "SQG K6 decoder test input is too large"); + const int count = static_cast(count64); + constexpr int threads = 256; + const int blocks = (count + threads - 1) / threads; + cudaStream_t stream = at::cuda::getCurrentCUDAStream(device).stream(); + decode_k6_sqg_codewords_kernel<<>>( + reinterpret_cast(codewords.data_ptr()), + reinterpret_cast(output.data_ptr()), count); + check_cuda(cudaPeekAtLastError(), "SQG K6 codeword decoder launch"); + return output; +} + +template +void launch_k6( const torch::Tensor& input, const torch::Tensor& trellis, torch::Tensor& output, @@ -104,6 +164,11 @@ void launch_k6_mcg( " N=", size_n); TORCH_CHECK(locks.numel() >= size_n / 16, "Trellis K6 lock workspace is too small"); + if constexpr (Codebook == kCodebookSqg) { + std::lock_guard guard(sqg_lut_mutex); + TORCH_CHECK(sqg_lut_devices.count(device) != 0, + "SQG K6 direct LUT is not configured on CUDA device ", device); + } int available_sms = 0; check_cuda(cudaDeviceGetAttribute(&available_sms, cudaDevAttrMultiProcessorCount, @@ -113,7 +178,7 @@ void launch_k6_mcg( int num_sms = requested_sms > 0 ? static_cast(requested_sms) : available_sms; num_sms = std::max(1, std::min(num_sms, std::min(available_sms, tiles))); - auto kernel = exl3_gemm_kernel; static std::once_flag smem_attribute_once; std::call_once(smem_attribute_once, [&] { @@ -144,9 +209,41 @@ void launch_k6_mcg( check_cuda(cudaPeekAtLastError(), "Trellis K6 kernel launch"); } +void launch_k6_mcg( + const torch::Tensor& input, + const torch::Tensor& trellis, + torch::Tensor& output, + const torch::Tensor& suh, + torch::Tensor& rotated_input, + const torch::Tensor& svh, + torch::Tensor& locks, + int64_t requested_sms) { + launch_k6(input, trellis, output, suh, rotated_input, svh, + locks, requested_sms); +} + +void launch_k6_sqg( + const torch::Tensor& input, + const torch::Tensor& trellis, + torch::Tensor& output, + const torch::Tensor& suh, + torch::Tensor& rotated_input, + const torch::Tensor& svh, + torch::Tensor& locks, + int64_t requested_sms) { + launch_k6(input, trellis, output, suh, rotated_input, svh, + locks, requested_sms); +} + } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("configure_k6_sqg_lut", &configure_k6_sqg_lut, + "Configure the process-lifetime SQG K6 T12 LUT"); + module.def("decode_k6_sqg_codewords", &decode_k6_sqg_codewords, + "Decode SQG K6 codewords for exhaustive validation"); module.def("launch_k6_mcg", &launch_k6_mcg, "B12X K6/MCG small-M dense GEMM"); + module.def("launch_k6_sqg", &launch_k6_sqg, + "B12X K6/SQG small-M dense GEMM"); } diff --git a/b12x/gemm/trellis_linear/csrc/vendor/quant/codebook.cuh b/b12x/gemm/trellis_linear/csrc/vendor/quant/codebook.cuh index e3f6bac25..3903d5102 100644 --- a/b12x/gemm/trellis_linear/csrc/vendor/quant/codebook.cuh +++ b/b12x/gemm/trellis_linear/csrc/vendor/quant/codebook.cuh @@ -53,6 +53,25 @@ __device__ inline half decode_3inst(uint32_t x) half_uint16 h((uint16_t) sum); return __hfma(h.as_half, k_inv_h, k_bias_h); } + if constexpr (cb == 3) + { + // K6 leaves ten history bits. At this rate the frozen XOR graph's + // two width-limited xorshifts are identities, so the exact T12 rank + // bucket needs one wrapping IMAD plus the six-bit branch reversal. + // The 4,096-entry FP16 table is the checkpoint-independent modal + // E4M3 staircase. It avoids random reads from a 65,536-entry direct + // table while preserving every codeword's reconstruction exactly. + const uint32_t codeword = x & 0xffffu; + const uint32_t history = codeword >> 6; + const uint32_t branch = codeword & 0x3fu; + const uint32_t product = history * 0x3fa7d929u + 0xc928fd8eu; + const uint32_t phase_bucket = (product & 0x3ffu) >> 4; + const uint32_t syndrome = product >> 26; + const uint32_t reversed_branch = __brev(branch) >> 26; + const uint32_t t12_index = + ((reversed_branch ^ syndrome) << 6) | phase_bucket; + return __ldg(b12x_sqg_k6_t12_lut + t12_index); + } } template @@ -101,6 +120,10 @@ __device__ inline half2 decode_3inst_2(uint32_t x0, uint32_t x1) half_uint16 h1((uint16_t) sum1); return __hfma2(__halves2half2(h0.as_half, h1.as_half), k_inv_h2, k_bias_h2); } + if constexpr (cb == 3) + { + return __halves2half2(decode_3inst(x0), decode_3inst(x1)); + } } __device__ inline half2 decode_mcg_product_2(uint32_t x0, uint32_t x1) diff --git a/b12x/gemm/trellis_linear/w4a8.py b/b12x/gemm/trellis_linear/w4a8.py new file mode 100644 index 000000000..eac220aab --- /dev/null +++ b/b12x/gemm/trellis_linear/w4a8.py @@ -0,0 +1,793 @@ +"""Direct E4M3 trellis decode into SM120 MXFP8 MMA. + +This module is the dense correctness/performance anchor for the routed port. +It consumes either a native uniform K3/K4 tensor or the compact P24/P33 pair +payload, reconstructs E4M3 weights in registers, and feeds those bytes +directly to ``m16n8k32`` block-scaled MMA. +No FP16/BF16 weight tile is materialized. SQG-Cheb uses the universal +2 KiB rank-to-E4M3 descriptor read through the read-only global cache; it is +runtime state, not checkpoint metadata. +""" + +from __future__ import annotations + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cutlass_dsl import Float32, Int32, Int64, Uint32 + +from b12x._lib.compiler import KernelCompileSpec, compile as b12x_compile +from b12x._lib.intrinsics import ( + get_ptr_as_int64, + mxfp8_mma_m16n8k32_f32_e4m3, + packed_decode_sqg_xor_cheb_t12_to_e4m3x8, +) +from b12x._lib.quant.sqg_e4m3 import ( + sqg_xor_cheb_t12_lut, +) +from b12x._lib.quant.mxfp8_rows import quantize_mxfp8_rows_cute +from b12x._lib.runtime_control import raise_if_kernel_resolution_frozen +from b12x._lib.utils import current_cuda_stream, make_ptr +from b12x.gemm._shared.wo_mxfp8 import ( + MXFP8Rows, + empty_mxfp8_rows_for_dense_gemm, +) +from b12x.moe._shared.kernels.w4a16.kernel import ( + _resolve_exl3_hadamard_128, + _trellis_dense_buffer, +) + + +_PAIR_KINDS = {"P24", "P33"} +_RATE_AXES = {"k", "n"} +_UNIFORM_BITS = {3, 4} +_W4A8_CODEBOOKS = {"sqg_xor_cheb_t12"} + + +def select_dense_w4a8_m_tile_rows(m: int) -> int: + """Choose the largest supported decode-reuse tile not exceeding M.""" + m = int(m) + if m <= 0: + raise ValueError("W4A8 dense M must be positive") + if m >= 128: + return 128 + if m >= 64: + return 64 + if m >= 32: + return 32 + return 16 + + +class _TrellisW4A8DenseLaunch: + """One-warp multi-MxN8 kernel with a direct native-tile B operand. + + Large prefill rows reuse each decoded B fragment across several M16 MMA + groups. Decode work therefore grows by M tiles while preserving the compact + weight representation and FP8 MMA arithmetic. + """ + + threads = 32 + + def __init__( + self, + *, + size_k: int, + size_n: int, + trellis_bits: int, + pair_kind: str | None, + rate_axis: str | None, + trellis_codebook: str, + m_tile_rows: int = 16, + ) -> None: + self.size_k = int(size_k) + self.size_n = int(size_n) + self.trellis_bits = int(trellis_bits) + self.pair_kind = None if pair_kind is None else str(pair_kind).upper() + self.rate_axis = None if rate_axis is None else str(rate_axis).lower() + self.trellis_codebook = str(trellis_codebook).lower() + self.m_tile_rows = int(m_tile_rows) + if self.m_tile_rows not in (16, 32, 64, 128): + raise ValueError("W4A8 dense M tile must be 16, 32, 64, or 128") + self.m_groups = self.m_tile_rows // 16 + if self.trellis_codebook not in _W4A8_CODEBOOKS: + raise ValueError( + "W4A8 trellis codebook must be one of " + f"{sorted(_W4A8_CODEBOOKS)}, got {self.trellis_codebook!r}" + ) + if (self.pair_kind is None) != (self.rate_axis is None): + raise ValueError("pair_kind and rate_axis must be supplied together") + if self.pair_kind is None: + if self.trellis_bits not in _UNIFORM_BITS: + raise ValueError( + "uniform W4A8 trellis bits must be one of " + f"{sorted(_UNIFORM_BITS)}, got {self.trellis_bits}" + ) + self.low_bits = self.high_bits = self.trellis_bits + average_bits = self.trellis_bits + else: + if self.pair_kind not in _PAIR_KINDS or self.rate_axis not in _RATE_AXES: + raise ValueError("W4A8 pair requires P24/P33 and rate axis k/n") + if self.trellis_bits != 3: + raise ValueError("P24/P33 W4A8 pair containers average three bits") + self.low_bits, self.high_bits = ( + (2, 4) if self.pair_kind == "P24" else (3, 3) + ) + average_bits = 3 + self.trellis_words = self.size_k * self.size_n * average_bits // 32 + + @cute.jit + def __call__( + self, + values_ptr: cute.Pointer, + scale_rows_ptr: cute.Pointer, + trellis_ptr: cute.Pointer, + rank_lut_ptr: cute.Pointer, + output_ptr: cute.Pointer, + m: Int32, + stream: cuda.CUstream, + ) -> None: + values = cute.make_tensor( + values_ptr, + cute.make_ordered_layout((m, self.size_k // 4), order=(1, 0)), + ) + scale_rows = cute.make_tensor( + scale_rows_ptr, + cute.make_ordered_layout((m, self.size_k // 32), order=(1, 0)), + ) + trellis = cute.make_tensor( + trellis_ptr, + cute.make_layout((self.trellis_words,)), + ) + rank_lut = cute.make_tensor( + rank_lut_ptr, + cute.make_layout((4096,)), + ) + output = cute.make_tensor( + output_ptr, + cute.make_ordered_layout((m, self.size_n), order=(1, 0)), + ) + self.kernel(values, scale_rows, trellis, rank_lut, output, m).launch( + grid=(cute.ceil_div(m, self.m_tile_rows), self.size_n // 8, 1), + block=[self.threads, 1, 1], + stream=stream, + ) + + @cute.jit + def _lane_geom(self, lane: Int32, bits: cutlass.Constexpr[int]): + bits_i32 = Int32(int(bits)) + ring_u32 = Int32(8 * int(bits)) + t_offset = Int32(8) * lane + b1 = (t_offset + Int32(257)) * bits_i32 + b0 = b1 - Int32(16) + b2 = b1 + Int32(7 * int(bits)) + i0 = b0 >> Int32(5) + i2 = (b2 - Int32(1)) >> Int32(5) + ia = i0 - ring_u32 * (i0 >= ring_u32).to(Int32) + ib = i2 - ring_u32 * (i2 >= ring_u32).to(Int32) + s2 = (i2 + Int32(1)) * Int32(32) - b2 + return ia, ib, s2 + + @cute.jit + def _tile_base( + self, + k16: Int32, + n16: Int32, + record: Int32, + bits: cutlass.Constexpr[int], + ) -> Int64: + if cutlass.const_expr(self.pair_kind is None): + n16_count = self.size_n // 16 + return (Int64(k16) * Int64(n16_count) + Int64(n16)) * Int64(8 * int(bits)) + if cutlass.const_expr(self.rate_axis == "n"): + pair_u32_per_k16 = 8 * 8 * (self.low_bits + self.high_bits) + high_base_u32 = 8 * 8 * self.low_bits + return ( + Int64(k16) * Int64(pair_u32_per_k16) + + Int64(record) * Int64(high_base_u32) + + Int64(n16) * Int64(8 * int(bits)) + ) + + n16_count = self.size_n // 16 + low_record_u32 = 8 * n16_count * 8 * self.low_bits + local_k16 = k16 - record * Int32(8) + return ( + Int64(record) * Int64(low_record_u32) + + Int64(local_k16) * Int64(n16_count * 8 * int(bits)) + + Int64(n16) * Int64(8 * int(bits)) + ) + + @cute.jit + def _decode_tile_at_base( + self, + trellis: cute.Tensor, + lane: Int32, + k16: Int32, + n16: Int32, + record: Int32, + n_high: Int32, + bits: cutlass.Constexpr[int], + rank_lut_addr: Int64, + tensor_base: Int64, + ) -> Uint32: + ia, ib, s2 = self._lane_geom(lane, bits) + base = tensor_base + self._tile_base(k16, n16, record, bits) + a = Uint32(trellis[base + Int64(ia)]) + b = Uint32(trellis[base + Int64(ib)]) + merged = (Int64(a) << Int64(32)) | Int64(b) + win_a = Uint32(merged >> Int64(s2)) + win_b = Uint32(merged >> Int64(s2 + Int32(4 * int(bits)))) + lo, hi = packed_decode_sqg_xor_cheb_t12_to_e4m3x8( + win_a, + win_b, + rank_lut_addr, + int(bits), + ) + value = lo + if n_high != Int32(0): + value = hi + return value + + @cute.jit + def _decode_tile( + self, + trellis: cute.Tensor, + lane: Int32, + k16: Int32, + n16: Int32, + record: Int32, + n_high: Int32, + bits: cutlass.Constexpr[int], + rank_lut_addr: Int64, + ) -> Uint32: + return self._decode_tile_at_base( + trellis, + lane, + k16, + n16, + record, + n_high, + bits, + rank_lut_addr, + Int64(0), + ) + + @cute.jit + def _decode_k32_bits_at_base( + self, + trellis: cute.Tensor, + lane: Int32, + k32: Int32, + n16: Int32, + record: Int32, + n_high: Int32, + bits: cutlass.Constexpr[int], + rank_lut_addr: Int64, + tensor_base: Int64, + ): + e0 = self._decode_tile_at_base( + trellis, + lane, + k32 * Int32(2), + n16, + record, + n_high, + bits, + rank_lut_addr, + tensor_base, + ) + e1 = self._decode_tile_at_base( + trellis, + lane, + k32 * Int32(2) + Int32(1), + n16, + record, + n_high, + bits, + rank_lut_addr, + tensor_base, + ) + c = lane & Int32(3) + own = e0 + send = e1 + if c >= Int32(2): + own = e1 + send = e0 + # See tests/moe/test_w4a8_fragment_probe.py: the opposite register is + # selected before xor-2 so every receiver gets the same K16 tile. + peer = Uint32(cute.arch.shuffle_sync_bfly(send, offset=2)) + return own, peer + + @cute.jit + def _decode_k32_bits( + self, + trellis: cute.Tensor, + lane: Int32, + k32: Int32, + n16: Int32, + record: Int32, + n_high: Int32, + bits: cutlass.Constexpr[int], + rank_lut_addr: Int64, + ): + return self._decode_k32_bits_at_base( + trellis, + lane, + k32, + n16, + record, + n_high, + bits, + rank_lut_addr, + Int64(0), + ) + + @cute.jit + def _decode_k32( + self, + trellis: cute.Tensor, + lane: Int32, + k32: Int32, + n_base: Int32, + rank_lut_addr: Int64, + ): + if cutlass.const_expr(self.pair_kind is None): + n16 = n_base // Int32(16) + if cutlass.const_expr(self.trellis_bits == 3): + return self._decode_k32_bits( + trellis, + lane, + k32, + n16, + Int32(0), + ((n_base & Int32(15)) >= Int32(8)).to(Int32), + 3, + rank_lut_addr, + ) + return self._decode_k32_bits( + trellis, + lane, + k32, + n16, + Int32(0), + ((n_base & Int32(15)) >= Int32(8)).to(Int32), + 4, + rank_lut_addr, + ) + n_record = n_base // Int32(128) + n_local = n_base - n_record * Int32(128) + n16 = n_local // Int32(16) + if cutlass.const_expr(self.rate_axis == "k"): + # K-axis containers concatenate two complete K128 records, so N + # remains the ordinary full-width native-tile coordinate. + n16 = n_base // Int32(16) + n_high = (n_base & Int32(15)) >= Int32(8) + if cutlass.const_expr(self.pair_kind == "P33"): + record = ( + n_record + if cutlass.const_expr(self.rate_axis == "n") + else (k32 // Int32(4)) + ) + return self._decode_k32_bits( + trellis, + lane, + k32, + n16, + record, + n_high.to(Int32), + 3, + rank_lut_addr, + ) + + b0 = Uint32(0) + b1 = Uint32(0) + if cutlass.const_expr(self.rate_axis == "n"): + if n_record == Int32(0): + b0, b1 = self._decode_k32_bits( + trellis, + lane, + k32, + n16, + Int32(0), + n_high.to(Int32), + 2, + rank_lut_addr, + ) + else: + b0, b1 = self._decode_k32_bits( + trellis, + lane, + k32, + n16, + Int32(1), + n_high.to(Int32), + 4, + rank_lut_addr, + ) + else: + if k32 < Int32(4): + b0, b1 = self._decode_k32_bits( + trellis, + lane, + k32, + n16, + Int32(0), + n_high.to(Int32), + 2, + rank_lut_addr, + ) + else: + b0, b1 = self._decode_k32_bits( + trellis, + lane, + k32, + n16, + Int32(1), + n_high.to(Int32), + 4, + rank_lut_addr, + ) + return b0, b1 + + @cute.kernel + def kernel( + self, + values: cute.Tensor, + scale_rows: cute.Tensor, + trellis: cute.Tensor, + rank_lut: cute.Tensor, + output: cute.Tensor, + m: Int32, + ) -> None: + lane = cute.arch.lane_idx() + rank_lut_addr = get_ptr_as_int64(rank_lut, Int32(0)) + bmx, bny, _ = cute.arch.block_idx() + c = lane & Int32(3) + g = lane >> Int32(2) + m_base = Int32(bmx) * Int32(self.m_tile_rows) + n_base = Int32(bny) * Int32(8) + accumulators = tuple( + cute.make_rmem_tensor((4,), Float32) for _ in range(self.m_groups) + ) + for group in cutlass.range_constexpr(self.m_groups): + accumulators[group].fill(0.0) + k32 = Int32(0) + while k32 < Int32(self.size_k // 32): + b0, b1 = self._decode_k32( + trellis, + lane, + k32, + n_base, + rank_lut_addr, + ) + word0 = k32 * Int32(8) + c * Int32(2) + for group in cutlass.range_constexpr(self.m_groups): + group_base = m_base + Int32(group * 16) + m_lo = group_base + g + m_hi = m_lo + Int32(8) + a0 = Uint32(0) + a1 = Uint32(0) + a2 = Uint32(0) + a3 = Uint32(0) + if m_lo < m: + a0 = Uint32(values[m_lo, word0]) + a2 = Uint32(values[m_lo, word0 + Int32(1)]) + if m_hi < m: + a1 = Uint32(values[m_hi, word0]) + a3 = Uint32(values[m_hi, word0 + Int32(1)]) + + sf_row = group_base + g + ((lane & Int32(1)) << Int32(3)) + sf = Uint32(127) + if sf_row < m: + sf = Uint32(scale_rows[sf_row, k32]) + sfa = sf * Uint32(0x01010101) + frag = accumulators[group] + d0, d1, d2, d3 = mxfp8_mma_m16n8k32_f32_e4m3( + frag[0], + frag[1], + frag[2], + frag[3], + a0, + a1, + a2, + a3, + b0, + b1, + sfa, + Uint32(0x7F7F7F7F), + ) + frag[0] = d0 + frag[1] = d1 + frag[2] = d2 + frag[3] = d3 + k32 += Int32(1) + + col = n_base + c * Int32(2) + for group in cutlass.range_constexpr(self.m_groups): + group_base = m_base + Int32(group * 16) + m_lo = group_base + g + m_hi = m_lo + Int32(8) + frag = accumulators[group] + if m_lo < m: + output[m_lo, col] = cutlass.Float16(frag[0]) + output[m_lo, col + Int32(1)] = cutlass.Float16(frag[1]) + if m_hi < m: + output[m_hi, col] = cutlass.Float16(frag[2]) + output[m_hi, col + Int32(1)] = cutlass.Float16(frag[3]) + + +@functools.cache +def _compile_dense( + size_k: int, + size_n: int, + trellis_bits: int, + pair_kind: str | None, + rate_axis: str | None, + trellis_codebook: str, + m_tile_rows: int, + device_index: int, +): + launch = _TrellisW4A8DenseLaunch( + size_k=size_k, + size_n=size_n, + trellis_bits=trellis_bits, + pair_kind=pair_kind, + rate_axis=rate_axis, + trellis_codebook=trellis_codebook, + m_tile_rows=m_tile_rows, + ) + key = ( + int(size_k), + int(size_n), + int(trellis_bits), + pair_kind, + rate_axis, + trellis_codebook, + int(m_tile_rows), + int(device_index), + ) + raise_if_kernel_resolution_frozen("cute.compile", target=launch, cache_key=key) + return b12x_compile( + launch, + make_ptr(cutlass.Uint32, 16, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, 16, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint32, 16, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Uint8, 16, cute.AddressSpace.gmem, assumed_align=16), + make_ptr(cutlass.Float16, 16, cute.AddressSpace.gmem, assumed_align=16), + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key( + "gemm.trellis_w4a8_dense", + 4, + key, + ), + ) + + +def _validate_quantized( + quantized: MXFP8Rows, + *, + m: int, + k: int, + device: torch.device, +) -> None: + if tuple(quantized.values.shape) != (m, k): + raise ValueError( + "quantized values must have shape " + f"{(m, k)}, got {tuple(quantized.values.shape)}" + ) + if tuple(quantized.scale_rows.shape) not in { + (m, k // 32), + (1, m, k // 32), + }: + raise ValueError( + "quantized scale_rows must have one-group shape " + f"{(m, k // 32)} or {(1, m, k // 32)}, got " + f"{tuple(quantized.scale_rows.shape)}" + ) + for name, tensor in ( + ("values", quantized.values), + ("scale_rows", quantized.scale_rows), + ): + if tensor.device != device or not tensor.is_contiguous(): + raise ValueError(f"quantized {name} must be contiguous on {device}") + # scale_mma is intentionally a permuted view over contiguous physical + # storage. The quantizer writes that storage through its raw base pointer; + # this dense kernel consumes scale_rows and never dereferences the view. + if quantized.scale_mma.device != device: + raise ValueError(f"quantized scale_mma must be on {device}") + + +def run_trellis256_dense_w4a8( + x: torch.Tensor, + prepared_dense, + *, + output: torch.Tensor | None = None, + input_f16: torch.Tensor | None = None, + rotated_f16: torch.Tensor | None = None, + quantized: MXFP8Rows | None = None, + gemm_output_f16: torch.Tensor | None = None, + output_f16: torch.Tensor | None = None, + hadamard_128=None, + m_tile_rows: int | None = None, +) -> torch.Tensor: + """Execute one native K3/K4 or compact P24/P33 linear through W4A8 MMA.""" + if not isinstance(x, torch.Tensor) or not x.is_cuda: + raise ValueError("x must be a CUDA tensor") + if x.ndim != 2 or int(x.shape[0]) <= 0 or not x.is_contiguous(): + raise ValueError("x must be a non-empty contiguous rank-2 tensor") + if x.dtype not in (torch.float16, torch.bfloat16): + raise TypeError(f"x must be fp16 or bf16, got {x.dtype}") + trellis_codebook = str(getattr(prepared_dense, "trellis_codebook", "")).lower() + if trellis_codebook not in _W4A8_CODEBOOKS: + raise NotImplementedError( + "W4A8 trellis decode requires one of " + f"{sorted(_W4A8_CODEBOOKS)}, got {trellis_codebook!r}" + ) + pair_kind_raw = getattr(prepared_dense, "trellis_pair_kind", None) + rate_axis_raw = getattr(prepared_dense, "trellis_rate_axis", None) + if (pair_kind_raw is None) != (rate_axis_raw is None): + raise ValueError("prepared W4A8 pair metadata is incomplete") + pair_kind = None if pair_kind_raw is None else str(pair_kind_raw).upper() + rate_axis = None if rate_axis_raw is None else str(rate_axis_raw).lower() + trellis_bits = int(getattr(prepared_dense, "trellis_bits", 0)) + if pair_kind is None: + if trellis_bits not in _UNIFORM_BITS: + raise ValueError( + "uniform W4A8 dense trellis requires native K3 or K4, got " + f"K{trellis_bits}" + ) + elif pair_kind not in _PAIR_KINDS or rate_axis not in _RATE_AXES: + raise ValueError("W4A8 dense trellis pair requires P24/P33 and axis k/n") + if int(getattr(prepared_dense, "num_experts", 0)) != 1: + raise ValueError("W4A8 dense trellis requires E=1 prepared weights") + if prepared_dense.trellis.device != x.device: + raise ValueError("x and prepared weights must share one CUDA device") + + m, size_k = (int(v) for v in x.shape) + size_n = int(prepared_dense.out_features) + if size_k != int(prepared_dense.in_features): + raise ValueError( + f"x has K={size_k}, prepared weight expects {prepared_dense.in_features}" + ) + device = x.device + if m_tile_rows is None: + m_tile_rows = select_dense_w4a8_m_tile_rows(m) + m_tile_rows = int(m_tile_rows) + if m_tile_rows not in (16, 32, 64, 128): + raise ValueError("m_tile_rows must be 16, 32, 64, or 128") + output = _trellis_dense_buffer( + "w4a8 output", output, shape=(m, size_n), dtype=x.dtype, device=device + ) + gemm_output_f16 = _trellis_dense_buffer( + "w4a8 gemm_output_f16", + gemm_output_f16, + shape=(m, size_n), + dtype=torch.float16, + device=device, + ) + hadamard_128 = _resolve_exl3_hadamard_128(hadamard_128) + if x.dtype == torch.float16: + x_f16 = x + else: + input_f16 = _trellis_dense_buffer( + "w4a8 input_f16", + input_f16, + shape=(m, size_k), + dtype=torch.float16, + device=device, + ) + input_f16.copy_(x) + x_f16 = input_f16 + rotated_f16 = _trellis_dense_buffer( + "w4a8 rotated_f16", + rotated_f16, + shape=(m, size_k), + dtype=torch.float16, + device=device, + ) + hadamard_128(x_f16, rotated_f16, prepared_dense.suh, None, 1.0) + + if quantized is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "W4A8 trellis quantization storage must be supplied during " + "graph capture" + ) + quantized = empty_mxfp8_rows_for_dense_gemm(m, size_k, device=device) + _validate_quantized(quantized, m=m, k=size_k, device=device) + quantize_mxfp8_rows_cute( + rotated_f16, + quantized.values, + quantized.scale_rows, + quantized.scale_mma, + value_order="trellis_native_mma", + ) + + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + compiled = _compile_dense( + size_k, + size_n, + trellis_bits, + pair_kind, + rate_axis, + trellis_codebook, + m_tile_rows, + int(device_index), + ) + rank_lut = sqg_xor_cheb_t12_lut(device) + compiled( + make_ptr( + cutlass.Uint32, + quantized.values.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + cutlass.Uint8, + quantized.scale_rows.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + cutlass.Uint32, + prepared_dense.trellis.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + cutlass.Uint8, + rank_lut.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + make_ptr( + cutlass.Float16, + gemm_output_f16.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + m, + current_cuda_stream(), + ) + + if output.dtype == torch.float16: + hadamard_128(gemm_output_f16, output, None, prepared_dense.svh, 1.0) + else: + output_f16 = _trellis_dense_buffer( + "w4a8 output_f16", + output_f16, + shape=(m, size_n), + dtype=torch.float16, + device=device, + ) + hadamard_128(gemm_output_f16, output_f16, None, prepared_dense.svh, 1.0) + output.copy_(output_f16) + return output + + +def run_trellis256_uniform_dense_w4a8( + x: torch.Tensor, + prepared_dense, + **kwargs, +) -> torch.Tensor: + """Execute a checkpoint-native uniform K3/K4 SQG tensor through W4A8. + + This strict entry point exists so GLM benchmarks cannot accidentally use + the Kimi P24/P33 pair-container contract. The compact trellis payload is + decoded inside the MMA loop; no dense E4M3 or FP16 weight is materialized. + """ + if getattr(prepared_dense, "trellis_pair_kind", None) is not None: + raise ValueError("uniform W4A8 entry point rejects P24/P33 pair weights") + return run_trellis256_dense_w4a8(x, prepared_dense, **kwargs) + + +__all__ = [ + "run_trellis256_dense_w4a8", + "run_trellis256_uniform_dense_w4a8", + "select_dense_w4a8_m_tile_rows", +] diff --git a/b12x/moe/__init__.py b/b12x/moe/__init__.py index 1ca2cbca0..7d168d090 100644 --- a/b12x/moe/__init__.py +++ b/b12x/moe/__init__.py @@ -6,6 +6,8 @@ ``qsrt_sqg_e4m3`` plus uniform-K5/K6 ``sqg_fp16_d3l`` W4A16 source formats. - ``ep_moe``: expert-parallel MoE (replicated input -> local partial; cross-rank reduction is the caller's job, typically ``comm.pcie``). +- ``glm_sqg_w4a8``: graph-safe mixed-K3/K4 W4A8 routed execution for the + ``glm52_sqg_atoms_v2`` checkpoint format on SM120/SM121. """ from __future__ import annotations @@ -13,7 +15,7 @@ import importlib from typing import Any -_OP_MODULES = ("fused_moe", "ep_moe") +_OP_MODULES = ("fused_moe", "ep_moe", "glm_sqg_w4a8") def __getattr__(name: str) -> Any: diff --git a/b12x/moe/_shared/kernels/glm_trellis_transform.py b/b12x/moe/_shared/kernels/glm_trellis_transform.py new file mode 100644 index 000000000..6af6cda4f --- /dev/null +++ b/b12x/moe/_shared/kernels/glm_trellis_transform.py @@ -0,0 +1,317 @@ +"""Exact GLM route transforms around compact trellis projections.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from b12x.moe._shared.kernels.w4a16.kernel import ( + _run_trellis_dense_hadamard128, +) + + +@triton.jit +def _glm_route_silu_kernel( + gate, + up, + route_experts, + gate_svh, + up_svh, + output, + width: tl.constexpr, + block_n: tl.constexpr, +): + route = tl.program_id(0) + block = tl.program_id(1) + offsets = block * block_n + tl.arange(0, block_n) + mask = offsets < width + expert = tl.load(route_experts + route).to(tl.int64) + gate_value = tl.load(gate + route * width + offsets, mask=mask, other=0.0) + up_value = tl.load(up + route * width + offsets, mask=mask, other=0.0) + gate_scale = tl.load(gate_svh + expert * width + offsets, mask=mask, other=0.0) + up_scale = tl.load(up_svh + expert * width + offsets, mask=mask, other=0.0) + gate_value = gate_value.to(tl.float32) * gate_scale.to(tl.float32) + up_value = up_value.to(tl.float32) * up_scale.to(tl.float32) + activated = gate_value * tl.sigmoid(gate_value) * up_value + tl.store(output + route * width + offsets, activated, mask=mask) + + +@triton.jit +def _glm_route_scale_kernel( + source, + route_experts, + scale_table, + output, + width: tl.constexpr, + block_n: tl.constexpr, +): + route = tl.program_id(0) + block = tl.program_id(1) + offsets = block * block_n + tl.arange(0, block_n) + mask = offsets < width + expert = tl.load(route_experts + route).to(tl.int64) + values = tl.load(source + route * width + offsets, mask=mask, other=0.0) + scales = tl.load(scale_table + expert * width + offsets, mask=mask, other=0.0) + tl.store(output + route * width + offsets, values * scales, mask=mask) + + +@triton.jit +def _glm_topk_weighted_sum_kernel( + routes, + topk_weights, + output, + width: tl.constexpr, + topk: tl.constexpr, + block_n: tl.constexpr, +): + token = tl.program_id(0) + block = tl.program_id(1) + offsets = block * block_n + tl.arange(0, block_n) + mask = offsets < width + accumulator = tl.zeros((block_n,), dtype=tl.float32) + for route_in_token in tl.static_range(topk): + route = token * topk + route_in_token + values = tl.load(routes + route * width + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + weight = tl.load(topk_weights + route).to(tl.float32) + accumulator += values * weight + tl.store(output + token * width + offsets, accumulator, mask=mask) + + +def run_glm_gate_up_output_transform_silu( + gate_transformed: torch.Tensor, + up_transformed: torch.Tensor, + route_experts: torch.Tensor, + gate_svh: torch.Tensor, + up_svh: torch.Tensor, + gate_hadamard: torch.Tensor, + up_hadamard: torch.Tensor, + output: torch.Tensor, + *, + ones: torch.Tensor, +) -> torch.Tensor: + """Invert gate/up H128 bases, apply expert scales, then exact SwiGLU.""" + + if gate_transformed.shape != up_transformed.shape or gate_transformed.ndim != 2: + raise ValueError("gate/up transformed outputs must be aligned rank-2 tensors") + routes, width = (int(value) for value in gate_transformed.shape) + device = gate_transformed.device + for name, tensor in ( + ("gate_transformed", gate_transformed), + ("up_transformed", up_transformed), + ("gate_hadamard", gate_hadamard), + ("up_hadamard", up_hadamard), + ("output", output), + ): + if ( + tensor.shape != gate_transformed.shape + or tensor.dtype != torch.float16 + or tensor.device != device + or not tensor.is_contiguous() + ): + raise TypeError( + f"{name} must be contiguous FP16 {tuple(gate_transformed.shape)} " + f"on {device}" + ) + if ( + route_experts.shape != (routes,) + or route_experts.dtype != torch.int32 + or route_experts.device != device + or not route_experts.is_contiguous() + ): + raise TypeError("route_experts must be contiguous int32 [routes]") + if gate_svh.shape != up_svh.shape or gate_svh.ndim != 2: + raise ValueError("gate/up svh tables must be aligned [experts, intermediate]") + for name, tensor in (("gate_svh", gate_svh), ("up_svh", up_svh)): + if ( + int(tensor.shape[1]) != width + or tensor.dtype != torch.float16 + or tensor.device != device + or not tensor.is_contiguous() + ): + raise TypeError(f"{name} must be contiguous FP16 [experts,{width}]") + if ( + ones.shape != (width,) + or ones.dtype != torch.float16 + or ones.device != device + or not ones.is_contiguous() + ): + raise TypeError(f"ones must be contiguous FP16 [{width}] on {device}") + + _run_trellis_dense_hadamard128( + gate_transformed, + gate_hadamard, + ones, + scale_before=False, + ) + _run_trellis_dense_hadamard128( + up_transformed, + up_hadamard, + ones, + scale_before=False, + ) + block_n = 256 + _glm_route_silu_kernel[(routes, triton.cdiv(width, block_n))]( + gate_hadamard, + up_hadamard, + route_experts, + gate_svh, + up_svh, + output, + width=width, + block_n=block_n, + num_warps=4, + ) + return output + + +def run_glm_down_input_transform( + activation: torch.Tensor, + route_experts: torch.Tensor, + down_suh: torch.Tensor, + scaled: torch.Tensor, + rotated: torch.Tensor, + *, + ones: torch.Tensor, +) -> torch.Tensor: + """Apply expert-private down ``suh`` followed by normalized H128.""" + + if activation.ndim != 2: + raise ValueError("down activation must be rank 2") + routes, width = (int(value) for value in activation.shape) + device = activation.device + for name, tensor in ( + ("activation", activation), + ("scaled", scaled), + ("rotated", rotated), + ): + if ( + tensor.shape != activation.shape + or tensor.dtype != torch.float16 + or tensor.device != device + or not tensor.is_contiguous() + ): + raise TypeError( + f"{name} must be contiguous FP16 {tuple(activation.shape)} on {device}" + ) + if ( + route_experts.shape != (routes,) + or route_experts.dtype != torch.int32 + or route_experts.device != device + or not route_experts.is_contiguous() + ): + raise TypeError("route_experts must be contiguous int32 [routes]") + if ( + down_suh.ndim != 2 + or int(down_suh.shape[1]) != width + or down_suh.dtype != torch.float16 + or down_suh.device != device + or not down_suh.is_contiguous() + ): + raise TypeError(f"down_suh must be contiguous FP16 [experts,{width}]") + if ( + ones.shape != (width,) + or ones.dtype != torch.float16 + or ones.device != device + or not ones.is_contiguous() + ): + raise TypeError(f"ones must be contiguous FP16 [{width}] on {device}") + + block_n = 256 + _glm_route_scale_kernel[(routes, triton.cdiv(width, block_n))]( + activation, + route_experts, + down_suh, + scaled, + width=width, + block_n=block_n, + num_warps=4, + ) + _run_trellis_dense_hadamard128( + scaled, + rotated, + ones, + scale_before=True, + ) + return rotated + + +def run_glm_down_output_transform_sum( + down_transformed: torch.Tensor, + topk_weights: torch.Tensor, + down_svh: torch.Tensor, + down_canonical: torch.Tensor, + output: torch.Tensor, + *, + topk: int, +) -> torch.Tensor: + """Invert down H128, apply shared ``svh``, and sum signed top-k routes.""" + + if down_transformed.ndim != 2: + raise ValueError("down transformed output must be rank 2") + routes, width = (int(value) for value in down_transformed.shape) + topk = int(topk) + if topk <= 0 or routes % topk: + raise ValueError("route count must be positive and divisible by topk") + tokens = routes // topk + device = down_transformed.device + if ( + down_transformed.dtype != torch.float16 + or not down_transformed.is_cuda + or not down_transformed.is_contiguous() + or down_canonical.shape != down_transformed.shape + or down_canonical.dtype != torch.float16 + or down_canonical.device != device + or not down_canonical.is_contiguous() + ): + raise TypeError("down buffers must be aligned contiguous CUDA FP16") + if ( + down_svh.shape != (width,) + or down_svh.dtype != torch.float16 + or down_svh.device != device + or not down_svh.is_contiguous() + ): + raise TypeError(f"down_svh must be contiguous FP16 [{width}]") + if ( + topk_weights.shape != (tokens, topk) + or topk_weights.dtype != torch.float32 + or topk_weights.device != device + or not topk_weights.is_contiguous() + ): + raise TypeError(f"topk_weights must be contiguous FP32 [{tokens},{topk}]") + if ( + output.shape != (tokens, width) + or output.dtype not in (torch.float16, torch.bfloat16) + or output.device != device + or not output.is_contiguous() + ): + raise TypeError( + f"output must be contiguous FP16/BF16 [{tokens},{width}] on {device}" + ) + + _run_trellis_dense_hadamard128( + down_transformed, + down_canonical, + down_svh, + scale_before=False, + ) + block_n = 256 + _glm_topk_weighted_sum_kernel[(tokens, triton.cdiv(width, block_n))]( + down_canonical, + topk_weights, + output, + width=width, + topk=topk, + block_n=block_n, + num_warps=4, + ) + return output + + +__all__ = [ + "run_glm_down_input_transform", + "run_glm_down_output_transform_sum", + "run_glm_gate_up_output_transform_silu", +] diff --git a/b12x/moe/_shared/kernels/glm_trellis_w4a8.py b/b12x/moe/_shared/kernels/glm_trellis_w4a8.py new file mode 100644 index 000000000..84140ff10 --- /dev/null +++ b/b12x/moe/_shared/kernels/glm_trellis_w4a8.py @@ -0,0 +1,1730 @@ +"""Route-packed GLM SQG W4A8 projection primitives. + +Unlike Kimi's coupled P24/P33 runtime profile, GLM keeps an independent K3/K4 +decision for every expert tensor. This module therefore stores one native +trellis pool per rate and one global-expert-to-pool-slot map per projection. +Route packing groups rows by expert; one warp then decodes a weight fragment +once and reuses it across a complete M64 route block. + +The primitive intentionally stops at the transformed projection output. The +caller owns the exact GLM ``suh``/Hadamard input transform and the expert-local +``svh`` output transform. Keeping those operations explicit lets the same +projection serve gate, up, and (if quality permits) down without inventing a +Kimi pair-mode contract for GLM. +""" + +from __future__ import annotations + +import functools +import os +from collections.abc import Sequence +from dataclasses import dataclass + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cutlass_dsl import Float32, Int32, Int64, Uint32 + +from b12x._lib.compiler import KernelCompileSpec, compile as b12x_compile +from b12x._lib.intrinsics import ( + cp_async4_shared_global, + cp_async_u32_shared_global, + get_ptr_as_int64, + ld_shared_u32, + ld_shared_v2_u32, + mxfp8_mma_m16n8k32_f32_e4m3, + packed_decode_sqg_xor_cheb_t12_to_e4m3x8, + shared_ptr_to_u32, +) +from b12x._lib.quant.sqg_e4m3 import sqg_xor_cheb_t12_lut +from b12x._lib.runtime_control import raise_if_kernel_resolution_frozen +from b12x._lib.utils import current_cuda_stream, make_ptr +from b12x.gemm._shared.wo_mxfp8 import MXFP8Rows +from b12x.gemm.trellis_linear.w4a8 import ( + _TrellisW4A8DenseLaunch, + _validate_quantized, +) + + +_GLM_ROUTE_BLOCK_ROWS = 128 +_GLM_TRELLIS_CODEBOOK = "sqg_xor_cheb_t12" + + +@dataclass(frozen=True) +class GLMRoutePackedW4A8Projection: + """Prepared native-trellis pools for one independently rated projection. + + Each pool concatenates complete native tensors in slot order. The maps + translate a logical expert ID to its rate-local slot, or ``-1`` when that + expert uses the other rate. Exactly one map must own every active expert; + the materializer is responsible for sealing that partition. + """ + + size_k: int + size_n: int + num_experts: int + trellis_k3: torch.Tensor + trellis_k4: torch.Tensor + expert_slots_k3: torch.Tensor + expert_slots_k4: torch.Tensor + trellis_codebook: str = _GLM_TRELLIS_CODEBOOK + + +def prepare_glm_route_packed_w4a8_projection( + trellis_by_expert: Sequence[torch.Tensor], + bits_by_expert: Sequence[int], + *, + size_k: int, + size_n: int, +) -> GLMRoutePackedW4A8Projection: + """Build rate-local native pools without coupling projection rate choices.""" + + tensors = tuple(trellis_by_expert) + rates = tuple(int(bits) for bits in bits_by_expert) + if not tensors or len(tensors) != len(rates): + raise ValueError("GLM projection preparation requires one rate per expert") + if any(bits not in (3, 4) for bits in rates): + raise ValueError("GLM projection preparation supports only K3 and K4") + device = tensors[0].device + expected_prefix = (int(size_k) // 16, int(size_n) // 16) + for expert, (tensor, bits) in enumerate(zip(tensors, rates, strict=True)): + expected_shape = (*expected_prefix, 16 * bits) + if ( + tensor.dtype != torch.int16 + or tensor.device != device + or not tensor.is_contiguous() + or tuple(tensor.shape) != expected_shape + ): + raise ValueError( + f"expert {expert} K{bits} trellis must be contiguous int16 " + f"{expected_shape} on {device}" + ) + + pools: dict[int, list[torch.Tensor]] = {3: [], 4: []} + slots = { + 3: torch.full((len(tensors),), -1, dtype=torch.int32, device=device), + 4: torch.full((len(tensors),), -1, dtype=torch.int32, device=device), + } + for expert, (tensor, bits) in enumerate(zip(tensors, rates, strict=True)): + slots[bits][expert] = len(pools[bits]) + pools[bits].append(tensor) + + def stack_rate(bits: int) -> torch.Tensor: + if pools[bits]: + return torch.stack(pools[bits]).contiguous() + return torch.empty((0,), dtype=torch.int16, device=device) + + return GLMRoutePackedW4A8Projection( + size_k=int(size_k), + size_n=int(size_n), + num_experts=len(tensors), + trellis_k3=stack_rate(3), + trellis_k4=stack_rate(4), + expert_slots_k3=slots[3], + expert_slots_k4=slots[4], + ) + + +class _GLMRoutePackedW4A8ProjectionLaunch(_TrellisW4A8DenseLaunch): + """One-warp M64xN8 projection over expert-homogeneous route blocks.""" + + def __init__( + self, + *, + size_k: int, + size_n: int, + trellis_bits: int, + topk: int, + shared_input: bool, + route_block_rows: int = _GLM_ROUTE_BLOCK_ROWS, + ) -> None: + super().__init__( + size_k=size_k, + size_n=size_n, + trellis_bits=trellis_bits, + pair_kind=None, + rate_axis=None, + trellis_codebook=_GLM_TRELLIS_CODEBOOK, + m_tile_rows=route_block_rows, + ) + self.topk = int(topk) + self.shared_input = bool(shared_input) + self.route_block_rows = int(route_block_rows) + if self.topk <= 0: + raise ValueError("GLM route-packed W4A8 topk must be positive") + if self.route_block_rows != _GLM_ROUTE_BLOCK_ROWS: + raise ValueError( + "GLM route-packed W4A8 requires fixed 128-row route blocks" + ) + + @cute.jit + def __call__( + self, + values_ptr: cute.Pointer, + scale_rows_ptr: cute.Pointer, + trellis_ptr: cute.Pointer, + rank_lut_ptr: cute.Pointer, + packed_route_indices_ptr: cute.Pointer, + block_expert_ids_ptr: cute.Pointer, + expert_slots_ptr: cute.Pointer, + output_ptr: cute.Pointer, + input_rows: Int32, + routes: Int32, + packed_routes: Int32, + route_blocks: Int32, + num_experts: Int32, + pool_experts: Int32, + stream: cuda.CUstream, + ) -> None: + values = cute.make_tensor( + values_ptr, + cute.make_ordered_layout((input_rows, self.size_k // 4), order=(1, 0)), + ) + scale_rows = cute.make_tensor( + scale_rows_ptr, + cute.make_ordered_layout((input_rows, self.size_k // 32), order=(1, 0)), + ) + trellis = cute.make_tensor( + trellis_ptr, + cute.make_layout((Int64(pool_experts) * Int64(self.trellis_words),)), + ) + rank_lut = cute.make_tensor(rank_lut_ptr, cute.make_layout((4096,))) + packed_route_indices = cute.make_tensor( + packed_route_indices_ptr, cute.make_layout((packed_routes,)) + ) + block_expert_ids = cute.make_tensor( + block_expert_ids_ptr, cute.make_layout((route_blocks,)) + ) + expert_slots = cute.make_tensor( + expert_slots_ptr, cute.make_layout((num_experts,)) + ) + output = cute.make_tensor( + output_ptr, + cute.make_ordered_layout((routes, self.size_n), order=(1, 0)), + ) + self.kernel( + values, + scale_rows, + trellis, + rank_lut, + packed_route_indices, + block_expert_ids, + expert_slots, + output, + routes, + packed_routes, + num_experts, + pool_experts, + ).launch( + grid=(route_blocks, self.size_n // 8, 1), + block=[self.threads, 1, 1], + stream=stream, + ) + + @cute.jit + def _run_rate( + self, + values: cute.Tensor, + scale_rows: cute.Tensor, + trellis: cute.Tensor, + rank_lut: cute.Tensor, + packed_route_indices: cute.Tensor, + output: cute.Tensor, + lane: Int32, + block: Int32, + n_base: Int32, + c: Int32, + g: Int32, + tensor_base: Int64, + routes: Int32, + packed_routes: Int32, + bits: cutlass.Constexpr[int], + ) -> None: + accumulators = tuple( + cute.make_rmem_tensor((4,), Float32) for _ in range(self.m_groups) + ) + for group in cutlass.range_constexpr(self.m_groups): + accumulators[group].fill(0.0) + + rank_lut_addr = get_ptr_as_int64(rank_lut, Int32(0)) + k32 = Int32(0) + while k32 < Int32(self.size_k // 32): + n16 = n_base // Int32(16) + n_high = ((n_base & Int32(15)) >= Int32(8)).to(Int32) + b0, b1 = self._decode_k32_bits_at_base( + trellis, + lane, + k32, + n16, + Int32(0), + n_high, + bits, + rank_lut_addr, + tensor_base, + ) + + word0 = k32 * Int32(8) + c * Int32(2) + packed_base = block * Int32(self.route_block_rows) + for group in cutlass.range_constexpr(self.m_groups): + group_base = packed_base + Int32(group * 16) + packed_lo = group_base + g + packed_hi = packed_lo + Int32(8) + route_lo = routes + route_hi = routes + if packed_lo < packed_routes: + route_lo = packed_route_indices[packed_lo].to(Int32) + if packed_hi < packed_routes: + route_hi = packed_route_indices[packed_hi].to(Int32) + valid_lo = (route_lo >= Int32(0)) & (route_lo < routes) + valid_hi = (route_hi >= Int32(0)) & (route_hi < routes) + input_lo = route_lo + input_hi = route_hi + if cutlass.const_expr(self.shared_input): + input_lo = route_lo // Int32(self.topk) + input_hi = route_hi // Int32(self.topk) + + a0 = Uint32(0) + a1 = Uint32(0) + a2 = Uint32(0) + a3 = Uint32(0) + if valid_lo: + a0 = Uint32(values[input_lo, word0]) + a2 = Uint32(values[input_lo, word0 + Int32(1)]) + if valid_hi: + a1 = Uint32(values[input_hi, word0]) + a3 = Uint32(values[input_hi, word0 + Int32(1)]) + + scale_route = route_lo + scale_valid = valid_lo + if (lane & Int32(1)) != Int32(0): + scale_route = route_hi + scale_valid = valid_hi + scale_input = scale_route + if cutlass.const_expr(self.shared_input): + scale_input = scale_route // Int32(self.topk) + sf = Uint32(127) + if scale_valid: + sf = Uint32(scale_rows[scale_input, k32]) + sfa = sf * Uint32(0x01010101) + + frag = accumulators[group] + d0, d1, d2, d3 = mxfp8_mma_m16n8k32_f32_e4m3( + frag[0], + frag[1], + frag[2], + frag[3], + a0, + a1, + a2, + a3, + b0, + b1, + sfa, + Uint32(0x7F7F7F7F), + ) + frag[0] = d0 + frag[1] = d1 + frag[2] = d2 + frag[3] = d3 + k32 += Int32(1) + + col = n_base + c * Int32(2) + packed_base = block * Int32(self.route_block_rows) + for group in cutlass.range_constexpr(self.m_groups): + group_base = packed_base + Int32(group * 16) + packed_lo = group_base + g + packed_hi = packed_lo + Int32(8) + route_lo = routes + route_hi = routes + if packed_lo < packed_routes: + route_lo = packed_route_indices[packed_lo].to(Int32) + if packed_hi < packed_routes: + route_hi = packed_route_indices[packed_hi].to(Int32) + frag = accumulators[group] + if route_lo >= Int32(0) and route_lo < routes: + output[route_lo, col] = cutlass.Float16(frag[0]) + output[route_lo, col + Int32(1)] = cutlass.Float16(frag[1]) + if route_hi >= Int32(0) and route_hi < routes: + output[route_hi, col] = cutlass.Float16(frag[2]) + output[route_hi, col + Int32(1)] = cutlass.Float16(frag[3]) + + @cute.kernel + def kernel( + self, + values: cute.Tensor, + scale_rows: cute.Tensor, + trellis: cute.Tensor, + rank_lut: cute.Tensor, + packed_route_indices: cute.Tensor, + block_expert_ids: cute.Tensor, + expert_slots: cute.Tensor, + output: cute.Tensor, + routes: Int32, + packed_routes: Int32, + num_experts: Int32, + pool_experts: Int32, + ) -> None: + lane = cute.arch.lane_idx() + block_idx, n_idx, _ = cute.arch.block_idx() + block = Int32(block_idx) + n_base = Int32(n_idx) * Int32(8) + c = lane & Int32(3) + g = lane >> Int32(2) + + expert = block_expert_ids[block].to(Int32) + slot = Int32(-1) + selected = Int32(0) + if expert >= Int32(0) and expert < num_experts: + slot = expert_slots[expert].to(Int32) + if slot >= Int32(0) and slot < pool_experts: + selected = Int32(1) + tensor_base = Int64(0) + if selected != Int32(0): + tensor_base = Int64(slot) * Int64(self.trellis_words) + + accumulators = tuple( + cute.make_rmem_tensor((4,), Float32) for _ in range(self.m_groups) + ) + for group in cutlass.range_constexpr(self.m_groups): + accumulators[group].fill(0.0) + + rank_lut_addr = get_ptr_as_int64(rank_lut, Int32(0)) + k32 = Int32(0) + if selected == Int32(0): + k32 = Int32(self.size_k // 32) + while k32 < Int32(self.size_k // 32): + n16 = n_base // Int32(16) + n_high = ((n_base & Int32(15)) >= Int32(8)).to(Int32) + if cutlass.const_expr(self.trellis_bits == 3): + b0, b1 = self._decode_k32_bits_at_base( + trellis, + lane, + k32, + n16, + Int32(0), + n_high, + 3, + rank_lut_addr, + tensor_base, + ) + else: + b0, b1 = self._decode_k32_bits_at_base( + trellis, + lane, + k32, + n16, + Int32(0), + n_high, + 4, + rank_lut_addr, + tensor_base, + ) + + word0 = k32 * Int32(8) + c * Int32(2) + packed_base = block * Int32(self.route_block_rows) + for group in cutlass.range_constexpr(self.m_groups): + group_base = packed_base + Int32(group * 16) + packed_lo = group_base + g + packed_hi = packed_lo + Int32(8) + route_lo = routes + route_hi = routes + if packed_lo < packed_routes: + route_lo = packed_route_indices[packed_lo].to(Int32) + if packed_hi < packed_routes: + route_hi = packed_route_indices[packed_hi].to(Int32) + valid_lo = (route_lo >= Int32(0)) & (route_lo < routes) + valid_hi = (route_hi >= Int32(0)) & (route_hi < routes) + input_lo = route_lo + input_hi = route_hi + if cutlass.const_expr(self.shared_input): + input_lo = route_lo // Int32(self.topk) + input_hi = route_hi // Int32(self.topk) + + a0 = Uint32(0) + a1 = Uint32(0) + a2 = Uint32(0) + a3 = Uint32(0) + if valid_lo: + a0 = Uint32(values[input_lo, word0]) + a2 = Uint32(values[input_lo, word0 + Int32(1)]) + if valid_hi: + a1 = Uint32(values[input_hi, word0]) + a3 = Uint32(values[input_hi, word0 + Int32(1)]) + + scale_route = route_lo + scale_valid = valid_lo + if (lane & Int32(1)) != Int32(0): + scale_route = route_hi + scale_valid = valid_hi + scale_input = scale_route + if cutlass.const_expr(self.shared_input): + scale_input = scale_route // Int32(self.topk) + sf = Uint32(127) + if scale_valid: + sf = Uint32(scale_rows[scale_input, k32]) + sfa = sf * Uint32(0x01010101) + + frag = accumulators[group] + d0, d1, d2, d3 = mxfp8_mma_m16n8k32_f32_e4m3( + frag[0], + frag[1], + frag[2], + frag[3], + a0, + a1, + a2, + a3, + b0, + b1, + sfa, + Uint32(0x7F7F7F7F), + ) + frag[0] = d0 + frag[1] = d1 + frag[2] = d2 + frag[3] = d3 + k32 += Int32(1) + + col = n_base + c * Int32(2) + packed_base = block * Int32(self.route_block_rows) + for group in cutlass.range_constexpr(self.m_groups): + group_base = packed_base + Int32(group * 16) + packed_lo = group_base + g + packed_hi = packed_lo + Int32(8) + route_lo = routes + route_hi = routes + if packed_lo < packed_routes: + route_lo = packed_route_indices[packed_lo].to(Int32) + if packed_hi < packed_routes: + route_hi = packed_route_indices[packed_hi].to(Int32) + frag = accumulators[group] + if selected != Int32(0): + if route_lo >= Int32(0) and route_lo < routes: + output[route_lo, col] = cutlass.Float16(frag[0]) + output[route_lo, col + Int32(1)] = cutlass.Float16(frag[1]) + if route_hi >= Int32(0) and route_hi < routes: + output[route_hi, col] = cutlass.Float16(frag[2]) + output[route_hi, col + Int32(1)] = cutlass.Float16(frag[3]) + + +class _GLMRoutePackedW4A8MixedProjectionLaunch(_GLMRoutePackedW4A8ProjectionLaunch): + """Dispatch K3/K4 once per expert-homogeneous route block. + + One grid reads both native pools and selects the owning decoder from the + block's expert identifier. This preserves independent tensor rates without + launching inactive CTAs for the non-owning bitrate. + """ + + def __init__( + self, + *, + size_k: int, + size_n: int, + topk: int, + shared_input: bool, + route_block_rows: int = _GLM_ROUTE_BLOCK_ROWS, + ) -> None: + super().__init__( + size_k=size_k, + size_n=size_n, + trellis_bits=3, + topk=topk, + shared_input=shared_input, + route_block_rows=route_block_rows, + ) + self.trellis_words_k3 = self.size_k * self.size_n * 3 // 32 + self.trellis_words_k4 = self.size_k * self.size_n * 4 // 32 + + @cute.jit + def __call__( + self, + values_ptr: cute.Pointer, + scale_rows_ptr: cute.Pointer, + trellis_k3_ptr: cute.Pointer, + trellis_k4_ptr: cute.Pointer, + rank_lut_ptr: cute.Pointer, + packed_route_indices_ptr: cute.Pointer, + block_expert_ids_ptr: cute.Pointer, + expert_slots_k3_ptr: cute.Pointer, + expert_slots_k4_ptr: cute.Pointer, + output_ptr: cute.Pointer, + input_rows: Int32, + routes: Int32, + packed_routes: Int32, + route_blocks: Int32, + num_experts: Int32, + pool_experts_k3: Int32, + pool_experts_k4: Int32, + stream: cuda.CUstream, + ) -> None: + values = cute.make_tensor( + values_ptr, + cute.make_ordered_layout((input_rows, self.size_k // 4), order=(1, 0)), + ) + scale_rows = cute.make_tensor( + scale_rows_ptr, + cute.make_ordered_layout((input_rows, self.size_k // 32), order=(1, 0)), + ) + trellis_k3 = cute.make_tensor( + trellis_k3_ptr, + cute.make_layout((Int64(pool_experts_k3) * Int64(self.trellis_words_k3),)), + ) + trellis_k4 = cute.make_tensor( + trellis_k4_ptr, + cute.make_layout((Int64(pool_experts_k4) * Int64(self.trellis_words_k4),)), + ) + rank_lut = cute.make_tensor(rank_lut_ptr, cute.make_layout((4096,))) + packed_route_indices = cute.make_tensor( + packed_route_indices_ptr, cute.make_layout((packed_routes,)) + ) + block_expert_ids = cute.make_tensor( + block_expert_ids_ptr, cute.make_layout((route_blocks,)) + ) + expert_slots_k3 = cute.make_tensor( + expert_slots_k3_ptr, cute.make_layout((num_experts,)) + ) + expert_slots_k4 = cute.make_tensor( + expert_slots_k4_ptr, cute.make_layout((num_experts,)) + ) + output = cute.make_tensor( + output_ptr, + cute.make_ordered_layout((routes, self.size_n), order=(1, 0)), + ) + self.kernel( + values, + scale_rows, + trellis_k3, + trellis_k4, + rank_lut, + packed_route_indices, + block_expert_ids, + expert_slots_k3, + expert_slots_k4, + output, + routes, + packed_routes, + num_experts, + pool_experts_k3, + pool_experts_k4, + ).launch( + grid=(route_blocks, self.size_n // 8, 1), + block=[self.threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + values: cute.Tensor, + scale_rows: cute.Tensor, + trellis_k3: cute.Tensor, + trellis_k4: cute.Tensor, + rank_lut: cute.Tensor, + packed_route_indices: cute.Tensor, + block_expert_ids: cute.Tensor, + expert_slots_k3: cute.Tensor, + expert_slots_k4: cute.Tensor, + output: cute.Tensor, + routes: Int32, + packed_routes: Int32, + num_experts: Int32, + pool_experts_k3: Int32, + pool_experts_k4: Int32, + ) -> None: + lane = cute.arch.lane_idx() + block_idx, n_idx, _ = cute.arch.block_idx() + block = Int32(block_idx) + n_base = Int32(n_idx) * Int32(8) + c = lane & Int32(3) + g = lane >> Int32(2) + + expert = block_expert_ids[block].to(Int32) + rate = Int32(0) + tensor_base_k3 = Int64(0) + tensor_base_k4 = Int64(0) + if expert >= Int32(0) and expert < num_experts: + slot_k3 = expert_slots_k3[expert].to(Int32) + slot_k4 = expert_slots_k4[expert].to(Int32) + if slot_k3 >= Int32(0) and slot_k3 < pool_experts_k3: + rate = Int32(3) + tensor_base_k3 = Int64(slot_k3) * Int64(self.trellis_words_k3) + elif slot_k4 >= Int32(0) and slot_k4 < pool_experts_k4: + rate = Int32(4) + tensor_base_k4 = Int64(slot_k4) * Int64(self.trellis_words_k4) + if rate == Int32(3): + self._run_rate( + values, + scale_rows, + trellis_k3, + rank_lut, + packed_route_indices, + output, + lane, + block, + n_base, + c, + g, + tensor_base_k3, + routes, + packed_routes, + 3, + ) + elif rate == Int32(4): + self._run_rate( + values, + scale_rows, + trellis_k4, + rank_lut, + packed_route_indices, + output, + lane, + block, + n_base, + c, + g, + tensor_base_k4, + routes, + packed_routes, + 4, + ) + + +_V2_TILE_N = 64 +_V2_N8_PER_WARP = 2 +_V2_THREADS = 128 +_V2_KTILE_K32 = 4 # k32 steps per staged A tile (K128) +_V2_A_BLOCK_WORDS = 16 * 32 # 16 rows x K128 bytes, u32 words +_V2_STAGES = 2 + + +class _GLMRoutePackedW4A8TileLaunch(_GLMRoutePackedW4A8MixedProjectionLaunch): + """Four-warp M64xN256 route-packed projection with staged A operands. + + The M64xN8 one-warp kernel re-read the quantized A rows from global + memory once per N8 column strip (256x amplification at N=2048) and had a + single warp to cover decode plus MMA latency. This launch keeps the + packed-route workspace and decode primitives unchanged and fixes the + schedule: + + - one CTA covers M64 x N256 with 128 threads; each warp owns eight n8 + strips, so a decoded K32xN8 fragment feeds four M16 MMAs instead of + being re-decoded per M16 pair; + - A bytes and UE8M0 row scales are cp.async double-buffered through + shared memory with the xor-swizzle from the fused W4A8 pipeline, so A + global traffic per route block falls from size_n/8 reads to size_n/256; + - B stays register-resident: the existing warp trellis decode already + produces the exact m16n8k32 fragment layout, so there is no B staging, + no second decode of a fragment inside a CTA, and the compact payload is + the only weight representation touched. + + Rate independence is preserved: the expert's K3 or K4 pool is selected + once per CTA exactly as in the mixed one-warp launch. + """ + + def __init__( + self, + *, + size_k: int, + size_n: int, + topk: int, + shared_input: bool, + blocks_per_cta: int = 8, + stages: int = 2, + ) -> None: + super().__init__( + size_k=size_k, + size_n=size_n, + topk=topk, + shared_input=shared_input, + ) + if int(stages) not in (2, 3, 4): + raise ValueError("v2 stages must be 2, 3, or 4") + self.stages = int(stages) + m16_blocks = self.route_block_rows // 16 + if int(blocks_per_cta) not in (2, 4, 8): + raise ValueError("v2 blocks_per_cta must be 2, 4, or 8") + if m16_blocks % int(blocks_per_cta): + raise ValueError("v2 blocks_per_cta must divide the route block") + if self.size_n % _V2_TILE_N: + raise ValueError("v2 requires size_n divisible by 64") + if self.size_k % (_V2_KTILE_K32 * 32): + raise ValueError("v2 requires size_k divisible by 128") + self.blocks_per_cta = int(blocks_per_cta) + # M128 blocks are split into M16 CTA parts when requested. + self.m_parts = m16_blocks // self.blocks_per_cta + self.a_words = self.blocks_per_cta * _V2_A_BLOCK_WORDS + self.asf_words = self.blocks_per_cta * 16 + # Maximum compact B payload for one K128 x N64 tile is K4: + # 8 K16 rows x 4 N16 columns x 32 u32 words per tile. + self.b_words = (_V2_KTILE_K32 * 2) * (_V2_TILE_N // 16) * (8 * 4) + self.b_offset_words = self.a_words + self.asf_words + self.stage_words = self.b_offset_words + self.b_words + self.stage_bytes = self.stage_words * 4 + self.lut_words = 4096 // 4 + self.k_tiles = self.size_k // (_V2_KTILE_K32 * 32) + + @cute.jit + def __call__( + self, + values_ptr: cute.Pointer, + scale_rows_ptr: cute.Pointer, + trellis_k3_ptr: cute.Pointer, + trellis_k4_ptr: cute.Pointer, + rank_lut_ptr: cute.Pointer, + packed_route_indices_ptr: cute.Pointer, + block_expert_ids_ptr: cute.Pointer, + expert_slots_k3_ptr: cute.Pointer, + expert_slots_k4_ptr: cute.Pointer, + output_ptr: cute.Pointer, + input_rows: Int32, + routes: Int32, + packed_routes: Int32, + route_blocks: Int32, + num_experts: Int32, + pool_experts_k3: Int32, + pool_experts_k4: Int32, + stream: cuda.CUstream, + ) -> None: + values = cute.make_tensor( + values_ptr, + cute.make_layout((Int64(input_rows) * Int64(self.size_k // 4),)), + ) + scale_rows = cute.make_tensor( + scale_rows_ptr, + cute.make_layout((Int64(input_rows) * Int64(self.size_k // 32),)), + ) + trellis_k3 = cute.make_tensor( + trellis_k3_ptr, + cute.make_layout((Int64(pool_experts_k3) * Int64(self.trellis_words_k3),)), + ) + trellis_k4 = cute.make_tensor( + trellis_k4_ptr, + cute.make_layout((Int64(pool_experts_k4) * Int64(self.trellis_words_k4),)), + ) + rank_lut = cute.make_tensor(rank_lut_ptr, cute.make_layout((4096,))) + packed_route_indices = cute.make_tensor( + packed_route_indices_ptr, cute.make_layout((packed_routes,)) + ) + block_expert_ids = cute.make_tensor( + block_expert_ids_ptr, cute.make_layout((route_blocks,)) + ) + expert_slots_k3 = cute.make_tensor( + expert_slots_k3_ptr, cute.make_layout((num_experts,)) + ) + expert_slots_k4 = cute.make_tensor( + expert_slots_k4_ptr, cute.make_layout((num_experts,)) + ) + output = cute.make_tensor( + output_ptr, + cute.make_ordered_layout((routes, self.size_n), order=(1, 0)), + ) + self.kernel( + values, + scale_rows, + trellis_k3, + trellis_k4, + rank_lut, + packed_route_indices, + block_expert_ids, + expert_slots_k3, + expert_slots_k4, + output, + routes, + packed_routes, + num_experts, + pool_experts_k3, + pool_experts_k4, + ).launch( + grid=( + route_blocks * Int32(self.m_parts), + self.size_n // _V2_TILE_N, + 1, + ), + block=[_V2_THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + values: cute.Tensor, + scale_rows: cute.Tensor, + trellis_k3: cute.Tensor, + trellis_k4: cute.Tensor, + rank_lut: cute.Tensor, + packed_route_indices: cute.Tensor, + block_expert_ids: cute.Tensor, + expert_slots_k3: cute.Tensor, + expert_slots_k4: cute.Tensor, + output: cute.Tensor, + routes: Int32, + packed_routes: Int32, + num_experts: Int32, + pool_experts_k3: Int32, + pool_experts_k4: Int32, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + tid = Int32(tidx) + + smem = cutlass.utils.SmemAllocator() + + @cute.struct + class Storage: + sData: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Uint32, + self.stages * self.stage_words + self.lut_words, + ], + 16, + ] + + storage = smem.allocate(Storage) + # Hoisted before any dynamic control flow (flattener trap). + s_base = shared_ptr_to_u32(storage.sData.data_ptr()) + + part = Int32(0) + block = Int32(bidx) + if cutlass.const_expr(self.m_parts > 1): + part = block % Int32(self.m_parts) + block = block // Int32(self.m_parts) + n_tile = Int32(bidy) + + expert = block_expert_ids[block].to(Int32) + rate = Int32(0) + tensor_base_k3 = Int64(0) + tensor_base_k4 = Int64(0) + if expert >= Int32(0) and expert < num_experts: + slot_k3 = expert_slots_k3[expert].to(Int32) + slot_k4 = expert_slots_k4[expert].to(Int32) + if slot_k3 >= Int32(0) and slot_k3 < pool_experts_k3: + rate = Int32(3) + tensor_base_k3 = Int64(slot_k3) * Int64(self.trellis_words_k3) + elif slot_k4 >= Int32(0) and slot_k4 < pool_experts_k4: + rate = Int32(4) + tensor_base_k4 = Int64(slot_k4) * Int64(self.trellis_words_k4) + if rate == Int32(3): + self._tile_body( + values, + scale_rows, + trellis_k3, + rank_lut, + packed_route_indices, + output, + s_base, + tid, + block, + part, + n_tile, + tensor_base_k3, + routes, + packed_routes, + 3, + ) + elif rate == Int32(4): + self._tile_body( + values, + scale_rows, + trellis_k4, + rank_lut, + packed_route_indices, + output, + s_base, + tid, + block, + part, + n_tile, + tensor_base_k4, + routes, + packed_routes, + 4, + ) + + @cute.jit + def _stage_ktile( + self, + values: cute.Tensor, + scale_rows: cute.Tensor, + trellis: cute.Tensor, + stage_base: Int32, + tid: Int32, + kt: Int32, + n_tile: Int32, + tensor_base: Int64, + a_src_words: cute.Tensor, # rmem [blocks]: -1 = padded row (skip) + a_dst_addr: Int32, + asf_src_byte: Int32, # -1 = padded row (skip) + asf_dst_addr: Int32, + bits: cutlass.Constexpr[int], + ) -> None: + """cp.async one compact (A, Asf, B) K128 tile into shared memory. + + A: blocks x 16 rows x 128B, xor-swizzled 16B units exactly like the + fused W4A8 pipeline so fragment reads are bank-conflict-free. Rows + whose packed route slot is padding are skipped; their smem bytes are + never consumed because the epilogue masks those routes out. + """ + for blk in cutlass.range_constexpr(self.blocks_per_cta): + src = Int32(a_src_words[blk]) + if src >= Int32(0): + cp_async4_shared_global( + stage_base + Int32(blk * _V2_A_BLOCK_WORDS * 4) + a_dst_addr, + get_ptr_as_int64(values, src + kt * Int32(32)), + ) + if asf_src_byte >= Int32(0): + cp_async_u32_shared_global( + stage_base + Int32(self.a_words * 4) + asf_dst_addr, + get_ptr_as_int64(scale_rows, asf_src_byte + kt * Int32(4)), + ) + + # Cooperatively stage the native compressed trellis payload. The + # source layout is [K16, N16, 8*bits u32]; the shared layout packs only + # this CTA's 8-by-4 tile. K3 uses 192 16-byte units, K4 uses 256. + units_per_trellis_tile = 2 * int(bits) + n16_per_cta = _V2_TILE_N // 16 + total_units = _V2_KTILE_K32 * 2 * n16_per_cta * units_per_trellis_tile + for copy in cutlass.range_constexpr(2): + unit = tid + Int32(copy * _V2_THREADS) + if unit < Int32(total_units): + tile = unit // Int32(units_per_trellis_tile) + unit_in_tile = unit % Int32(units_per_trellis_tile) + k16_local = tile // Int32(n16_per_cta) + n16_local = tile % Int32(n16_per_cta) + global_k16 = kt * Int32(_V2_KTILE_K32 * 2) + k16_local + global_n16 = n_tile * Int32(n16_per_cta) + n16_local + src_word = ( + tensor_base + + (Int64(global_k16) * Int64(self.size_n // 16) + Int64(global_n16)) + * Int64(8 * int(bits)) + + Int64(unit_in_tile) * Int64(4) + ) + cp_async4_shared_global( + stage_base + Int32(self.b_offset_words * 4) + unit * Int32(16), + get_ptr_as_int64(trellis, src_word), + ) + + @cute.jit + def _decode_tile_at_base_shared( + self, + b_base: Int32, + lane: Int32, + k16_local: Int32, + n16_local: Int32, + n_high: Int32, + bits: cutlass.Constexpr[int], + rank_lut_addr: Int32, + ) -> Uint32: + ia, ib, s2 = self._lane_geom(lane, bits) + tile_word = (k16_local * Int32(_V2_TILE_N // 16) + n16_local) * Int32( + 8 * int(bits) + ) + a = ld_shared_u32(b_base + (tile_word + ia) * Int32(4)) + b = ld_shared_u32(b_base + (tile_word + ib) * Int32(4)) + merged = (Int64(a) << Int64(32)) | Int64(b) + win_a = Uint32(merged >> Int64(s2)) + win_b = Uint32(merged >> Int64(s2 + Int32(4 * int(bits)))) + lo, hi = packed_decode_sqg_xor_cheb_t12_to_e4m3x8( + win_a, + win_b, + rank_lut_addr, + int(bits), + t12_in_shared=True, + ) + value = lo + if n_high != Int32(0): + value = hi + return value + + @cute.jit + def _decode_k32_pair_shared( + self, + b_base: Int32, + lane: Int32, + kb: Int32, + n16_local: Int32, + bits: cutlass.Constexpr[int], + rank_lut_addr: Int32, + ): + """Decode BOTH n8 halves of one staged n16 tile for one K32 step. + + ``_decode_tile_at_base_shared`` computes the low and high n8 halves + and discards one by ``n_high``, so both halves of every staged tile + were decoded twice. This performs the same two K16 tile reads once, + keeps both halves, and reproduces the two-call fragment values and + lane placement bit-exactly with half the shared-memory reads, SQG + hashing, and LUT gathers. + """ + ia, ib, s2 = self._lane_geom(lane, bits) + n16_words = Int32(_V2_TILE_N // 16) * Int32(8 * int(bits)) + tile_word0 = (kb * Int32(2) * Int32(_V2_TILE_N // 16) + n16_local) * Int32( + 8 * int(bits) + ) + a0 = ld_shared_u32(b_base + (tile_word0 + ia) * Int32(4)) + b0w = ld_shared_u32(b_base + (tile_word0 + ib) * Int32(4)) + merged0 = (Int64(a0) << Int64(32)) | Int64(b0w) + lo0, hi0 = packed_decode_sqg_xor_cheb_t12_to_e4m3x8( + Uint32(merged0 >> Int64(s2)), + Uint32(merged0 >> Int64(s2 + Int32(4 * int(bits)))), + rank_lut_addr, + int(bits), + t12_in_shared=True, + ) + tile_word1 = tile_word0 + n16_words + a1 = ld_shared_u32(b_base + (tile_word1 + ia) * Int32(4)) + b1w = ld_shared_u32(b_base + (tile_word1 + ib) * Int32(4)) + merged1 = (Int64(a1) << Int64(32)) | Int64(b1w) + lo1, hi1 = packed_decode_sqg_xor_cheb_t12_to_e4m3x8( + Uint32(merged1 >> Int64(s2)), + Uint32(merged1 >> Int64(s2 + Int32(4 * int(bits)))), + rank_lut_addr, + int(bits), + t12_in_shared=True, + ) + c = lane & Int32(3) + own_lo = lo0 + send_lo = lo1 + own_hi = hi0 + send_hi = hi1 + if c >= Int32(2): + own_lo = lo1 + send_lo = lo0 + own_hi = hi1 + send_hi = hi0 + peer_lo = Uint32(cute.arch.shuffle_sync_bfly(send_lo, offset=2)) + peer_hi = Uint32(cute.arch.shuffle_sync_bfly(send_hi, offset=2)) + return own_lo, peer_lo, own_hi, peer_hi + + @cute.jit + def _decode_k32_bits_at_base_shared( + self, + b_base: Int32, + lane: Int32, + kb: Int32, + n16_local: Int32, + n_high: Int32, + bits: cutlass.Constexpr[int], + rank_lut_addr: Int32, + ): + e0 = self._decode_tile_at_base_shared( + b_base, + lane, + kb * Int32(2), + n16_local, + n_high, + bits, + rank_lut_addr, + ) + e1 = self._decode_tile_at_base_shared( + b_base, + lane, + kb * Int32(2) + Int32(1), + n16_local, + n_high, + bits, + rank_lut_addr, + ) + c = lane & Int32(3) + own = e0 + send = e1 + if c >= Int32(2): + own = e1 + send = e0 + peer = Uint32(cute.arch.shuffle_sync_bfly(send, offset=2)) + return own, peer + + @cute.jit + def _tile_body( + self, + values: cute.Tensor, + scale_rows: cute.Tensor, + trellis: cute.Tensor, + rank_lut: cute.Tensor, + packed_route_indices: cute.Tensor, + output: cute.Tensor, + s_base: Int32, + tid: Int32, + block: Int32, + part: Int32, + n_tile: Int32, + tensor_base: Int64, + routes: Int32, + packed_routes: Int32, + bits: cutlass.Constexpr[int], + ) -> None: + warp = tid >> Int32(5) + lane = tid & Int32(31) + q = lane >> Int32(2) + c = lane & Int32(3) + + packed_base = block * Int32(self.route_block_rows) + part * Int32( + self.blocks_per_cta * 16 + ) + a_row_words = Int32(self.size_k // 4) + asf_row_bytes = Int32(self.size_k // 32) + + # ---- Per-thread staging source addresses, resolved once. ---- + # A copy role: row = tid>>3 in [0,16), 16B unit v = tid&7, swizzled + # destination unit p = v ^ (row & 7). + stage_row = tid >> Int32(3) + stage_v = tid & Int32(7) + stage_p = stage_v ^ (stage_row & Int32(7)) + a_dst_addr = (stage_row << Int32(7)) + (stage_p << Int32(4)) + a_src_words = cute.make_rmem_tensor((self.blocks_per_cta,), cutlass.Int32) + for blk in cutlass.range_constexpr(self.blocks_per_cta): + src_words = Int32(-1) + slot = packed_base + Int32(blk * 16) + stage_row + if slot < packed_routes: + route = packed_route_indices[slot].to(Int32) + if route >= Int32(0) and route < routes: + src_row = route + if cutlass.const_expr(self.shared_input): + src_row = route // Int32(self.topk) + src_words = src_row * a_row_words + (stage_v << Int32(2)) + a_src_words[blk] = src_words + # Asf copy role: tid < blocks*16 loads one u32 (4 UE8M0 bytes = one + # K128 tile) for packed row tid. + asf_src_byte = Int32(-1) + asf_dst_addr = tid << Int32(2) + if tid < Int32(self.blocks_per_cta * 16): + slot = packed_base + tid + if slot < packed_routes: + route = packed_route_indices[slot].to(Int32) + if route >= Int32(0) and route < routes: + src_row = route + if cutlass.const_expr(self.shared_input): + src_row = route // Int32(self.topk) + asf_src_byte = src_row * asf_row_bytes + # The decoder performs many byte lookups per fragment. Stage the + # immutable 4 KiB T12 table once per CTA to replace long-scoreboard + # global byte loads with shared-memory lookups. + lut_base = s_base + Int32(self.stages * self.stage_bytes) + for copy in cutlass.range_constexpr(2): + lut_src = tid * Int32(32) + Int32(copy * 16) + cp_async4_shared_global( + lut_base + lut_src, + get_ptr_as_int64(rank_lut, lut_src), + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + rank_lut_addr = Int32(lut_base) + + # f32 accumulators [m-block][n8 strip][fragment], full K sweep. + facc = cute.make_rmem_tensor((self.blocks_per_cta, _V2_N8_PER_WARP, 4), Float32) + facc.fill(0.0) + + # ---- cp.async pipeline: prefetch the first stages-1 k-tiles. ---- + for p in cutlass.range_constexpr(self.stages - 1): + if Int32(p) < Int32(self.k_tiles): + self._stage_ktile( + values, + scale_rows, + trellis, + s_base + Int32(p % self.stages) * Int32(self.stage_bytes), + tid, + Int32(p), + n_tile, + tensor_base, + a_src_words, + a_dst_addr, + asf_src_byte, + asf_dst_addr, + bits, + ) + cute.arch.cp_async_commit_group() + + kt = Int32(0) + k_tiles = Int32(self.k_tiles) + while kt < k_tiles: + nxt = kt + Int32(self.stages - 1) + if nxt < k_tiles: + self._stage_ktile( + values, + scale_rows, + trellis, + s_base + (nxt % Int32(self.stages)) * Int32(self.stage_bytes), + tid, + nxt, + n_tile, + tensor_base, + a_src_words, + a_dst_addr, + asf_src_byte, + asf_dst_addr, + bits, + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(self.stages - 1) + cute.arch.sync_threads() + + cur = s_base + (kt % Int32(self.stages)) * Int32(self.stage_bytes) + sasf_base = cur + Int32(self.a_words * 4) + sb_base = cur + Int32(self.b_offset_words * 4) + + # Lane's SFA row follows the one-warp kernel's parity rule: even + # lanes carry row q, odd lanes row q+8. + asf_row = q + ((lane & Int32(1)) << Int32(3)) + asc = cute.make_rmem_tensor((self.blocks_per_cta,), cutlass.Uint32) + for blk in cutlass.range_constexpr(self.blocks_per_cta): + asc[blk] = ld_shared_u32( + sasf_base + Int32(blk * 64) + (asf_row << Int32(2)) + ) + + for kb in cutlass.range_constexpr(_V2_KTILE_K32): + # Swizzle-aware A fragment loads (rows q and q+8). + u_phys = (Int32(kb * 2) + (c >> Int32(1))) ^ q + a_lo = ( + cur + + (q << Int32(7)) + + (u_phys << Int32(4)) + + ((c & Int32(1)) << Int32(3)) + ) + a_frag = cute.make_rmem_tensor((self.blocks_per_cta, 4), cutlass.Uint32) + for blk in cutlass.range_constexpr(self.blocks_per_cta): + blk_off = Int32(blk * _V2_A_BLOCK_WORDS * 4) + f0, f2 = ld_shared_v2_u32(a_lo + blk_off) + f1, f3 = ld_shared_v2_u32(a_lo + blk_off + Int32(8 * 128)) + a_frag[blk, 0] = f0 + a_frag[blk, 1] = f1 + a_frag[blk, 2] = f2 + a_frag[blk, 3] = f3 + sfa = cute.make_rmem_tensor((self.blocks_per_cta,), cutlass.Uint32) + for blk in cutlass.range_constexpr(self.blocks_per_cta): + sf = (Uint32(asc[blk]) >> Uint32(8 * kb)) & Uint32(0xFF) + sfa[blk] = sf * Uint32(0x01010101) + + for t in cutlass.range_constexpr(_V2_N8_PER_WARP // 2): + n_base_lo = ( + n_tile * Int32(_V2_TILE_N) + + warp * Int32(_V2_N8_PER_WARP * 8) + + Int32(t * 16) + ) + n16_local = (n_base_lo >> Int32(4)) - n_tile * Int32( + _V2_TILE_N // 16 + ) + b0_lo, b1_lo, b0_hi, b1_hi = self._decode_k32_pair_shared( + sb_base, + lane, + Int32(kb), + n16_local, + bits, + rank_lut_addr, + ) + for h in cutlass.range_constexpr(2): + i = t * 2 + h + if cutlass.const_expr(h == 0): + b0 = b0_lo + b1 = b1_lo + else: + b0 = b0_hi + b1 = b1_hi + for blk in cutlass.range_constexpr(self.blocks_per_cta): + frag = facc + d0, d1, d2, d3 = mxfp8_mma_m16n8k32_f32_e4m3( + frag[blk, i, 0], + frag[blk, i, 1], + frag[blk, i, 2], + frag[blk, i, 3], + Uint32(a_frag[blk, 0]), + Uint32(a_frag[blk, 1]), + Uint32(a_frag[blk, 2]), + Uint32(a_frag[blk, 3]), + b0, + b1, + Uint32(sfa[blk]), + Uint32(0x7F7F7F7F), + ) + frag[blk, i, 0] = d0 + frag[blk, i, 1] = d1 + frag[blk, i, 2] = d2 + frag[blk, i, 3] = d3 + cute.arch.sync_threads() + kt += Int32(1) + + # ---- epilogue: scatter valid routes, identical to the one-warp + # kernel's mapping (rows q/q+8 of each M16 block, columns 2c/2c+1 of + # each n8 strip). + for blk in cutlass.range_constexpr(self.blocks_per_cta): + group_base = packed_base + Int32(blk * 16) + packed_lo = group_base + q + packed_hi = packed_lo + Int32(8) + route_lo = routes + route_hi = routes + if packed_lo < packed_routes: + route_lo = packed_route_indices[packed_lo].to(Int32) + if packed_hi < packed_routes: + route_hi = packed_route_indices[packed_hi].to(Int32) + for i in cutlass.range_constexpr(_V2_N8_PER_WARP): + col = ( + n_tile * Int32(_V2_TILE_N) + + warp * Int32(_V2_N8_PER_WARP * 8) + + Int32(i * 8) + + c * Int32(2) + ) + if route_lo >= Int32(0) and route_lo < routes: + output[route_lo, col] = cutlass.Float16(facc[blk, i, 0]) + output[route_lo, col + Int32(1)] = cutlass.Float16(facc[blk, i, 1]) + if route_hi >= Int32(0) and route_hi < routes: + output[route_hi, col] = cutlass.Float16(facc[blk, i, 2]) + output[route_hi, col + Int32(1)] = cutlass.Float16(facc[blk, i, 3]) + + +def _ptr(dtype, address: int): + return make_ptr(dtype, address, cute.AddressSpace.gmem, assumed_align=16) + + +@functools.cache +def _compile_glm_route_packed_projection( + size_k: int, + size_n: int, + trellis_bits: int, + topk: int, + shared_input: bool, + device_index: int, +): + launch = _GLMRoutePackedW4A8ProjectionLaunch( + size_k=size_k, + size_n=size_n, + trellis_bits=trellis_bits, + topk=topk, + shared_input=shared_input, + ) + key = ( + int(size_k), + int(size_n), + int(trellis_bits), + int(topk), + bool(shared_input), + int(device_index), + ) + raise_if_kernel_resolution_frozen("cute.compile", target=launch, cache_key=key) + return b12x_compile( + launch, + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint8, 16), + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint8, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Float16, 16), + 1, + 1, + 1, + 1, + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key( + "moe.glm_route_packed_trellis_w4a8_projection", + 1, + key, + ), + ) + + +@functools.cache +def _compile_glm_route_packed_mixed_projection( + size_k: int, + size_n: int, + topk: int, + shared_input: bool, + device_index: int, +): + launch = _GLMRoutePackedW4A8MixedProjectionLaunch( + size_k=size_k, + size_n=size_n, + topk=topk, + shared_input=shared_input, + ) + key = ( + int(size_k), + int(size_n), + int(topk), + bool(shared_input), + int(device_index), + ) + raise_if_kernel_resolution_frozen("cute.compile", target=launch, cache_key=key) + return b12x_compile( + launch, + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint8, 16), + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint8, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Float16, 16), + 1, + 1, + 1, + 1, + 1, + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key( + "moe.glm_route_packed_trellis_w4a8_mixed_projection", + 1, + key, + ), + ) + + +@functools.cache +def _compile_glm_route_packed_tile_projection( + size_k: int, + size_n: int, + topk: int, + shared_input: bool, + blocks_per_cta: int, + stages: int, + device_index: int, +): + launch = _GLMRoutePackedW4A8TileLaunch( + size_k=size_k, + size_n=size_n, + topk=topk, + shared_input=shared_input, + blocks_per_cta=blocks_per_cta, + stages=stages, + ) + key = ( + int(size_k), + int(size_n), + int(topk), + bool(shared_input), + int(blocks_per_cta), + int(stages), + int(device_index), + ) + raise_if_kernel_resolution_frozen("cute.compile", target=launch, cache_key=key) + return b12x_compile( + launch, + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint8, 16), + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint32, 16), + _ptr(cutlass.Uint8, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Int32, 16), + _ptr(cutlass.Float16, 16), + 1, + 1, + 1, + 1, + 1, + 1, + 1, + current_cuda_stream(), + compile_spec=KernelCompileSpec.from_key( + "moe.glm_route_packed_trellis_w4a8_tile_projection", + 1, + key, + ), + ) + + +def _selected_w4a8_kernel() -> str: + kernel = os.environ.get("B12X_GLM_W4A8_KERNEL", "m128n64").strip().lower() + if kernel not in ("m64n8", "m128n64"): + raise ValueError("B12X_GLM_W4A8_KERNEL must be 'm64n8' (one-warp) or 'm128n64'") + return kernel + + +def _v2_stages() -> int: + stages = int(os.environ.get("B12X_GLM_W4A8_V2_STAGES", "2")) + if stages not in (2, 3, 4): + raise ValueError("B12X_GLM_W4A8_V2_STAGES must be 2, 3, or 4") + return stages + + +def _v2_blocks_per_cta() -> int: + blocks = int(os.environ.get("B12X_GLM_W4A8_V2_BLOCKS", "8")) + if blocks not in (2, 4, 8): + raise ValueError("B12X_GLM_W4A8_V2_BLOCKS must be 2, 4, or 8") + return blocks + + +def _validate_pool( + name: str, + pool: torch.Tensor, + *, + bits: int, + size_k: int, + size_n: int, + device: torch.device, +) -> int: + if pool.dtype != torch.int16 or pool.device != device or not pool.is_contiguous(): + raise TypeError(f"{name} must be contiguous int16 on {device}") + values_per_expert = int(size_k) * int(size_n) + encoded_i16_per_expert = values_per_expert * int(bits) // 16 + if encoded_i16_per_expert <= 0 or int(pool.numel()) % encoded_i16_per_expert: + raise ValueError(f"{name} does not contain complete native K{bits} tensors") + return int(pool.numel()) // encoded_i16_per_expert + + +def _validate_slot_map( + name: str, + slots: torch.Tensor, + *, + num_experts: int, + device: torch.device, +) -> None: + if ( + slots.dtype != torch.int32 + or slots.device != device + or not slots.is_contiguous() + or tuple(slots.shape) != (int(num_experts),) + ): + raise TypeError( + f"{name} must be contiguous int32 [{int(num_experts)}] on {device}" + ) + + +def run_glm_route_packed_w4a8_projection( + quantized: MXFP8Rows, + prepared: GLMRoutePackedW4A8Projection, + packed_route_indices: torch.Tensor, + block_expert_ids: torch.Tensor, + output: torch.Tensor, + *, + topk: int, + shared_input: bool, + clear_output: bool = True, +) -> torch.Tensor: + """Execute one mixed-K3/K4 GLM projection over packed expert routes.""" + + size_k = int(prepared.size_k) + size_n = int(prepared.size_n) + num_experts = int(prepared.num_experts) + topk = int(topk) + if topk <= 0: + raise ValueError("topk must be positive") + if str(prepared.trellis_codebook).lower() != _GLM_TRELLIS_CODEBOOK: + raise ValueError("GLM route-packed W4A8 requires SQG-XOR-Cheb-T12") + if output.ndim != 2 or output.dtype != torch.float16 or not output.is_cuda: + raise TypeError("output must be contiguous CUDA FP16 [routes, N]") + if not output.is_contiguous() or int(output.shape[1]) != size_n: + raise ValueError("output must be contiguous with the prepared N dimension") + device = output.device + routes = int(output.shape[0]) + if routes <= 0 or routes % topk: + raise ValueError("output routes must be positive and divisible by topk") + if size_k <= 0 or size_k % 32 or size_n <= 0 or size_n % 8: + raise ValueError("GLM W4A8 projection dimensions must close K32 and N8") + input_rows = routes // topk if shared_input else routes + _validate_quantized(quantized, m=input_rows, k=size_k, device=device) + + for name, value in ( + ("packed_route_indices", packed_route_indices), + ("block_expert_ids", block_expert_ids), + ): + if ( + value.dtype != torch.int32 + or value.device != device + or not value.is_contiguous() + ): + raise TypeError(f"{name} must be contiguous int32 on {device}") + route_blocks = min( + int(block_expert_ids.numel()), + int(packed_route_indices.numel()) // _GLM_ROUTE_BLOCK_ROWS, + ) + if route_blocks <= 0: + raise ValueError("packed route workspace contains no complete M64 block") + + pool_counts = { + 3: _validate_pool( + "trellis_k3", + prepared.trellis_k3, + bits=3, + size_k=size_k, + size_n=size_n, + device=device, + ), + 4: _validate_pool( + "trellis_k4", + prepared.trellis_k4, + bits=4, + size_k=size_k, + size_n=size_n, + device=device, + ), + } + _validate_slot_map( + "expert_slots_k3", + prepared.expert_slots_k3, + num_experts=num_experts, + device=device, + ) + _validate_slot_map( + "expert_slots_k4", + prepared.expert_slots_k4, + num_experts=num_experts, + device=device, + ) + if clear_output: + output.zero_() + + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + rank_lut = sqg_xor_cheb_t12_lut(device) + # Shapes without N256/K128 closure use the one-warp kernel, whose geometry + # accepts the remaining aligned dimensions. + use_tile_kernel = ( + _selected_w4a8_kernel() == "m128n64" + and size_n % _V2_TILE_N == 0 + and size_k % (_V2_KTILE_K32 * 32) == 0 + ) + if use_tile_kernel: + compiled = _compile_glm_route_packed_tile_projection( + size_k, + size_n, + topk, + bool(shared_input), + _v2_blocks_per_cta(), + _v2_stages(), + int(device_index), + ) + else: + compiled = _compile_glm_route_packed_mixed_projection( + size_k, + size_n, + topk, + bool(shared_input), + int(device_index), + ) + compiled( + _ptr(cutlass.Uint32, quantized.values.data_ptr()), + _ptr(cutlass.Uint8, quantized.scale_rows.data_ptr()), + _ptr(cutlass.Uint32, prepared.trellis_k3.data_ptr()), + _ptr(cutlass.Uint32, prepared.trellis_k4.data_ptr()), + _ptr(cutlass.Uint8, rank_lut.data_ptr()), + _ptr(cutlass.Int32, packed_route_indices.data_ptr()), + _ptr(cutlass.Int32, block_expert_ids.data_ptr()), + _ptr(cutlass.Int32, prepared.expert_slots_k3.data_ptr()), + _ptr(cutlass.Int32, prepared.expert_slots_k4.data_ptr()), + _ptr(cutlass.Float16, output.data_ptr()), + input_rows, + routes, + int(packed_route_indices.numel()), + route_blocks, + num_experts, + pool_counts[3], + pool_counts[4], + current_cuda_stream(), + ) + return output + + +__all__ = [ + "GLMRoutePackedW4A8Projection", + "prepare_glm_route_packed_w4a8_projection", + "run_glm_route_packed_w4a8_projection", +] diff --git a/b12x/moe/_shared/kernels/w4a16/host.py b/b12x/moe/_shared/kernels/w4a16/host.py index 831db31e0..c815c3f67 100644 --- a/b12x/moe/_shared/kernels/w4a16/host.py +++ b/b12x/moe/_shared/kernels/w4a16/host.py @@ -267,6 +267,32 @@ def packed_gemm_scratch_elements( return max(elements, 1) +def dense_trellis_gemm_scratch_elements_upper_bound( + *, + rows: int, + size_n: int, + sms: int, +) -> int: + """Bound dense-Trellis scratch for every supported row-block schedule. + + Dense launch policy may select a different row block for the same matrix + geometry at different row counts. Caller-owned CUDA-graph storage must + cover every supported selection rather than one observed warmup decision. + """ + + if int(rows) <= 0 or int(size_n) <= 0 or int(sms) <= 0: + raise ValueError("rows, size_n, and sms must be positive") + return max( + packed_gemm_scratch_elements( + size_n=int(size_n), + route_slots=((int(rows) + block_size - 1) // block_size) * block_size, + moe_block_size=block_size, + sms=int(sms), + ) + for block_size in _W4A16_ALLOWED_ROUTED_SIZES + ) + + def plan_w4a16_buffers( prepared, *, diff --git a/b12x/moe/_shared/kernels/w4a16/kernel.py b/b12x/moe/_shared/kernels/w4a16/kernel.py index 46938eb1a..ae6b44af8 100644 --- a/b12x/moe/_shared/kernels/w4a16/kernel.py +++ b/b12x/moe/_shared/kernels/w4a16/kernel.py @@ -186,8 +186,8 @@ def _sqg_xor_cheb_t12_smem_enabled() -> bool: def _validate_trellis256_codebook_bits(codebook: str, bits: int) -> None: - if codebook == "sqg_xor_cheb_t12" and bits not in (2, 3, 4): - raise ValueError("sqg_xor_cheb_t12 is defined only for K2/K3/K4") + if codebook == "sqg_xor_cheb_t12" and bits not in (2, 3, 4, 6): + raise ValueError("sqg_xor_cheb_t12 is defined for K2/K3/K4/K6") if codebook == SQG_FP16_D3L and bits not in (5, 6): raise ValueError("sqg_fp16_d3l is defined only for uniform K5/K6") @@ -4096,6 +4096,7 @@ def _scaled_dequant_b_fragment_trellis256_bits( trellis_lut_addr, int(bits), t12_in_shared=self.sqg_xor_cheb_t12_smem, + stream_layout=int(bits) == 6, ) if cutlass.const_expr(self.is_fp16): o0, o1 = fp8x4_e4m3_to_half2x2(e_lo) @@ -12016,7 +12017,7 @@ def _trellis_dense_buffer( return buffer -def _use_k6_mcg_small( +def _use_k6_small( *, device: torch.device, m: int, @@ -12027,13 +12028,13 @@ def _use_k6_mcg_small( external_hadamard_128, explicit_launch_config: bool, ) -> bool: - """Select the capture-safe K6/MCG kernel only on its compiled target.""" + """Select a capture-safe K6 small-M kernel on its compiled target.""" return ( not explicit_launch_config and tuple(torch.cuda.get_device_capability(device)) == (12, 0) and m <= 128 and trellis_bits == 6 - and trellis_codebook == "mcg" + and trellis_codebook in ("mcg", "sqg_xor_cheb_t12") and trellis_pair_kind is None and compute_dtype == torch.float16 and external_hadamard_128 is None @@ -12120,11 +12121,10 @@ def _run_trellis256_dense_current_device( None if hadamard_128 is None else _resolve_exl3_hadamard_128(hadamard_128) ) - # Keep the established K6/MCG decode path independent from the generic - # Trellis scheduler. It owns both H128 rotations, needs no GEMM scratch, - # and is safe to capture with only caller-owned output/rotation storage. - # Compact pair payloads and the newer SQG codebooks use the generic path. - use_k6_mcg_small = _use_k6_mcg_small( + # K6 small-M kernels own both H128 rotations and need no GEMM scratch. + # Each accepted codebook has a distinct decoder; compact pair payloads and + # larger row counts remain on the generic Trellis scheduler. + use_k6_small = _use_k6_small( device=x.device, m=m, trellis_bits=trellis_bits, @@ -12136,7 +12136,7 @@ def _run_trellis256_dense_current_device( _moe_block_size is not None or _force_tile_config is not None ), ) - if use_k6_mcg_small: + if use_k6_small: if x.dtype == torch.float16: x_f16 = x else: @@ -12167,14 +12167,15 @@ def _run_trellis256_dense_current_device( device=x.device, ) small_output = output_f16 - from b12x.gemm.trellis_linear._small_m import run_k6_mcg + from b12x.gemm.trellis_linear._small_m import run_k6_mcg, run_k6_sqg trellis_i16 = prepared_dense.trellis.view(torch.int16).view( size_k // 16, size_n // 16, trellis_bits * 16, ) - run_k6_mcg( + run_small = run_k6_mcg if trellis_codebook == "mcg" else run_k6_sqg + run_small( x_f16, trellis_i16, small_output, @@ -12187,6 +12188,13 @@ def _run_trellis256_dense_current_device( output.copy_(small_output) return output + if ( + _moe_block_size is None + and trellis_bits == 6 + and trellis_codebook == "sqg_xor_cheb_t12" + ): + _moe_block_size = 64 + gemm_output = _trellis_dense_buffer( "gemm_output", gemm_output, diff --git a/b12x/moe/_shared/kernels/w4a16/prepare.py b/b12x/moe/_shared/kernels/w4a16/prepare.py index 8bd20ecf8..3550bdd71 100644 --- a/b12x/moe/_shared/kernels/w4a16/prepare.py +++ b/b12x/moe/_shared/kernels/w4a16/prepare.py @@ -1191,9 +1191,7 @@ def prepare_w4a16_fc2_e8m0_weights( size_k=intermediate_size, size_n=hidden_size, ) - global_scale = torch.ones( - (num_experts,), dtype=torch.float32, device=w2_fp4.device - ) + global_scale = torch.ones((num_experts,), dtype=torch.float32, device=w2_fp4.device) return W4A16FC2Weights( w2=w2_fp4, w2_scale=packed_scale, @@ -1361,9 +1359,7 @@ def prepare_w4a16_x4t_weights( from b12x._lib.quant.x4t_scales import X4TScaleBatch - if not isinstance(w13_x4t, X4TScaleBatch) or not isinstance( - w2_x4t, X4TScaleBatch - ): + if not isinstance(w13_x4t, X4TScaleBatch) or not isinstance(w2_x4t, X4TScaleBatch): raise TypeError("X4T preparation requires X4TScaleBatch scale planes") w13_x4t.validate() w2_x4t.validate() @@ -1515,9 +1511,7 @@ def make_w4a16_packed_buffers( def _normalize_trellis256_codebook(codebook: str | int) -> str: if isinstance(codebook, int): - normalized = _TRELLIS256_CODEBOOK_SENTINELS.get( - int(codebook) & 0xFFFFFFFF - ) + normalized = _TRELLIS256_CODEBOOK_SENTINELS.get(int(codebook) & 0xFFFFFFFF) if normalized is None: raise ValueError( "unsupported trellis256 codebook sentinel " @@ -1535,8 +1529,8 @@ def _normalize_trellis256_codebook(codebook: str | int) -> str: def _validate_trellis256_codebook_bits(codebook: str, bits: int) -> None: - if codebook == "sqg_xor_cheb_t12" and bits not in (2, 3, 4): - raise ValueError("sqg_xor_cheb_t12 is defined only for K2/K3/K4") + if codebook == "sqg_xor_cheb_t12" and bits not in (2, 3, 4, 6): + raise ValueError("sqg_xor_cheb_t12 is defined for K2/K3/K4/K6") if codebook == SQG_FP16_D3L and bits not in (5, 6): raise ValueError("sqg_fp16_d3l is defined only for uniform K5/K6") @@ -1973,9 +1967,7 @@ def prepare_trellis256_moe_weights( "trellis3_t256 tile_config K dimensions must divide the model geometry" ) - _validate_trellis256_codebook_bits( - normalized_codebook, resolved_trellis_bits - ) + _validate_trellis256_codebook_bits(normalized_codebook, resolved_trellis_bits) if dummy_scale is None: dummy_scale = torch.zeros(4, dtype=torch.uint8, device=resolved_device) else: @@ -2090,6 +2082,8 @@ def prepare_trellis256_dense_weight( codebook: str | None = None, params_dtype: torch.dtype = torch.float16, dummy_scale: torch.Tensor | None = None, + global_scale: torch.Tensor | None = None, + workspace: torch.Tensor | None = None, ) -> PreparedTrellis256DenseWeight: """Prepare one native EXL3 linear for the dense trellis256 entry point. @@ -2174,17 +2168,45 @@ def prepare_trellis256_dense_weight( "trellis3_t256 dense dummy_scale must be a contiguous, 16-byte-" "aligned four-byte uint8 tensor on the weight device" ) + if global_scale is None: + global_scale = torch.ones(1, dtype=torch.float32, device=device) + elif ( + global_scale.device != device + or global_scale.dtype != torch.float32 + or tuple(global_scale.shape) != (1,) + or not global_scale.is_contiguous() + ): + raise ValueError( + "trellis3_t256 dense global_scale must be contiguous float32[1] " + "on the weight device" + ) + workspace_elements = max( + int(torch.cuda.get_device_properties(device).multi_processor_count) * 4 + 2, + out_features // 16, + ) + if workspace is None: + workspace = _make_workspace( + device, + max_blocks_per_sm=4, + min_elements=out_features // 16, + ) + elif ( + workspace.device != device + or workspace.dtype != torch.int32 + or not workspace.is_contiguous() + or workspace.numel() < workspace_elements + ): + raise ValueError( + "trellis3_t256 dense workspace must be contiguous int32 storage " + f"with at least {workspace_elements} elements on the weight device" + ) return PreparedTrellis256DenseWeight( trellis=packed, suh=suh, svh=svh, scale=dummy_scale, - global_scale=torch.ones(1, dtype=torch.float32, device=device), - workspace=_make_workspace( - device, - max_blocks_per_sm=4, - min_elements=out_features // 16, - ), + global_scale=global_scale, + workspace=workspace, in_features=in_features, out_features=out_features, params_dtype=params_dtype, @@ -2224,21 +2246,31 @@ def prepare_trellis256_pair_dense_weight( raise ValueError(f"trellis pair_kind must be P24 or P33, got {pair_kind!r}") rate_axis = str(rate_axis).lower() if rate_axis not in {"k", "n"}: - raise ValueError(f"trellis pair rate_axis must be 'k' or 'n', got {rate_axis!r}") + raise ValueError( + f"trellis pair rate_axis must be 'k' or 'n', got {rate_axis!r}" + ) if params_dtype not in (torch.float16, torch.bfloat16): raise ValueError("trellis3_t256 pair compute requires fp16 or bf16 MMA inputs") if payload.dtype != torch.int16: - raise TypeError(f"trellis pair payload must use torch.int16, got {payload.dtype}") + raise TypeError( + f"trellis pair payload must use torch.int16, got {payload.dtype}" + ) if payload.ndim != 1 or not payload.is_contiguous(): - raise ValueError("trellis pair payload must be a contiguous one-dimensional tensor") + raise ValueError( + "trellis pair payload must be a contiguous one-dimensional tensor" + ) if payload.device.type != "cuda": - raise ValueError(f"trellis pair payload requires CUDA storage, got {payload.device}") + raise ValueError( + f"trellis pair payload requires CUDA storage, got {payload.device}" + ) device = payload.device for name, scale in (("suh", suh), ("svh", svh)): if scale.device != device: raise ValueError(f"trellis pair {name} must be on {device}") if scale.dtype != torch.float16: - raise TypeError(f"trellis pair {name} must be torch.float16, got {scale.dtype}") + raise TypeError( + f"trellis pair {name} must be torch.float16, got {scale.dtype}" + ) if scale.ndim != 1 or not scale.is_contiguous(): raise ValueError(f"trellis pair {name} must be a contiguous vector") if not bool(torch.all(torch.isfinite(scale))): @@ -2280,9 +2312,7 @@ def prepare_trellis256_pair_dense_weight( # time, so retain the same bytes while interleaving the two complete # record spans at K16 granularity. low = payload[:low_words].reshape(orthogonal_tiles, 8 * 16 * low_bits) - high = payload[low_words:].reshape( - orthogonal_tiles, 8 * 16 * high_bits - ) + high = payload[low_words:].reshape(orthogonal_tiles, 8 * 16 * high_bits) prepared_i16 = torch.cat((low, high), dim=1).contiguous().reshape(-1) else: prepared_i16 = payload @@ -2305,9 +2335,7 @@ def prepare_trellis256_pair_dense_weight( else ( None if mul1_e4m3 is None - else torch.tensor( - mul1_e4m3, dtype=torch.uint32, device=device - ) + else torch.tensor(mul1_e4m3, dtype=torch.uint32, device=device) ) ), codebook=codebook, @@ -2339,9 +2367,7 @@ def prepare_trellis256_pair_dense_weight( trellis_bits=3, trellis_codebook=normalized_codebook, mcg=mcg if isinstance(mcg, torch.Tensor) else None, - mul1_e4m3=( - mul1_e4m3 if isinstance(mul1_e4m3, torch.Tensor) else None - ), + mul1_e4m3=(mul1_e4m3 if isinstance(mul1_e4m3, torch.Tensor) else None), trellis_pair_kind=pair_kind, trellis_rate_axis=rate_axis, ) @@ -2451,9 +2477,7 @@ def _modes(name: str, value: torch.Tensor) -> torch.Tensor: high = selected[..., low_words:].reshape( 2, ids.numel(), hidden_tiles, 8 * 16 * high_bits ) - swizzled = torch.cat((low, high), dim=-1).reshape( - 2, ids.numel(), pair_words - ) + swizzled = torch.cat((low, high), dim=-1).reshape(2, ids.numel(), pair_words) prepared_w13.index_copy_(1, ids, swizzled) for name, scale, shapes in ( @@ -2668,9 +2692,7 @@ def _prepare_qsrt_p33_p43_moe_weights( f"{total_u32} uint32 values" ) if w2_payload.numel() != 2 * total_u32: - raise ValueError( - f"w2_payload must contain {total_u32} compact uint32 values" - ) + raise ValueError(f"w2_payload must contain {total_u32} compact uint32 values") descriptors = ((pair_offsets_u32 << 1) | modes).contiguous() for name, scale, shapes in ( @@ -2782,18 +2804,10 @@ def _qsrt_coupled_rotation_signs( return torch.ones(length, dtype=torch.float32) generator = torch.Generator(device="cpu") generator.manual_seed( - ( - 0x6A09E667F3BCC909 * int(draw) - + 0xBB67AE8584CAA73B * int(axis) - ) + (0x6A09E667F3BCC909 * int(draw) + 0xBB67AE8584CAA73B * int(axis)) & ((1 << 63) - 1) ) - return ( - torch.randint(0, 2, (length,), generator=generator) - .mul_(2) - .sub_(1) - .float() - ) + return torch.randint(0, 2, (length,), generator=generator).mul_(2).sub_(1).float() def _prepare_qsrt_p22_atom_v2_moe_weights( @@ -2830,10 +2844,7 @@ def _prepare_qsrt_p22_atom_v2_moe_weights( "preactivation halves" ) source_intermediate_size = atom_count * 32 - if ( - intermediate_size < source_intermediate_size - or intermediate_size % 128 - ): + if intermediate_size < source_intermediate_size or intermediate_size % 128: raise ValueError( "pure-K2 runtime intermediate size must contain its atom extent " "and close 128-channel postactivation blocks" @@ -2894,9 +2905,8 @@ def _prepare_qsrt_p22_atom_v2_moe_weights( .reshape(atom_count, count, 2, hidden_tiles, 32) ) if matrix_index < 2: - restored = ( - values.permute(1, 3, 0, 2, 4) - .reshape(count, hidden_tiles, source_local_tiles, 32) + restored = values.permute(1, 3, 0, 2, 4).reshape( + count, hidden_tiles, source_local_tiles, 32 ) w13[ matrix_index, @@ -2905,9 +2915,8 @@ def _prepare_qsrt_p22_atom_v2_moe_weights( :source_local_tiles, ].copy_(restored) else: - restored = ( - values.permute(1, 0, 2, 3, 4) - .reshape(count, source_local_tiles, hidden_tiles, 32) + restored = values.permute(1, 0, 2, 3, 4).reshape( + count, source_local_tiles, hidden_tiles, 32 ) w2[ first_expert : first_expert + count, @@ -2929,8 +2938,7 @@ def _prepare_qsrt_p22_atom_v2_moe_weights( ) intermediate_rotations[ first_expert : first_expert + count, - matrix_index * intermediate_size : matrix_index - * intermediate_size + matrix_index * intermediate_size : matrix_index * intermediate_size + source_intermediate_size, ].copy_(scales) @@ -2940,9 +2948,7 @@ def _prepare_qsrt_p22_atom_v2_moe_weights( pre_begin = 2 * first_atom_slot * 32 pre_count = 2 * source_intermediate_size post_begin = first_atom_slot * 32 - signs = torch.ones( - (num_experts, 3 * intermediate_size), dtype=torch.float16 - ) + signs = torch.ones((num_experts, 3 * intermediate_size), dtype=torch.float16) for draw in sorted(set(int(value) for value in rotation_draws.tolist())): rows = torch.nonzero(rotation_draws == draw, as_tuple=False).flatten() pre = _qsrt_coupled_rotation_signs(2 * 3072, draw=draw, axis=1)[ @@ -2954,13 +2960,10 @@ def _prepare_qsrt_p22_atom_v2_moe_weights( draw_signs = torch.ones( (rows.numel(), 3 * intermediate_size), dtype=torch.float16 ) - draw_signs[:, :pre_count].copy_( - pre.to(torch.float16).expand(rows.numel(), -1) - ) + draw_signs[:, :pre_count].copy_(pre.to(torch.float16).expand(rows.numel(), -1)) draw_signs[ :, - 2 * intermediate_size : 2 * intermediate_size - + source_intermediate_size, + 2 * intermediate_size : 2 * intermediate_size + source_intermediate_size, ].copy_(post.to(torch.float16).expand(rows.numel(), -1)) signs.index_copy_(0, rows, draw_signs) coupled_signs = signs.to(device=device, non_blocking=True).contiguous() @@ -3204,8 +3207,7 @@ def restore_matrix( ) intermediate_scales[ first_expert : first_expert + count, - matrix_index - * intermediate_size : (matrix_index + 1) + matrix_index * intermediate_size : (matrix_index + 1) * intermediate_size, ].copy_(restored) @@ -3430,7 +3432,7 @@ def _restore_group_matrix_into( p43: bool, fc1: bool, ) -> None: - low_bits, high_bits = ((4, 3) if p43 else (3, 3)) + low_bits, high_bits = (4, 3) if p43 else (3, 3) matrix_bytes = ( _QSRT_V2_P43_MATRIX_TRELLIS_BYTES if p43 @@ -3486,12 +3488,8 @@ def _restore_group_matrix_into( ).copy_(high.permute(0, 2, 1, 3)) else: low_words_per_expert = low.numel() // chunk_count - target[:, :low_words_per_expert].copy_( - low.reshape(chunk_count, -1) - ) - target[:, low_words_per_expert:].copy_( - high.reshape(chunk_count, -1) - ) + target[:, :low_words_per_expert].copy_(low.reshape(chunk_count, -1)) + target[:, low_words_per_expert:].copy_(high.reshape(chunk_count, -1)) group_matrix_words = [] for ids, _source, p43, _bundle in groups: @@ -3510,7 +3508,7 @@ def _restore_group_matrix_into( for matrix_index in range(2): group_offset = matrix_index * matrix_words for group_words, (_ids, source, p43, _bundle) in zip( - group_matrix_words, groups + group_matrix_words, groups, strict=True ): _restore_group_matrix_into( source, @@ -3522,7 +3520,7 @@ def _restore_group_matrix_into( group_offset += group_words group_offset = 0 for group_words, (_ids, source, p43, _bundle) in zip( - group_matrix_words, groups + group_matrix_words, groups, strict=True ): _restore_group_matrix_into( source, @@ -3666,9 +3664,10 @@ def prepare_qsrt_atom_moe_weights( f"{tuple(atom_payload.shape)}" ) expected_inner_strides = (_QSRT_ATOM_BUNDLE_BYTES, 1) - if tuple(atom_payload.stride()[1:]) != expected_inner_strides or int( - atom_payload.stride(0) - ) < num_experts * _QSRT_ATOM_BUNDLE_BYTES: + if ( + tuple(atom_payload.stride()[1:]) != expected_inner_strides + or int(atom_payload.stride(0)) < num_experts * _QSRT_ATOM_BUNDLE_BYTES + ): raise ValueError( "QSRT atom payloads must be expert-major within each atom row; " "the row stride may include checkpoint alignment padding" @@ -3677,9 +3676,7 @@ def prepare_qsrt_atom_moe_weights( first_atom_slot = int(first_atom_slot) layer_index = int(layer_index) if not 0 <= first_atom_slot < _QSRT_ATOMS_PER_EXPERT: - raise ValueError( - f"first_atom_slot must be in 0..{_QSRT_ATOMS_PER_EXPERT - 1}" - ) + raise ValueError(f"first_atom_slot must be in 0..{_QSRT_ATOMS_PER_EXPERT - 1}") if first_atom_slot % _QSRT_ATOMS_PER_PAIR: raise ValueError("the current QSRT kernel requires a pair-aligned atom extent") if not 1 <= layer_index <= 92: @@ -3689,9 +3686,7 @@ def _normalize_vector(name: str, value: torch.Tensor) -> torch.Tensor: if not isinstance(value, torch.Tensor): raise TypeError(f"{name} must be a tensor") if value.device != device or tuple(value.shape) != (num_experts,): - raise ValueError( - f"{name} must have shape {(num_experts,)} on {device}" - ) + raise ValueError(f"{name} must have shape {(num_experts,)} on {device}") if value.dtype not in { torch.uint8, torch.int8, @@ -3711,9 +3706,7 @@ def _normalize_vector(name: str, value: torch.Tensor) -> torch.Tensor: if not bool(torch.all((r13 >= 0) & (r13 <= 2) & (r2 >= 0) & (r2 <= 2))): raise ValueError("compressed QSRT format codes must encode R0/R1/R2") physical_pair = first_atom_slot // _QSRT_ATOMS_PER_PAIR - rotation = ( - _QSRT_EXPERT_ROTATION_MULTIPLIER * expert_ids_i32 + layer_index - ) % 12 + rotation = (_QSRT_EXPERT_ROTATION_MULTIPLIER * expert_ids_i32 + layer_index) % 12 logical_pair = (physical_pair - rotation) % 12 fc1_pair_modes = (logical_pair < r13).to(dtype=torch.int32).contiguous() fc2_pair_modes = (logical_pair < r2).to(dtype=torch.int32).contiguous() @@ -3729,9 +3722,11 @@ def _matrix_words(matrix_index: int) -> torch.Tensor: raw = atom_payload.narrow( 2, begin, _QSRT_MATRIX_ATOM_TRELLIS_BYTES ).contiguous() - return raw.view(torch.int16).reshape( - _QSRT_ATOMS_PER_PAIR, num_experts, words_per_atom - ).permute(1, 0, 2) + return ( + raw.view(torch.int16) + .reshape(_QSRT_ATOMS_PER_PAIR, num_experts, words_per_atom) + .permute(1, 0, 2) + ) def _restore_matrix( matrix_index: int, modes: torch.Tensor, *, fc1: bool @@ -3744,9 +3739,7 @@ def _restore_matrix( ids = torch.nonzero(modes == mode, as_tuple=False).flatten() if int(ids.numel()) == 0: continue - selected = source.index_select(0, ids).narrow( - 2, 0, hidden_tiles * 16 * 6 - ) + selected = source.index_select(0, ids).narrow(2, 0, hidden_tiles * 16 * 6) low_words = hidden_tiles * 16 * low_bits low = selected[..., :low_words].reshape( -1, _QSRT_ATOMS_PER_PAIR, hidden_tiles, 16 * low_bits @@ -3755,16 +3748,10 @@ def _restore_matrix( -1, _QSRT_ATOMS_PER_PAIR, hidden_tiles, 16 * high_bits ) if fc1: - low = low.permute(0, 2, 1, 3).reshape( - ids.numel(), hidden_tiles, -1 - ) - high = high.permute(0, 2, 1, 3).reshape( - ids.numel(), hidden_tiles, -1 - ) + low = low.permute(0, 2, 1, 3).reshape(ids.numel(), hidden_tiles, -1) + high = high.permute(0, 2, 1, 3).reshape(ids.numel(), hidden_tiles, -1) # FC1 places both 128-channel records under each K16 tile. - restored = torch.cat((low, high), dim=-1).reshape( - ids.numel(), -1 - ) + restored = torch.cat((low, high), dim=-1).reshape(ids.numel(), -1) else: # FC2 retains its K-major low-plane/high-plane ordering. restored = torch.cat( @@ -3783,18 +3770,16 @@ def _restore_matrix( _restore_matrix(1, fc1_pair_modes, fc1=True), ) ).reshape(-1) - prepared_w2_i16 = _restore_matrix( - 2, fc2_pair_modes, fc1=False - ).reshape(-1) + prepared_w2_i16 = _restore_matrix(2, fc2_pair_modes, fc1=False).reshape(-1) def _local_scale(matrix_index: int) -> torch.Tensor: begin = _QSRT_MATRIX_SCALE_OFFSETS[matrix_index] - raw = atom_payload.narrow( - 2, begin, _QSRT_MATRIX_ATOM_SCALE_BYTES - ).contiguous() - values = raw.view(torch.float16).reshape( - _QSRT_ATOMS_PER_PAIR, num_experts, _QSRT_ATOM_CHANNELS - ).permute(1, 0, 2) + raw = atom_payload.narrow(2, begin, _QSRT_MATRIX_ATOM_SCALE_BYTES).contiguous() + values = ( + raw.view(torch.float16) + .reshape(_QSRT_ATOMS_PER_PAIR, num_experts, _QSRT_ATOM_CHANNELS) + .permute(1, 0, 2) + ) return torch.cat( ( values[..., :16].reshape(num_experts, -1), diff --git a/b12x/moe/glm_sqg_w4a8.py b/b12x/moe/glm_sqg_w4a8.py new file mode 100644 index 000000000..13001ac10 --- /dev/null +++ b/b12x/moe/glm_sqg_w4a8.py @@ -0,0 +1,747 @@ +"""Topology-neutral GLM SQG atoms-v2 W4A8 MoE execution. + +GLM assigns K3/K4 independently to gate, up, and down for every expert. The +prepared layer therefore owns six compact trellis pools (three projections by +two rates) and three independent expert-to-pool partitions. No TP rank is +encoded in the checkpoint contract: vLLM slices the intermediate axis before +calling :func:`prepare_weights`, and this module records that local extent. + +The runtime owns every mutable route, MXFP8, transform, and output buffer. A +prepared runtime can consequently be reused by eager execution and CUDA graph +capture without allocating in :func:`run`. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +import math +import os +import weakref + +import torch + +from b12x._lib.quant.mxfp8_rows import quantize_mxfp8_rows_cute +from b12x.gemm._shared.wo_mxfp8 import ( + MXFP8Rows, + empty_mxfp8_rows_for_dense_gemm, +) +from b12x.moe._shared.kernels.glm_trellis_transform import ( + run_glm_down_input_transform, + run_glm_down_output_transform_sum, + run_glm_gate_up_output_transform_silu, +) +from b12x.moe._shared.kernels.glm_trellis_w4a8 import ( + GLMRoutePackedW4A8Projection, + prepare_glm_route_packed_w4a8_projection, + run_glm_route_packed_w4a8_projection, +) +from b12x.moe._shared.kernels.w4a16.host import route_pack_capacity +from b12x.moe._shared.kernels.w4a16.kernel import ( + _run_trellis_dense_hadamard128, + pack_topk_routes_by_expert, +) + + +_ROUTE_BLOCK_ROWS = 128 +_CODEBOOK = "sqg_xor_cheb_t12" +_WEIGHT_REGISTRY: weakref.WeakValueDictionary[int, GLMSQGW4A8Weights] +_RUNTIME_REGISTRY: weakref.WeakValueDictionary[int, GLMSQGW4A8Runtime] + + +def glm_route_packed_w4a8_kernel_contract() -> dict[str, int | str]: + """Return the resolved route-packed schedule for serving provenance.""" + + return { + "kernel": os.environ.get("B12X_GLM_W4A8_KERNEL", "m128n64").strip().lower(), + "blocks_per_cta": int(os.environ.get("B12X_GLM_W4A8_V2_BLOCKS", "8")), + "stages": int(os.environ.get("B12X_GLM_W4A8_V2_STAGES", "2")), + "route_block_rows": _ROUTE_BLOCK_ROWS, + "tile_n": 64, + } + + +# Qualified SM120/SM121 schedule for the GLM-5.2 SQG atoms-v2 geometry. +_SM12X_M128N64_SCHEDULE = { + "kernel": "m128n64", + "blocks_per_cta": 8, + "stages": 2, + "route_block_rows": 128, + "tile_n": 64, +} + +_ACCEPTANCE_SCHEDULE_BY_ARCH = { + (12, 0): dict(_SM12X_M128N64_SCHEDULE), + (12, 1): dict(_SM12X_M128N64_SCHEDULE), +} + + +def _acceptance_arch() -> tuple[int, int]: + """Resolve the architecture key for the acceptance schedule.""" + forced = os.environ.get("B12X_GLM_W4A8_ACCEPT_ARCH", "").strip() + if forced: + text = forced.lower().removeprefix("sm_").removeprefix("sm") + if not text.isdigit() or len(text) < 2: + raise ValueError( + "B12X_GLM_W4A8_ACCEPT_ARCH must look like 'sm_120' or '120', " + f"got {forced!r}" + ) + return (int(text[:-1]), int(text[-1])) + import torch + + if not torch.cuda.is_available(): + raise RuntimeError( + "GLM SQG full-W4A8 acceptance needs a CUDA device to resolve the " + "architecture; set B12X_GLM_W4A8_ACCEPT_ARCH for offline checks" + ) + return tuple(torch.cuda.get_device_capability(torch.cuda.current_device())) + + +def validate_glm_route_packed_w4a8_acceptance_kernel() -> None: + """Require the qualified schedule on a supported SM120/SM121 device.""" + + arch = _acceptance_arch() + expected = _ACCEPTANCE_SCHEDULE_BY_ARCH.get(arch) + if expected is None: + raise RuntimeError( + "GLM SQG full-W4A8 supports SM120/SM121; device reports " + f"SM{arch[0]}{arch[1]}" + ) + resolved = glm_route_packed_w4a8_kernel_contract() + if resolved != expected: + raise RuntimeError( + f"GLM SQG full-W4A8 requires the qualified " + f"sm_{arch[0]}{arch[1]} schedule {expected}, got {resolved}" + ) + + +@dataclass(frozen=True) +class GLMSQGW4A8Weights: + """One TP-local view of a topology-neutral GLM atoms-v2 layer.""" + + gate: GLMRoutePackedW4A8Projection + up: GLMRoutePackedW4A8Projection + down: GLMRoutePackedW4A8Projection + gate_up_suh: torch.Tensor + gate_svh: torch.Tensor + up_svh: torch.Tensor + down_suh: torch.Tensor + down_svh: torch.Tensor + hidden_size: int + intermediate_size: int + global_intermediate_size: int + num_experts: int + tp_rank: int + tp_size: int + derived_down_target_id: str | None + down_target_beta: float | None + codebook: str = _CODEBOOK + direct_e4m3_weights: bool = True + allow_a16_fallback: bool = False + + +@dataclass(frozen=True) +class GLMSQGW4A8Runtime: + """Fixed-capacity caller-owned storage for one GLM W4A8 execution scope.""" + + max_tokens: int + topk: int + hidden_size: int + intermediate_size: int + num_experts: int + input_f16: torch.Tensor + hidden_rotated: torch.Tensor + hidden_quantized: MXFP8Rows + gate_transformed: torch.Tensor + up_transformed: torch.Tensor + gate_hadamard: torch.Tensor + up_hadamard: torch.Tensor + activated: torch.Tensor + down_scaled: torch.Tensor + down_rotated: torch.Tensor + down_quantized: MXFP8Rows + down_transformed: torch.Tensor + down_canonical: torch.Tensor + output: torch.Tensor + packed_route_indices: torch.Tensor + block_expert_ids: torch.Tensor + packed_route_count: torch.Tensor + expert_offsets: torch.Tensor + expert_counts: torch.Tensor + ones_intermediate: torch.Tensor + + +def _validate_scale( + name: str, + value: torch.Tensor, + *, + shape: tuple[int, ...], + device: torch.device, +) -> None: + if ( + value.dtype != torch.float16 + or value.device != device + or tuple(value.shape) != shape + or not value.is_contiguous() + ): + raise ValueError( + f"{name} must be contiguous FP16 {shape} on {device}, got " + f"{tuple(value.shape)}/{value.dtype}/{value.device}" + ) + if not bool(torch.all(torch.isfinite(value))): + raise ValueError(f"{name} contains non-finite values") + + +def _validate_projection_partition( + name: str, + projection: GLMRoutePackedW4A8Projection, +) -> None: + slots3 = projection.expert_slots_k3.detach().cpu() + slots4 = projection.expert_slots_k4.detach().cpu() + owns3 = slots3 >= 0 + owns4 = slots4 >= 0 + if not bool(torch.all(owns3 ^ owns4)): + raise ValueError(f"{name} K3/K4 slot maps do not partition all experts") + for bits, slots, owns, pool in ( + (3, slots3, owns3, projection.trellis_k3), + (4, slots4, owns4, projection.trellis_k4), + ): + live = slots[owns] + expected = torch.arange(live.numel(), dtype=torch.int32) + if not torch.equal(torch.sort(live).values, expected): + raise ValueError(f"{name} K{bits} slots are not dense and unique") + expected_numel = ( + live.numel() * int(projection.size_k) * int(projection.size_n) * bits // 16 + ) + if int(pool.numel()) != expected_numel: + raise ValueError( + f"{name} K{bits} pool has {pool.numel()} int16 values; " + f"expected {expected_numel}" + ) + + +def prepare_weights( + *, + gate_trellis: Sequence[torch.Tensor], + gate_bits: Sequence[int], + up_trellis: Sequence[torch.Tensor], + up_bits: Sequence[int], + down_trellis: Sequence[torch.Tensor], + down_bits: Sequence[int], + gate_up_suh: torch.Tensor, + gate_svh: torch.Tensor, + up_svh: torch.Tensor, + down_suh: torch.Tensor, + down_svh: torch.Tensor, + hidden_size: int, + intermediate_size: int, + global_intermediate_size: int | None = None, + tp_rank: int = 0, + tp_size: int = 1, + derived_down_target_id: str | None = None, + down_target_beta: float | None = None, +) -> GLMSQGW4A8Weights: + """Prepare six independent K3/K4 pools from one TP-local atoms-v2 extent. + + ``intermediate_size`` is the local TP width. Rate choices remain per + logical tensor and are never coupled across gate/up/down or rewritten by + topology. + """ + + hidden_size = int(hidden_size) + intermediate_size = int(intermediate_size) + tp_rank = int(tp_rank) + tp_size = int(tp_size) + if global_intermediate_size is None: + global_intermediate_size = intermediate_size * tp_size + global_intermediate_size = int(global_intermediate_size) + if hidden_size <= 0 or hidden_size % 128: + raise ValueError("GLM SQG W4A8 hidden_size must be a multiple of 128") + if intermediate_size <= 0 or intermediate_size % 128: + raise ValueError( + "GLM SQG W4A8 local intermediate_size must be a multiple of 128" + ) + if tp_size <= 0 or not 0 <= tp_rank < tp_size: + raise ValueError("tp_rank must identify one rank in tp_size") + if global_intermediate_size != intermediate_size * tp_size: + raise ValueError( + "topology-neutral intermediate extent does not close exactly: " + f"global={global_intermediate_size}, local={intermediate_size}, " + f"tp={tp_size}" + ) + if (derived_down_target_id is None) != (down_target_beta is None): + raise ValueError( + "derived_down_target_id and down_target_beta must be supplied together" + ) + if down_target_beta is not None and ( + isinstance(down_target_beta, bool) or not math.isfinite(float(down_target_beta)) + ): + raise ValueError("down_target_beta must be a finite real value") + num_experts = len(gate_trellis) + if num_experts <= 0 or any( + len(values) != num_experts + for values in ( + gate_bits, + up_trellis, + up_bits, + down_trellis, + down_bits, + ) + ): + raise ValueError("every GLM projection must describe every expert once") + device = gate_up_suh.device + + gate = prepare_glm_route_packed_w4a8_projection( + gate_trellis, + gate_bits, + size_k=hidden_size, + size_n=intermediate_size, + ) + up = prepare_glm_route_packed_w4a8_projection( + up_trellis, + up_bits, + size_k=hidden_size, + size_n=intermediate_size, + ) + down = prepare_glm_route_packed_w4a8_projection( + down_trellis, + down_bits, + size_k=intermediate_size, + size_n=hidden_size, + ) + for name, projection in (("gate", gate), ("up", up), ("down", down)): + if ( + projection.trellis_k3.device != device + or projection.trellis_k4.device != device + ): + raise ValueError(f"{name} trellis pools must be on {device}") + _validate_projection_partition(name, projection) + + _validate_scale("gate_up_suh", gate_up_suh, shape=(hidden_size,), device=device) + _validate_scale( + "gate_svh", + gate_svh, + shape=(num_experts, intermediate_size), + device=device, + ) + _validate_scale( + "up_svh", + up_svh, + shape=(num_experts, intermediate_size), + device=device, + ) + _validate_scale( + "down_suh", + down_suh, + shape=(num_experts, intermediate_size), + device=device, + ) + _validate_scale("down_svh", down_svh, shape=(hidden_size,), device=device) + + prepared = GLMSQGW4A8Weights( + gate=gate, + up=up, + down=down, + gate_up_suh=gate_up_suh, + gate_svh=gate_svh, + up_svh=up_svh, + down_suh=down_suh, + down_svh=down_svh, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + global_intermediate_size=global_intermediate_size, + num_experts=num_experts, + tp_rank=tp_rank, + tp_size=tp_size, + derived_down_target_id=derived_down_target_id, + down_target_beta=( + None if down_target_beta is None else float(down_target_beta) + ), + ) + _WEIGHT_REGISTRY[id(prepared)] = prepared + return prepared + + +def prepare_runtime( + weights: GLMSQGW4A8Weights, + *, + max_tokens: int, + topk: int, + output_dtype: torch.dtype, +) -> GLMSQGW4A8Runtime: + """Allocate fixed graph-stable storage for one serving execution scope.""" + + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("GLM SQG W4A8 runtime must be prepared before capture") + validate_glm_route_packed_w4a8_acceptance_kernel() + max_tokens = int(max_tokens) + topk = int(topk) + if max_tokens <= 0 or topk <= 0: + raise ValueError("max_tokens and topk must be positive") + if output_dtype not in (torch.float16, torch.bfloat16): + raise ValueError("GLM SQG W4A8 output must be FP16 or BF16") + device = weights.gate_up_suh.device + if device.type != "cuda": + raise ValueError("GLM SQG W4A8 execution requires CUDA weights") + capability = torch.cuda.get_device_capability(device) + if capability not in ((12, 0), (12, 1)): + raise ValueError( + "GLM SQG direct-E4M3 W4A8 requires SM120/SM121, got " + f"SM{capability[0]}{capability[1]}" + ) + hidden = int(weights.hidden_size) + intermediate = int(weights.intermediate_size) + max_routes = max_tokens * topk + _, packed_capacity, block_capacity = route_pack_capacity( + max_routes, + _ROUTE_BLOCK_ROWS, + weights.num_experts, + topk=topk, + ) + + def f16(rows: int, cols: int) -> torch.Tensor: + return torch.empty((rows, cols), dtype=torch.float16, device=device) + + runtime = GLMSQGW4A8Runtime( + max_tokens=max_tokens, + topk=topk, + hidden_size=hidden, + intermediate_size=intermediate, + num_experts=weights.num_experts, + input_f16=f16(max_tokens, hidden), + hidden_rotated=f16(max_tokens, hidden), + hidden_quantized=empty_mxfp8_rows_for_dense_gemm( + max_tokens, hidden, device=device + ), + gate_transformed=f16(max_routes, intermediate), + up_transformed=f16(max_routes, intermediate), + gate_hadamard=f16(max_routes, intermediate), + up_hadamard=f16(max_routes, intermediate), + activated=f16(max_routes, intermediate), + down_scaled=f16(max_routes, intermediate), + down_rotated=f16(max_routes, intermediate), + down_quantized=empty_mxfp8_rows_for_dense_gemm( + max_routes, intermediate, device=device + ), + down_transformed=f16(max_routes, hidden), + down_canonical=f16(max_routes, hidden), + output=torch.empty((max_tokens, hidden), dtype=output_dtype, device=device), + packed_route_indices=torch.empty( + packed_capacity, dtype=torch.int32, device=device + ), + block_expert_ids=torch.empty(block_capacity, dtype=torch.int32, device=device), + packed_route_count=torch.empty(1, dtype=torch.int32, device=device), + expert_offsets=torch.empty( + weights.num_experts + 1, dtype=torch.int32, device=device + ), + expert_counts=torch.empty( + weights.num_experts, dtype=torch.int32, device=device + ), + ones_intermediate=torch.ones(intermediate, dtype=torch.float16, device=device), + ) + _RUNTIME_REGISTRY[id(runtime)] = runtime + return runtime + + +def _active_quantized(value: MXFP8Rows, rows: int) -> MXFP8Rows: + scale_rows = ( + value.scale_rows[:, :rows] + if value.scale_rows.ndim == 3 + else value.scale_rows[:rows] + ) + return MXFP8Rows( + values=value.values[:rows], + scale_rows=scale_rows, + scale_mma=value.scale_mma, + ) + + +def _run_impl( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + weights: GLMSQGW4A8Weights, + runtime: GLMSQGW4A8Runtime, +) -> torch.Tensor: + """Run the complete route-packed GLM SQG W4A8 expert layer.""" + + if hidden_states.ndim != 2 or not hidden_states.is_contiguous(): + raise ValueError("hidden_states must be contiguous [tokens, hidden]") + tokens, hidden = (int(value) for value in hidden_states.shape) + if tokens <= 0 or tokens > runtime.max_tokens: + raise ValueError( + f"token count {tokens} exceeds planned capacity {runtime.max_tokens}" + ) + if ( + runtime.hidden_size != weights.hidden_size + or runtime.intermediate_size != weights.intermediate_size + or runtime.num_experts != weights.num_experts + ): + raise ValueError("prepared weights do not match the runtime geometry") + if hidden != weights.hidden_size: + raise ValueError(f"hidden width {hidden} != {weights.hidden_size}") + if hidden_states.dtype not in (torch.float16, torch.bfloat16): + raise TypeError("hidden_states must be FP16 or BF16") + if hidden_states.device != weights.gate_up_suh.device: + raise ValueError("hidden_states and prepared weights must share a device") + expected_route_shape = (tokens, runtime.topk) + if ( + tuple(topk_ids.shape) != expected_route_shape + or topk_ids.dtype != torch.int32 + or topk_ids.device != hidden_states.device + or not topk_ids.is_contiguous() + ): + raise TypeError(f"topk_ids must be contiguous int32 {expected_route_shape}") + if ( + tuple(topk_weights.shape) != expected_route_shape + or topk_weights.dtype != torch.float32 + or topk_weights.device != hidden_states.device + or not topk_weights.is_contiguous() + ): + raise TypeError(f"topk_weights must be contiguous FP32 {expected_route_shape}") + + routes = tokens * runtime.topk + packed, block_experts, _ = pack_topk_routes_by_expert( + topk_ids, + _ROUTE_BLOCK_ROWS, + weights.num_experts, + packed_route_indices=runtime.packed_route_indices, + block_expert_ids=runtime.block_expert_ids, + packed_route_count=runtime.packed_route_count, + expert_offsets=runtime.expert_offsets, + expert_counts=runtime.expert_counts, + ) + input_f16 = runtime.input_f16[:tokens] + input_f16.copy_(hidden_states) + hidden_rotated = runtime.hidden_rotated[:tokens] + _run_trellis_dense_hadamard128( + input_f16, + hidden_rotated, + weights.gate_up_suh, + scale_before=True, + ) + hidden_quantized = _active_quantized(runtime.hidden_quantized, tokens) + quantize_mxfp8_rows_cute( + hidden_rotated, + hidden_quantized.values, + hidden_quantized.scale_rows, + hidden_quantized.scale_mma, + value_order="trellis_native_mma", + ) + + gate = runtime.gate_transformed[:routes] + up = runtime.up_transformed[:routes] + run_glm_route_packed_w4a8_projection( + hidden_quantized, + weights.gate, + packed, + block_experts, + gate, + topk=runtime.topk, + shared_input=True, + ) + run_glm_route_packed_w4a8_projection( + hidden_quantized, + weights.up, + packed, + block_experts, + up, + topk=runtime.topk, + shared_input=True, + ) + activated = runtime.activated[:routes] + route_experts = topk_ids.reshape(-1) + run_glm_gate_up_output_transform_silu( + gate, + up, + route_experts, + weights.gate_svh, + weights.up_svh, + runtime.gate_hadamard[:routes], + runtime.up_hadamard[:routes], + activated, + ones=runtime.ones_intermediate, + ) + + down_rotated = runtime.down_rotated[:routes] + run_glm_down_input_transform( + activated, + route_experts, + weights.down_suh, + runtime.down_scaled[:routes], + down_rotated, + ones=runtime.ones_intermediate, + ) + down_quantized = _active_quantized(runtime.down_quantized, routes) + quantize_mxfp8_rows_cute( + down_rotated, + down_quantized.values, + down_quantized.scale_rows, + down_quantized.scale_mma, + value_order="trellis_native_mma", + ) + down_transformed = runtime.down_transformed[:routes] + run_glm_route_packed_w4a8_projection( + down_quantized, + weights.down, + packed, + block_experts, + down_transformed, + topk=runtime.topk, + shared_input=False, + ) + return run_glm_down_output_transform_sum( + down_transformed, + topk_weights, + weights.down_svh, + runtime.down_canonical[:routes], + runtime.output[:tokens], + topk=runtime.topk, + ) + + +def _weight_tensors(weights: GLMSQGW4A8Weights) -> list[torch.Tensor]: + tensors: list[torch.Tensor] = [] + for projection in (weights.gate, weights.up, weights.down): + tensors.extend( + ( + projection.trellis_k3, + projection.trellis_k4, + projection.expert_slots_k3, + projection.expert_slots_k4, + ) + ) + tensors.extend( + ( + weights.gate_up_suh, + weights.gate_svh, + weights.up_svh, + weights.down_suh, + weights.down_svh, + ) + ) + return tensors + + +def _runtime_tensors(runtime: GLMSQGW4A8Runtime) -> list[torch.Tensor]: + return [ + runtime.input_f16, + runtime.hidden_rotated, + runtime.hidden_quantized.values, + runtime.hidden_quantized.scale_rows, + runtime.hidden_quantized.scale_mma, + runtime.gate_transformed, + runtime.up_transformed, + runtime.gate_hadamard, + runtime.up_hadamard, + runtime.activated, + runtime.down_scaled, + runtime.down_rotated, + runtime.down_quantized.values, + runtime.down_quantized.scale_rows, + runtime.down_quantized.scale_mma, + runtime.down_transformed, + runtime.down_canonical, + runtime.output, + runtime.packed_route_indices, + runtime.block_expert_ids, + runtime.packed_route_count, + runtime.expert_offsets, + runtime.expert_counts, + runtime.ones_intermediate, + ] + + +@torch.library.custom_op( + "b12x::glm_sqg_w4a8_moe", + mutates_args="unknown", + device_types="cuda", +) +def _run_op( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + weight_tensors: list[torch.Tensor], + runtime_tensors: list[torch.Tensor], + weights_id: int, + runtime_id: int, +) -> None: + weights = _WEIGHT_REGISTRY.get(int(weights_id)) + runtime = _RUNTIME_REGISTRY.get(int(runtime_id)) + if weights is None or runtime is None: + raise RuntimeError("GLM SQG W4A8 prepared owner is no longer live") + expected_weights = _weight_tensors(weights) + expected_runtime = _runtime_tensors(runtime) + if len(weight_tensors) != len(expected_weights): + raise RuntimeError("GLM SQG W4A8 weight inventory changed after planning") + if len(runtime_tensors) != len(expected_runtime): + raise RuntimeError("GLM SQG W4A8 scratch inventory changed after planning") + if any( + actual.data_ptr() != expected.data_ptr() + for actual, expected in zip(weight_tensors, expected_weights, strict=True) + ): + raise RuntimeError("GLM SQG W4A8 custom op received different weight storage") + if any( + actual.data_ptr() != expected.data_ptr() + for actual, expected in zip(runtime_tensors, expected_runtime, strict=True) + ): + raise RuntimeError("GLM SQG W4A8 custom op received different scratch storage") + _run_impl(hidden_states, topk_weights, topk_ids, weights, runtime) + + +@_run_op.register_fake +def _run_op_fake( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + weight_tensors: list[torch.Tensor], + runtime_tensors: list[torch.Tensor], + weights_id: int, + runtime_id: int, +) -> None: + del ( + hidden_states, + topk_weights, + topk_ids, + weight_tensors, + runtime_tensors, + weights_id, + runtime_id, + ) + + +def run( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + weights: GLMSQGW4A8Weights, + runtime: GLMSQGW4A8Runtime, +) -> torch.Tensor: + """Run through an opaque compile-safe op and return its stable output view.""" + + tokens = int(hidden_states.shape[0]) + _run_op( + hidden_states, + topk_weights, + topk_ids, + _weight_tensors(weights), + _runtime_tensors(runtime), + id(weights), + id(runtime), + ) + return runtime.output[:tokens] + + +_WEIGHT_REGISTRY = weakref.WeakValueDictionary() +_RUNTIME_REGISTRY = weakref.WeakValueDictionary() + + +__all__ = [ + "GLMSQGW4A8Runtime", + "GLMSQGW4A8Weights", + "glm_route_packed_w4a8_kernel_contract", + "prepare_runtime", + "prepare_weights", + "run", + "validate_glm_route_packed_w4a8_acceptance_kernel", +] diff --git a/tests/gemm/test_sqg_k6_w6a16.py b/tests/gemm/test_sqg_k6_w6a16.py new file mode 100644 index 000000000..637c9f175 --- /dev/null +++ b/tests/gemm/test_sqg_k6_w6a16.py @@ -0,0 +1,338 @@ +from __future__ import annotations + +import math + +import pytest +import torch + +from b12x._lib.quant.sqg_e4m3 import ( + sqg_xor_cheb_t12_direct_lut_cpu, + sqg_xor_cheb_t12_lut_cpu, +) +from b12x.gemm import trellis_linear +from b12x.gemm.trellis_linear._small_m import ( + _configure_sqg_k6_lut, + _extension, +) +from b12x.moe._shared.kernels.w4a16.host import ( + dense_trellis_gemm_scratch_elements_upper_bound, + packed_gemm_scratch_elements, +) +from b12x.moe._shared.kernels.w4a16.kernel import ( + _run_trellis_dense_hadamard128, +) + + +def _sm12x_available() -> bool: + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability() in ((12, 0), (12, 1)) + + +def _pack_edges(edges: torch.Tensor, bits: int) -> torch.Tensor: + tiles = edges.reshape(-1, 256).to(torch.int64) & ((1 << bits) - 1) + symbol_shifts = torch.arange(bits - 1, -1, -1) + word_shifts = torch.arange(15, -1, -1) + spans = tiles.reshape(-1, 16, 16) + bitstream = ((spans[..., None] >> symbol_shifts) & 1).reshape(-1, 16, bits * 16) + words = (bitstream.reshape(-1, 16, bits, 16) << word_shifts).sum(dim=-1) + flat = words.reshape(-1, 16 * bits) + packed = flat.reshape(flat.shape[0], -1, 2).flip(-1).reshape(flat.shape) + return packed.to(torch.int16).reshape(*edges.shape[:-1], 16 * bits).contiguous() + + +def _reconstruct_states(edges: torch.Tensor, bits: int) -> torch.Tensor: + values = edges.to(torch.int64) & ((1 << bits) - 1) + states = torch.zeros_like(values) + for lag in range(math.ceil(16 / bits)): + states |= torch.roll(values, shifts=lag, dims=-1) << (lag * bits) + return (states & 0xFFFF).to(torch.int16) + + +def _tensor_core_permutation() -> torch.Tensor: + permutation = [0] * 256 + for thread in range(32): + rows = ( + (thread % 4) * 2, + (thread % 4) * 2 + 1, + (thread % 4) * 2 + 8, + (thread % 4) * 2 + 9, + ) + columns = (thread // 4, thread // 4 + 8) + permutation[thread * 8 : thread * 8 + 8] = [ + rows[0] * 16 + columns[0], + rows[1] * 16 + columns[0], + rows[2] * 16 + columns[0], + rows[3] * 16 + columns[0], + rows[0] * 16 + columns[1], + rows[1] * 16 + columns[1], + rows[2] * 16 + columns[1], + rows[3] * 16 + columns[1], + ] + return torch.tensor(permutation, dtype=torch.long) + + +def _decode_reference(edges: torch.Tensor, bits: int) -> torch.Tensor: + states = _reconstruct_states(edges, bits) + codebook = sqg_xor_cheb_t12_direct_lut_cpu().reshape(5, 1 << 16)[bits - 2] + values = codebook.view(torch.float8_e4m3fn).float() + decoded = values.index_select(0, (states.to(torch.int64) & 0xFFFF).flatten()) + decoded = decoded.reshape_as(states) + decoded = decoded.index_select(-1, torch.argsort(_tensor_core_permutation())) + k_tiles, n_tiles, _ = decoded.shape + return ( + decoded.reshape(k_tiles, n_tiles, 16, 16) + .permute(0, 2, 1, 3) + .reshape(k_tiles * 16, n_tiles * 16) + .contiguous() + ) + + +def _identity_hadamard( + source: torch.Tensor, + destination: torch.Tensor, + _left_scale, + _right_scale, + _scale: float, +) -> None: + destination.copy_(source) + + +def _separate_hadamard( + source: torch.Tensor, + destination: torch.Tensor, + left_scale, + right_scale, + _scale: float, +) -> None: + _run_trellis_dense_hadamard128( + source, + destination, + left_scale if left_scale is not None else right_scale, + scale_before=left_scale is not None, + ) + + +def test_sqg_k6_direct_table_matches_codebook_definition() -> None: + bits = 6 + width = 16 - bits + history_mask = (1 << width) - 1 + branch_mask = (1 << bits) - 1 + codeword = torch.arange(1 << 16, dtype=torch.int64) + history = codeword >> bits + branch = codeword & branch_mask + + mixed = history ^ (history >> 11) + mixed ^= (mixed << 11) & history_mask + product = (0x3FA7D929 * mixed + 0xC928FD8E) & 0xFFFFFFFF + phase = product & history_mask + syndrome = product >> (32 - bits) + reversed_branch = torch.zeros_like(branch) + for index in range(bits): + reversed_branch |= ((branch >> index) & 1) << (bits - 1 - index) + rank = ((reversed_branch ^ syndrome) << width) | phase + + expected = sqg_xor_cheb_t12_lut_cpu()[rank >> 4] + actual = sqg_xor_cheb_t12_direct_lut_cpu().reshape(5, 1 << 16)[bits - 2] + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120 GPU") +def test_sqg_k6_device_decoder_matches_every_direct_table_entry() -> None: + device = torch.device("cuda", torch.cuda.current_device()) + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + _configure_sqg_k6_lut(int(device_index)) + codewords = torch.arange(1 << 16, dtype=torch.int32).to( + device=device, dtype=torch.int16 + ) + actual = _extension().decode_k6_sqg_codewords(codewords) + labels = sqg_xor_cheb_t12_direct_lut_cpu().reshape(5, 1 << 16)[4] + expected = labels.view(torch.float8_e4m3fn).to(device=device, dtype=torch.float16) + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("rows", (1, 7, 32, 129, 8192)) +def test_dense_trellis_scratch_bound_covers_every_schedule(rows: int) -> None: + size_n = 6144 + sms = 104 + capacity = dense_trellis_gemm_scratch_elements_upper_bound( + rows=rows, + size_n=size_n, + sms=sms, + ) + for block_size in (8, 16, 32, 48, 64): + route_slots = ((rows + block_size - 1) // block_size) * block_size + required = packed_gemm_scratch_elements( + size_n=size_n, + route_slots=route_slots, + moe_block_size=block_size, + sms=sms, + ) + assert capacity >= required + + +@pytest.mark.skipif(not _sm12x_available(), reason="requires SM120 or SM121") +def test_sqg_k6_w6a16_matches_reference_and_replays_in_cuda_graph() -> None: + device = torch.device("cuda", torch.cuda.current_device()) + generator = torch.Generator().manual_seed(0x53514706) + bits = 6 + features = 128 + rows = 7 + edges = torch.randint( + 0, + 1 << bits, + (features // 16, features // 16, 256), + dtype=torch.int16, + generator=generator, + ) + packed = _pack_edges(edges, bits).to(device) + scale = torch.ones(features, dtype=torch.float16, device=device) + prepared = trellis_linear.prepare_weight( + packed, + scale, + scale.clone(), + codebook="sqg_xor_cheb_t12", + params_dtype=torch.float16, + ) + x = (torch.randn((rows, features), generator=generator) * 0.03125).to( + device=device, dtype=torch.float16 + ) + output = torch.empty((rows, features), dtype=torch.float16, device=device) + gemm_output = torch.empty_like(output) + rotated_f16 = torch.empty_like(x) + c_tmp = torch.empty( + trellis_linear.sqg_k6_w6a16_scratch_elements( + rows, + features, + device=device, + ), + dtype=torch.float32, + device=device, + ) + + actual = trellis_linear.run_sqg_k6_w6a16( + x, + prepared, + output=output, + gemm_output=gemm_output, + rotated_f16=rotated_f16, + c_tmp=c_tmp, + hadamard_128=_identity_hadamard, + ).clone() + expected_weight = _decode_reference(edges, bits).to(device) + expected = (x.float() @ expected_weight).to(torch.float16) + relative_error = (actual - expected).float().norm() / expected.float().norm() + cosine = torch.nn.functional.cosine_similarity( + actual.float().flatten(), expected.float().flatten(), dim=0 + ) + assert float(relative_error) <= 2.0e-2 + assert float(cosine) >= 0.999 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = trellis_linear.run_sqg_k6_w6a16( + x, + prepared, + output=output, + gemm_output=gemm_output, + rotated_f16=rotated_f16, + c_tmp=c_tmp, + hadamard_128=_identity_hadamard, + ) + graph.replay() + torch.cuda.synchronize(device) + assert torch.equal(captured, actual) + + +@pytest.mark.parametrize("rows", (1, 7, 32, 128)) +@pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120 GPU") +def test_sqg_k6_small_m_matches_generic_pipeline_and_replays( + rows: int, +) -> None: + device = torch.device("cuda", torch.cuda.current_device()) + generator = torch.Generator().manual_seed(0x53514760 + rows) + features = 128 + packed = torch.randint( + -32768, + 32767, + (features // 16, features // 16, 96), + dtype=torch.int16, + generator=generator, + ).to(device) + suh = torch.randn(features, dtype=torch.float16, device=device) + svh = torch.randn(features, dtype=torch.float16, device=device) + prepared = trellis_linear.prepare_weight( + packed, + suh, + svh, + codebook="sqg_xor_cheb_t12", + params_dtype=torch.float16, + ) + x = torch.randn( + (rows, features), + dtype=torch.float16, + generator=generator, + ).to(device) + generic_output = torch.empty_like(x) + generic_gemm = torch.empty_like(x) + actual_output = torch.empty_like(x) + actual_gemm = torch.empty_like(x) + rotated_f16 = torch.empty_like(x) + c_tmp = torch.empty( + trellis_linear.sqg_k6_w6a16_scratch_elements( + rows, + features, + device=device, + ), + dtype=torch.float32, + device=device, + ) + + expected = trellis_linear.run_sqg_k6_w6a16( + x, + prepared, + output=generic_output, + gemm_output=generic_gemm, + rotated_f16=rotated_f16, + c_tmp=c_tmp, + hadamard_128=_separate_hadamard, + ).clone() + prepared.workspace.zero_() + actual = trellis_linear.run_sqg_k6_w6a16( + x, + prepared, + output=actual_output, + gemm_output=actual_gemm, + rotated_f16=rotated_f16, + c_tmp=c_tmp, + ).clone() + torch.cuda.synchronize(device) + + delta = actual.float() - expected.float() + relative_l2 = delta.norm() / expected.float().norm().clamp_min(1e-12) + cosine = torch.nn.functional.cosine_similarity( + actual.float().flatten(), expected.float().flatten(), dim=0 + ) + max_relative_to_range = delta.abs().max() / expected.float().abs().max().clamp_min( + 1e-12 + ) + assert float(relative_l2) < 1.5e-3 + assert float(cosine) > 0.999998 + assert float(max_relative_to_range) < 2e-3 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = trellis_linear.run_sqg_k6_w6a16( + x, + prepared, + output=actual_output, + gemm_output=actual_gemm, + rotated_f16=rotated_f16, + c_tmp=c_tmp, + ) + actual_output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize(device) + assert torch.equal(captured, actual) diff --git a/tests/gemm/test_trellis_linear.py b/tests/gemm/test_trellis_linear.py index e00d49038..4a4e1e64b 100644 --- a/tests/gemm/test_trellis_linear.py +++ b/tests/gemm/test_trellis_linear.py @@ -18,10 +18,9 @@ from b12x.moe._shared.kernels.w4a16.kernel import ( _run_trellis_dense_hadamard128, _trellis256_dense_launch_geometry, - _use_k6_mcg_small, + _use_k6_small, ) - _MCG = np.uint64(0xCBAC1FED) _MUL1 = np.uint64(0x83DCD12D) _MASK = np.uint32(0x8FFF8FFF) @@ -326,6 +325,73 @@ def fake_run(*args, **kwargs): assert seen["kwargs"]["_force_tile_config"] is None +def test_prepare_weight_reuses_caller_owned_persistent_storage(monkeypatch) -> None: + trellis = torch.empty(0) + suh = torch.empty(0) + svh = torch.empty(0) + dummy_scale = torch.empty(0) + global_scale = torch.empty(0) + workspace = torch.empty(0) + prepared = object() + seen = {} + + def fake_prepare(*args, **kwargs): + seen["args"] = args + seen["kwargs"] = kwargs + return prepared + + monkeypatch.setattr(api, "prepare_trellis256_dense_weight", fake_prepare) + + actual = trellis_linear.prepare_weight( + trellis, + suh, + svh, + codebook="sqg_xor_cheb_t12", + params_dtype=torch.float16, + dummy_scale=dummy_scale, + global_scale=global_scale, + workspace=workspace, + ) + + assert actual is prepared + assert seen["args"] == (trellis, suh, svh) + assert seen["kwargs"]["dummy_scale"] is dummy_scale + assert seen["kwargs"]["global_scale"] is global_scale + assert seen["kwargs"]["workspace"] is workspace + + +def test_sqg_k6_endpoint_rejects_other_weight_contracts(monkeypatch) -> None: + calls = [] + + def fake_run(x, weight, **kwargs): + calls.append((x, weight, kwargs)) + return "output" + + monkeypatch.setattr(api, "run_trellis256_dense", fake_run) + x = object() + weight = SimpleNamespace( + trellis_codebook="sqg_xor_cheb_t12", + trellis_bits=6, + trellis_pair_kind=None, + mcg=None, + mul1_e4m3=None, + ) + + assert trellis_linear.run_sqg_k6_w6a16(x, weight) == "output" + for override in ( + {"trellis_codebook": "mcg"}, + {"trellis_bits": 4}, + {"trellis_pair_kind": "P33"}, + {"mcg": object()}, + {"mul1_e4m3": object()}, + ): + invalid = SimpleNamespace(**(vars(weight) | override)) + with pytest.raises(ValueError): + trellis_linear.run_sqg_k6_w6a16(x, invalid) + + assert len(calls) == 1 + + def test_is_supported_uses_standard_sm12x_gate(monkeypatch) -> None: seen = {} @@ -346,6 +412,8 @@ def fake_gate(device, *, requires): (2048, 4096, 120, 120), (6144, 1024, 188, 64), (6144, 1024, 48, 48), + (6144, 512, 188, 48), + (6144, 512, 32, 32), (512, 6144, 188, 96), (512, 6144, 80, 80), # The unsharded dimensions are deliberately not inferred from TP4. @@ -365,7 +433,11 @@ def test_k6_small_m_default_sms_preserves_glm_decode_overlap( assert _default_num_sms(size_k, size_n, available_sms) == expected -def test_k6_small_m_rejects_unsupported_arch_before_jit(monkeypatch) -> None: +@pytest.mark.parametrize("runner", (_small_m.run_k6_mcg, _small_m.run_k6_sqg)) +def test_k6_small_m_rejects_unsupported_arch_before_jit( + monkeypatch, + runner, +) -> None: monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0)) monkeypatch.setattr( _small_m, @@ -374,21 +446,25 @@ def test_k6_small_m_rejects_unsupported_arch_before_jit(monkeypatch) -> None: ) with pytest.raises(NotImplementedError, match="built for sm_120 only"): - _small_m.run_k6_mcg(*(torch.empty(0) for _ in range(7))) + runner(*(torch.empty(0) for _ in range(7))) @pytest.mark.parametrize( - ("capability", "explicit_launch_config", "expected"), + ("capability", "codebook", "rows", "explicit_launch_config", "expected"), [ - ((12, 0), False, True), - ((12, 0), True, False), - ((12, 1), False, False), - ((9, 0), False, False), + ((12, 0), "mcg", 128, False, True), + ((12, 0), "sqg_xor_cheb_t12", 128, False, True), + ((12, 0), "sqg_xor_cheb_t12", 129, False, False), + ((12, 0), "sqg_xor_cheb_t12", 128, True, False), + ((12, 1), "mcg", 128, False, False), + ((9, 0), "mcg", 128, False, False), ], ) def test_k6_small_m_dispatch_requires_compiled_target( monkeypatch, capability: tuple[int, int], + codebook: str, + rows: int, explicit_launch_config: bool, expected: bool, ) -> None: @@ -399,11 +475,11 @@ def test_k6_small_m_dispatch_requires_compiled_target( ) assert ( - _use_k6_mcg_small( + _use_k6_small( device=torch.device("cuda"), - m=128, + m=rows, trellis_bits=6, - trellis_codebook="mcg", + trellis_codebook=codebook, trellis_pair_kind=None, compute_dtype=torch.float16, external_hadamard_128=None, diff --git a/tests/moe/test_glm_sqg_w4a8.py b/tests/moe/test_glm_sqg_w4a8.py new file mode 100644 index 000000000..a8ef77866 --- /dev/null +++ b/tests/moe/test_glm_sqg_w4a8.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import math + +import pytest +import torch + +from b12x._lib.quant.mxfp8_rows import quantize_mxfp8_rows_cute +from b12x._lib.quant.sqg_e4m3 import sqg_xor_cheb_t12_direct_lut_cpu +from b12x.gemm._shared.wo_mxfp8 import ( + dequantize_mxfp8_rows_torch, + empty_mxfp8_rows_for_dense_gemm, + quantize_mxfp8_rows_torch, +) +from b12x.moe import glm_sqg_w4a8 +from b12x.moe._shared.kernels.glm_trellis_w4a8 import ( + prepare_glm_route_packed_w4a8_projection, + run_glm_route_packed_w4a8_projection, +) +from b12x.moe._shared.kernels.w4a16.kernel import pack_topk_routes_by_expert + + +def _sm12x_available() -> bool: + if not torch.cuda.is_available(): + return False + return torch.cuda.get_device_capability() in ((12, 0), (12, 1)) + + +def _pack_edges(edges: torch.Tensor, bits: int) -> torch.Tensor: + tiles = edges.reshape(-1, 256).to(torch.int64) & ((1 << bits) - 1) + symbol_shifts = torch.arange(bits - 1, -1, -1) + word_shifts = torch.arange(15, -1, -1) + spans = tiles.reshape(-1, 16, 16) + bitstream = ((spans[..., None] >> symbol_shifts) & 1).reshape(-1, 16, bits * 16) + words = (bitstream.reshape(-1, 16, bits, 16) << word_shifts).sum(dim=-1) + flat = words.reshape(-1, 16 * bits) + packed = flat.reshape(flat.shape[0], -1, 2).flip(-1).reshape(flat.shape) + return packed.to(torch.int16).reshape(*edges.shape[:-1], 16 * bits).contiguous() + + +def _tensor_core_permutation() -> torch.Tensor: + permutation = [0] * 256 + for thread in range(32): + rows = ( + (thread % 4) * 2, + (thread % 4) * 2 + 1, + (thread % 4) * 2 + 8, + (thread % 4) * 2 + 9, + ) + columns = (thread // 4, thread // 4 + 8) + permutation[thread * 8 : thread * 8 + 8] = [ + rows[0] * 16 + columns[0], + rows[1] * 16 + columns[0], + rows[2] * 16 + columns[0], + rows[3] * 16 + columns[0], + rows[0] * 16 + columns[1], + rows[1] * 16 + columns[1], + rows[2] * 16 + columns[1], + rows[3] * 16 + columns[1], + ] + return torch.tensor(permutation, dtype=torch.long) + + +def _decode_reference(edges: torch.Tensor, bits: int) -> torch.Tensor: + values = edges.to(torch.int64) & ((1 << bits) - 1) + states = torch.zeros_like(values) + for lag in range(math.ceil(16 / bits)): + states |= torch.roll(values, shifts=lag, dims=-1) << (lag * bits) + states &= 0xFFFF + codebook = sqg_xor_cheb_t12_direct_lut_cpu().reshape(5, 1 << 16)[bits - 2] + decoded = ( + codebook.view(torch.float8_e4m3fn).float().index_select(0, states.flatten()) + ) + decoded = decoded.reshape_as(states) + decoded = decoded.index_select(-1, torch.argsort(_tensor_core_permutation())) + k_tiles, n_tiles, _ = decoded.shape + return ( + decoded.reshape(k_tiles, n_tiles, 16, 16) + .permute(0, 2, 1, 3) + .reshape(k_tiles * 16, n_tiles * 16) + .contiguous() + ) + + +def test_acceptance_schedule_is_qualified_for_sm120(monkeypatch) -> None: + monkeypatch.setenv("B12X_GLM_W4A8_ACCEPT_ARCH", "sm_120") + monkeypatch.delenv("B12X_GLM_W4A8_KERNEL", raising=False) + monkeypatch.delenv("B12X_GLM_W4A8_V2_BLOCKS", raising=False) + monkeypatch.delenv("B12X_GLM_W4A8_V2_STAGES", raising=False) + + glm_sqg_w4a8.validate_glm_route_packed_w4a8_acceptance_kernel() + + +def test_acceptance_schedule_rejects_unqualified_override(monkeypatch) -> None: + monkeypatch.setenv("B12X_GLM_W4A8_ACCEPT_ARCH", "120") + monkeypatch.setenv("B12X_GLM_W4A8_V2_BLOCKS", "4") + + with pytest.raises(RuntimeError, match="qualified sm_120 schedule"): + glm_sqg_w4a8.validate_glm_route_packed_w4a8_acceptance_kernel() + + +@pytest.mark.parametrize("arch", ("sm_90", "sm_103", "sm_130")) +def test_acceptance_schedule_rejects_other_architectures(monkeypatch, arch) -> None: + monkeypatch.setenv("B12X_GLM_W4A8_ACCEPT_ARCH", arch) + + with pytest.raises(RuntimeError, match="supports SM120/SM121"): + glm_sqg_w4a8.validate_glm_route_packed_w4a8_acceptance_kernel() + + +@pytest.mark.skipif(not _sm12x_available(), reason="requires SM120 or SM121") +def test_mixed_k3_k4_projection_matches_reference_and_replays_in_cuda_graph() -> None: + device = torch.device("cuda", torch.cuda.current_device()) + generator = torch.Generator().manual_seed(0x53514734) + size_k = 128 + size_n = 128 + topk = 2 + tokens = 4 + bits_by_expert = (3, 4) + edges = tuple( + torch.randint( + 0, + 1 << bits, + (size_k // 16, size_n // 16, 256), + dtype=torch.int16, + generator=generator, + ) + for bits in bits_by_expert + ) + packed = tuple( + _pack_edges(expert_edges, bits).to(device) + for expert_edges, bits in zip(edges, bits_by_expert, strict=True) + ) + prepared = prepare_glm_route_packed_w4a8_projection( + packed, + bits_by_expert, + size_k=size_k, + size_n=size_n, + ) + topk_ids = torch.tensor( + ((0, 1), (1, 0), (0, 1), (1, 0)), + dtype=torch.int32, + device=device, + ) + packed_routes, block_experts, _ = pack_topk_routes_by_expert( + topk_ids, + 128, + len(bits_by_expert), + ) + source = (torch.randn((tokens, size_k), generator=generator) * 0.03125).to( + device=device, dtype=torch.float16 + ) + quantized = empty_mxfp8_rows_for_dense_gemm( + tokens, + size_k, + device=device, + ) + output = torch.empty((tokens * topk, size_n), dtype=torch.float16, device=device) + + def run() -> torch.Tensor: + quantize_mxfp8_rows_cute( + source, + quantized.values, + quantized.scale_rows, + quantized.scale_mma, + value_order="trellis_native_mma", + ) + return run_glm_route_packed_w4a8_projection( + quantized, + prepared, + packed_routes, + block_experts, + output, + topk=topk, + shared_input=True, + ) + + actual = run().clone() + torch.cuda.synchronize(device) + canonical_quantized = quantize_mxfp8_rows_torch(source) + source_reference = dequantize_mxfp8_rows_torch( + canonical_quantized.values, + canonical_quantized.scale_rows, + ) + weights = tuple( + _decode_reference(expert_edges, bits).to(device) + for expert_edges, bits in zip(edges, bits_by_expert, strict=True) + ) + expected = torch.stack( + tuple( + source_reference[route // topk] @ weights[int(expert)] + for route, expert in enumerate(topk_ids.flatten().tolist()) + ) + ).to(torch.float16) + relative_error = (actual - expected).float().norm() / expected.float().norm() + cosine = torch.nn.functional.cosine_similarity( + actual.float().flatten(), expected.float().flatten(), dim=0 + ) + assert float(relative_error) <= 2.0e-2 + assert float(cosine) >= 0.999 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize(device) + assert torch.equal(captured, actual) diff --git a/tests/quantization/test_sqg_e4m3.py b/tests/quantization/test_sqg_e4m3.py index dc1fff4e1..083d1f1d0 100644 --- a/tests/quantization/test_sqg_e4m3.py +++ b/tests/quantization/test_sqg_e4m3.py @@ -9,6 +9,7 @@ SQG_E4M3_RANK_LUT_ENTRIES, SQG_E4M3_STATE_ENTRIES, SQG_E4M3_STATE_LUT_ENTRIES, + SQG_XOR_CHEB_T12_DIRECT_LUT_ENTRIES, _sqg_state_for_histories, decode_sqg_cheb_normal_e4m3_ranks_torch, sqg_xor_cheb_t12_direct_lut_cpu, @@ -23,10 +24,10 @@ def test_sqg_xor_cheb_t12_direct_lut_is_finite_and_bijective_by_rank() -> None: direct = sqg_xor_cheb_t12_direct_lut_cpu() assert direct.dtype == torch.uint8 - assert direct.shape == (SQG_E4M3_DIRECT_LUT_ENTRIES,) + assert direct.shape == (SQG_XOR_CHEB_T12_DIRECT_LUT_ENTRIES,) t12 = sqg_xor_cheb_t12_lut_cpu() expected_histogram = torch.bincount(t12.to(torch.int64), minlength=256) * 16 - for rate_index in range(3): + for rate_index in range(5): labels = direct[rate_index << 16 : (rate_index + 1) << 16] assert not bool(torch.any(labels == 0x80)) assert not bool(torch.any((labels & 0x7F) == 0x7F)) @@ -48,13 +49,9 @@ def test_sqg_cheb_rank_descriptor_closes_all_ranks() -> None: ranks = torch.arange(1 << 16, dtype=torch.int64) negative = ranks < 0x8000 magnitude_rank = ranks & 0x7FFF - magnitude_rank = torch.where( - negative, magnitude_rank ^ 0x7FFF, magnitude_rank - ) + magnitude_rank = torch.where(negative, magnitude_rank ^ 0x7FFF, magnitude_rank) entry = table[(magnitude_rank >> 5)].to(torch.int64) - magnitude = (entry & 0xFF) + ( - (magnitude_rank & 31) >= (entry >> 8) - ).to(torch.int64) + magnitude = (entry & 0xFF) + ((magnitude_rank & 31) >= (entry >> 8)).to(torch.int64) magnitude += (magnitude_rank >= 32764).to(torch.int64) magnitude += (magnitude_rank >= 32767).to(torch.int64) actual = magnitude | ((negative & (magnitude != 0)).to(torch.int64) << 7) @@ -129,6 +126,4 @@ def test_sqg_cheb_k2_q8h4_w2_direct_lut_matches_profile_5() -> None: native_k2 = native[: 1 << 16] profile_k2 = profile[: 1 << 16] assert not torch.equal(profile_k2, native_k2) - assert torch.equal( - torch.sort(profile_k2).values, torch.sort(native_k2).values - ) + assert torch.equal(torch.sort(profile_k2).values, torch.sort(native_k2).values)