Kernels:

kernels-community
/

finegrained-kernels

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on 2 days ago

Commit

fbf55f6

verified ·

1 Parent(s): d91aa5e

Uploaded using `kernel-builder`.

Browse files

Files changed (12) hide show

README.md +2 -0
build/torch-xpu/__init__.py +6 -2
build/torch-xpu/_ops.py +1 -1
build/torch-xpu/batched.py +12 -2
build/torch-xpu/grouped.py +7 -0
build/torch-xpu/loading/scales.py +263 -0
build/torch-xpu/matmul.py +14 -5
build/torch-xpu/metadata.json +11 -11
build/torch-xpu/metadata.json.sigstore +1 -1
build/torch-xpu/moe.py +83 -18
build/torch-xpu/scheduling.py +44 -13
build/torch-xpu/symbols.json +170 -20

README.md CHANGED Viewed

@@ -30,6 +30,8 @@ matmul_2d(...)
 - `rms_norm_rows`
 - `swizzle_mx_scales`
 - `unswizzle_mx_scales`
 - `mxfp8_act_quant`
 - `mxfp4_act_quant`
 - `nvfp4_act_quant`

 - `rms_norm_rows`
 - `swizzle_mx_scales`
 - `unswizzle_mx_scales`
+- `prepare_nvfp4_csf_scales`
+- `decompress_nvfp4_csf_scales`
 - `mxfp8_act_quant`
 - `mxfp4_act_quant`
 - `nvfp4_act_quant`

build/torch-xpu/__init__.py CHANGED Viewed

@@ -22,6 +22,7 @@ from . import backward
 from .formats import get_supported_act_fns
 from .norm import get_supported_norms, rms_norm_rows
 from .swizzle import swizzle_mx_scales, unswizzle_mx_scales
 from .quant import mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant
 __all__ = [
@@ -41,10 +42,13 @@ __all__ = [
     # the per-expert output norm a model can hand the forwards by name, and its standalone op
     "get_supported_norms",
     "rms_norm_rows",
-    # load-time helpers: the swizzled scale layout the SM100 scaled-MMA reads, and the row-wise
-    # quantizers a loader uses to quantize weights into the group formats
     "swizzle_mx_scales",
     "unswizzle_mx_scales",
     "mxfp8_act_quant",
     "mxfp4_act_quant",
     "nvfp4_act_quant",

 from .formats import get_supported_act_fns
 from .norm import get_supported_norms, rms_norm_rows
 from .swizzle import swizzle_mx_scales, unswizzle_mx_scales
+from .loading.scales import prepare_nvfp4_csf_scales, decompress_nvfp4_csf_scales
 from .quant import mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant
 __all__ = [
     # the per-expert output norm a model can hand the forwards by name, and its standalone op
     "get_supported_norms",
     "rms_norm_rows",
+    # load-time helpers: the swizzled scale layout the SM100 scaled-MMA reads, CSF scale-stream
+    # preparation/decoding, and the row-wise quantizers a loader uses to quantize weights into the
+    # group formats
     "swizzle_mx_scales",
     "unswizzle_mx_scales",
+    "prepare_nvfp4_csf_scales",
+    "decompress_nvfp4_csf_scales",
     "mxfp8_act_quant",
     "mxfp4_act_quant",
     "nvfp4_act_quant",

build/torch-xpu/_ops.py CHANGED Viewed

@@ -25,7 +25,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "finegrained_kernels"
-    unique_id = "78a2a29"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

 def _find_ops_name() -> str:
     kernel_name = "finegrained_kernels"
+    unique_id = "bbd13a8"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

build/torch-xpu/batched.py CHANGED Viewed

@@ -31,9 +31,9 @@ from .formats import check_activation_format, global_scale_stride, gate_out_cols
 from .epilogue import fused_glu
 from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, static_expert_act_operands, tensor_wide_act_operands
 from .swizzle import swizzled_scale_descriptor
-from .scheduling import is_off_rank
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
-from .loading.scales import weight_scale_rows
 from .loading.tiles import (
     advance_ptrs,
     load_act_block_dynamic,
@@ -1788,6 +1788,8 @@ def matmul_batched(
     scatter_idx: torch.Tensor | None = None,
     a_global_scale: torch.Tensor | None = None,
     b_global_scale: torch.Tensor | None = None,
     output_global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one
@@ -1825,7 +1827,15 @@ def matmul_batched(
     - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →
       ``w8a8_tensor_dynamic_fp8_matmul_batched``.
     - block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``.
     """
     check_activation_format(activation_format, quantize_output)
     assert act_fn is not None or not quantize_output, "the raw gate|up output (act_fn=None) is not requantized"
     assert (a_global_scale is None and b_global_scale is None) or (

 from .epilogue import fused_glu
 from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, static_expert_act_operands, tensor_wide_act_operands
 from .swizzle import swizzled_scale_descriptor
+from .scheduling import is_off_rank, routed_expert_start
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
+from .loading.scales import csf_weight_scales, weight_scale_rows
 from .loading.tiles import (
     advance_ptrs,
     load_act_block_dynamic,
     scatter_idx: torch.Tensor | None = None,
     a_global_scale: torch.Tensor | None = None,
     b_global_scale: torch.Tensor | None = None,
+    b_scale_exceptions: torch.Tensor | None = None,
+    b_scale_exception_offsets: torch.Tensor | None = None,
     output_global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one
     - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →
       ``w8a8_tensor_dynamic_fp8_matmul_batched``.
     - block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``.
+    - CSF-compressed NVFP4 weight scales (``prepare_nvfp4_csf_scales``): ``Bs`` is the fixed
+      stream, ``b_scale_exceptions``/``b_scale_exception_offsets`` its exceptions and expert
+      offsets; the op decodes them before the GEMM (only the routed experts when routing allows).
     """
+    if b_scale_exceptions is not None:
+        Bs = csf_weight_scales(
+            Bs, b_scale_exceptions, b_scale_exception_offsets,
+            routed_expert_start(expert_ids, B.shape[0]),
+        )
     check_activation_format(activation_format, quantize_output)
     assert act_fn is not None or not quantize_output, "the raw gate|up output (act_fn=None) is not requantized"
     assert (a_global_scale is None and b_global_scale is None) or (

build/torch-xpu/grouped.py CHANGED Viewed

@@ -30,6 +30,7 @@ from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, mx_act_quant_group
 from .swizzle import swizzled_scale_descriptor
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
 from .scheduling import build_tile_layout, expand_gather_below_parity, expand_regime, build_packed_schedule, load_packed_schedule, prefetch_packed_entry, resolve_grouped_tile, resolve_grouped_tile_packed
 from .loading.tiles import (
     load_act_block_dynamic,
     load_act_mx,
@@ -2227,6 +2228,8 @@ def matmul_grouped(
     scatter_idx: torch.Tensor | None = None,
     a_global_scale: torch.Tensor | None = None,
     b_global_scale: torch.Tensor | None = None,
     output_global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """Grouped matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision).
@@ -2268,7 +2271,11 @@ def matmul_grouped(
     - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →
       ``w8a8_tensor_dynamic_fp8_matmul_grouped``.
     - block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_grouped``.
     """
     check_activation_format(activation_format, quantize_output)
     assert act_fn is not None or not quantize_output, "the raw gate|up output (act_fn=None) is not requantized"
     if gather_idx is not None and activation_format == "bf16" and As is None:

 from .swizzle import swizzled_scale_descriptor
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
 from .scheduling import build_tile_layout, expand_gather_below_parity, expand_regime, build_packed_schedule, load_packed_schedule, prefetch_packed_entry, resolve_grouped_tile, resolve_grouped_tile_packed
+from .loading.scales import csf_weight_scales
 from .loading.tiles import (
     load_act_block_dynamic,
     load_act_mx,
     scatter_idx: torch.Tensor | None = None,
     a_global_scale: torch.Tensor | None = None,
     b_global_scale: torch.Tensor | None = None,
+    b_scale_exceptions: torch.Tensor | None = None,
+    b_scale_exception_offsets: torch.Tensor | None = None,
     output_global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """Grouped matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision).
     - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →
       ``w8a8_tensor_dynamic_fp8_matmul_grouped``.
     - block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_grouped``.
+    - CSF-compressed NVFP4 weight scales (``prepare_nvfp4_csf_scales``): ``Bs`` is the fixed
+      stream, ``b_scale_exceptions``/``b_scale_exception_offsets`` its exceptions and expert
+      offsets; the op decodes them before the GEMM (only the routed experts when routing allows).
     """
+    Bs = csf_weight_scales(Bs, b_scale_exceptions, b_scale_exception_offsets, expert_start)
     check_activation_format(activation_format, quantize_output)
     assert act_fn is not None or not quantize_output, "the raw gate|up output (act_fn=None) is not requantized"
     if gather_idx is not None and activation_format == "bf16" and As is None:

build/torch-xpu/loading/scales.py CHANGED Viewed

@@ -12,9 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import triton
 import triton.language as tl
 from ..quant import fp8_act_quant_inline, mx_act_quant_inline
@@ -337,3 +344,259 @@ def mx_2d_scale_ptrs(
         as_mask = offs_am_lin < M
         bs_mask = offs_bn_lin < (2 * N if GATE else N)
     return as_ptrs, bs_ptrs, as_mask, bs_mask

 # See the License for the specific language governing permissions and
 # limitations under the License.
+from collections.abc import Sequence
+import torch
 import triton
 import triton.language as tl
+from ..compat import (
+    add_op_namespace_prefix, compile_time_only_triton_op, compile_time_only_triton_wrap,
+    device_context, sm_count,
+)
 from ..quant import fp8_act_quant_inline, mx_act_quant_inline
         as_mask = offs_am_lin < M
         bs_mask = offs_bn_lin < (2 * N if GATE else N)
     return as_ptrs, bs_ptrs, as_mask, bs_mask
+# CSF byte-window scales: stream preparation and staged decoding before ordinary GEMMs.
+def prepare_nvfp4_csf_scales(
+    fixed: torch.Tensor, exceptions: torch.Tensor | Sequence[torch.Tensor]
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Prepare one matrix or an expert stack without expanding or moving its scales.
+    CSF stores each row of E4M3 scale bytes as a base plus 4-bit deltas, which is
+    lossless whenever a row's bytes span at most 16 codes; the rest are exceptions.
+    - `fixed` is uint8 `[rows/16, 16*(1+cols/2)]`, or the same shape with a leading
+      expert dimension. Each 16-row slab holds 16 row bases (bytes 0..240), then those
+      rows' deltas, two per byte (even column in the low nibble). Scale `(r, c)` decodes
+      to byte `base[r] + delta[r, c]`.
+    - `exceptions` is one uint32 `[count]` tensor for a matrix, or one per expert, for
+      the scales outside their row's window. Each word is `position | byte << 24`:
+      the low 24 bits are the row-major index `r * cols + c` (sorted and unique), the
+      high 8 bits are the exact E4M3 byte, which overwrites the decoded value.
+    Returns, on the input device:
+    - `fixed`: contiguous `[experts, rows/16, 16*(1+cols/2)]` (a 2D input gains `experts=1`).
+    - `exceptions`: every expert's words concatenated into one uint32 vector.
+    - `expert_offsets`: uint32 `[experts + 1]` prefix sums; expert `e` owns
+      `exceptions[expert_offsets[e]:expert_offsets[e + 1]]`.
+    These are `decompress_nvfp4_csf_scales`' inputs. Decode on the accelerator at load time
+    for dense residency, or into shared scratch immediately before GEMM. The checks here
+    run once per checkpoint load, off the forward path, so the decoder does not repeat them.
+    """
+    if fixed.dtype != torch.uint8 or fixed.ndim not in (2, 3):
+        raise ValueError("CSF fixed stream must be a 2D or 3D uint8 tensor")
+    fixed = fixed.unsqueeze(0) if fixed.ndim == 2 else fixed
+    experts, slabs, slab_bytes = fixed.shape
+    if experts == 0 or slabs == 0 or slab_bytes <= 16 or slab_bytes % 16:
+        raise ValueError("Invalid CSF slab shape")
+    rows, cols = slabs * 16, (slab_bytes // 16 - 1) * 2
+    if rows * cols > 1 << 24:
+        raise ValueError(
+            "CSF exception positions support at most 2**24 scales per matrix"
+        )
+    if bool((fixed[:, :, :16] > 240).any()):
+        raise ValueError("CSF row bases must be in 0..240")
+    streams = [exceptions] if isinstance(exceptions, torch.Tensor) else list(exceptions)
+    if len(streams) != experts:
+        raise ValueError("CSF needs one exception stream per expert")
+    offsets, words = [0], []
+    for stream in streams:
+        if (
+            stream.dtype != torch.uint32
+            or stream.ndim != 1
+            or stream.device != fixed.device
+        ):
+            raise ValueError(
+                "CSF exception streams must be 1D uint32 tensors on the fixed stream's device"
+            )
+        values = stream.to(torch.int64)
+        positions = values & 0xFFFFFF
+        if positions.numel() and (
+            int(positions[-1]) >= rows * cols
+            or bool((positions[1:] <= positions[:-1]).any())
+        ):
+            raise ValueError(
+                "CSF exception positions must be sorted, unique and in range"
+            )
+        offsets.append(offsets[-1] + stream.numel())
+        words.append(stream)
+    if offsets[-1] >= 1 << 32:
+        raise ValueError("CSF exception directory exceeds uint32 offsets")
+    return (
+        fixed.contiguous(),
+        torch.cat(words),
+        torch.tensor(offsets, dtype=torch.uint32, device=fixed.device),
+    )
+@triton.jit
+def _decode_fixed_kernel(
+    Fixed,
+    Out,
+    ExpertStart,
+    E: tl.constexpr,
+    ROWS: tl.constexpr,
+    COLS: tl.constexpr,
+    BLOCK: tl.constexpr,
+    ALIGNED_OUTPUT: tl.constexpr,
+):
+    pairs: tl.constexpr = ROWS * (COLS // 2)
+    tiles: tl.constexpr = tl.cdiv(pairs, BLOCK)
+    offsets = tl.arange(0, BLOCK)
+    # Contiguous pairs avoid padding each row to a power-of-two tile width.
+    # Keep indices local to an expert; only the expert base needs 64-bit addressing.
+    for tile in range(tl.program_id(0), E * tiles, tl.num_programs(0)):
+        expert = tile // tiles
+        active = True
+        if ExpertStart is not None:
+            active = tl.load(ExpertStart + expert + 1) > tl.load(ExpertStart + expert)
+        if active:
+            pair = (tile % tiles) * BLOCK + offsets
+            row = pair // (COLS // 2)
+            slab = expert.to(tl.int64) * (ROWS // 16) + row // 16
+            slab_bytes: tl.constexpr = 16 * (1 + COLS // 2)
+            base = tl.load(Fixed + slab * slab_bytes + row % 16, pair < pairs, other=0).to(
+                tl.uint16
+            )
+            packed = tl.load(
+                Fixed + slab * slab_bytes + 16 + (row % 16) * (COLS // 2) + pair % (COLS // 2),
+                pair < pairs,
+                other=0,
+            ).to(tl.uint16)
+            low = base + (packed & 15)
+            high = base + (packed >> 4)
+            # A contiguous sliced workspace can start at an odd byte.
+            if ALIGNED_OUTPUT:
+                output = Out.to(tl.pointer_type(tl.uint16))
+                tl.store(
+                    output + expert.to(tl.int64) * pairs + pair,
+                    low | (high << 8),
+                    pair < pairs,
+                )
+            else:
+                values = tl.interleave(low, high).to(tl.uint8).to(tl.float8e4nv, bitcast=True)
+                output_offsets = (tile % tiles) * BLOCK * 2 + tl.arange(0, BLOCK * 2)
+                tl.store(
+                    Out + expert.to(tl.int64) * pairs * 2 + output_offsets,
+                    values,
+                    output_offsets < pairs * 2,
+                )
+@triton.jit
+def _apply_exceptions_kernel(
+    Exceptions,
+    ExpertOffsets,
+    Out,
+    ExpertStart,
+    PLANE: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    expert = tl.program_id(0)
+    if ExpertStart is not None:
+        if tl.load(ExpertStart + expert + 1) == tl.load(ExpertStart + expert):
+            return
+    start = tl.load(ExpertOffsets + expert).to(tl.int64)
+    end = tl.load(ExpertOffsets + expert + 1).to(tl.int64)
+    for base in range(start, end, BLOCK):
+        index = base + tl.arange(0, BLOCK)
+        word = tl.load(Exceptions + index, index < end, other=0)
+        position = word & 0xFFFFFF
+        value = (word >> 24).to(tl.uint8).to(tl.float8e4nv, bitcast=True)
+        tl.store(Out + expert.to(tl.int64) * PLANE + position, value, index < end)
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("decompress_nvfp4_csf_scales_into"),
+    mutates_args={"out"},
+    opaque=True,
+)
+def _decompress_into(
+    fixed: torch.Tensor,
+    exceptions: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    out: torch.Tensor,
+    expert_start: torch.Tensor | None,
+) -> None:
+    experts, rows, cols = out.shape
+    with device_context(fixed.device):
+        programs = min(
+            experts * triton.cdiv(rows * cols // 2, 4096),
+            16 * sm_count(fixed.device.index),
+        )
+        compile_time_only_triton_wrap(_decode_fixed_kernel)[(programs,)](
+            fixed,
+            out,
+            expert_start,
+            experts,
+            rows,
+            cols,
+            4096,
+            ALIGNED_OUTPUT=out.storage_offset() % 2 == 0,
+        )
+        # The fixed stream decodes every scale branch-free as base + delta, including the
+        # few outside their row's 16-code window. This second, small pass overwrites those
+        # positions with their exact bytes (stream order runs after the first kernel).
+        if exceptions.numel():
+            compile_time_only_triton_wrap(_apply_exceptions_kernel)[(experts,)](
+                exceptions, expert_offsets, out, expert_start, rows * cols, 256
+            )
+def decompress_nvfp4_csf_scales(
+    fixed: torch.Tensor,
+    exceptions: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+    expert_start: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Decode prepared CSF streams to dense E4M3 scales for an ordinary NVFP4 GEMM.
+    With no `out`, allocate dense scales on the accelerator at checkpoint loading.
+    Pass reusable contiguous [experts, rows, cols] workspace to keep the streams
+    compressed between layers. Consume the workspace before overwriting it; it
+    must not be retained for backward or shared by concurrent consumers.
+    A shared dense workspace saves memory across layers, not within one layer.
+    Pass the grouped GEMM's cumulative `expert_start` to decode only experts with
+    routed tokens. Inactive workspace planes stay untouched and must not be read.
+    The streams must come from `prepare_nvfp4_csf_scales`, which validates them; only
+    `out` and `expert_start`, which are new here, are checked.
+    """
+    if fixed.device.type == "cpu":
+        raise ValueError("CSF decompression requires an accelerator")
+    experts, slabs, slab_bytes = fixed.shape
+    rows, cols = slabs * 16, (slab_bytes // 16 - 1) * 2
+    shape = (experts, rows, cols)
+    if out is None:
+        out = torch.empty(shape, dtype=torch.float8_e4m3fn, device=fixed.device)
+    if (
+        out.shape != shape
+        or out.dtype != torch.float8_e4m3fn
+        or not out.is_contiguous()
+    ):
+        raise ValueError(f"CSF workspace must be contiguous E4M3 with shape {shape}")
+    if any(t.device != out.device for t in (fixed, exceptions, expert_offsets)):
+        raise ValueError("CSF buffers and workspace must share a device")
+    if expert_start is not None and (
+        expert_start.shape != (experts + 1,)
+        or expert_start.dtype not in (torch.int32, torch.int64)
+        or not expert_start.is_contiguous()
+        or expert_start.device != out.device
+    ):
+        raise ValueError("CSF expert_start must be a contiguous device int32/int64 vector of length experts + 1")
+    _decompress_into(fixed, exceptions, expert_offsets, out, expert_start)
+    return out
+def csf_weight_scales(
+    scales: torch.Tensor | None,
+    exceptions: torch.Tensor | None,
+    exception_offsets: torch.Tensor | None,
+    expert_start: torch.Tensor | None = None,
+) -> torch.Tensor | None:
+    """A public op's weight scales: with ``exceptions`` given, ``scales`` is the prepared CSF fixed
+    stream and decodes here (only the experts with routed rows when ``expert_start`` is given) into
+    a fresh dense E4M3 buffer. The caching allocator recycles that buffer across calls and layers,
+    so callers keep only the compressed streams resident. Dense scales pass through."""
+    if exceptions is None:
+        return scales
+    return decompress_nvfp4_csf_scales(
+        scales, exceptions, exception_offsets, expert_start=expert_start
+    )

build/torch-xpu/matmul.py CHANGED Viewed

@@ -27,7 +27,7 @@ from .descriptors import maybe_descriptor, rebind_bd_descriptors, rebind_mx_desc
 from .formats import check_activation_format, global_scale_stride, gate_out_cols, normalize_global_scale, e2m1_as_uint8, is_mx, mx_scale_family, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_2d_operands, weight_format
 from .swizzle import swizzle_mx_scales, swizzled_scale_descriptor
 from .quant import MX_ACT_QUANT, fp8_act_quant_block_dynamic, maybe_act_quant, mxfp8_act_quant, nvfp4_act_quant, quantize_rows_static, tensor_wide_act_operands
-from .loading.scales import apply_global_scale, mx_2d_scale_ptrs, weight_scale_rows
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
 from .loading.tiles import (
     advance_ptrs,
@@ -1912,6 +1912,8 @@ def matmul_2d(
     output_dtype: torch.dtype | None = None,
     a_global_scale: torch.Tensor | None = None,
     b_global_scale: torch.Tensor | None = None,
     output_global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor | list[torch.Tensor]:
     """Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of
@@ -1937,11 +1939,18 @@ def matmul_2d(
     ``[..., N]`` intermediate (``act_fn=None``: the raw ``[..., 2N]`` gate|up output instead); ``quantize_output=True`` (MX weights only) stores that intermediate in
     ``activation_format`` and returns ``(C, Cs)`` for the next op.
-    Routes by weight dtype + inferred block: ``Bs`` None → full-precision; MX/NVFP4 (int8/E4M3 weights
-    with UE8M0 group-32 or E4M3 group-16 ``Bs``) → ``mx_dynamic_matmul``; tensor-wide FP8 (scalar
-    ``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``; block FP8 → the static variant when ``As`` is a
-    per-tensor scalar, else dynamic.
     """
     check_activation_format(activation_format, quantize_output)
     assert act_fn is not None or not quantize_output, "the raw gate|up output (act_fn=None) is not requantized"
     assert (a_global_scale is None and b_global_scale is None) or (

 from .formats import check_activation_format, global_scale_stride, gate_out_cols, normalize_global_scale, e2m1_as_uint8, is_mx, mx_scale_family, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_2d_operands, weight_format
 from .swizzle import swizzle_mx_scales, swizzled_scale_descriptor
 from .quant import MX_ACT_QUANT, fp8_act_quant_block_dynamic, maybe_act_quant, mxfp8_act_quant, nvfp4_act_quant, quantize_rows_static, tensor_wide_act_operands
+from .loading.scales import apply_global_scale, csf_weight_scales, mx_2d_scale_ptrs, weight_scale_rows
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
 from .loading.tiles import (
     advance_ptrs,
     output_dtype: torch.dtype | None = None,
     a_global_scale: torch.Tensor | None = None,
     b_global_scale: torch.Tensor | None = None,
+    b_scale_exceptions: torch.Tensor | None = None,
+    b_scale_exception_offsets: torch.Tensor | None = None,
     output_global_scale: torch.Tensor | None = None,
 ) -> torch.Tensor | list[torch.Tensor]:
     """Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of
     ``[..., N]`` intermediate (``act_fn=None``: the raw ``[..., 2N]`` gate|up output instead); ``quantize_output=True`` (MX weights only) stores that intermediate in
     ``activation_format`` and returns ``(C, Cs)`` for the next op.
+    Routes by weight dtype + inferred block:
+    - ``Bs`` None → full-precision.
+    - MX/NVFP4 (int8/E4M3 weights with UE8M0 group-32 or E4M3 group-16 ``Bs``) →
+      ``mx_dynamic_matmul``.
+    - tensor-wide FP8 (scalar ``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``.
+    - block FP8 → the static variant when ``As`` is a per-tensor scalar, else dynamic.
+    - CSF-compressed NVFP4 weight scales (``prepare_nvfp4_csf_scales``): ``Bs`` is the fixed
+      stream, ``b_scale_exceptions``/``b_scale_exception_offsets`` its exceptions and expert
+      offsets; the op decodes them before the GEMM.
     """
+    if b_scale_exceptions is not None:  # one prepared matrix decodes as a 1-expert stack
+        Bs = csf_weight_scales(Bs, b_scale_exceptions, b_scale_exception_offsets)[0]
     check_activation_format(activation_format, quantize_output)
     assert act_fn is not None or not quantize_output, "the raw gate|up output (act_fn=None) is not requantized"
     assert (a_global_scale is None and b_global_scale is None) or (

build/torch-xpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "finegrained-kernels",
-  "id": "_finegrained_kernels_xpu_78a2a29",
   "version": 0,
   "kernels-minver": "0.14.0",
   "license": "Apache-2.0",
@@ -12,30 +12,30 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "xFRwaLTont8ZLTxS5ajuiDqM9gWwXAU4xap2prqdhrY=",
-      "_ops.py": "6ajnsgIjvPzGZBsQyo6hy2ukYkhjeD9B6q+AogINgPY=",
       "autotune_config.py": "d5fb2Vqss0fgC9TZbarrmQIInL1n1L7rHz6o3Gqv6cI=",
       "backward.py": "K1S6KD3J/12SthXiqBcmNO9kUikctQpccsbxsO8lpRg=",
-      "batched.py": "EA/eS6jhAvLinodFgT6A6Mz5hWJQjIwOrKvFWUMgz7c=",
       "bayesian_autotuner.py": "NDRUQrhk+bEEvrOkIilWJPfqdKpcsd4E4QQ9AHzylKI=",
       "compat.py": "K3RCq1ETi80duzv3zSldIwrL5e46vZAWtLxLlg1LXvM=",
       "descriptors.py": "4acRMazGPFlbydNt1Hbo8RpAqb3fOtnOQBv5yb4onqw=",
       "embedding.py": "k7jiz2vtEJuYKo+w6TuS/4WO1D4Bv96O5DMFjEpYUGo=",
       "epilogue.py": "AmE9Mhs0iWhnS5CneW3XGkg2Vvrj7+Ej2RmoorIM5gs=",
       "formats.py": "TT9zs05HzUMOVF5DG1OTRg8AjXssXHGoJoRONRJSyrg=",
-      "grouped.py": "sNQ5L3DzTk6BFMyiCWbbkIwAWX2VEAG3SEpE+2gKHAQ=",
       "loading/__init__.py": "so7kYmlVoiASfoOzemdU5fBTiP8yA85gVPNQExIq/KU=",
-      "loading/scales.py": "tcenA3+yIqrfCMttlMJNWgGJDUIwWWpUGltox6qaRxQ=",
       "loading/tiles.py": "p1go4zzcZ6nog6JCXgNbjFyhp5cL8iu6Is96gUV4+ks=",
-      "matmul.py": "GFc0Olq/UV3xLJxuYuQNEGvbVL8TJ9oGbWFCKqjBUM8=",
       "mma.py": "gLkCU5QjhQVMacixNR24Yi07aNelQyN+I0pJOuRxE4E=",
-      "moe.py": "kmGfWaq+hx++ri1Xt1X2aIuafGM96zOK4EDAO9n+Jyw=",
       "norm.py": "VvE6GUE0Oztf7PDcviLCcB4fhnXt7/zgdE+aApv6DBM=",
       "pruners.py": "AWpe/H5N3mXQnqxsU9AMkX8KPbwJ5qo4vu+fCM/KFlE=",
       "quant.py": "1fNRvb63DAOcFNLfKw3cSwMqSwLua4smI+xvhlsUk3k=",
-      "scheduling.py": "cxwoL5NTJxARVxLStDsNalm/jqbkCObm8uyLRFL/yYc=",
       "swizzle.py": "ueIJ9FsO8QZhVt3nMFlA+8b+U9/ZHabHFlhufsONV04=",
-      "symbols.json": "wHUHJ8vYH5uj+XDbZv4GynNezpqBAuAVg5r4/l35UsA="
     }
   },
   "provenance": {
@@ -45,7 +45,7 @@
       "dirty": false
     },
     "kernel": {
-      "commit": "78a2a29b886fcaaa5ee6a8cda4cbadf28e622a37",
       "dirty": false
     }
   }

 {
   "name": "finegrained-kernels",
+  "id": "_finegrained_kernels_xpu_bbd13a8",
   "version": 0,
   "kernels-minver": "0.14.0",
   "license": "Apache-2.0",
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "YGdTqU34klTI43VQVxTlk3qo0iJPSFaECyKN7nbwgqM=",
+      "_ops.py": "A9CLOqiJNpGlsP+gfdW8KzwepufQlyrG2qGK1YRVoog=",
       "autotune_config.py": "d5fb2Vqss0fgC9TZbarrmQIInL1n1L7rHz6o3Gqv6cI=",
       "backward.py": "K1S6KD3J/12SthXiqBcmNO9kUikctQpccsbxsO8lpRg=",
+      "batched.py": "R6swbEtGGx/h+766rP4wsSsfmr4KfYzRF0cq7goLcyI=",
       "bayesian_autotuner.py": "NDRUQrhk+bEEvrOkIilWJPfqdKpcsd4E4QQ9AHzylKI=",
       "compat.py": "K3RCq1ETi80duzv3zSldIwrL5e46vZAWtLxLlg1LXvM=",
       "descriptors.py": "4acRMazGPFlbydNt1Hbo8RpAqb3fOtnOQBv5yb4onqw=",
       "embedding.py": "k7jiz2vtEJuYKo+w6TuS/4WO1D4Bv96O5DMFjEpYUGo=",
       "epilogue.py": "AmE9Mhs0iWhnS5CneW3XGkg2Vvrj7+Ej2RmoorIM5gs=",
       "formats.py": "TT9zs05HzUMOVF5DG1OTRg8AjXssXHGoJoRONRJSyrg=",
+      "grouped.py": "5F+eYL+8eSr3BVRwX9R3d3HhmvKj4YdVw6u3W6MlNWg=",
       "loading/__init__.py": "so7kYmlVoiASfoOzemdU5fBTiP8yA85gVPNQExIq/KU=",
+      "loading/scales.py": "Ywe/tL94/S2YQI2T/d0ywfEUAkSYVpQTxpzGlb4cu9A=",
       "loading/tiles.py": "p1go4zzcZ6nog6JCXgNbjFyhp5cL8iu6Is96gUV4+ks=",
+      "matmul.py": "TdCy8rzuCtYEzowgQiPzbaFXAqANzUxuG62d6ZEYc0g=",
       "mma.py": "gLkCU5QjhQVMacixNR24Yi07aNelQyN+I0pJOuRxE4E=",
+      "moe.py": "9X93HFPen8gm7cB66X9EIxnOxGAmF4o1SBcBAZ6h+sE=",
       "norm.py": "VvE6GUE0Oztf7PDcviLCcB4fhnXt7/zgdE+aApv6DBM=",
       "pruners.py": "AWpe/H5N3mXQnqxsU9AMkX8KPbwJ5qo4vu+fCM/KFlE=",
       "quant.py": "1fNRvb63DAOcFNLfKw3cSwMqSwLua4smI+xvhlsUk3k=",
+      "scheduling.py": "A67jzd628qAVOqthPsRQQqbLfBybsZzOMO1e7wmZyT0=",
       "swizzle.py": "ueIJ9FsO8QZhVt3nMFlA+8b+U9/ZHabHFlhufsONV04=",
+      "symbols.json": "49uaUpO+Ntuu6o54Da6UAHvpFY9MgyjV/seN/pk8KlI="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "commit": "bbd13a828401b2fc6232fb13ccc122b0d3442fc5",
       "dirty": false
     }
   }

build/torch-xpu/metadata.json.sigstore CHANGED Viewed

@@ -1 +1 @@

- {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrOgAwIBAgIUaJFSxhpSWMO+Yia4J1eh8hBr+R0wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDIwMTMxNTU5WhcNMjYxMDIwMTMyNTU5WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEWt0ciJhHthUmSorEYVFfA9osevZEL54NMFNHGraSNn8VcvFNeOL2M9+LKrahPxsmsuwl0u3lcxKCtdKhoqmcrKOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg3OGEyYTI5Yjg4NmZjYWFhNWVlNmE4Y2RhNGNiYWRmMjhlNjIyYTM3MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg3OGEyYTI5Yjg4NmZjYWFhNWVlNmE4Y2RhNGNiYWRmMjhlNjIyYTM3MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg3OGEyYTI5Yjg4NmZjYWFhNWVlNmE4Y2RhNGNiYWRmMjhlNjIyYTM3MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNzhhMmEyOWI4ODZmY2FhYTVlZTZhOGNkYTRjYmFkZjI4ZTYyMmEzNzAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTQyODI2MzM5My9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhWXRZDAAABAMASDBGAiEA4kZJuKrCfhK+U2gZyDuFFQdDzpYkksvJQQ5Oi3N1l0YCIQDM+Br8zIDbZyQuT4pwho2jHSiSaDMgxjSp4evZhVWnBzAKBggqhkjOPQQDAwNoADBlAjEAzxxTvKkxYWujsdRVN16qq9u9FFrMWBuN0Ft70vFcDaHYokP8/1DtU6kUJyP8OkZmAjATzE+uSBijKKOFUipYtktXOV5Om7S0VgcA8zpEvhJTUdQZ9n5YifVl41u7dNhskds="}, "tlogEntries":[{"logIndex":"3337654267", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792502159", "inclusionPromise":{"signedEntryTimestamp":"MEUCIBDRBfn0OEr/3FI257Q8J6VNl0wPOC4C9GtEHCLvNMTiAiEAnA1frW7ZaIw2jb3g+xvOtbVeES3WKa4S0gjEkCnqI6s="}, "inclusionProof":{"logIndex":"3215750005", "rootHash":"EpuX/pRAzmWjhxJJhokpLRrLM4SW3jmWkE5KChIWqqs=", "treeSize":"3215750022", "hashes":["+ZeLGpISpSAYliOcVIzDBZa7/lRtEiq0zYgusOE8i1M=", "bTc4oPq97OYTzCwCBDAdAPlPOzw4/tVYm6M5A2Wv3Yk=", "0Xs6asm0G6/DH7iTugZlSR07Imsx+672Wz5bMCB7h2U=", "reyrj33h3Dm3n+/sj7mWmRDYNL6sZZpBbYnYKs03u3k=", "RJx/cVf+3XKI/JqI1krMKi6OZW7LQW+zgfDdarPNl4I=", "9/LkgO5mZRBXdGtH+UTxlIc5W1oBba0dcd9HRKcRnxs=", "LHM96E7LWMo97lKFxzwhu+bu4iqAgL9T08PQziX/ILw=", "BvQoQxsxgHD/VUQH9EPYOE3ZPJFndtBpjVYPSNDrKkM=", "27lVQFnJ0RKrJBEQGk4oUznEo3pFrggEC9gXDrQPiec=", "kys4WGMTuPqGaZc79EMJO/inn7y58MLx7inDyoMzO7I=", "kOrL2VZ9Dyg9SYQbyjbCK6l98qmY0xZiHeM68R2bcj8=", "sV2HYjtjXhxymU+V7/5mR942xjHqAOkHVQC7t+zu2us=", "XailVEBW00TuSppkvc+4KYSUGmNa81YudM5FIt3nvSE=", "zK7qW++aMer3MkBNGlcXVS4wV7x70Epwzh1KB82PWiQ=", "wc+TLATo5lnmbuwIr4/sQmH0osErTOAAiMahYD4f4sE=", "HxMxUZSh35Y+k1xWQyJAxe5BvkFVx4UmVeFrI31tchA=", "6Rb7uOgk2v5YKcFSbAPOpLlWAIN1ciDPdlp4J/CUu5k=", "HovTPBWzoErvUg1u5sIi3x1ZonZ7yJekVsBrmpomCIA=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3215750022\nEpuX/pRAzmWjhxJJhokpLRrLM4SW3jmWkE5KChIWqqs=\n\n— rekor.sigstore.dev wNI9ajBFAiEAq2OMQrs8uxumi2h0nA5ZrHtYbTWATE2oUPVNlEPL++YCIDyHeEb1qqierZe6+ZukfeIOtSSZn9JBOILcjnoPz4AG\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI5ZjRhNGU4YmUxYzQ5OWMxOTJlMWRiZDVkMDczYmE1NTJhMDZiZmMxMDRhNjA2MTRmZmUwNDVjZWMxZDU5MDg5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNWTnNNL3JRTHFoeUVZcXBIZVdtWHRBRWdzcUNGc0NtT3JXRW0rT2IwYmdBSWdFb3hlU3pXM0ZzL280ZXRLZTZKS3JZeDBvaEluRHRLQmxoYkNBV2s0ZlBRPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KUFowRjNTVUpCWjBsVllVcEdVM2hvY0ZOWFRVOHJXV2xoTkVveFpXZzRhRUp5SzFJd2QwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1hkTlZFMTRUbFJWTlZkb1kwNU5hbGw0VFVSSmQwMVVUWGxPVkZVMVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZYZERCamFVcG9TSFJvVlcxVGIzSkZXVlpHWmtFNWIzTmxkbHBGVERVMFRrMUdUa2dLUjNKaFUwNXVPRlpqZGtaT1pVOU1NazA1SzB4TGNtRm9VSGh6YlhOMWQyd3dkVE5zWTNoTFEzUmtTMmh2Y1cxamNrdFBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN6VDBkRmVRcFpWRWsxV1dwbk5FNXRXbXBaVjBab1RsZFdiRTV0UlRSWk1sSm9Ua2RPYVZsWFVtMU5hbWhzVG1wSmVWbFVUVE5OUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpOUFIwVjVXVlJKTlZscVp6Uk9iVnBxQ2xsWFJtaE9WMVpzVG0xRk5Ga3lVbWhPUjA1cFdWZFNiVTFxYUd4T2FrbDVXVlJOTTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpOUFIwVjVDbGxVU1RWWmFtYzBUbTFhYWxsWFJtaE9WMVpzVG0xRk5Ga3lVbWhPUjA1cFdWZFNiVTFxYUd4T2FrbDVXVlJOTTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPZW1ob1RXMUZlVTlYU1RSUFJGcHRXVEpHYUZsVVZteGFWRnBvVDBkT2ExbFVVbXBaYlVacldtcEpORnBVV1hrS1RXMUZlazU2UVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUlJlVTlFU1RKTmVrMDFUWGs1YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVjFoU1drUkJRVUZDUVUxQlUwUkNSMEZwUlVFMGExcEtkVXR5UTJab1N5dFZNbWRhZVVSMVJrWlJaRVI2Y0ZscmEzTjJTZ3BSVVRWUGFUTk9NV3d3V1VOSlVVUk5LMEp5T0hwSlJHSmFlVkYxVkRSd2QyaHZNbXBJVTJsVFlVUk5aM2hxVTNBMFpYWmFhRlpYYmtKNlFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2IwRkVRbXhCYWtWQmVuaDRWSFpMYTNoWlYzVnFjMlJTVms0eE5uRnhPWFU1UmtaeVRWZENkVTR3Um5RM01IWkdZMFJoU0ZrS2IydFFPQzh4UkhSVk5tdFZTbmxRT0U5cldtMUJha0ZVZWtVcmRWTkNhV3BMUzA5R1ZXbHdXWFJyZEZoUFZqVlBiVGRUTUZablkwRTRlbkJGZG1oS1ZBcFZaRkZhT1c0MVdXbG1WbXcwTVhVM1pFNW9jMnRrY3owS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgGvvpyM5daNrIiQ9KigW43lvD1woe+c5H7EvPv83/0KoCFQD3L1oK+wRy6+QrK/QSKMNEGRsNihgPMjAyNjEwMjAxMzE1NTlaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAyMDEzMTU1OVowLwYJKoZIhvcNAQkEMSIEINyAkk8G8TfnuJF2m0ryaVdbp85BgA5mEd3C4/QuT0sMMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQCFPPrAYXsCL4pef9zarV04KhKP9RSFgnvoqbDig1IrdvGpytpT4hGVzQEV3zEWLSgCMQD84JuotindjTYofPx8aHpL1yG50cu1URXGLfjdP8M+i47x/8ztaAJY3noPr9t9zzE="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"n0pOi+HEmcGS4dvV0HO6VSoGv8EEpgYU/+BFzsHVkIk="}, "signature":"MEUCIQCVNsM/rQLqhyEYqpHeWmXtAEgsqCFsCmOrWEm+Ob0bgAIgEoxeSzW3Fs/o4etKe6JKrYx0ohInDtKBlhbCAWk4fPQ="}}

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrKgAwIBAgIUHC8wcl1kX/JbMbKiE4FVYO+prL8wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDIwMTUxMDQwWhcNMjYxMDIwMTUyMDQwWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEsQRLUB6xoFYle2fcSUsPjnEcCgZmLp9bz2P/5j0EQp46NcS5Y2UddimgWh/Ip1oF4Zepi60nKaK6Lz0rRTiCRqOCBdEwggXNMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChiYmQxM2E4Mjg0MDFiMmZjNjIzMmZiMTNjY2MxMjJiMGQzNDQyZmM1MBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChiYmQxM2E4Mjg0MDFiMmZjNjIzMmZiMTNjY2MxMjJiMGQzNDQyZmM1MBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChiYmQxM2E4Mjg0MDFiMmZjNjIzMmZiMTNjY2MxMjJiMGQzNDQyZmM1MB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoYmJkMTNhODI4NDAxYjJmYzYyMzJmYjEzY2NjMTIyYjBkMzQ0MmZjNTAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTQ0NTI1MTI3My9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGKBgorBgEEAdZ5AgQCBHwEegB4AHYA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhWd1ZQQAABAMARzBFAiAHypITgd6CL+AzKRIiyXZbN8RJQggLIhzqSO4CVV9xbwIhAPtbH3/e4kudi0vclPLu2lu8bJitW8nWdh0hKIuWxZe0MAoGCCqGSM49BAMDA2cAMGQCMFs+unF1TyFOoXIcBvSskpYRG8e4CldJD4Gt1YlVO/ssvGyiDJF5uRDrOxJCLX6M+QIwLaEYxRj2nI7caflU2T8sLAmBqZhZMlxClfIqry2w1//2R9SqqKNsQGy/594prJOD"}, "tlogEntries":[{"logIndex":"3339130426", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792509041", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQCvyujPC0C5/M5HySh4ke5DSmM05dlGmXdqKnY3R51RJwIhAIeOpUMWWZmVeFfbti/ubzk+v9AYz518in38Z97187jh"}, "inclusionProof":{"logIndex":"3217226164", "rootHash":"aAfmDK/36rgiINi+u3H43jsdR+KvBbhhl0Nc0r2I9vQ=", "treeSize":"3217226274", "hashes":["ca5GNalktTsbGurshy30GhzGIP6QDbSSc6K35t2YdKQ=", "tNeFN0AKToulcumZEd+LgxYyuo5Px6DJMPZR8u5NEOw=", "iuAmI+PATWE4twmkvPDrSb9TWpfs+p9SUt6E1ZMOwHM=", "ehzDMfhx2swpET186mXA9gC2zZhApiyG0dB9i2UYw+Y=", "2Oryv3UdAONZselufvuadwJTSQyCZPhkxX5U3A4wFxY=", "DlFBBDkOhxmAmp8EZEwOdsg8PI+V/KLQr8qTu6uX81c=", "nhXquwdTMVAMDps1UKG2FEQt0npG20a+lqdGUM9PQuQ=", "8k4n9R6Kw6arfKtVv1/fuqlusNDOMFt9XEAU/VsBzAg=", "ihjw7czfJ6AUHTN3VRWC7CZIu8lz90V6ecA5BCETGxQ=", "yad744SnxxtV0Yyd6qv47P5yFF+UFzYSVmN+mh8jtVc=", "09pyo+G5mF68+pa7t1psiPGzJhglifTSaw9CbIKMqMU=", "aPYknwzqrIDM5LVynpVZH33H93HdgyR50BOomht+QP8=", "0RgfIOKppEB265lEdj/u+COTsb06TcDX0G42G5Lwe7o=", "dabD48Wgka/nV9/Ur20Wl2IqgXu7sesULelMO7ruTBk=", "blZpLEAd3NmIOAwBvBlK4Ldv3f/OCZsIFySRzXPzpAM=", "bb50n/N0zb4c/IMphDa1DnrDkBgUgubxyoMvGlURzSs=", "XqpeS/BHHOExU96Hvg5L+dHwq0RPJDHCmKTx/jXAg50=", "6Rb7uOgk2v5YKcFSbAPOpLlWAIN1ciDPdlp4J/CUu5k=", "HovTPBWzoErvUg1u5sIi3x1ZonZ7yJekVsBrmpomCIA=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3217226274\naAfmDK/36rgiINi+u3H43jsdR+KvBbhhl0Nc0r2I9vQ=\n\n— rekor.sigstore.dev wNI9ajBEAiB/ah2bQraH05wHYTAzmTtOS/3rOKjmWrjiBU5D2VY/swIgWud4RW8qpqy3L/uGRPONDi7UIPFrrCPEdZVa4jRs+IU=\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiIwOTk0YTYzNDAwNjM3NTFmMjQ0ZThkMmM2OTIxNzNkMzNmNGYwOTMyYTk1ODg1ZDJkYzBmN2Y5YTYwNTUyNjUwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FUUNJSHhvWUcxeklPcldWK1ppL0l2NGdXT3pxWmxUV3dWU29xVWozSWxweERyd0FpQjc4VlFLNDRrM2VBQlVIMUc2WStVOFBzdE1NcWt6K1k0MlFFdDNNUkV6S1E9PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KTFowRjNTVUpCWjBsVlNFTTRkMk5zTVd0WUwwcGlUV0pMYVVVMFJsWlpUeXR3Y2t3NGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1hkTlZGVjRUVVJSZDFkb1kwNU5hbGw0VFVSSmQwMVVWWGxOUkZGM1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZ6VVZKTVZVSTJlRzlHV1d4bE1tWmpVMVZ6VUdwdVJXTkRaMXB0VEhBNVlub3lVQzhLTldvd1JWRndORFpPWTFNMVdUSlZaR1JwYldkWGFDOUpjREZ2UmpSYVpYQnBOakJ1UzJGTE5reDZNSEpTVkdsRFVuRlBRMEprUlhkbloxaE9UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmhwV1cxUmVBcE5Na1UwVFdwbk1FMUVSbWxOYlZwcVRtcEplazF0V21sTlZFNXFXVEpOZUUxcVNtbE5SMUY2VGtSUmVWcHRUVEZOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUdsWmJWRjRUVEpGTkUxcVp6Qk5SRVpwQ2sxdFdtcE9ha2w2VFcxYWFVMVVUbXBaTWsxNFRXcEthVTFIVVhwT1JGRjVXbTFOTVUxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUdsWmJWRjRDazB5UlRSTmFtY3dUVVJHYVUxdFdtcE9ha2w2VFcxYWFVMVVUbXBaTWsxNFRXcEthVTFIVVhwT1JGRjVXbTFOTVUxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlaYlVwclRWUk9hRTlFU1RST1JFRjRXV3BLYlZsNldYbE5la3B0V1dwRmVsa3lUbXBOVkVsNVdXcENhMDE2VVRBS1RXMWFhazVVUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUlJNRTVVU1RGTlZFa3pUWGs1YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTENrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NIZEZaV2RDTkVGSVdVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVjJReFdsRlJRVUZDUVUxQlVucENSa0ZwUVVoNWNFbFVaMlEyUTB3clFYcExVa2xwZVZoYVlrNDRVa3BSWjJkTVNXaDZjUXBUVHpSRFZsWTVlR0ozU1doQlVIUmlTRE12WlRScmRXUnBNSFpqYkZCTWRUSnNkVGhpU21sMFZ6aHVWMlJvTUdoTFNYVlhlRnBsTUUxQmIwZERRM0ZIQ2xOTk5EbENRVTFFUVRKalFVMUhVVU5OUm5NcmRXNUdNVlI1Ums5dldFbGpRblpUYzJ0d1dWSkhPR1UwUTJ4a1NrUTBSM1F4V1d4V1R5OXpjM1pIZVdrS1JFcEdOWFZTUkhKUGVFcERURmcyVFN0UlNYZE1ZVVZaZUZKcU1tNUpOMk5oWm14Vk1sUTRjMHhCYlVKeFdtaGFUV3g0UTJ4bVNYRnllVEozTVM4dk1ncFNPVk54Y1V0T2MxRkhlUzgxT1RSd2NrcFBSQW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyzADAgEAMIICwgYJKoZIhvcNAQcCoIICszCCAq8CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgvs31/7DRRm4Abken+kCChhxSnCcdrFfS+RhRBNwyTe4CFQD29NJRSvRnC2TNE2dAR1RU1ya64RgPMjAyNjEwMjAxNTEwNDFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHcMIIB2AIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAyMDE1MTA0MVowLwYJKoZIhvcNAQkEMSIEIKLJxH8rfpuw1xi5MT7pRVIDUGbqpHyi2bDmoevTMIOqMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRoMGYCMQDxhU45xIFeO/VwOLA1gA4oONgxXs+6y9+N+ogLV+cf69Tn8knNkEyRO+4X3UcscEICMQCCpEa6UXRjq0I/+2CFCZLXTSzLGuhSoMDkvHFWZtcbqYrF8xw7RaDqK7xcJ4cxpLE="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"CZSmNABjdR8kTo0saSFz0z9PCTKpWIXS3A9/mmBVJlA="}, "signature":"MEQCIHxoYG1zIOrWV+Zi/Iv4gWOzqZlTWwVSoqUj3IlpxDrwAiB78VQK44k3eABUH1G6Y+U8PstMMqkz+Y42QEt3MREzKQ=="}}

build/torch-xpu/moe.py CHANGED Viewed

@@ -58,8 +58,9 @@ from .compat import (
 )
 from .formats import get_supported_act_fns, is_mx, is_mxfp4, is_per_expert_global, mx_scale_family, normalize_global_scale, ue8m0_as_uint8, weight_format
 from .norm import norm_column_factor, rms_inv_rows, rms_norm_rows
 from .quant import _launch_act_quant, MX_ACT_QUANT, mx_act_quant_routed_scales
-from .scheduling import compute_grouped_scheduling, is_off_rank, off_rank
 from .epilogue import fused_glu
@@ -245,6 +246,23 @@ def _torch_weighted_reduce(down_out, top_k_index, top_k_weights, num_experts):
     return contrib.view(num_tokens, num_top_k, down_out.size(1)).sum(dim=1)
 def _host_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit, gate):
     """The activation on the host, for the arms the epilogue does not fuse: a callable takes the
     raw GEMM output (interleaved gate|up columns when ``gate``) and returns the intermediate —
@@ -345,6 +363,10 @@ def moe_fused_grouped(
     down_proj: torch.Tensor,  # (E, H, I)
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
@@ -379,17 +401,24 @@ def moe_fused_grouped(
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
-    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
-    fmt = _block_format(
-        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
-    )
     num_top_k = top_k_index.size(-1)
-    glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate)
-    fused = glu["act_fn"] is not None
     NUM_EXPERTS = gate_up_proj.size(0)
     expert_start, gather_idx, scatter_idx = compute_grouped_scheduling(
         top_k_index, NUM_EXPERTS, num_top_k
     )
     # Phase 1: gate_up + SiLU + requant in the block format -> expert-ordered quantized
     # intermediate. scatter_idx=None: the down reads the intermediate in place. (C, Cs) under a
@@ -462,6 +491,10 @@ def moe_fused_batched(
     down_proj: torch.Tensor,  # (E, H, I)
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
@@ -498,14 +531,23 @@ def moe_fused_batched(
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
-    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
     fmt = _block_format(
         gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
     )
     glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate)
     fused = glu["act_fn"] is not None
-    NUM_EXPERTS = gate_up_proj.size(0)
-    expert_ids = top_k_index.reshape(-1)
     gather_idx = _gather_idx(top_k_index)
     # Phase 1: gate_up + SiLU + requant in the block format -> per-row quantized
@@ -587,6 +629,10 @@ def moe_unfused_grouped(
     down_proj: torch.Tensor,
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
@@ -617,16 +663,22 @@ def moe_unfused_grouped(
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
-    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
-    fmt = _block_format(
-        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
-    )
     num_top_k = top_k_index.size(-1)
     NUM_EXPERTS = gate_up_proj.size(0)
     expert_start, gather_idx, scatter_idx = compute_grouped_scheduling(
         top_k_index, NUM_EXPERTS, num_top_k
     )
     # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> expert-ordered (S, 2I).
     gate_up_in, gate_up_in_scale = _gate_up_activations(
@@ -866,6 +918,10 @@ def moe_unfused_batched(
     down_proj: torch.Tensor,
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
@@ -895,12 +951,21 @@ def moe_unfused_batched(
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
-    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
     fmt = _block_format(
         gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
     )
-    NUM_EXPERTS = gate_up_proj.size(0)
-    expert_ids = top_k_index.reshape(-1)
     gather_idx = _gather_idx(top_k_index)
     # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> (S, 2I).

 )
 from .formats import get_supported_act_fns, is_mx, is_mxfp4, is_per_expert_global, mx_scale_family, normalize_global_scale, ue8m0_as_uint8, weight_format
 from .norm import norm_column_factor, rms_inv_rows, rms_norm_rows
+from .loading.scales import csf_weight_scales
 from .quant import _launch_act_quant, MX_ACT_QUANT, mx_act_quant_routed_scales
+from .scheduling import compute_grouped_scheduling, is_off_rank, off_rank, routed_expert_start
 from .epilogue import fused_glu
     return contrib.view(num_tokens, num_top_k, down_out.size(1)).sum(dim=1)
+def _csf_weight_scales(
+    gate_up_proj_scale_inv, gate_up_proj_scale_exceptions, gate_up_proj_scale_exception_offsets,
+    down_proj_scale_inv, down_proj_scale_exceptions, down_proj_scale_exception_offsets, expert_start,
+):
+    """Both projections' weight scales for the GEMMs: a CSF fixed stream (given its exceptions)
+    decodes, only the experts with routed rows in ``expert_start``; dense scales pass through."""
+    return (
+        csf_weight_scales(
+            gate_up_proj_scale_inv, gate_up_proj_scale_exceptions, gate_up_proj_scale_exception_offsets,
+            expert_start,
+        ),
+        csf_weight_scales(
+            down_proj_scale_inv, down_proj_scale_exceptions, down_proj_scale_exception_offsets, expert_start
+        ),
+    )
 def _host_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit, gate):
     """The activation on the host, for the arms the epilogue does not fuse: a callable takes the
     raw GEMM output (interleaved gate|up columns when ``gate``) and returns the intermediate —
     down_proj: torch.Tensor,  # (E, H, I)
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_scale_exceptions: torch.Tensor | None = None,
+    gate_up_proj_scale_exception_offsets: torch.Tensor | None = None,
+    down_proj_scale_exceptions: torch.Tensor | None = None,
+    down_proj_scale_exception_offsets: torch.Tensor | None = None,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.
+    ``*_proj_scale_exceptions``/``*_proj_scale_exception_offsets`` mark that projection's
+    ``*_proj_scale_inv`` as a CSF-compressed NVFP4 fixed stream (``prepare_nvfp4_csf_scales``);
+    the block decodes it, only the routed experts, before its GEMMs."""
     num_top_k = top_k_index.size(-1)
     NUM_EXPERTS = gate_up_proj.size(0)
     expert_start, gather_idx, scatter_idx = compute_grouped_scheduling(
         top_k_index, NUM_EXPERTS, num_top_k
     )
+    gate_up_proj_scale_inv, down_proj_scale_inv = _csf_weight_scales(
+        gate_up_proj_scale_inv, gate_up_proj_scale_exceptions, gate_up_proj_scale_exception_offsets,
+        down_proj_scale_inv, down_proj_scale_exceptions, down_proj_scale_exception_offsets, expert_start,
+    )
+    fmt = _block_format(
+        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
+    )
+    glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate)
+    fused = glu["act_fn"] is not None
     # Phase 1: gate_up + SiLU + requant in the block format -> expert-ordered quantized
     # intermediate. scatter_idx=None: the down reads the intermediate in place. (C, Cs) under a
     down_proj: torch.Tensor,  # (E, H, I)
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_scale_exceptions: torch.Tensor | None = None,
+    gate_up_proj_scale_exception_offsets: torch.Tensor | None = None,
+    down_proj_scale_exceptions: torch.Tensor | None = None,
+    down_proj_scale_exception_offsets: torch.Tensor | None = None,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.
+    ``*_proj_scale_exceptions``/``*_proj_scale_exception_offsets`` mark that projection's
+    ``*_proj_scale_inv`` as a CSF-compressed NVFP4 fixed stream (``prepare_nvfp4_csf_scales``);
+    the block decodes it, only the routed experts, before its GEMMs."""
+    NUM_EXPERTS = gate_up_proj.size(0)
+    expert_ids = top_k_index.reshape(-1)
+    if gate_up_proj_scale_exceptions is not None or down_proj_scale_exceptions is not None:
+        gate_up_proj_scale_inv, down_proj_scale_inv = _csf_weight_scales(
+            gate_up_proj_scale_inv, gate_up_proj_scale_exceptions, gate_up_proj_scale_exception_offsets,
+            down_proj_scale_inv, down_proj_scale_exceptions, down_proj_scale_exception_offsets,
+            routed_expert_start(expert_ids, NUM_EXPERTS),  # decode only the routed experts
+        )
     fmt = _block_format(
         gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
     )
     glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate)
     fused = glu["act_fn"] is not None
     gather_idx = _gather_idx(top_k_index)
     # Phase 1: gate_up + SiLU + requant in the block format -> per-row quantized
     down_proj: torch.Tensor,
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_scale_exceptions: torch.Tensor | None = None,
+    gate_up_proj_scale_exception_offsets: torch.Tensor | None = None,
+    down_proj_scale_exceptions: torch.Tensor | None = None,
+    down_proj_scale_exception_offsets: torch.Tensor | None = None,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.
+    ``*_proj_scale_exceptions``/``*_proj_scale_exception_offsets`` mark that projection's
+    ``*_proj_scale_inv`` as a CSF-compressed NVFP4 fixed stream (``prepare_nvfp4_csf_scales``);
+    the block decodes it, only the routed experts, before its GEMMs."""
     num_top_k = top_k_index.size(-1)
     NUM_EXPERTS = gate_up_proj.size(0)
     expert_start, gather_idx, scatter_idx = compute_grouped_scheduling(
         top_k_index, NUM_EXPERTS, num_top_k
     )
+    gate_up_proj_scale_inv, down_proj_scale_inv = _csf_weight_scales(
+        gate_up_proj_scale_inv, gate_up_proj_scale_exceptions, gate_up_proj_scale_exception_offsets,
+        down_proj_scale_inv, down_proj_scale_exceptions, down_proj_scale_exception_offsets, expert_start,
+    )
+    fmt = _block_format(
+        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
+    )
     # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> expert-ordered (S, 2I).
     gate_up_in, gate_up_in_scale = _gate_up_activations(
     down_proj: torch.Tensor,
     gate_up_proj_scale_inv: torch.Tensor,
     down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_scale_exceptions: torch.Tensor | None = None,
+    gate_up_proj_scale_exception_offsets: torch.Tensor | None = None,
+    down_proj_scale_exceptions: torch.Tensor | None = None,
+    down_proj_scale_exception_offsets: torch.Tensor | None = None,
     gate_up_proj_weight_global_scale: torch.Tensor | None = None,
     down_proj_weight_global_scale: torch.Tensor | None = None,
     gate_up_proj_input_global_scale: torch.Tensor | None = None,
     ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
     callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
     projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine.
+    ``*_proj_scale_exceptions``/``*_proj_scale_exception_offsets`` mark that projection's
+    ``*_proj_scale_inv`` as a CSF-compressed NVFP4 fixed stream (``prepare_nvfp4_csf_scales``);
+    the block decodes it, only the routed experts, before its GEMMs."""
+    NUM_EXPERTS = gate_up_proj.size(0)
+    expert_ids = top_k_index.reshape(-1)
+    if gate_up_proj_scale_exceptions is not None or down_proj_scale_exceptions is not None:
+        gate_up_proj_scale_inv, down_proj_scale_inv = _csf_weight_scales(
+            gate_up_proj_scale_inv, gate_up_proj_scale_exceptions, gate_up_proj_scale_exception_offsets,
+            down_proj_scale_inv, down_proj_scale_exceptions, down_proj_scale_exception_offsets,
+            routed_expert_start(expert_ids, NUM_EXPERTS),  # decode only the routed experts
+        )
     fmt = _block_format(
         gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
     )
     gather_idx = _gather_idx(top_k_index)
     # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> (S, 2I).

build/torch-xpu/scheduling.py CHANGED Viewed

@@ -86,13 +86,15 @@ def _exclusive_offsets_kernel(
     ExpertFreq, ExpertStart, Counters, NUM_EXPERTS: tl.constexpr
 ):
     """Exclusive cumsum of per-expert token counts → ``expert_start`` (leading 0, trailing the
-    local count), and zero the scatter counters, the EP-sentinel bucket's included — one launch."""
     offs = tl.arange(0, NUM_EXPERTS)
     incl = tl.cumsum(tl.load(ExpertFreq + offs), 0)
     tl.store(ExpertStart, 0)
     tl.store(ExpertStart + 1 + offs, incl)
-    tl.store(Counters + offs, tl.zeros([NUM_EXPERTS], tl.int32))
-    tl.store(Counters + NUM_EXPERTS, 0)
 @triton.jit
@@ -179,20 +181,25 @@ def compute_grouped_scheduling(
     return expert_start, gather_idx, scatter_idx
-@compile_time_only_triton_op(
-    add_op_namespace_prefix("compute_grouped_scheduling"), mutates_args=(), opaque=True
-)
-def _compute_grouped_scheduling(
-    expert_ids: torch.Tensor, num_experts: int, num_top_k: int
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     device = expert_ids.device
-    expert_ids = expert_ids.int().contiguous()  # routing kernels index with unit stride
     num_routed_tokens = expert_ids.numel()  # S = num_tokens * num_top_k
     expert_freq = torch.zeros(num_experts, dtype=torch.int32, device=device)
     expert_start = torch.empty(num_experts + 1, dtype=torch.int32, device=device)
-    counters = torch.empty(num_experts + 1, dtype=torch.int32, device=device)
-    perm = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
-    perm_token = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
     with device_context(device):
         compile_time_only_triton_wrap(_count_kernel)[
             (triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)
@@ -209,6 +216,30 @@ def _compute_grouped_scheduling(
             counters,
             NUM_EXPERTS=num_experts,
         )
         compile_time_only_triton_wrap(_scatter_kernel)[
             (triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)
         ](

     ExpertFreq, ExpertStart, Counters, NUM_EXPERTS: tl.constexpr
 ):
     """Exclusive cumsum of per-expert token counts → ``expert_start`` (leading 0, trailing the
+    local count), and zero the scatter counters, the EP-sentinel bucket's included — one launch.
+    ``Counters`` None (no scatter follows) skips them."""
     offs = tl.arange(0, NUM_EXPERTS)
     incl = tl.cumsum(tl.load(ExpertFreq + offs), 0)
     tl.store(ExpertStart, 0)
     tl.store(ExpertStart + 1 + offs, incl)
+    if Counters is not None:
+        tl.store(Counters + offs, tl.zeros([NUM_EXPERTS], tl.int32))
+        tl.store(Counters + NUM_EXPERTS, 0)
 @triton.jit
     return expert_start, gather_idx, scatter_idx
+def routed_expert_start(expert_ids: torch.Tensor, num_experts: int) -> torch.Tensor | None:
+    """``expert_start`` of an unsorted (batched) routing: its consumers only need which experts
+    have routed rows, to skip per-expert work on the rest, so only the count and offset launches
+    run (no sort). ``None`` (every expert) when the expert count is not a power of 2, which the
+    scheduling kernels cannot hold."""
+    if num_experts & (num_experts - 1):
+        return None
+    return _compute_expert_start(expert_ids, num_experts)
+def _launch_expert_start(
+    expert_ids: torch.Tensor, num_experts: int, counters: torch.Tensor | None = None
+) -> torch.Tensor:
+    """The count + exclusive-offset launches → ``expert_start``; ``counters`` (the scatter's
+    per-expert cursors, when a sort follows) is zeroed in the same launch."""
     device = expert_ids.device
     num_routed_tokens = expert_ids.numel()  # S = num_tokens * num_top_k
     expert_freq = torch.zeros(num_experts, dtype=torch.int32, device=device)
     expert_start = torch.empty(num_experts + 1, dtype=torch.int32, device=device)
     with device_context(device):
         compile_time_only_triton_wrap(_count_kernel)[
             (triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)
             counters,
             NUM_EXPERTS=num_experts,
         )
+    return expert_start
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("compute_expert_start"), mutates_args=(), opaque=True
+)
+def _compute_expert_start(expert_ids: torch.Tensor, num_experts: int) -> torch.Tensor:
+    return _launch_expert_start(expert_ids.int().contiguous(), num_experts)
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("compute_grouped_scheduling"), mutates_args=(), opaque=True
+)
+def _compute_grouped_scheduling(
+    expert_ids: torch.Tensor, num_experts: int, num_top_k: int
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = expert_ids.device
+    expert_ids = expert_ids.int().contiguous()  # routing kernels index with unit stride
+    num_routed_tokens = expert_ids.numel()  # S = num_tokens * num_top_k
+    counters = torch.empty(num_experts + 1, dtype=torch.int32, device=device)
+    expert_start = _launch_expert_start(expert_ids, num_experts, counters)
+    perm = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
+    perm_token = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
+    with device_context(device):
         compile_time_only_triton_wrap(_scatter_kernel)[
             (triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)
         ](

build/torch-xpu/symbols.json CHANGED Viewed

@@ -1,13 +1,13 @@
 {
   "schema_version": 1,
-  "module": "_finegrained_kernels_xpu_78a2a29",
   "functions": [
     {
       "name": "matmul_2d",
       "qualname": "matmul_2d",
-      "module": "_finegrained_kernels_xpu_78a2a29.matmul",
       "kind": "function",
-      "doc": "Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of\n``matmul_grouped``/``matmul_batched``, taking the same ``As``/``Bs`` scale spec,\n``activation_format`` / gate|up fusion kwargs, and inferred quant block (no ``block_size`` argument — it\nfalls out of the weight-scale shape, so the data can't disagree with a parameter).\n\n``As`` is the activation block-scale spec, same as the routed ops: ``None`` → the op quantizes raw\n``A`` dynamically; a per-tensor scalar → static (calibrated) FP8 activation scale; block scales →\npre-quantized ``A``. ``Bs`` is the weight block scale. ``a_global_scale``/``b_global_scale`` are the\ntwo-level NVFP4 per-tensor second-level scales (calibrated ``input_scale`` ``g_a`` and\n``weight_scale_2`` ``g_b``) — provided, never computed; the op folds ``g_a·g_b`` onto the\naccumulator (``a_global_scale`` rides raw ``A`` or a pre-quantized ``As`` alike).\n``output_global_scale`` (NVFP4 requant only) is the NEXT proj's provided ``input_scale``:\nthe fused requant normalizes the GLU intermediate by it, returning ``(C, Cs)`` the down-proj\nconsumes as ``As=Cs, a_global_scale=output_global_scale``.\n\n``activation_format`` is the activations' quantization: ``None`` (default) = the weights' own\nformat, ``\"bf16\"`` = weight-only (no activation quant), or an explicit ``\"fp8\"`` / ``\"mxfp8\"`` /\n``\"mxfp4\"`` / ``\"nvfp4\"`` (``\"mxfp8\"`` on MXFP4 weights = the W4A8 chain; ``check_activation_format``\nhas the per-family matrix). ``gate=True`` reads ``B`` as the interleaved ``(2N, K)`` gate|up weight and\nfuses the ``act_fn`` GLU (``swiglu_alpha``/``swiglu_limit`` its clamped variant), returning the\n``[..., N]`` intermediate (``act_fn=None``: the raw ``[..., 2N]`` gate|up output instead); ``quantize_output=True`` (MX weights only) stores that intermediate in\n``activation_format`` and returns ``(C, Cs)`` for the next op.\n\nRoutes by weight dtype + inferred block: ``Bs`` None → full-precision; MX/NVFP4 (int8/E4M3 weights\nwith UE8M0 group-32 or E4M3 group-16 ``Bs``) → ``mx_dynamic_matmul``; tensor-wide FP8 (scalar\n``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``; block FP8 → the static variant when ``As`` is a\nper-tensor scalar, else dynamic.",
       "signature": {
         "parameters": [
           {
@@ -100,6 +100,18 @@
             "default": "None",
             "annotation": "torch.Tensor | None"
           },
           {
             "name": "output_global_scale",
             "kind": "KEYWORD_ONLY",
@@ -113,9 +125,9 @@
     {
       "name": "matmul_batched",
       "qualname": "matmul_batched",
-      "module": "_finegrained_kernels_xpu_78a2a29.batched",
       "kind": "function",
-      "doc": "Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one\nprogram per routed row (``expert_ids`` gives its expert).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3 offline for bd/tensor, inline for MX; ``\"bf16\"`` = weight-only;\n``\"mxfp4\"`` = packed E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue; ``act_fn=None`` with ``gate`` skips it, returning the raw interleaved gate|up\noutput for a host-side activation.\n``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4\nsecond-level scales ride the separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor,\nweights per-expert ``(E,)``; from ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b``\nonto the accumulator. The activation global ``g_a`` is CALIBRATED (the checkpoint's\n``input_scale``): ``a_global_scale=g_a`` with a raw ``A`` has the op quantize ``A / g_a`` per\nblock, and rides a pre-quantized ``As`` the same way. Under an NVFP4 requant the\nfused requant normalizes the GLU intermediate by the PROVIDED ``output_global_scale`` (the next\nproj's calibrated ``input_scale``) before the block quant and returns ``[C, Cs]``; the down\nconsumes it as ``As=Cs, a_global_scale=output_global_scale``. ``gather_idx``/``scatter_idx`` (each None or a ``(S,)`` map)\naddress the source row of ``A`` / destination row of ``C`` per program — None means row\n``s``; the gather lets the gate_up read unexpanded activations (source row\n``s // num_top_k``) with no copy.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_batched`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` → ``mx_dynamic_matmul_batched``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_batched``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``.",
       "signature": {
         "parameters": [
           {
@@ -226,6 +238,18 @@
             "default": "None",
             "annotation": "torch.Tensor | None"
           },
           {
             "name": "output_global_scale",
             "kind": "KEYWORD_ONLY",
@@ -239,9 +263,9 @@
     {
       "name": "matmul_grouped",
       "qualname": "matmul_grouped",
-      "module": "_finegrained_kernels_xpu_78a2a29.grouped",
       "kind": "function",
-      "doc": "Grouped matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision).\n``expert_start`` is the ``(E+1,)`` tiling schedule from one ``compute_grouped_scheduling``\npass, shared by every grouped GEMM of the layer (the expert sort is virtual — nothing is\nphysically permuted).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3, or the fp4 chains; ``\"bf16\"`` = weight-only; ``\"mxfp4\"`` =\npacked E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue; ``act_fn=None`` with ``gate`` skips it, returning the raw interleaved gate|up\noutput for a host-side activation.\n``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4 second-level scales ride\nthe separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor, weights per-expert ``(E,)``;\nfrom ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b`` onto the accumulator. The\nactivation global ``g_a`` is CALIBRATED (the checkpoint's ``input_scale``): ``a_global_scale=g_a``\nwith a raw ``A`` has the op quantize ``A / g_a`` per block, and rides a pre-quantized ``As`` the\nsame way. Under an NVFP4 requant the fused requant normalizes the GLU intermediate by the\nPROVIDED ``output_global_scale`` (the next proj's calibrated ``input_scale``) before the block\nquant and returns ``[C, Cs]``; the down consumes it as ``As=Cs, a_global_scale=output_global_scale``.\nRow order is carried by the standalone maps: ``gather_idx`` gathers ``A`` (``None`` ->\nalready expert-ordered), ``scatter_idx`` scatters the output. The fused MoE chain is one\nscheduling pass: gate_up with ``scatter_idx=None`` + ``gate=True`` + ``quantize_output=True``,\nthen down with ``gather_idx=None`` and the\nintermediate's scales as ``As``. EP-sentinel routes sit past ``expert_start[-1]``, which\nno tile reaches.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_grouped`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` ��� ``mx_dynamic_matmul_grouped``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_grouped``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_grouped``.",
       "signature": {
         "parameters": [
           {
@@ -352,6 +376,18 @@
             "default": "None",
             "annotation": "torch.Tensor | None"
           },
           {
             "name": "output_global_scale",
             "kind": "KEYWORD_ONLY",
@@ -365,9 +401,9 @@
     {
       "name": "moe_fused_batched",
       "qualname": "moe_fused_batched",
-      "module": "_finegrained_kernels_xpu_78a2a29.moe",
       "kind": "function",
-      "doc": "Fused batched MoE (decode): gate_up + SiLU + requant epilogue → per-row quantized\nintermediate → batched down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4 — decode runs the software/swap arms\nbelow the native mxf4nvf4 M=128 staging); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"`` runs the all-fp4 W4A4 chain, ``\"bf16\"`` is weight-only); ``None``\nfollows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each\nstep through the activation dtype so the output matches the unfused reference to\nreduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.",
       "signature": {
         "parameters": [
           {
@@ -412,6 +448,30 @@
             "default": null,
             "annotation": "torch.Tensor"
           },
           {
             "name": "gate_up_proj_weight_global_scale",
             "kind": "POSITIONAL_OR_KEYWORD",
@@ -521,9 +581,9 @@
     {
       "name": "moe_fused_grouped",
       "qualname": "moe_fused_grouped",
-      "module": "_finegrained_kernels_xpu_78a2a29.moe",
       "kind": "function",
-      "doc": "Fused grouped MoE (prefill): gather gate_up + SiLU + requant epilogue → quantized\nexpert-ordered intermediate → grouped down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"``/``\"nvfp4\"`` run all-fp4 W4A4 chains, ``\"bf16\"`` is weight-only);\n``None`` follows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each step through\nthe activation dtype so the output matches the unfused reference to reduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.",
       "signature": {
         "parameters": [
           {
@@ -568,6 +628,30 @@
             "default": null,
             "annotation": "torch.Tensor"
           },
           {
             "name": "gate_up_proj_weight_global_scale",
             "kind": "POSITIONAL_OR_KEYWORD",
@@ -677,7 +761,7 @@
     {
       "name": "embedding_lookup",
       "qualname": "embedding_lookup",
-      "module": "_finegrained_kernels_xpu_78a2a29.embedding",
       "kind": "function",
       "doc": "Gather and dequantize selected rows without expanding the embedding table.\n\n``weight_format`` selects ``fp8``, ``nvfp4``, ``mxfp8``, ``mxfp4``, ``int8``, or ``int4``.\nBy default, a global scale selects NVFP4, signed int8 weights select INT8, and otherwise\nFP8 is used. FP8 accepts scalar, per-row, or square block scales (``block_size``).\nNVFP4 uses E4M3 group-16 scales and one FP32 multiplier; MX formats use UE8M0 group-32\nscales. FP4 packs adjacent columns, low nibble first. INT8 uses per-row scales; INT4\npacks adjacent vocabulary rows and uses per-row column groups inferred from its scale\ngrid. Integer values are signed, symmetric, and have no zero point.\n\nOutputs default to the floating scale dtype, or bf16 for NVFP4/MX. Table/grid strides\nmay describe TP slices. Indices follow ``torch.nn.functional.embedding``: negative and\nout-of-range values fail; padding indices look up their stored row. Empty, repeated,\nand noncontiguous indices work.",
       "signature": {
@@ -731,7 +815,7 @@
     {
       "name": "get_supported_act_fns",
       "qualname": "get_supported_act_fns",
-      "module": "_finegrained_kernels_xpu_78a2a29.formats",
       "kind": "function",
       "doc": "The activation names the gate|up epilogue fuses in-kernel; any other activation is applied\non the host by the MoE forwards (``act_fn`` as a callable).",
       "signature": {
@@ -742,7 +826,7 @@
     {
       "name": "get_supported_norms",
       "qualname": "get_supported_norms",
-      "module": "_finegrained_kernels_xpu_78a2a29.norm",
       "kind": "function",
       "doc": "The norm names the kernels fuse (``post_expert_norm``), each ``x * factor * rsqrt(mean +\neps)`` over the row: ``\"rms_norm\"`` scales the normalized row by ``weight``,\n``\"centered_rms_norm\"`` by ``1 + weight`` (the zero-centered parameterization), and\n``\"input_scaled_rms_norm\"`` scales by ``1 + weight`` BEFORE normalizing, so the mean square\nis taken on the scaled row. Anything else stays a host callable.",
       "signature": {
@@ -753,7 +837,7 @@
     {
       "name": "rms_norm_rows",
       "qualname": "rms_norm_rows",
-      "module": "_finegrained_kernels_xpu_78a2a29.norm",
       "kind": "function",
       "doc": "RMS-normalize ``(S, H)`` rows in one launch, in the ``get_supported_norms()`` form named by\n``norm`` (fp32 accumulate, input dtype out). The MoE chain folds this into\n``weighted_reduce`` instead; this is the standalone form.",
       "signature": {
@@ -789,7 +873,7 @@
     {
       "name": "swizzle_mx_scales",
       "qualname": "swizzle_mx_scales",
-      "module": "_finegrained_kernels_xpu_78a2a29.swizzle",
       "kind": "function",
       "doc": "Reorder a block-scale tensor into the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05\nscaled-MMA consumes, one triton launch per matrix (``_swizzle_scales_kernel`` — no torch\npermute/transpose copies). The scale values are unchanged — this only rearranges them from\nplain row-major into the swizzled order the tensor core reads directly (the same layout\ncuBLAS/CUTLASS require for MXFP8/NVFP4 ``scaled_mm``). Plain row-major forces a gather that\ncaps the scaled dot below the fp8/fp4 peak; the swizzle removes it. Apply to weight scales\nonce at model load; the ops accept the result directly.\n\n``scale`` is the per-block scale grid — UE8M0 (``float8_e8m0fnu``, group-32 MXFP8/MXFP4) or\nE4M3 (``float8_e4m3fn``, group-16 NVFP4); the reorder is dtype-agnostic (moves bytes):\n\n- 2D ``(rows, K // group)``: one matrix (a dense weight's scale).\n- 3D ``(E, rows, K // group)``: an expert stack, each expert swizzled independently so its\n  blocks stay ``ceil(rows/128)``-aligned (the reader indexes expert ``e`` at\n  ``e * ceil(rows/128)``; byte-identical to a flat swizzle when ``rows % 128 == 0``).\n\n``gather_idx`` (2D only): a 1-D ``(padded_rows,)`` index mapping each output (sorted) row to\nits source row in ``scale`` (``-1`` = padding → zero row), folded into the kernel's load — the\nrouted/expert-sorted, per-tile-padded layout a grouped GEMM reads affine per BM=128 tile.\n``padded_rows`` must be a multiple of 128.\n\n``rows``/``cols`` are zero-padded to (128, 4) multiples; returns the 5D\n``(G, row_blocks, ceil(cols/4), 2, 256)`` artifact the ops read — ``G`` the expert count (1 for\na matrix), so an expert stack shards/gathers on its leading dim. Bit-identical to CUTLASS's\npacker (verified).",
       "signature": {
@@ -813,7 +897,7 @@
     {
       "name": "unswizzle_mx_scales",
       "qualname": "unswizzle_mx_scales",
-      "module": "_finegrained_kernels_xpu_78a2a29.swizzle",
       "kind": "function",
       "doc": "Inverse of ``swizzle_mx_scales``: a SWIZZLE_32_4_4 artifact back to the plain row-major\n``(rows, cols)`` grid — or ``(num_experts, rows, cols)`` for an expert stack — with the\n(128, 4) padding trimmed. Pure reshape/permute on the byte view (no kernel); it runs at\nsave/dequantize time, never on a forward, so the torch-level cost is irrelevant.\n\n``rows``/``cols`` are the LOGICAL dims of ONE matrix (for a gate|up slab, the full ``2N``);\nboth are trimmed back out of the (128, 4) padding. ``num_experts`` selects the stacked\n``(E, rows, cols)`` result and must match the artifact's leading dim.",
       "signature": {
@@ -846,10 +930,76 @@
         "return_annotation": "torch.Tensor"
       }
     },
     {
       "name": "mxfp8_act_quant",
       "qualname": "mxfp8_act_quant",
-      "module": "_finegrained_kernels_xpu_78a2a29.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MX once (E4M3 values + UE8M0 group-32 uint8 scales)\ninstead of inline per weight-tile — the fused gate_up re-ran the inline quant per N-tile\n(16x redundant amax/convert ALU + 2x act bytes), which held it at ~380 TFLOPS while the\npre-quantized down kernel ran at ~1080. One pass costs ~50µs at 8k tokens. Bit-exact with\nthe inline form (same group boundaries). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path (same dense grid, per-element store).",
       "signature": {
@@ -873,7 +1023,7 @@
     {
       "name": "mxfp4_act_quant",
       "qualname": "mxfp4_act_quant",
-      "module": "_finegrained_kernels_xpu_78a2a29.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MXFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8, two codes per byte, first value in the low nibble) + UE8M0\ngroup-32 uint8 scales (``(T, K//32)``, amax/6 ceil'd to a power of two). Bit-identical\nto the fused epilogues' inline form (shared ``mx_act_quant_inline`` arm). Feeds the\nW4A4 arm of the MX matmul ops; quantizing activations to fp4 at runtime is an\naccuracy call the caller owns — the ops never do it implicitly. ``swizzled=True`` emits the\nscale directly in SWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {
@@ -897,7 +1047,7 @@
     {
       "name": "nvfp4_act_quant",
       "qualname": "nvfp4_act_quant",
-      "module": "_finegrained_kernels_xpu_78a2a29.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to NVFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8) + E4M3 group-16 block scales (``(T, K//16)`` — ``amax/6`` rounded to\nE4M3, NOT a power of two; values divide by the DECODED scale before the E2M1 grid,\nthe standard two-step). ``global_scale`` is the CALIBRATED per-tensor second level\n(fp32, the checkpoint's ``input_scale``): values are normalized by it before\nthe block quant, so block scales stay in e4m3 range regardless of the activation's\ndynamic range — the canonical two-level format. The GEMM folds ``g_a·g_b`` back onto\nthe accumulator (pass ``As=scales`` with ``a_global_scale=global_scale``). Per-expert globals\nneed each row's expert, which ``expert_index`` gives where the caller already holds that map\n(the batched op's ``expert_ids``, sentinels and all) — it is read only when the globals are\nper expert, and ``mx_act_quant_grouped`` needs none at all, its tiles carrying an expert.\n``None`` = single-level (``g_a = 1``). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {

 {
   "schema_version": 1,
+  "module": "_finegrained_kernels_xpu_bbd13a8",
   "functions": [
     {
       "name": "matmul_2d",
       "qualname": "matmul_2d",
+      "module": "_finegrained_kernels_xpu_bbd13a8.matmul",
       "kind": "function",
+      "doc": "Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of\n``matmul_grouped``/``matmul_batched``, taking the same ``As``/``Bs`` scale spec,\n``activation_format`` / gate|up fusion kwargs, and inferred quant block (no ``block_size`` argument — it\nfalls out of the weight-scale shape, so the data can't disagree with a parameter).\n\n``As`` is the activation block-scale spec, same as the routed ops: ``None`` → the op quantizes raw\n``A`` dynamically; a per-tensor scalar → static (calibrated) FP8 activation scale; block scales →\npre-quantized ``A``. ``Bs`` is the weight block scale. ``a_global_scale``/``b_global_scale`` are the\ntwo-level NVFP4 per-tensor second-level scales (calibrated ``input_scale`` ``g_a`` and\n``weight_scale_2`` ``g_b``) — provided, never computed; the op folds ``g_a·g_b`` onto the\naccumulator (``a_global_scale`` rides raw ``A`` or a pre-quantized ``As`` alike).\n``output_global_scale`` (NVFP4 requant only) is the NEXT proj's provided ``input_scale``:\nthe fused requant normalizes the GLU intermediate by it, returning ``(C, Cs)`` the down-proj\nconsumes as ``As=Cs, a_global_scale=output_global_scale``.\n\n``activation_format`` is the activations' quantization: ``None`` (default) = the weights' own\nformat, ``\"bf16\"`` = weight-only (no activation quant), or an explicit ``\"fp8\"`` / ``\"mxfp8\"`` /\n``\"mxfp4\"`` / ``\"nvfp4\"`` (``\"mxfp8\"`` on MXFP4 weights = the W4A8 chain; ``check_activation_format``\nhas the per-family matrix). ``gate=True`` reads ``B`` as the interleaved ``(2N, K)`` gate|up weight and\nfuses the ``act_fn`` GLU (``swiglu_alpha``/``swiglu_limit`` its clamped variant), returning the\n``[..., N]`` intermediate (``act_fn=None``: the raw ``[..., 2N]`` gate|up output instead); ``quantize_output=True`` (MX weights only) stores that intermediate in\n``activation_format`` and returns ``(C, Cs)`` for the next op.\n\nRoutes by weight dtype + inferred block:\n- ``Bs`` None → full-precision.\n- MX/NVFP4 (int8/E4M3 weights with UE8M0 group-32 or E4M3 group-16 ``Bs``) →\n  ``mx_dynamic_matmul``.\n- tensor-wide FP8 (scalar ``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``.\n- block FP8 → the static variant when ``As`` is a per-tensor scalar, else dynamic.\n- CSF-compressed NVFP4 weight scales (``prepare_nvfp4_csf_scales``): ``Bs`` is the fixed\n  stream, ``b_scale_exceptions``/``b_scale_exception_offsets`` its exceptions and expert\n  offsets; the op decodes them before the GEMM.",
       "signature": {
         "parameters": [
           {
             "default": "None",
             "annotation": "torch.Tensor | None"
           },
+          {
+            "name": "b_scale_exceptions",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "b_scale_exception_offsets",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
           {
             "name": "output_global_scale",
             "kind": "KEYWORD_ONLY",
     {
       "name": "matmul_batched",
       "qualname": "matmul_batched",
+      "module": "_finegrained_kernels_xpu_bbd13a8.batched",
       "kind": "function",
+      "doc": "Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one\nprogram per routed row (``expert_ids`` gives its expert).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3 offline for bd/tensor, inline for MX; ``\"bf16\"`` = weight-only;\n``\"mxfp4\"`` = packed E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue; ``act_fn=None`` with ``gate`` skips it, returning the raw interleaved gate|up\noutput for a host-side activation.\n``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4\nsecond-level scales ride the separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor,\nweights per-expert ``(E,)``; from ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b``\nonto the accumulator. The activation global ``g_a`` is CALIBRATED (the checkpoint's\n``input_scale``): ``a_global_scale=g_a`` with a raw ``A`` has the op quantize ``A / g_a`` per\nblock, and rides a pre-quantized ``As`` the same way. Under an NVFP4 requant the\nfused requant normalizes the GLU intermediate by the PROVIDED ``output_global_scale`` (the next\nproj's calibrated ``input_scale``) before the block quant and returns ``[C, Cs]``; the down\nconsumes it as ``As=Cs, a_global_scale=output_global_scale``. ``gather_idx``/``scatter_idx`` (each None or a ``(S,)`` map)\naddress the source row of ``A`` / destination row of ``C`` per program — None means row\n``s``; the gather lets the gate_up read unexpanded activations (source row\n``s // num_top_k``) with no copy.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_batched`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` → ``mx_dynamic_matmul_batched``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_batched``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``.\n- CSF-compressed NVFP4 weight scales (``prepare_nvfp4_csf_scales``): ``Bs`` is the fixed\n  stream, ``b_scale_exceptions``/``b_scale_exception_offsets`` its exceptions and expert\n  offsets; the op decodes them before the GEMM (only the routed experts when routing allows).",
       "signature": {
         "parameters": [
           {
             "default": "None",
             "annotation": "torch.Tensor | None"
           },
+          {
+            "name": "b_scale_exceptions",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "b_scale_exception_offsets",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
           {
             "name": "output_global_scale",
             "kind": "KEYWORD_ONLY",
     {
       "name": "matmul_grouped",
       "qualname": "matmul_grouped",
+      "module": "_finegrained_kernels_xpu_bbd13a8.grouped",
       "kind": "function",
+      "doc": "Grouped matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision).\n``expert_start`` is the ``(E+1,)`` tiling schedule from one ``compute_grouped_scheduling``\npass, shared by every grouped GEMM of the layer (the expert sort is virtual — nothing is\nphysically permuted).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3, or the fp4 chains; ``\"bf16\"`` = weight-only; ``\"mxfp4\"`` =\npacked E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue; ``act_fn=None`` with ``gate`` skips it, returning the raw interleaved gate|up\noutput for a host-side activation.\n``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4 second-level scales ride\nthe separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor, weights per-expert ``(E,)``;\nfrom ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b`` onto the accumulator. The\nactivation global ``g_a`` is CALIBRATED (the checkpoint's ``input_scale``): ``a_global_scale=g_a``\nwith a raw ``A`` has the op quantize ``A / g_a`` per block, and rides a pre-quantized ``As`` the\nsame way. Under an NVFP4 requant the fused requant normalizes the GLU intermediate by the\nPROVIDED ``output_global_scale`` (the next proj's calibrated ``input_scale``) before the block\nquant and returns ``[C, Cs]``; the down consumes it as ``As=Cs, a_global_scale=output_global_scale``.\nRow order is carried by the standalone maps: ``gather_idx`` gathers ``A`` (``None`` ->\nalready expert-ordered), ``scatter_idx`` scatters the output. The fused MoE chain is one\nscheduling pass: gate_up with ``scatter_idx=None`` + ``gate=True`` + ``quantize_output=True``,\nthen down with ``gather_idx=None`` and the\nintermediate's scales as ``As``. EP-sentinel routes sit past ``expert_start[-1]``, which\nno tile reaches.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_grouped`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` → ``mx_dynamic_matmul_grouped``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_grouped``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_grouped``.\n- CSF-compressed NVFP4 weight scales (``prepare_nvfp4_csf_scales``): ``Bs`` is the fixed\n  stream, ``b_scale_exceptions``/``b_scale_exception_offsets`` its exceptions and expert\n  offsets; the op decodes them before the GEMM (only the routed experts when routing allows).",
       "signature": {
         "parameters": [
           {
             "default": "None",
             "annotation": "torch.Tensor | None"
           },
+          {
+            "name": "b_scale_exceptions",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "b_scale_exception_offsets",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
           {
             "name": "output_global_scale",
             "kind": "KEYWORD_ONLY",
     {
       "name": "moe_fused_batched",
       "qualname": "moe_fused_batched",
+      "module": "_finegrained_kernels_xpu_bbd13a8.moe",
       "kind": "function",
+      "doc": "Fused batched MoE (decode): gate_up + SiLU + requant epilogue → per-row quantized\nintermediate → batched down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4 — decode runs the software/swap arms\nbelow the native mxf4nvf4 M=128 staging); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"`` runs the all-fp4 W4A4 chain, ``\"bf16\"`` is weight-only); ``None``\nfollows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each\nstep through the activation dtype so the output matches the unfused reference to\nreduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.\n``*_proj_scale_exceptions``/``*_proj_scale_exception_offsets`` mark that projection's\n``*_proj_scale_inv`` as a CSF-compressed NVFP4 fixed stream (``prepare_nvfp4_csf_scales``);\nthe block decodes it, only the routed experts, before its GEMMs.",
       "signature": {
         "parameters": [
           {
             "default": null,
             "annotation": "torch.Tensor"
           },
+          {
+            "name": "gate_up_proj_scale_exceptions",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "gate_up_proj_scale_exception_offsets",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "down_proj_scale_exceptions",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "down_proj_scale_exception_offsets",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
           {
             "name": "gate_up_proj_weight_global_scale",
             "kind": "POSITIONAL_OR_KEYWORD",
     {
       "name": "moe_fused_grouped",
       "qualname": "moe_fused_grouped",
+      "module": "_finegrained_kernels_xpu_bbd13a8.moe",
       "kind": "function",
+      "doc": "Fused grouped MoE (prefill): gather gate_up + SiLU + requant epilogue → quantized\nexpert-ordered intermediate → grouped down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"``/``\"nvfp4\"`` run all-fp4 W4A4 chains, ``\"bf16\"`` is weight-only);\n``None`` follows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each step through\nthe activation dtype so the output matches the unfused reference to reduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.\n``*_proj_scale_exceptions``/``*_proj_scale_exception_offsets`` mark that projection's\n``*_proj_scale_inv`` as a CSF-compressed NVFP4 fixed stream (``prepare_nvfp4_csf_scales``);\nthe block decodes it, only the routed experts, before its GEMMs.",
       "signature": {
         "parameters": [
           {
             "default": null,
             "annotation": "torch.Tensor"
           },
+          {
+            "name": "gate_up_proj_scale_exceptions",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "gate_up_proj_scale_exception_offsets",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "down_proj_scale_exceptions",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "down_proj_scale_exception_offsets",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
           {
             "name": "gate_up_proj_weight_global_scale",
             "kind": "POSITIONAL_OR_KEYWORD",
     {
       "name": "embedding_lookup",
       "qualname": "embedding_lookup",
+      "module": "_finegrained_kernels_xpu_bbd13a8.embedding",
       "kind": "function",
       "doc": "Gather and dequantize selected rows without expanding the embedding table.\n\n``weight_format`` selects ``fp8``, ``nvfp4``, ``mxfp8``, ``mxfp4``, ``int8``, or ``int4``.\nBy default, a global scale selects NVFP4, signed int8 weights select INT8, and otherwise\nFP8 is used. FP8 accepts scalar, per-row, or square block scales (``block_size``).\nNVFP4 uses E4M3 group-16 scales and one FP32 multiplier; MX formats use UE8M0 group-32\nscales. FP4 packs adjacent columns, low nibble first. INT8 uses per-row scales; INT4\npacks adjacent vocabulary rows and uses per-row column groups inferred from its scale\ngrid. Integer values are signed, symmetric, and have no zero point.\n\nOutputs default to the floating scale dtype, or bf16 for NVFP4/MX. Table/grid strides\nmay describe TP slices. Indices follow ``torch.nn.functional.embedding``: negative and\nout-of-range values fail; padding indices look up their stored row. Empty, repeated,\nand noncontiguous indices work.",
       "signature": {
     {
       "name": "get_supported_act_fns",
       "qualname": "get_supported_act_fns",
+      "module": "_finegrained_kernels_xpu_bbd13a8.formats",
       "kind": "function",
       "doc": "The activation names the gate|up epilogue fuses in-kernel; any other activation is applied\non the host by the MoE forwards (``act_fn`` as a callable).",
       "signature": {
     {
       "name": "get_supported_norms",
       "qualname": "get_supported_norms",
+      "module": "_finegrained_kernels_xpu_bbd13a8.norm",
       "kind": "function",
       "doc": "The norm names the kernels fuse (``post_expert_norm``), each ``x * factor * rsqrt(mean +\neps)`` over the row: ``\"rms_norm\"`` scales the normalized row by ``weight``,\n``\"centered_rms_norm\"`` by ``1 + weight`` (the zero-centered parameterization), and\n``\"input_scaled_rms_norm\"`` scales by ``1 + weight`` BEFORE normalizing, so the mean square\nis taken on the scaled row. Anything else stays a host callable.",
       "signature": {
     {
       "name": "rms_norm_rows",
       "qualname": "rms_norm_rows",
+      "module": "_finegrained_kernels_xpu_bbd13a8.norm",
       "kind": "function",
       "doc": "RMS-normalize ``(S, H)`` rows in one launch, in the ``get_supported_norms()`` form named by\n``norm`` (fp32 accumulate, input dtype out). The MoE chain folds this into\n``weighted_reduce`` instead; this is the standalone form.",
       "signature": {
     {
       "name": "swizzle_mx_scales",
       "qualname": "swizzle_mx_scales",
+      "module": "_finegrained_kernels_xpu_bbd13a8.swizzle",
       "kind": "function",
       "doc": "Reorder a block-scale tensor into the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05\nscaled-MMA consumes, one triton launch per matrix (``_swizzle_scales_kernel`` — no torch\npermute/transpose copies). The scale values are unchanged — this only rearranges them from\nplain row-major into the swizzled order the tensor core reads directly (the same layout\ncuBLAS/CUTLASS require for MXFP8/NVFP4 ``scaled_mm``). Plain row-major forces a gather that\ncaps the scaled dot below the fp8/fp4 peak; the swizzle removes it. Apply to weight scales\nonce at model load; the ops accept the result directly.\n\n``scale`` is the per-block scale grid — UE8M0 (``float8_e8m0fnu``, group-32 MXFP8/MXFP4) or\nE4M3 (``float8_e4m3fn``, group-16 NVFP4); the reorder is dtype-agnostic (moves bytes):\n\n- 2D ``(rows, K // group)``: one matrix (a dense weight's scale).\n- 3D ``(E, rows, K // group)``: an expert stack, each expert swizzled independently so its\n  blocks stay ``ceil(rows/128)``-aligned (the reader indexes expert ``e`` at\n  ``e * ceil(rows/128)``; byte-identical to a flat swizzle when ``rows % 128 == 0``).\n\n``gather_idx`` (2D only): a 1-D ``(padded_rows,)`` index mapping each output (sorted) row to\nits source row in ``scale`` (``-1`` = padding → zero row), folded into the kernel's load — the\nrouted/expert-sorted, per-tile-padded layout a grouped GEMM reads affine per BM=128 tile.\n``padded_rows`` must be a multiple of 128.\n\n``rows``/``cols`` are zero-padded to (128, 4) multiples; returns the 5D\n``(G, row_blocks, ceil(cols/4), 2, 256)`` artifact the ops read — ``G`` the expert count (1 for\na matrix), so an expert stack shards/gathers on its leading dim. Bit-identical to CUTLASS's\npacker (verified).",
       "signature": {
     {
       "name": "unswizzle_mx_scales",
       "qualname": "unswizzle_mx_scales",
+      "module": "_finegrained_kernels_xpu_bbd13a8.swizzle",
       "kind": "function",
       "doc": "Inverse of ``swizzle_mx_scales``: a SWIZZLE_32_4_4 artifact back to the plain row-major\n``(rows, cols)`` grid — or ``(num_experts, rows, cols)`` for an expert stack — with the\n(128, 4) padding trimmed. Pure reshape/permute on the byte view (no kernel); it runs at\nsave/dequantize time, never on a forward, so the torch-level cost is irrelevant.\n\n``rows``/``cols`` are the LOGICAL dims of ONE matrix (for a gate|up slab, the full ``2N``);\nboth are trimmed back out of the (128, 4) padding. ``num_experts`` selects the stacked\n``(E, rows, cols)`` result and must match the artifact's leading dim.",
       "signature": {
         "return_annotation": "torch.Tensor"
       }
     },
+    {
+      "name": "prepare_nvfp4_csf_scales",
+      "qualname": "prepare_nvfp4_csf_scales",
+      "module": "_finegrained_kernels_xpu_bbd13a8.loading.scales",
+      "kind": "function",
+      "doc": "Prepare one matrix or an expert stack without expanding or moving its scales.\n\nCSF stores each row of E4M3 scale bytes as a base plus 4-bit deltas, which is\nlossless whenever a row's bytes span at most 16 codes; the rest are exceptions.\n\n- `fixed` is uint8 `[rows/16, 16*(1+cols/2)]`, or the same shape with a leading\n  expert dimension. Each 16-row slab holds 16 row bases (bytes 0..240), then those\n  rows' deltas, two per byte (even column in the low nibble). Scale `(r, c)` decodes\n  to byte `base[r] + delta[r, c]`.\n- `exceptions` is one uint32 `[count]` tensor for a matrix, or one per expert, for\n  the scales outside their row's window. Each word is `position | byte << 24`:\n  the low 24 bits are the row-major index `r * cols + c` (sorted and unique), the\n  high 8 bits are the exact E4M3 byte, which overwrites the decoded value.\n\nReturns, on the input device:\n\n- `fixed`: contiguous `[experts, rows/16, 16*(1+cols/2)]` (a 2D input gains `experts=1`).\n- `exceptions`: every expert's words concatenated into one uint32 vector.\n- `expert_offsets`: uint32 `[experts + 1]` prefix sums; expert `e` owns\n  `exceptions[expert_offsets[e]:expert_offsets[e + 1]]`.\n\nThese are `decompress_nvfp4_csf_scales`' inputs. Decode on the accelerator at load time\nfor dense residency, or into shared scratch immediately before GEMM. The checks here\nrun once per checkpoint load, off the forward path, so the decoder does not repeat them.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "fixed",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "exceptions",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor | collections.abc.Sequence[torch.Tensor]"
+          }
+        ],
+        "return_annotation": "tuple[torch.Tensor, torch.Tensor, torch.Tensor]"
+      }
+    },
+    {
+      "name": "decompress_nvfp4_csf_scales",
+      "qualname": "decompress_nvfp4_csf_scales",
+      "module": "_finegrained_kernels_xpu_bbd13a8.loading.scales",
+      "kind": "function",
+      "doc": "Decode prepared CSF streams to dense E4M3 scales for an ordinary NVFP4 GEMM.\n\nWith no `out`, allocate dense scales on the accelerator at checkpoint loading.\nPass reusable contiguous [experts, rows, cols] workspace to keep the streams\ncompressed between layers. Consume the workspace before overwriting it; it\nmust not be retained for backward or shared by concurrent consumers.\nA shared dense workspace saves memory across layers, not within one layer.\nPass the grouped GEMM's cumulative `expert_start` to decode only experts with\nrouted tokens. Inactive workspace planes stay untouched and must not be read.\n\nThe streams must come from `prepare_nvfp4_csf_scales`, which validates them; only\n`out` and `expert_start`, which are new here, are checked.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "fixed",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "exceptions",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "expert_offsets",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "out",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "expert_start",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          }
+        ],
+        "return_annotation": "torch.Tensor"
+      }
+    },
     {
       "name": "mxfp8_act_quant",
       "qualname": "mxfp8_act_quant",
+      "module": "_finegrained_kernels_xpu_bbd13a8.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MX once (E4M3 values + UE8M0 group-32 uint8 scales)\ninstead of inline per weight-tile — the fused gate_up re-ran the inline quant per N-tile\n(16x redundant amax/convert ALU + 2x act bytes), which held it at ~380 TFLOPS while the\npre-quantized down kernel ran at ~1080. One pass costs ~50µs at 8k tokens. Bit-exact with\nthe inline form (same group boundaries). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path (same dense grid, per-element store).",
       "signature": {
     {
       "name": "mxfp4_act_quant",
       "qualname": "mxfp4_act_quant",
+      "module": "_finegrained_kernels_xpu_bbd13a8.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MXFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8, two codes per byte, first value in the low nibble) + UE8M0\ngroup-32 uint8 scales (``(T, K//32)``, amax/6 ceil'd to a power of two). Bit-identical\nto the fused epilogues' inline form (shared ``mx_act_quant_inline`` arm). Feeds the\nW4A4 arm of the MX matmul ops; quantizing activations to fp4 at runtime is an\naccuracy call the caller owns — the ops never do it implicitly. ``swizzled=True`` emits the\nscale directly in SWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {
     {
       "name": "nvfp4_act_quant",
       "qualname": "nvfp4_act_quant",
+      "module": "_finegrained_kernels_xpu_bbd13a8.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to NVFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8) + E4M3 group-16 block scales (``(T, K//16)`` — ``amax/6`` rounded to\nE4M3, NOT a power of two; values divide by the DECODED scale before the E2M1 grid,\nthe standard two-step). ``global_scale`` is the CALIBRATED per-tensor second level\n(fp32, the checkpoint's ``input_scale``): values are normalized by it before\nthe block quant, so block scales stay in e4m3 range regardless of the activation's\ndynamic range — the canonical two-level format. The GEMM folds ``g_a·g_b`` back onto\nthe accumulator (pass ``As=scales`` with ``a_global_scale=global_scale``). Per-expert globals\nneed each row's expert, which ``expert_index`` gives where the caller already holds that map\n(the batched op's ``expert_ids``, sentinels and all) — it is read only when the globals are\nper expert, and ``mx_act_quant_grouped`` needs none at all, its tiles carrying an expert.\n``None`` = single-level (``g_a = 1``). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {