# isort: off
# fmt: off
import triton
import triton.language as tl
from triton_kernels.tensor_details.layout_details.blackwell_scale import (
    unswizzle_mx_scale_bw,
)
from triton_kernels.tensor_details.layout_details.hopper_scale import unswizzle_mxfp4_scale_hopper
from triton_kernels.tensor_details.layout_details.hopper_value import mxfp4_to_bf16_triton
from triton_kernels.tensor_details.layout_details.cdna4_scale import unswizzle_mx_scale_cdna4
from triton_kernels.tensor_details.layout_details.gfx1250_scale import unswizzle_mx_scale_gfx1250
from triton_kernels.numerics_details.flexpoint import float_to_flex, load_scale
from triton_kernels.numerics_details.mxfp_details._downcast_to_mxfp import MXFP_BLOCK_SIZE, NVFP_BLOCK_SIZE
from triton_kernels.numerics_details.mxfp_details._upcast_from_mxfp import upcast_mxfp4_tile, upcast_nvfp4_tile
from triton_kernels.target_info import cuda_capability_geq
from ._common import (
    compute_offsets,
    get_scaled_dot_format_string,
    make_matmul_repr,
    matmul_launch_metadata,
    compute_pids,
    output_mx_scale_store_ptr,
)


@triton.jit
def round_f32_to_tf32(x: tl.tensor):
    # use cvt.rn on Hopper+ to match the rounding of TMA.
    ASM: tl.constexpr = "cvt.rn.tf32.f32 $0, $1;" if cuda_capability_geq(9, 0) else "cvt.rna.tf32.f32 $0, $1;"
    return tl.inline_asm_elementwise(ASM, "=r, r", [x], dtype=tl.float32, is_pure=True, pack=1)


@triton.jit
def _compute_packed_n_w(N, W_N_DIVISOR: tl.constexpr, SWIZZLE_MX_VALUE: tl.constexpr):
    packed_n_w = tl.cdiv(N, W_N_DIVISOR)
    if SWIZZLE_MX_VALUE == "HOPPER_VALUE":
        packed_n_w = tl.cdiv(packed_n_w, 64) * 64
    return packed_n_w

_matmul_repr = make_matmul_repr("_matmul", [0, 1, 2])
@triton.jit(do_not_specialize=["TOKENS_PER_EXPT_FOR_ANNOTATION"],
            repr=_matmul_repr, launch_metadata=matmul_launch_metadata)
def _matmul(
             Y, YPtr, stride_y_k, stride_y_z, stride_y_m, stride_y_n,
             YExpectedScale, YActualScale, YChecksumScale,
             stride_y_mx_k, stride_y_mx_z, stride_y_mx_m, stride_y_mx_n,
             X, XPtr, stride_x_z, stride_x_m, stride_x_k, X_TRANSPOSE: tl.constexpr,
             XScale,
             XMxScale, stride_x_mx_z, stride_x_mx_m, stride_x_mx_k,
             XTensorScale, stride_x_tensor_scale_z, stride_x_tensor_scale_m,
             W, WPtr, stride_w_e, stride_w_k, stride_w_n, W_TRANSPOSE: tl.constexpr,
             WScale,
             WMxScale, stride_w_mx_e, stride_w_mx_k, stride_w_mx_n,
             WTensorScale, stride_w_tensor_scale_e, stride_w_tensor_scale_n,
             OutAcc, stride_acc_z, stride_acc_m, stride_acc_n,
             OutAccScale, Y_ACC_IS_Y: tl.constexpr,
             B, stride_b_e, # Bias
             M, N, K, K_W, # shapes
             # expt data
             Betas, Gammas,
             GatherIndx,
             WriteBackIndx, writeback_size,
             RAGGED_DIMENSION: tl.constexpr,
             XSliceSizes, XSliceOffs, XBlockOffs, XBlockSchedule, X_EXPECTED_SLICE_SIZE: tl.constexpr, X_SLICE_SIZES_DIVISIBILITY: tl.constexpr,
             XOutputScaleBlockOffs,
             WSliceSizes, WSliceOffs, WBlockOffs, WBlockSchedule, W_EXPECTED_SLICE_SIZE: tl.constexpr, _W_SLICE_SIZES_DIVISIBILITY: tl.constexpr,
             # true grid size
             batch_size, grid_m, grid_n,
             # Out scale
             out_alpha,
             # fused activation function
             ACTIVATION_FN: tl.constexpr, activation_fn_args, ACTIVATION_REDUCTION_N: tl.constexpr,
             # epilogue transform
             EPILOGUE_FN: tl.constexpr, epilogue_fn_args,
             # MoE config
             N_EXPTS_TOT: tl.constexpr,
             # precision config
             MAX_NUM_IMPRECISE_ACC: tl.constexpr, ALLOW_TF32: tl.constexpr,
             FLEXPOINT_SATURATE_INF: tl.constexpr,
             PER_BATCH_W_SCALE: tl.constexpr,
             PER_BATCH_OUT_SCALE: tl.constexpr,
             PER_BATCH_ACC_SCALE: tl.constexpr,
             # optimization config
             BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
             GROUP_M: tl.constexpr, XCD_SWIZZLE: tl.constexpr,
             # One of ["HOPPER", "BLACKWELL", None]
             SWIZZLE_MX_VALUE: tl.constexpr,
             # One of ["HOPPER", "BLACKWELL", None]
             SWIZZLE_MX_SCALE: tl.constexpr,
             MX_BLOCK_SIZE: tl.constexpr,
             EPILOGUE_SUBTILE: tl.constexpr,
             EVEN_K: tl.constexpr, SPLIT_K: tl.constexpr,
             W_CACHE_MODIFIER: tl.constexpr,
             NUM_SMS: tl.constexpr,
             X_TMA_MODE: tl.constexpr,
             Y_TMA_MODE: tl.constexpr,
             Y_MX_SCALE_LAYOUT: tl.constexpr = None,
             OUT_N_TILE_ALIGNED: tl.constexpr = False,
             TOKENS_PER_EXPT_FOR_ANNOTATION=None,
             UPCAST_INDICES: tl.constexpr = False,
             SWAP_XW: tl.constexpr = False,
             IS_EPILOGUE_QUANT_MX: tl.constexpr = False,
             Y_VALUE_PACK_FACTOR: tl.constexpr = 1,
             FLATTEN_LOOPS: tl.constexpr = True, # Only relevant to persistent kernel
             pYPtrs=None,
             map_dst_coord=None,
             all_writes_issued=None,
             reduce_rank = 0,
             n_reduce_shards: tl.constexpr = 1,
             ):
    tl.assume(stride_y_k >= 0)
    tl.assume(stride_y_z >= 0)
    tl.assume(stride_y_m >= 0)
    tl.assume(stride_y_n >= 0)
    tl.assume(stride_x_z >= 0)
    tl.assume(stride_x_m >= 0)
    tl.assume(stride_x_k >= 0)
    tl.assume(stride_w_e >= 0)
    tl.assume(stride_w_k >= 0)
    tl.assume(stride_w_n >= 0)
    if stride_w_mx_e is not None:
        tl.assume(stride_w_mx_e >= 0)
    if stride_w_mx_k is not None:
        tl.assume(stride_w_mx_k >= 0)
    if stride_w_mx_n is not None:
        tl.assume(stride_w_mx_n >= 0)
    if B is not None:
        tl.assume(stride_b_e >= 0)
    tl.assume(batch_size >= 0)
    tl.assume(grid_m >= 0)
    tl.assume(grid_n >= 0)

    x_type: tl.constexpr = X.dtype.element_ty
    w_type: tl.constexpr = W.dtype.element_ty
    is_x_microscaled: tl.constexpr = XMxScale is not None
    is_w_microscaled: tl.constexpr = WMxScale is not None
    has_x_tensor_scale: tl.constexpr = XTensorScale is not None
    has_w_tensor_scale: tl.constexpr = WTensorScale is not None
    is_w_mxfp4: tl.constexpr = w_type == tl.uint8 and is_w_microscaled
    MX_PACK_DIVISOR: tl.constexpr = MX_BLOCK_SIZE
    if is_x_microscaled or is_w_microscaled:
        MX_SCALE_BLOCK_K: tl.constexpr = BLOCK_K // MX_PACK_DIVISOR

    if is_w_microscaled:
        tl.static_assert(MX_BLOCK_SIZE == NVFP_BLOCK_SIZE or MX_BLOCK_SIZE == MXFP_BLOCK_SIZE,
                         "Unsupported microscale factor")
        tl.static_assert(w_type == tl.uint8 or (w_type == tl.float8e4nv or w_type == tl.float8e5),
                         "mx_weight_ptr must be uint8 or fp8")
        # NOTE: uint8 scale means OCP E8M0 here. Direct NVFP-style scales stay float8e4nv.
        tl.static_assert(
            WMxScale.dtype.element_ty == tl.uint8
            or WMxScale.dtype.element_ty == tl.float8e4nv,
            "mx_scale_ptr must be uint8 or float8e4nv",
        )
        tl.static_assert(BLOCK_K % MX_PACK_DIVISOR == 0, f"{BLOCK_K=} must be a multiple of {MX_PACK_DIVISOR=}")
        tl.static_assert(SWIZZLE_MX_VALUE == "HOPPER_VALUE" or SWIZZLE_MX_VALUE == "STRIDED", "Only Hopper swizzling is supported for values")

        if SWIZZLE_MX_VALUE == "HOPPER_VALUE":
            tl.static_assert(is_w_mxfp4, "Only mxfp4 is supported for HOPPER swizzling")
            # We have pack 2 fp4 values in a byte but we divide the dimension by 2
            # when swizzling
            W_K_DIVISOR: tl.constexpr = 1
            W_K_MULTIPLIER: tl.constexpr = 2
            W_N_DIVISOR: tl.constexpr = 4
        else:
            # We have pack 2 fp4 values in a  byte
            W_K_DIVISOR: tl.constexpr = 2 if is_w_mxfp4 else 1
            W_K_MULTIPLIER: tl.constexpr = 1
            W_N_DIVISOR: tl.constexpr = 1

        if W_TRANSPOSE:
            # When weight is transposed, 2 fp4 values are packed per Byte along
            # the contiguous dimension, K.
            PACKED_BLOCK_K_W: tl.constexpr = (BLOCK_K // W_K_DIVISOR) * W_K_MULTIPLIER
            PACKED_BLOCK_N_W: tl.constexpr = BLOCK_N // W_N_DIVISOR
        else:
            # When weight is not transposed, fp4 values are *not* packed along
            # the contiguous dimension, N.
            PACKED_BLOCK_K_W: tl.constexpr = BLOCK_K
            PACKED_BLOCK_N_W: tl.constexpr = BLOCK_N // W_K_DIVISOR
    else:
        W_K_DIVISOR: tl.constexpr = 1
        W_K_MULTIPLIER: tl.constexpr = 1
        W_N_DIVISOR: tl.constexpr = 1
        PACKED_BLOCK_K_W: tl.constexpr = BLOCK_K
        PACKED_BLOCK_N_W: tl.constexpr = BLOCK_N
        tl.static_assert(SWIZZLE_MX_VALUE == "STRIDED")
        tl.static_assert(SWIZZLE_MX_SCALE == "STRIDED")
    if is_x_microscaled:
        is_x_fp4: tl.constexpr = x_type == tl.uint8
        tl.static_assert(x_type == tl.float8e4nv or x_type == tl.uint8, "mx_act_ptr must be float8e4nv or uint8")
        # NOTE: uint8 scale means OCP E8M0 here. Direct NVFP-style scales stay float8e4nv.
        tl.static_assert(
            XMxScale.dtype.element_ty == tl.uint8
            or XMxScale.dtype.element_ty == tl.float8e4nv,
            "mx_scale_ptr must be uint8 or float8e4nv",
        )
        tl.static_assert(BLOCK_K % MX_PACK_DIVISOR == 0, "BLOCK_K must be a multiple of MX_PACK_DIVISOR")
    else:
        is_x_fp4: tl.constexpr = False
    is_out_microscaled: tl.constexpr = stride_y_mx_z is not None
    is_out_fp4: tl.constexpr = is_out_microscaled and Y_VALUE_PACK_FACTOR == 2

    if _W_SLICE_SIZES_DIVISIBILITY is None:
        W_SLICE_SIZES_DIVISIBILITY: tl.constexpr = 1
    else:
        if PACKED_BLOCK_K_W > BLOCK_K:
            W_SLICE_SIZES_DIVISIBILITY: tl.constexpr =  _W_SLICE_SIZES_DIVISIBILITY * (PACKED_BLOCK_K_W // BLOCK_K)
        else:
            W_SLICE_SIZES_DIVISIBILITY: tl.constexpr =  _W_SLICE_SIZES_DIVISIBILITY // (BLOCK_K // PACKED_BLOCK_K_W)

    OUT_BLOCK_N: tl.constexpr = BLOCK_N // ACTIVATION_REDUCTION_N
    yN = N // ACTIVATION_REDUCTION_N

    pid = tl.program_id(0)
    if RAGGED_DIMENSION == "M":
        padding_m = grid_m - tl.load(XBlockOffs + N_EXPTS_TOT)
    else:
        padding_m: tl.constexpr = 0

    index_type: tl.constexpr = tl.int64 if UPCAST_INDICES else tl.int32

    unpadded_m = grid_m - padding_m
    tl.assume(unpadded_m >= 0)
    total_actual_tiles = batch_size * unpadded_m * grid_n * SPLIT_K

    if padding_m > 0 and pid >= total_actual_tiles:
        if pYPtrs is not None:
            all_writes_issued.fn(*all_writes_issued.captured)
        return

    pid_s, pid_m, pid_n, pid_k = compute_pids(pid, unpadded_m, grid_n, total_actual_tiles, XCD_SWIZZLE, GROUP_M, SPLIT_K)

    (
        expt_id, start_z, start_z_out,
        start_m, slice_block_off_m, off_m,
        off_k_x, off_k_w
    ) = compute_offsets(
            pid_s, pid_m, pid_k,
            XBlockSchedule, XSliceOffs, XBlockOffs, X_SLICE_SIZES_DIVISIBILITY,
            WBlockSchedule, WSliceOffs, W_SLICE_SIZES_DIVISIBILITY,
            RAGGED_DIMENSION,
            BLOCK_M, BLOCK_K, PACKED_BLOCK_K_W, SPLIT_K
        )
    if X_SLICE_SIZES_DIVISIBILITY is not None:
        off_k_x = off_k_x // X_SLICE_SIZES_DIVISIBILITY * X_SLICE_SIZES_DIVISIBILITY
    if W_SLICE_SIZES_DIVISIBILITY is not None:
        off_k_w = off_k_w // W_SLICE_SIZES_DIVISIBILITY * W_SLICE_SIZES_DIVISIBILITY


    if RAGGED_DIMENSION == "M":
        eM = tl.multiple_of(tl.load(XSliceSizes + expt_id), X_SLICE_SIZES_DIVISIBILITY)
    else:
        eM = M


    if RAGGED_DIMENSION == "K":
        K_W = tl.multiple_of(tl.load(WSliceOffs + pid_s + 1), W_SLICE_SIZES_DIVISIBILITY)
        K_X = tl.multiple_of(tl.load(XSliceOffs + pid_s + 1), X_SLICE_SIZES_DIVISIBILITY)
    else:
        if SWIZZLE_MX_VALUE == "HOPPER_VALUE":
            # Hopper value swizzling physically pads K to a full 64-row tile.
            K_W = tl.cdiv(K, 64) * 64
        else:
            K_W = K
        K_X = K
    if PACKED_BLOCK_K_W > BLOCK_K:
        K_W = K_W * (PACKED_BLOCK_K_W // BLOCK_K)
    else:
        K_W = K_W // (BLOCK_K // PACKED_BLOCK_K_W)

    loop_k = tl.multiple_of(tl.load(XSliceSizes + pid_s), X_SLICE_SIZES_DIVISIBILITY) if RAGGED_DIMENSION == "K" else K - off_k_x
    k_tiles = tl.cdiv(loop_k, BLOCK_K * SPLIT_K)

    # For split-k, advance to the output k slice
    if SPLIT_K > 1:
        Y += pid_k.to( index_type) * stride_y_k
        if is_out_microscaled:
            YActualScale += pid_k.to(index_type) * stride_y_mx_k

    expt_id, off_m = expt_id.to(index_type), off_m.to(index_type)
    start_m, start_z = start_m.to(index_type), start_z.to(index_type)
    pid_n, pid_k = pid_n.to(index_type), pid_k.to(index_type)
    # A pointers
    offs_x_m = off_m + tl.arange(0, BLOCK_M)
    offs_x_m = tl.max_contiguous(tl.multiple_of(offs_x_m % eM, BLOCK_M), BLOCK_M)
    X += start_z * stride_x_z
    if GatherIndx is None:
        X += start_m * stride_x_m
    else:
        GatherIndx += start_m
        # no needs to bounds-check here because `offs_x_m` wraps around M dim
        offs_x_m = tl.load(GatherIndx + offs_x_m)
    if is_x_fp4:
        offs_k_x = off_k_x // 2 + tl.arange(0, BLOCK_K // 2)
    else:
        offs_k_x = off_k_x + tl.arange(0, BLOCK_K)
    XPtrs = X + offs_x_m.to(index_type)[:, None] * stride_x_m + offs_k_x.to(index_type)[None, :] * stride_x_k


    # TODO: refactor if/else when triton front end improves
    if is_w_microscaled:
        WMxScale += expt_id * stride_w_mx_e

        if SWIZZLE_MX_SCALE == "BLACKWELL_SCALE":
            # TODO: support non W_TRANSPOSE with blackwell swizzling
            tl.static_assert(W_TRANSPOSE)
            tl.static_assert(BLOCK_N % 128 == 0)
            tl.static_assert(MX_SCALE_BLOCK_K % 4 == 0)
            PACKED_MX_BLOCK: tl.constexpr = (MX_SCALE_BLOCK_K // 4) * 32 * 4 * 4
            SCALE_BLOCK_N: tl.constexpr = BLOCK_N // 128
            stride_scale_k: tl.constexpr = 1
        elif SWIZZLE_MX_SCALE == "HOPPER_SCALE":
            # TODO: support non W_TRANSPOSE with Hopper swizzling
            tl.static_assert(W_TRANSPOSE)
            n_warps: tl.constexpr = tl.extra.cuda.num_warps()
            tl.static_assert(n_warps == 8 or n_warps == 4)
            tl.static_assert(BLOCK_N % (2 * n_warps * 2 * 8) == 0)
            tl.static_assert(MX_SCALE_BLOCK_K % 2 == 0)
            PACKED_MX_BLOCK: tl.constexpr = MX_SCALE_BLOCK_K * 32
            SCALE_BLOCK_N: tl.constexpr = BLOCK_N // 32
            stride_scale_k = stride_w_mx_k
        elif SWIZZLE_MX_SCALE == "CDNA4_SCALE":
            tl.static_assert(stride_w_mx_k is not None)
            tl.static_assert(stride_w_mx_n is not None)
            NON_K_PRESHUFFLE_BLOCK_SIZE: tl.constexpr = 32
            PACKED_MX_BLOCK: tl.constexpr = MX_SCALE_BLOCK_K * NON_K_PRESHUFFLE_BLOCK_SIZE
            SCALE_BLOCK_N: tl.constexpr = BLOCK_N // NON_K_PRESHUFFLE_BLOCK_SIZE
            stride_scale_k = stride_w_mx_k
        elif SWIZZLE_MX_SCALE == "GFX1250_SCALE":
            tl.static_assert(stride_w_mx_k is not None)
            tl.static_assert(stride_w_mx_n is not None)
            NON_K_PRESHUFFLE_BLOCK_SIZE: tl.constexpr = 128
            PACKED_MX_BLOCK: tl.constexpr = MX_SCALE_BLOCK_K * NON_K_PRESHUFFLE_BLOCK_SIZE
            SCALE_BLOCK_N: tl.constexpr = BLOCK_N // NON_K_PRESHUFFLE_BLOCK_SIZE
            stride_scale_k = stride_w_mx_k
        else:
            PACKED_MX_BLOCK: tl.constexpr = MX_SCALE_BLOCK_K
            SCALE_BLOCK_N: tl.constexpr = BLOCK_N
            stride_scale_k = stride_w_mx_k
        offs_n_scale = (pid_n * SCALE_BLOCK_N + tl.arange(0, SCALE_BLOCK_N)) % N
        offs_n_scale = tl.max_contiguous(tl.multiple_of(offs_n_scale, SCALE_BLOCK_N), SCALE_BLOCK_N)
        # K dimension must be the last dimension for the scales
        offs_k_scale = off_k_w // PACKED_BLOCK_K_W * PACKED_MX_BLOCK + tl.arange(0, PACKED_MX_BLOCK)
        WMxScalePtrs = WMxScale + offs_k_scale.to(index_type)[None, :] * stride_scale_k + offs_n_scale.to(index_type)[:, None] * stride_w_mx_n
    else:
        WMxScalePtrs = None
        offs_k_scale = None

    # B pointers
    offs_w_n = pid_n * PACKED_BLOCK_N_W + tl.arange(0, PACKED_BLOCK_N_W)
    packed_n_w = _compute_packed_n_w(N, W_N_DIVISOR, SWIZZLE_MX_VALUE)
    offs_w_n = tl.max_contiguous(tl.multiple_of(offs_w_n % packed_n_w, PACKED_BLOCK_N_W), PACKED_BLOCK_N_W)

    if is_x_microscaled:
        XMxScale += start_z.to(index_type) * stride_x_mx_z
        if GatherIndx is None:
            XMxScale += start_m * stride_x_mx_m
        offs_x_k_scale = off_k_x // MX_BLOCK_SIZE + tl.arange(0, MX_SCALE_BLOCK_K)
        XMxScalePtrs = XMxScale + offs_x_m.to(index_type)[:, None] * stride_x_mx_m + offs_x_k_scale.to(index_type)[None, :] * stride_x_mx_k
    else:
        XMxScalePtrs = None
    if has_x_tensor_scale:
        XTensorScale += start_z.to(index_type) * stride_x_tensor_scale_z
        if GatherIndx is None:
            XTensorScale += start_m * stride_x_tensor_scale_m
        XTensorScalePtrs = XTensorScale + offs_x_m.to(index_type) * stride_x_tensor_scale_m
    else:
        XTensorScalePtrs = None

    offs_w_k = off_k_w + tl.arange(0, PACKED_BLOCK_K_W)
    W += expt_id * stride_w_e
    WPtrs = W + (offs_w_k.to(index_type)[:, None] * stride_w_k + offs_w_n.to(index_type)[None, :] * stride_w_n)
    if has_w_tensor_scale:
        offs_w_tensor_scale_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
        WTensorScalePtrs = (
            WTensorScale
            + expt_id * stride_w_tensor_scale_e
            + offs_w_tensor_scale_n.to(index_type) * stride_w_tensor_scale_n
        )
    else:
        WTensorScalePtrs = None
    # compute output
    acc_dtype: tl.constexpr = tl.float64 if x_type == tl.float64 and w_type == tl.float64 else tl.float32
    acc = tl.zeros((BLOCK_N, BLOCK_M) if SWAP_XW else (BLOCK_M, BLOCK_N), dtype=acc_dtype)
    x_k_limit = K_X + BLOCK_K * SPLIT_K
    w_k_limit = K_W + PACKED_BLOCK_K_W * SPLIT_K

    for ki in range(k_tiles):
        x_k_limit -= BLOCK_K * SPLIT_K
        w_k_limit -= PACKED_BLOCK_K_W * SPLIT_K
        if EVEN_K:
            if is_x_fp4:
                mask_k_x = tl.full([BLOCK_K // 2], True, dtype=tl.int1)
            else:
                mask_k_x = tl.full([BLOCK_K], True, dtype=tl.int1)
            mask_k_w = tl.full([PACKED_BLOCK_K_W], True, dtype=tl.int1)
            if is_w_microscaled and SWIZZLE_MX_SCALE == "STRIDED":
                mask_k_scale = tl.full([PACKED_MX_BLOCK], True, dtype=tl.int1)
            if is_x_microscaled:
                mask_x_k_scale = tl.full([MX_SCALE_BLOCK_K], True, dtype=tl.int1)
        else:
            if is_x_fp4:
                mask_k_x = offs_k_x * 2 < x_k_limit
            else:
                mask_k_x = offs_k_x < x_k_limit
            mask_k_w = offs_w_k < w_k_limit
            if is_w_microscaled and SWIZZLE_MX_SCALE == "STRIDED":
                # dividing by W_K_DIVISOR because w_k_limit is also already
                # divided by W_K_DIVISOR (2 for mxfp4 wehre 2 fp4 values are
                # packed per Byte along K)
                mask_k_scale = offs_k_scale * (MX_PACK_DIVISOR // W_K_DIVISOR) < w_k_limit
            if is_x_microscaled:
                mask_x_k_scale = offs_x_k_scale * MX_PACK_DIVISOR < x_k_limit

        x = tl.load(XPtrs, mask=mask_k_x[None, :], other=0.0)
        w = tl.load(WPtrs, mask=mask_k_w[:, None], other=0.0, cache_modifier=W_CACHE_MODIFIER)
        if cuda_capability_geq(8, 0):
            if x.dtype == tl.float32 and ALLOW_TF32:
                x = round_f32_to_tf32(x)
            if w.dtype == tl.float32 and ALLOW_TF32:
                w = round_f32_to_tf32(w)
        if is_x_microscaled or is_w_microscaled:
            x_format: tl.constexpr = get_scaled_dot_format_string(x.dtype)
            w_format: tl.constexpr = get_scaled_dot_format_string(w.dtype)

            if is_x_microscaled:
                x_scales = tl.load(XMxScalePtrs, mask=mask_x_k_scale[None, :], other=0.0)
            elif x_format == "fp16" or x_format == "bf16":
                x_scales: tl.constexpr = None
            else:
                if WMxScale.dtype.element_ty == tl.uint8:
                    x_scales = tl.full((BLOCK_M, MX_SCALE_BLOCK_K), 127, dtype=tl.uint8)
                else:
                    x_scales = tl.full((BLOCK_M, MX_SCALE_BLOCK_K), 1.0, dtype=tl.float8e4nv)

            if is_w_microscaled:
                if SWIZZLE_MX_SCALE == "BLACKWELL_SCALE":
                    w_scales = unswizzle_mx_scale_bw(tl.load(WMxScalePtrs))
                elif SWIZZLE_MX_SCALE == "HOPPER_SCALE":
                    # Handshake with the swizzling code
                    num_warps: tl.constexpr = tl.extra.cuda.num_warps()
                    tl.static_assert(num_warps == 8 or num_warps == 4)
                    w_scales = unswizzle_mxfp4_scale_hopper(tl.load(WMxScalePtrs), mx_axis=1, num_warps=num_warps)
                    mask_k_scale = off_k_x + tl.arange(0, MX_SCALE_BLOCK_K) * MX_PACK_DIVISOR < x_k_limit
                    scale_zero = tl.full(w_scales.shape, 0, dtype=w_scales.dtype)
                    w_scales = tl.where(mask_k_scale[None, :], w_scales, scale_zero)
                elif SWIZZLE_MX_SCALE == "CDNA4_SCALE":
                    w_scales = unswizzle_mx_scale_cdna4(tl.load(WMxScalePtrs), BLOCK_N, MX_SCALE_BLOCK_K)
                elif SWIZZLE_MX_SCALE == "GFX1250_SCALE":
                    w_scales = unswizzle_mx_scale_gfx1250(tl.load(WMxScalePtrs), BLOCK_N, MX_SCALE_BLOCK_K)
                else:
                    w_scales = tl.load(WMxScalePtrs, mask=mask_k_scale[None, :], other=0.0)
            else:
                w_scales: tl.constexpr = None

            if is_x_fp4 and not is_w_microscaled and not cuda_capability_geq(10, 0):
                tl.static_assert(w_format == "fp16" or w_format == "bf16")
                x_dense = upcast_mxfp4_tile(x, x_scales, w.dtype)
                if SWAP_XW:
                    acc = tl.dot(w.T, x_dense.T, acc, max_num_imprecise_acc=MAX_NUM_IMPRECISE_ACC, allow_tf32=ALLOW_TF32)
                else:
                    acc = tl.dot(x_dense, w, acc, max_num_imprecise_acc=MAX_NUM_IMPRECISE_ACC, allow_tf32=ALLOW_TF32)
            elif SWIZZLE_MX_VALUE == "HOPPER_VALUE":
                # Handshake with the swizzling code
                tl.static_assert(w_format == "e2m1")
                tl.static_assert(SWAP_XW)
                wT = mxfp4_to_bf16_triton(w.T, w_scales, mx_axis=1)
                tl.static_assert(wT.dtype == tl.bfloat16)
                if is_x_microscaled:
                    acc = tl.dot_scaled(wT, None, "bf16", x.T, x_scales, x_format, acc=acc, fast_math=True)
                else:
                    tl.static_assert(x_format == "bf16")
                    acc = tl.dot(wT, x.T, acc, max_num_imprecise_acc=MAX_NUM_IMPRECISE_ACC, allow_tf32=ALLOW_TF32)
            elif (is_w_mxfp4 and WMxScale.dtype.element_ty == tl.float8e4nv and not is_x_microscaled
                  and x_format != "fp16" and x_format != "bf16"):
                # This fallback upcasts only local values and block scales.
                # Row/fiber scales are applied once after the K loop.
                w_dense = upcast_nvfp4_tile(w.T, w_scales, tl.bfloat16).T
                if SWAP_XW:
                    acc = tl.dot_scaled(w_dense.T, None, "bf16", x.T, x_scales, x_format, acc=acc, fast_math=True)
                else:
                    acc = tl.dot_scaled(x, x_scales, x_format, w_dense, None, "bf16", acc=acc, fast_math=True)
            else:
                rhs_k_pack: tl.constexpr = W_TRANSPOSE or not is_w_microscaled or W_K_DIVISOR != 2
                if SWAP_XW:
                    acc = tl.dot_scaled(w.T, w_scales, w_format, x.T, x_scales, x_format, acc=acc, fast_math=True)
                else:
                    acc = tl.dot_scaled(x, x_scales, x_format, w, w_scales, w_format, acc=acc, fast_math=True, rhs_k_pack=rhs_k_pack)
            if is_w_microscaled:
                if SWIZZLE_MX_SCALE == "BLACKWELL_SCALE":
                    WMxScalePtrs += (MX_SCALE_BLOCK_K // 4 * SPLIT_K) * stride_w_mx_k
                else:
                    WMxScalePtrs += (PACKED_MX_BLOCK * SPLIT_K) * stride_w_mx_k
            if is_x_microscaled:
                XMxScalePtrs += (MX_SCALE_BLOCK_K * SPLIT_K) * stride_x_mx_k
        else:
            if SWAP_XW:
                acc = tl.dot(w.T, x.T, acc, max_num_imprecise_acc=MAX_NUM_IMPRECISE_ACC, allow_tf32=ALLOW_TF32)
            else:
                acc = tl.dot(x, w, acc, max_num_imprecise_acc=MAX_NUM_IMPRECISE_ACC, allow_tf32=ALLOW_TF32)
        if is_x_fp4:
            XPtrs += ((BLOCK_K // 2) * SPLIT_K) * stride_x_k
        else:
            XPtrs += (BLOCK_K * SPLIT_K) * stride_x_k
        WPtrs += (PACKED_BLOCK_K_W * SPLIT_K) * stride_w_k
    if has_x_tensor_scale or has_w_tensor_scale:
        x_tensor_scales = tl.full((BLOCK_M,), 1.0, tl.float32)
        w_tensor_scales = tl.full((BLOCK_N,), 1.0, tl.float32)
        if has_x_tensor_scale:
            x_tensor_scales = tl.load(XTensorScalePtrs)
        if has_w_tensor_scale:
            w_tensor_scales = tl.load(WTensorScalePtrs, mask=offs_w_tensor_scale_n < N, other=0.0)
        if SWAP_XW:
            acc *= w_tensor_scales[:, None] * x_tensor_scales[None, :]
        else:
            acc *= x_tensor_scales[:, None] * w_tensor_scales[None, :]
    # bias + scale
    offs_m = off_m + tl.arange(0, BLOCK_M)
    offs_y_n = BLOCK_N * pid_n + tl.arange(0, BLOCK_N)
    mask_m = offs_m < eM
    mask_n = offs_y_n < N
    if B is not None:
        BPtrs = B + expt_id * stride_b_e + offs_y_n
        if pid_k == 0:
            bias = tl.load(BPtrs, mask=mask_n, other=0)
        else:
            bias = tl.full([BLOCK_N], 0, dtype=tl.float32)
    else:
        bias = tl.full([BLOCK_N], 0, dtype=tl.float32)
    if Betas is not None:
        betas = tl.load(Betas + start_m + offs_m, mask=mask_m, other=0.0)
    else:
        betas = tl.full([BLOCK_M], 1, dtype=tl.float32)
    if Gammas is not None:
        gammas = tl.load(Gammas + start_m + offs_m, mask=mask_m, other=0.0)
    else:
        gammas = tl.full([BLOCK_M], 1, dtype=tl.float32)
    # flexpoint
    x_scale = load_scale(XScale)
    if PER_BATCH_W_SCALE:
        w_scale = load_scale(WScale + expt_id)
    else:
        w_scale = load_scale(WScale)
    if SWAP_XW:
        acc = acc.trans()
    acc *= x_scale * w_scale
    acc = acc + bias[None, :] * betas[:, None]
    if out_alpha is not None:
        acc *= out_alpha
    if ACTIVATION_FN is not None:
        out = ACTIVATION_FN(acc, *activation_fn_args)
        tl.static_assert(out.shape[1] == OUT_BLOCK_N, f"Activation fn out.shape[1] ({out.shape[1]}) doesn't match computed OUT_BLOCK_N ({OUT_BLOCK_N})")
        offs_y_n = OUT_BLOCK_N * pid_n + tl.arange(0, OUT_BLOCK_N)
        mask_n = offs_y_n < yN
    else:
        tl.static_assert(ACTIVATION_REDUCTION_N == 1, "Activation reduction must be 1 if no activation fn is provided")
        out = acc
    out *= gammas[:, None]
    # write-back
    Y += start_z_out.to(index_type) * stride_y_z
    if WriteBackIndx is not None:
        WriteBackIndx += start_m
        dst_idx = tl.load(WriteBackIndx + offs_m, mask=start_m + offs_m < writeback_size, other=-1)
        mask_m = mask_m & (dst_idx != -1)
        offs_y_m = dst_idx
    else:
        Y += start_m * stride_y_m
        offs_y_m = offs_m

    if is_out_fp4:
        tl.static_assert(Y_TMA_MODE is None, "FP4 outputs are only supported without output TMA")
        offs_y_n_store = pid_n * (OUT_BLOCK_N // 2) + tl.arange(0, OUT_BLOCK_N // 2)
        YPtrs = Y + offs_y_m.to(index_type)[:, None] * stride_y_m + offs_y_n_store.to(index_type)[None, :] * stride_y_n
        mask_store = mask_m[:, None] if OUT_N_TILE_ALIGNED else mask_m[:, None] & (offs_y_n_store[None, :] < tl.cdiv(yN, 2))
    else:
        YPtrs = Y + offs_y_m.to(index_type)[:, None] * stride_y_m + offs_y_n.to(index_type)[None, :] * stride_y_n
        mask_store = mask_m[:, None] if OUT_N_TILE_ALIGNED else mask_m[:, None] & mask_n[None, :]
    mask = mask_m[:, None] if OUT_N_TILE_ALIGNED else mask_m[:, None] & mask_n[None, :]

    if OutAcc is not None:
        if PER_BATCH_ACC_SCALE:
            ScalePtr = OutAccScale + start_z_out
        else:
            ScalePtr = OutAccScale

        if Y_ACC_IS_Y:
            AccPtrs = YPtrs
        else:
            AccPtrs = OutAcc + start_z_out.to(index_type) * stride_acc_z + offs_y_m.to(index_type)[:, None] * stride_acc_m + offs_y_n.to(index_type)[None, :] * stride_acc_n
        out += tl.load(AccPtrs, mask=mask, other=0.0) * load_scale(ScalePtr)

    if is_out_microscaled:
        MX_SCALE_BLOCK_N: tl.constexpr = OUT_BLOCK_N // MX_BLOCK_SIZE
        N_MX_BLOCK = tl.cdiv(N, MX_BLOCK_SIZE)
        tl.static_assert(EPILOGUE_FN is not None)
        if PER_BATCH_OUT_SCALE:
            YExpectedScale = YExpectedScale + start_z_out
        # OCP MX outputs leave YExpectedScale unset, so this is an identity there.
        # NVFP4 uses YExpectedScale to precondition the dense output before the
        # microscaling epilogue writes direct e4m3 block scales.
        out = float_to_flex(out, YExpectedScale, None, None, mask, Y, False)
        out, out_scale = EPILOGUE_FN(out, mask, *epilogue_fn_args)
        tl.static_assert(BLOCK_N % MX_SCALE_BLOCK_N == 0, "")
        offs_y_n_scale = MX_SCALE_BLOCK_N * pid_n + tl.arange(0, MX_SCALE_BLOCK_N)
        output_scale_rows = offs_y_m if WriteBackIndx is not None else start_m + offs_y_m
        YActualScalePtrs = output_mx_scale_store_ptr(
            YActualScale,
            offs_m,
            output_scale_rows,
            offs_y_n_scale,
            start_z_out,
            start_m,
            M,
            N_MX_BLOCK,
            XOutputScaleBlockOffs,
            expt_id,
            pid_k - pid_k,
            pid_k - pid_k,
            batch_size,
            stride_y_mx_k,
            stride_y_mx_z,
            stride_y_mx_m,
            stride_y_mx_n,
            WriteBackIndx is not None,
            False,
            Y_TMA_MODE,
            RAGGED_DIMENSION,
            Y_MX_SCALE_LAYOUT,
            INDEX_TYPE=index_type,
        )
        mask_n_scale = offs_y_n_scale < N_MX_BLOCK
        if Y_MX_SCALE_LAYOUT == "BLACKWELL_ACT_SCALE" and not OUT_N_TILE_ALIGNED:
            out_scale = tl.where(mask_n_scale[None, :], out_scale, 0)
            mask_n_scale = offs_y_n_scale < tl.cdiv(N_MX_BLOCK, 4) * 4
        scale_store_mask = mask_m[:, None] if OUT_N_TILE_ALIGNED else mask_m[:, None] & mask_n_scale[None, :]
        tl.store(YActualScalePtrs, out_scale, mask=scale_store_mask)
    else:
        if PER_BATCH_OUT_SCALE:
            YExpectedScale = YExpectedScale + start_z_out
            YActualScale = YActualScale + start_z_out
        out = float_to_flex(out, YExpectedScale, YActualScale, YChecksumScale, mask, Y, FLEXPOINT_SATURATE_INF)
        if EPILOGUE_FN is not None and not IS_EPILOGUE_QUANT_MX:
            out = EPILOGUE_FN(out, *epilogue_fn_args, target_dtype=YPtrs.dtype.element_ty)
    if pYPtrs is None:
        tl.store(YPtrs, out, mask=mask_store)
    else:
        tl.static_assert(not is_out_fp4, "FP4 outputs are not supported with fused comms")
        tl.static_assert(Y_TMA_MODE is None, "TMA is not supported with fused comms")
        dst_shard_idx, dst_y_m, dst_y_n = map_dst_coord.fn(
            off_m if WriteBackIndx is None else None, offs_y_m,
            OUT_BLOCK_N * pid_n, offs_y_n,
            *map_dst_coord.captured)
        offs_mn = (
            dst_y_m.to(index_type)[:, None] * stride_y_m * n_reduce_shards + reduce_rank * stride_y_m +
            dst_y_n[None, :] * stride_y_n
        )
        for i in tl.static_range(n_reduce_shards):
            if dst_shard_idx is not None:
                peer = dst_shard_idx * n_reduce_shards + (reduce_rank + i) % n_reduce_shards
            else:
                peer = (reduce_rank + i) % n_reduce_shards
            peer_Y_ptr = tl.load(pYPtrs + peer).to(tl.pointer_type(YPtr.type.element_ty))
            if len(peer_Y_ptr.shape) == 0:
                tl.multiple_of(peer_Y_ptr, 16)
            else:
                tl.multiple_of(peer_Y_ptr, [16, 16])
            tl.store(peer_Y_ptr + offs_mn, out, mask=mask)


    if pYPtrs is not None:
        all_writes_issued.fn(*all_writes_issued.captured)
