vllm.model_executor.determinism.batch_invariant

Functions:

• all_reduce_batch_invariant –

  Sum all-reduce whose result does not depend on the message size.

• bmm_kernel –

  Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N).

• log_softmax –

  Compute log_softmax using Triton kernel.

• matmul_descriptor_persistent –

  Persistent matmul using tensor descriptors (Intel XPU fast path).

• matmul_kernel_descriptor_persistent –

  Persistent matmul using tensor descriptors for 2D block I/O.

• mean_dim –

  Triton implementation of torch.mean with single dimension reduction.

• mean_kernel –

  Kernel for computing mean along a single dimension.

• reduce_scatter_batch_invariant –

  Sum reduce-scatter over dim 0 whose result does not depend on the size.

• rms_norm_batch_invariant –

  Compute RMS normalization using Triton kernel.

• softmax –

  Compute softmax along the last dimension using a Triton kernel.

_fixed_order_sum_kernel(src_ptr, out_ptr, numel, stride_rank, WORLD_SIZE, BLOCK_SIZE, LARGE)

Sum WORLD_SIZE contributions in ascending rank order, fp32 accumulator.

An output element sums exactly WORLD_SIZE values in a compile-time constant order and rounds once, so the result depends only on the rank ordering and never on how many elements are being reduced.

Source code in vllm/model_executor/determinism/batch_invariant.py

@triton.jit
def _fixed_order_sum_kernel(
    src_ptr,  # (WORLD_SIZE, numel) gathered contributions, rank-major
    out_ptr,  # (numel,)
    numel,
    stride_rank,
    WORLD_SIZE: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    LARGE: tl.constexpr,
):
    """Sum ``WORLD_SIZE`` contributions in ascending rank order, fp32 accumulator.

    An output element sums exactly ``WORLD_SIZE`` values in a compile-time
    constant order and rounds once, so the result depends only on the rank
    ordering and never on how many elements are being reduced.
    """
    pid = tl.program_id(0)
    if LARGE:
        pid = pid.to(tl.int64)
        stride_rank = stride_rank.to(tl.int64)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < numel

    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    for r in tl.static_range(WORLD_SIZE):
        vals = tl.load(src_ptr + r * stride_rank + offs, mask=mask, other=0.0)
        acc += vals.to(tl.float32)

    tl.store(out_ptr + offs, acc.to(out_ptr.dtype.element_ty), mask=mask)

_resolve_block_sizes()

Size the tiles to the shared memory, to avoid triton OutOfResources.

Runs when the mode is enabled, or on first use by a caller that never enables it, such as a kernel unit test.

Source code in vllm/model_executor/determinism/batch_invariant.py

def _resolve_block_sizes():
    """Size the tiles to the shared memory, to avoid triton OutOfResources.

    Runs when the mode is enabled, or on first use by a caller that never
    enables it, such as a kernel unit test.
    """
    global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages
    global _BLOCK_SIZES_RESOLVED
    _BLOCK_SIZES_RESOLVED = True
    if not current_platform.is_cuda_alike():
        return
    if get_max_shared_memory_bytes() > 106496:
        _fp16_block_size_n = 256
        _fp32_block_size_n = 128
        _fp32_num_stages = 3
    else:
        _fp16_block_size_n = 128
        # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376).
        # Only N=1 OORs; keep 128/stages=3 for every other N.
        _fp32_block_size_n = 32
        _fp32_num_stages = 2

_rms_norm_kernel(input_ptr, weight_ptr, output_ptr, input_row_stride, output_row_stride, n_cols, eps, BLOCK_SIZE, HAS_WEIGHT)

Compute RMS normalization along the last dimension of a 2D tensor. RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight Each block handles one row of the input tensor.

Source code in vllm/model_executor/determinism/batch_invariant.py

@triton.jit
def _rms_norm_kernel(
    input_ptr,
    weight_ptr,
    output_ptr,
    input_row_stride,
    output_row_stride,
    n_cols,
    eps,
    BLOCK_SIZE: tl.constexpr,
    HAS_WEIGHT: tl.constexpr,
):
    """Compute RMS normalization along the last dimension of a 2D tensor.
    RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight
    Each block handles one row of the input tensor.
    """
    row_idx = tl.program_id(0).to(tl.int64)
    row_start_ptr = input_ptr + row_idx * input_row_stride
    output_row_start_ptr = output_ptr + row_idx * output_row_stride

    # Step 1: Compute sum of squares in float32 to avoid overflow
    sum_sq = tl.zeros([1], dtype=tl.float32)
    for col_offset in range(0, n_cols, BLOCK_SIZE):
        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
        mask = col_idx < n_cols

        vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
        # Convert to float32 for accumulation to prevent overflow
        vals_f32 = vals.to(tl.float32)
        sq_vals = vals_f32 * vals_f32
        sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0))

    # Step 2: Compute RMS (root mean square) in float32
    mean_sq = sum_sq / n_cols
    rms = tl.sqrt(mean_sq + eps)
    inv_rms = 1.0 / rms

    # Step 3: Normalize and apply weight
    for col_offset in range(0, n_cols, BLOCK_SIZE):
        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
        mask = col_idx < n_cols
        vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
        # Compute in float32 then convert back to input dtype
        vals_f32 = vals.to(tl.float32)
        output_f32 = vals_f32 * inv_rms
        if HAS_WEIGHT:
            weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0)
            output_f32 = output_f32 * weight.to(tl.float32)
        output = output_f32.to(vals.dtype)
        tl.store(output_row_start_ptr + col_idx, output, mask=mask)

_softmax_kernel(input_ptr, output_ptr, input_row_stride, input_col_stride, output_row_stride, n_cols, BLOCK_SIZE, LOG=False)

Compute softmax, or log_softmax when LOG, along the last dimension of a 2D tensor. Each block handles one row of the input tensor, in a fixed block order, so the result never depends on the row count.

Source code in vllm/model_executor/determinism/batch_invariant.py

@triton.jit
def _softmax_kernel(
    input_ptr,
    output_ptr,
    input_row_stride,
    input_col_stride,
    output_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
    LOG: tl.constexpr = False,
):
    """Compute softmax, or log_softmax when ``LOG``, along the last dimension of a
    2D tensor. Each block handles one row of the input tensor, in a fixed block
    order, so the result never depends on the row count.
    """
    # Get the row index for this block
    row_idx = tl.program_id(0).to(tl.int64)

    # Compute base pointers for input and output rows
    row_start_ptr = input_ptr + row_idx * input_row_stride
    output_row_start_ptr = output_ptr + row_idx * output_row_stride

    # Step 1: Find maximum value in the row for numerical stability
    max_val = -float("inf")
    for col_offset in range(0, n_cols, BLOCK_SIZE):
        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
        mask = col_idx < n_cols

        # Load values
        vals = tl.load(
            row_start_ptr + col_idx * input_col_stride, mask=mask, other=-float("inf")
        )

        # Update maximum
        max_val = tl.max(tl.maximum(vals, max_val))

    # Step 2: Compute sum of exp(x - max_val)
    sum_exp = 0.0
    for col_offset in range(0, n_cols, BLOCK_SIZE):
        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
        mask = col_idx < n_cols

        # Load values
        vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask, other=0.0)

        # Compute exp(x - max_val) and accumulate
        exp_vals = tl.exp(vals - max_val)
        sum_exp += tl.sum(tl.where(mask, exp_vals, 0.0))

    log_sum_exp = tl.log(sum_exp) if LOG else 0.0

    # Step 3: normalise
    for col_offset in range(0, n_cols, BLOCK_SIZE):
        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
        mask = col_idx < n_cols

        # Load values
        vals = tl.load(row_start_ptr + col_idx * input_col_stride, mask=mask)

        if LOG:
            output = vals - max_val - log_sum_exp
        else:
            output = tl.exp(vals - max_val) / sum_exp

        # Store results
        tl.store(output_row_start_ptr + col_idx, output, mask=mask)

all_reduce_batch_invariant(input_, group=None)

Sum all-reduce whose result does not depend on the message size.

Library all-reduces (NCCL/RCCL) pick their algorithm, channel count and chunk boundaries from the message size, so the order in which a given element's contributions are summed changes with the number of tokens in the batch.

Instead, all-gather the contributions -- pure data movement, so bitwise reproducible at any size -- and reduce them with _fixed_order_sum_kernel.

Source code in vllm/model_executor/determinism/batch_invariant.py

def all_reduce_batch_invariant(
    input_: torch.Tensor, group: "torch.distributed.ProcessGroup | None" = None
) -> torch.Tensor:
    """Sum all-reduce whose result does not depend on the message size.

    Library all-reduces (NCCL/RCCL) pick their algorithm, channel count and chunk
    boundaries from the message size, so the order in which a given element's
    contributions are summed changes with the number of tokens in the batch.

    Instead, all-gather the contributions -- pure data movement, so bitwise
    reproducible at any size -- and reduce them with ``_fixed_order_sum_kernel``.
    """
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return input_

    x = input_.contiguous()
    numel = x.numel()
    gathered = torch.empty((world_size, numel), dtype=x.dtype, device=x.device)
    dist.all_gather_into_tensor(gathered, x.view(-1), group=group)

    out = torch.empty_like(x)
    BLOCK_SIZE = 1024
    grid = (triton.cdiv(numel, BLOCK_SIZE),)
    _fixed_order_sum_kernel[grid](
        gathered,
        out.view(-1),
        numel,
        numel,
        WORLD_SIZE=world_size,
        BLOCK_SIZE=BLOCK_SIZE,
        LARGE=gathered.numel() > 2**31,
    )
    return out

bmm_kernel(a_ptr, b_ptr, c_ptr, B, M, N, K, stride_ab, stride_am, stride_ak, stride_bb, stride_bk, stride_bn, stride_cb, stride_cm, stride_cn, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, A_LARGE, B_LARGE, C_LARGE)

Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N).

Each program computes one (batch_idx, tile_m, tile_n) tile, accumulating along K in a fixed order to preserve batch invariance.

Source code in vllm/model_executor/determinism/batch_invariant.py

@triton.jit
def bmm_kernel(
    a_ptr,  # (*, ) pointer to A, (B, M, K)
    b_ptr,  # (*, ) pointer to B, (B, K, N)
    c_ptr,  # (*, ) pointer to C, (B, M, N)
    B,  # int, batch size
    M,  # int, output rows
    N,  # int, output cols
    K,  # int, reduction dim
    stride_ab,
    stride_am,
    stride_ak,
    stride_bb,
    stride_bk,
    stride_bn,
    stride_cb,
    stride_cm,
    stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    A_LARGE: tl.constexpr,
    B_LARGE: tl.constexpr,
    C_LARGE: tl.constexpr,
):
    """Batched GEMM: (B, M, K) x (B, K, N) -> (B, M, N).

    Each program computes one (batch_idx, tile_m, tile_n) tile, accumulating
    along K in a fixed order to preserve batch invariance.
    """
    pid_b = tl.program_id(0)
    pid = tl.program_id(1)

    if pid_b >= B:
        return

    # number of tiles along M / N
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    if pid_m >= num_pid_m or pid_n >= num_pid_n:
        return

    # offs_m / offs_n: raw global row/col indices for this tile
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    # masks for valid logical rows/cols within (M, N)
    mask_m = offs_m < M  # [BLOCK_SIZE_M]
    mask_n = offs_n < N  # [BLOCK_SIZE_N]

    if A_LARGE or B_LARGE or C_LARGE:
        offs_m = offs_m.to(tl.int64)
        offs_n = offs_n.to(tl.int64)

    offs_m = tl.where(mask_m, offs_m, 0)
    offs_n = tl.where(mask_n, offs_n, 0)

    # hint for triton contiguous memory
    offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_SIZE_M), BLOCK_SIZE_M)
    offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_SIZE_N), BLOCK_SIZE_N)

    # base pointers for current batch, shape-wise:
    #   a_batch_ptr points to A[pid_b, 0, 0]
    #   b_batch_ptr points to B[pid_b, 0, 0]
    #   c_batch_ptr points to C[pid_b, 0, 0]
    a_batch_ptr = a_ptr + pid_b * stride_ab
    b_batch_ptr = b_ptr + pid_b * stride_bb
    c_batch_ptr = c_ptr + pid_b * stride_cb

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    # number of K-blocks this tile iterates over
    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
    offs_k_mask = tl.arange(0, BLOCK_SIZE_K)

    for ki in range(k_tiles):
        if A_LARGE or B_LARGE:
            # offs_k: [BLOCK_SIZE_K], global K indices
            offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64)
        else:
            offs_k = ki * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)

        # a_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_K]
        #   element (i, j) points to A[pid_b, offs_m[i], offs_k[j]]
        a_ptrs = a_batch_ptr + (
            offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        )
        # b_ptrs: [BLOCK_SIZE_K, BLOCK_SIZE_N]
        #   element (i, j) points to B[pid_b, offs_k[i], offs_n[j]]
        b_ptrs = b_batch_ptr + (
            offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
        )

        # valid K lanes for this block
        k_valid = offs_k_mask < (K - ki * BLOCK_SIZE_K)
        # A mask within (M, K): [BLOCK_SIZE_M, BLOCK_SIZE_K]
        a_mask = mask_m[:, None] & k_valid[None, :]
        # B mask within (K, N): [BLOCK_SIZE_K, BLOCK_SIZE_N]
        b_mask = k_valid[:, None] & mask_n[None, :]

        # a: [BLOCK_SIZE_M, BLOCK_SIZE_K] from A[offs_m, offs_k]
        a = tl.load(
            a_ptrs,
            mask=a_mask,
            other=0.0,
        )
        # b: [BLOCK_SIZE_K, BLOCK_SIZE_N] from B[offs_k, offs_n]
        b = tl.load(
            b_ptrs,
            mask=b_mask,
            other=0.0,
        )
        accumulator = tl.dot(a, b, accumulator)

    # c_m / c_n: [BLOCK_SIZE_M] / [BLOCK_SIZE_N], row/col indices for C
    c_m = offs_m
    c_n = offs_n
    if C_LARGE:
        c_m = c_m.to(tl.int64)
        c_n = c_n.to(tl.int64)

    # c_ptrs: [BLOCK_SIZE_M, BLOCK_SIZE_N]
    #   element (i, j) points to C[pid_b, c_m[i], c_n[j]]
    c_ptrs = c_batch_ptr + stride_cm * c_m[:, None] + stride_cn * c_n[None, :]
    # mask out elements that fall outside logical (M, N) range
    c_mask = mask_m[:, None] & mask_n[None, :]
    # cast FP32 accumulator back to original dtype of C
    c = accumulator.to(c_ptr.dtype.element_ty)
    tl.store(c_ptrs, c, mask=c_mask)

log_softmax(input, dim=-1)

Compute log_softmax using Triton kernel.

Parameters:

• input

  (Tensor) –

  Input tensor

• dim

  (int, default: -1 ) –

  Dimension along which to compute log_softmax (only -1 or last dim supported)

Returns:

• Tensor –

  Tensor with log_softmax applied along the specified dimension

Source code in vllm/model_executor/determinism/batch_invariant.py

def log_softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """Compute log_softmax using Triton kernel.

    Args:
        input: Input tensor
        dim: Dimension along which to compute log_softmax
             (only -1 or last dim supported)

    Returns:
        Tensor with log_softmax applied along the specified dimension

    """
    if dim != -1 and dim != input.ndim - 1:
        raise ValueError(
            "This implementation only supports log_softmax along the last dimension"
        )

    # Flatten all dimensions except the last one
    original_shape = input.shape
    input_2d = input.reshape(-1, input.shape[-1])

    n_rows, n_cols = input_2d.shape

    # Allocate output tensor
    output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device)

    # Choose block size based on the number of columns
    BLOCK_SIZE = 1024

    # Launch kernel with one block per row
    grid = (n_rows,)
    _softmax_kernel[grid](
        input_2d,
        output,
        input_2d.stride(0),
        input_2d.stride(1),
        output.stride(0),
        n_cols,
        BLOCK_SIZE=BLOCK_SIZE,
        LOG=True,
    )
    # Reshape output back to original shape
    return output.reshape(original_shape)

matmul_descriptor_persistent(a, b, bias=None)

Persistent matmul using tensor descriptors (Intel XPU fast path).

Parameters:

• a

  (Tensor) –

  Input matrix [M, K], must be contiguous.

• b

  (Tensor) –

  Weight matrix [K, N] (standard layout — transposed internally).

• bias

  (Tensor | None, default: None ) –

  Optional 1D bias vector [N].

Returns:

• Tensor –

  Output matrix [M, N] with dtype matching the inputs.

Source code in vllm/model_executor/determinism/batch_invariant.py

def matmul_descriptor_persistent(
    a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
) -> torch.Tensor:
    """Persistent matmul using tensor descriptors (Intel XPU fast path).

    Args:
        a: Input matrix [M, K], must be contiguous.
        b: Weight matrix [K, N] (standard layout — transposed internally).
        bias: Optional 1D bias vector [N].

    Returns:
        Output matrix [M, N] with dtype matching the inputs.

    """
    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
    assert a.dtype == b.dtype, "Incompatible dtypes"
    assert bias is None or bias.dim() == 1, "bias must be 1D"

    M, K = a.shape
    _, N = b.shape
    dtype = a.dtype

    # Tensor descriptors require contiguous row-major layout.
    a = a.contiguous()
    # Descriptor kernel expects B in [N, K] layout (K-contiguous).
    b_t = b.t().contiguous()

    c = torch.empty((M, N), device=a.device, dtype=dtype)
    # This path is XPU-only. _NUM_SMS is set by enable_batch_invariant_mode()
    # before torch.compile traces this function, so Dynamo can lift it
    # as a compile-time constant.
    NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(a.device.index)

    def grid(META):
        return (
            min(
                NUM_SMS,
                triton.cdiv(M, META["BLOCK_SIZE_M"])
                * triton.cdiv(N, META["BLOCK_SIZE_N"]),
            ),
        )

    cfg = _get_descriptor_matmul_config(M, N, K, dtype)

    matmul_kernel_descriptor_persistent[grid](
        a,
        b_t,
        c,
        bias,
        M,
        N,
        K,
        NUM_SMS=NUM_SMS,
        HAS_BIAS=bias is not None,
        **cfg,
    )
    return c

matmul_kernel_descriptor_persistent(a_ptr, b_ptr, c_ptr, bias_ptr, M, N, K, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M, NUM_SMS, HAS_BIAS)

Persistent matmul using tensor descriptors for 2D block I/O.

Expects b_ptr to point to a transposed B matrix of shape [N, K] with row-major (K-contiguous) layout. The dot product transposes each loaded B-tile back: dot(A_tile, B_tile.T).

~3x faster than the pointer-based persistent kernel on Intel XPU because tensor descriptors leverage hardware 2D block load/store with automatic bounds checking (no explicit masks needed).

Source code in vllm/model_executor/determinism/batch_invariant.py

@triton.jit(launch_metadata=_matmul_launch_metadata)
def matmul_kernel_descriptor_persistent(
    a_ptr,
    b_ptr,
    c_ptr,  #
    bias_ptr,
    M,
    N,
    K,  #
    BLOCK_SIZE_M: tl.constexpr,  #
    BLOCK_SIZE_N: tl.constexpr,  #
    BLOCK_SIZE_K: tl.constexpr,  #
    GROUP_SIZE_M: tl.constexpr,  #
    NUM_SMS: tl.constexpr,  #
    HAS_BIAS: tl.constexpr,
):
    """Persistent matmul using tensor descriptors for 2D block I/O.

    Expects b_ptr to point to a transposed B matrix of shape [N, K] with
    row-major (K-contiguous) layout.  The dot product transposes each loaded
    B-tile back: dot(A_tile, B_tile.T).

    ~3x faster than the pointer-based persistent kernel on Intel XPU because
    tensor descriptors leverage hardware 2D block load/store with automatic
    bounds checking (no explicit masks needed).
    """
    dtype = c_ptr.dtype.element_ty
    start_pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
    num_tiles = num_pid_m * num_pid_n

    a_desc = tl.make_tensor_descriptor(
        a_ptr,
        shape=[M, K],
        strides=[K, 1],
        block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],
    )
    b_desc = tl.make_tensor_descriptor(
        b_ptr,
        shape=[N, K],
        strides=[K, 1],
        block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],
    )
    c_desc = tl.make_tensor_descriptor(
        c_ptr,
        shape=[M, N],
        strides=[N, 1],
        block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N],
    )

    tile_id_c = start_pid - NUM_SMS
    num_pid_in_group = GROUP_SIZE_M * num_pid_n

    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
        pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
        offs_am = pid_m * BLOCK_SIZE_M
        offs_bn = pid_n * BLOCK_SIZE_N

        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for ki in range(k_tiles):
            offs_k = ki * BLOCK_SIZE_K
            a = a_desc.load([offs_am, offs_k])
            b = b_desc.load([offs_bn, offs_k])
            accumulator = tl.dot(a, b.T, accumulator)

        tile_id_c += NUM_SMS
        pid_m, pid_n = _compute_pid(
            tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M
        )
        offs_cm = pid_m * BLOCK_SIZE_M
        offs_cn = pid_n * BLOCK_SIZE_N

        if HAS_BIAS:
            bias_offs = offs_cn + tl.arange(0, BLOCK_SIZE_N)
            bias_vals = tl.load(bias_ptr + bias_offs, mask=bias_offs < N, other=0.0).to(
                tl.float32
            )
            accumulator += bias_vals[None, :]

        c = accumulator.to(dtype)
        c_desc.store([offs_cm, offs_cn], c)

mean_dim(input, dim, keepdim=False, dtype=None)

Triton implementation of torch.mean with single dimension reduction.

Parameters:

• input

  (Tensor) –

  Input tensor

• dim

  (int) –

  Single dimension along which to compute mean

• keepdim

  (bool, default: False ) –

  Whether to keep the reduced dimension

• dtype

  (dtype | None, default: None ) –

  Output dtype. If None, uses input dtype (or float32 for integer inputs)

Returns:

• Tensor –

  Tensor with mean values along specified dimension

Source code in vllm/model_executor/determinism/batch_invariant.py

def mean_dim(
    input: torch.Tensor,
    dim: int,
    keepdim: bool = False,
    dtype: torch.dtype | None = None,
) -> torch.Tensor:
    """Triton implementation of torch.mean with single dimension reduction.

    Args:
        input: Input tensor
        dim: Single dimension along which to compute mean
        keepdim: Whether to keep the reduced dimension
        dtype: Output dtype. If None, uses input dtype
               (or float32 for integer inputs)

    Returns:
        Tensor with mean values along specified dimension

    """
    # Validate inputs
    assert -input.ndim <= dim < input.ndim, (
        f"Invalid dimension {dim} for tensor with {input.ndim} dimensions"
    )

    # Handle negative dim
    if dim < 0:
        dim = dim + input.ndim

    # Handle dtype
    if dtype is None:
        if input.dtype in [torch.int8, torch.int16, torch.int32, torch.int64]:
            dtype = torch.float32
        else:
            dtype = input.dtype

    # Convert input to appropriate dtype if needed
    if input.dtype != dtype:
        input = input.to(dtype)

    # Get input shape and strides
    shape = list(input.shape)

    # Calculate dimensions for kernel
    M = 1
    for i in range(dim):
        M *= shape[i]

    N = shape[dim]

    K = 1
    for i in range(dim + 1, len(shape)):
        K *= shape[i]

    # Reshape input to 3D view (M, N, K)
    input_3d = input.reshape(M, N, K)

    # Create output shape
    if keepdim:
        output_shape = shape.copy()
        output_shape[dim] = 1
    else:
        output_shape = shape[:dim] + shape[dim + 1 :]

    # Create output tensor
    output = torch.empty(output_shape, dtype=dtype, device=input.device)

    # Reshape output for kernel
    output_2d = output.reshape(M, 1, K).squeeze(1) if keepdim else output.reshape(M, K)

    # Launch kernel
    grid = (M * K,)
    BLOCK_SIZE = 1024

    mean_kernel[grid](
        input_3d,
        output_2d,
        input_3d.stride(0),
        input_3d.stride(1),
        input_3d.stride(2),
        output_2d.stride(0),
        output_2d.stride(1) if output_2d.ndim > 1 else 0,
        M,
        N,
        K,
        BLOCK_SIZE,
    )

    return output

mean_kernel(input_ptr, output_ptr, input_stride0, input_stride1, input_stride2, output_stride0, output_stride1, M, N, K, BLOCK_SIZE)

Kernel for computing mean along a single dimension. Input is viewed as (M, N, K) where N is the dimension being reduced.

Source code in vllm/model_executor/determinism/batch_invariant.py

@triton.jit
def mean_kernel(
    input_ptr,
    output_ptr,
    input_stride0,
    input_stride1,
    input_stride2,
    output_stride0,
    output_stride1,
    M,  # size before reduction dim
    N,  # size of reduction dim
    K,  # size after reduction dim
    BLOCK_SIZE: tl.constexpr,
):
    """Kernel for computing mean along a single dimension.
    Input is viewed as (M, N, K) where N is the dimension being reduced.
    """
    # Program ID gives us which output element we're computing
    pid = tl.program_id(0)

    # Compute output indices
    m_idx = pid // K
    k_idx = pid % K

    # Bounds check
    if m_idx >= M or k_idx >= K:
        return

    # Accumulate sum across reduction dimension
    acc = 0.0
    for n_start in range(0, N, BLOCK_SIZE):
        n_offsets = n_start + tl.arange(0, BLOCK_SIZE)
        mask = n_offsets < N

        # Calculate input indices
        input_idx = (
            m_idx * input_stride0 + n_offsets * input_stride1 + k_idx * input_stride2
        )

        # Load and accumulate
        vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0)
        acc += tl.sum(vals)

    # Compute mean and store
    mean_val = acc / N
    output_idx = m_idx * output_stride0 + k_idx * output_stride1
    tl.store(output_ptr + output_idx, mean_val)

reduce_scatter_batch_invariant(input_, group=None, sizes=None)

Sum reduce-scatter over dim 0 whose result does not depend on the size.

ncclReduceScatter picks its chunking from the message size just as ncclAllReduce does. Route around it the same way all_reduce_batch_invariant does: an all-to-all sends rank d exactly the rows it owns from every rank -- pure data movement, so bitwise reproducible however the library chunks it -- and lands them rank-major, which is the layout _fixed_order_sum_kernel already consumes.

sizes (the reduce-scatterv case) changes only which rows a rank receives, never how the received contributions are summed, so variable shard sizes are invariant for the same reason.

Parameters:

• input_

  (Tensor) –

  Contiguous full-size input; dim 0 is the scattered axis.

• group

  (ProcessGroup | None, default: None ) –

  Process group, defaults to the world group.

• sizes

  (list[int] | None, default: None ) –

  Per-rank row counts. Uniform split when None.

Returns:

• Tensor –

  This rank's shard of the sum.

Source code in vllm/model_executor/determinism/batch_invariant.py

def reduce_scatter_batch_invariant(
    input_: torch.Tensor,
    group: "torch.distributed.ProcessGroup | None" = None,
    sizes: list[int] | None = None,
) -> torch.Tensor:
    """Sum reduce-scatter over dim 0 whose result does not depend on the size.

    ``ncclReduceScatter`` picks its chunking from the message size just as
    ``ncclAllReduce`` does. Route around it the same way
    ``all_reduce_batch_invariant`` does: an all-to-all sends rank ``d`` exactly
    the rows it owns from every rank -- pure data movement, so bitwise
    reproducible however the library chunks it -- and lands them rank-major,
    which is the layout ``_fixed_order_sum_kernel`` already consumes.

    ``sizes`` (the reduce-scatterv case) changes only *which* rows a rank
    receives, never how the received contributions are summed, so variable
    shard sizes are invariant for the same reason.

    Args:
        input_: Contiguous full-size input; dim 0 is the scattered axis.
        group: Process group, defaults to the world group.
        sizes: Per-rank row counts. Uniform split when ``None``.

    Returns:
        This rank's shard of the sum.

    """
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return input_

    x = input_.contiguous()
    if sizes is None:
        assert x.shape[0] % world_size == 0
        sizes = [x.shape[0] // world_size] * world_size
    else:
        assert len(sizes) == world_size
        assert x.shape[0] == sum(sizes)

    rows = sizes[dist.get_rank(group)]
    out = torch.empty((rows, *x.shape[1:]), dtype=x.dtype, device=x.device)
    recv = torch.empty(
        (world_size * rows, *x.shape[1:]), dtype=x.dtype, device=x.device
    )
    # Every rank takes part even when its own shard is empty: it still has rows
    # to send to the others.
    dist.all_to_all_single(
        recv,
        x,
        output_split_sizes=[rows] * world_size,
        input_split_sizes=sizes,
        group=group,
    )

    numel = out.numel()
    if numel == 0:
        return out
    BLOCK_SIZE = 1024
    grid = (triton.cdiv(numel, BLOCK_SIZE),)
    _fixed_order_sum_kernel[grid](
        recv.view(-1),
        out.view(-1),
        numel,
        numel,
        WORLD_SIZE=world_size,
        BLOCK_SIZE=BLOCK_SIZE,
        LARGE=recv.numel() > 2**31,
    )
    return out

rms_norm_batch_invariant(input, weight, eps=1e-06, residual=None)

Compute RMS normalization using Triton kernel.

Parameters:

• input

  (Tensor) –

  Input tensor of shape (..., hidden_size)

• weight

  (Tensor | None) –

  Weight tensor of shape (hidden_size,), or None to skip the per-channel multiply (RMSNorm(has_weight=False))

• eps

  (float, default: 1e-06 ) –

  Small constant for numerical stability

• residual

  (Tensor | None, default: None ) –

  Optional residual tensor fused into the normalization path

Returns:

• Tensor | tuple[Tensor, Tensor] –

  RMS normalized tensor, or (output, residual_out) when residual

• Tensor | tuple[Tensor, Tensor] –

  is provided

Source code in vllm/model_executor/determinism/batch_invariant.py

def rms_norm_batch_invariant(
    input: torch.Tensor,
    weight: torch.Tensor | None,
    eps: float = 1e-6,
    residual: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Compute RMS normalization using Triton kernel.

    Args:
        input: Input tensor of shape (..., hidden_size)
        weight: Weight tensor of shape (hidden_size,), or None to skip the
            per-channel multiply (``RMSNorm(has_weight=False)``)
        eps: Small constant for numerical stability
        residual: Optional residual tensor fused into the normalization path

    Returns:
        RMS normalized tensor, or ``(output, residual_out)`` when ``residual``
        is provided

    """
    if residual is not None:
        assert input.shape == residual.shape, (
            f"Input shape {input.shape} must match residual shape {residual.shape}"
        )
        import vllm._custom_ops as ops

        ops.fused_add_rms_norm(input, residual, weight, eps)
        return input, residual

    if weight is not None:
        assert weight.dim() == 1, "Weight must be 1-dimensional"
        assert input.shape[-1] == weight.shape[0], (
            f"Input last dimension ({input.shape[-1]}) must match "
            f"weight dimension ({weight.shape[0]})"
        )
        weight = weight.contiguous()

    # Flatten all dimensions except the last one
    original_shape = input.shape
    input_2d = input.reshape(-1, input.shape[-1])
    input_2d = input_2d.contiguous()

    n_rows, n_cols = input_2d.shape

    output = torch.empty_like(input_2d)
    BLOCK_SIZE = 1024
    grid = (n_rows,)
    _rms_norm_kernel[grid](
        input_2d,
        weight if weight is not None else input_2d,
        output,
        input_2d.stride(0),
        output.stride(0),
        n_cols,
        eps,
        BLOCK_SIZE=BLOCK_SIZE,
        HAS_WEIGHT=weight is not None,
    )
    return output.reshape(original_shape)

softmax(input, dim=-1)

Compute softmax along the last dimension using a Triton kernel.

Source code in vllm/model_executor/determinism/batch_invariant.py

def softmax(input: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """Compute softmax along the last dimension using a Triton kernel."""
    if dim != -1 and dim != input.ndim - 1:
        raise ValueError(
            "This implementation only supports softmax along the last dimension"
        )

    original_shape = input.shape
    input_2d = input.reshape(-1, input.shape[-1])
    n_rows, n_cols = input_2d.shape
    output = torch.empty((n_rows, n_cols), dtype=input_2d.dtype, device=input_2d.device)

    _softmax_kernel[(n_rows,)](
        input_2d,
        output,
        input_2d.stride(0),
        input_2d.stride(1),
        output.stride(0),
        n_cols,
        BLOCK_SIZE=1024,
    )
    return output.reshape(original_shape)