[Common] MOE Split dBias

[Oleg-Goncharov]

Description

This PR adds a new kernel that computes dbias separately for each tensor in a group and outputs a grouped dbias tensor containing per-tensor dbias values.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Added the grouped dbias kernel

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Overview

Greptile Summary

This PR implements grouped dbias computation for MXFP8 quantization, allowing per-tensor bias gradients to be computed separately and stored in a grouped tensor format.

Key Changes:

• Added group_reduce_dbias_kernel in common.cuh to reduce workspace data separately for each tensor in a group
• Modified group_quantize to accept GroupedTensor* dbias instead of single Tensor*
• Updated all API signatures across activation functions (GeLU, ReLU, SiLU, etc.) to use grouped dbias
• Tests now validate per-tensor dbias outputs and include additional test cases

Important Notes:

• Grouped dbias is only supported for tensors with constant last dimension (SAME_BOTH_DIMS and VARYING_FIRST_DIM shape representations)
• The workspace offset computation for VARYING_FIRST_DIM relies on the validation that each tensor's first dimension is divisible by 128 (enforced in get_tensor_rows_num)
• Expected dbias output shape is now [num_tensors, last_logical_dim] instead of [last_logical_dim]

Confidence Score: 4/5

• This PR is safe to merge with minor considerations about the workspace offset computation
• The implementation is well-structured with proper validation, updated tests, and clear documentation. The workspace offset computation for VARYING_FIRST_DIM is correct due to the 128-divisibility validation, but relies on this invariant being maintained. Test coverage includes multiple shape configurations and the API changes are consistent across all activation functions.
• Pay close attention to transformer_engine/common/cast/core/common.cuh (workspace offset computation logic)

Important Files Changed

Filename	Overview
transformer_engine/common/cast/core/common.cuh	Added group_reduce_dbias_kernel to reduce dbias separately for each tensor in a group, outputting grouped dbias tensor
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Modified group_quantize to accept grouped dbias tensor and call grouped_reduce_dbias instead of single-tensor reduction
transformer_engine/common/include/transformer_engine/cast.h	Updated API signatures to accept NVTEGroupedTensor dbias parameter and added documentation about grouped dbias limitations
tests/cpp/operator/test_cast_mxfp8_grouped.cu	Updated tests to use grouped dbias tensors, compute per-tensor dbias references, and added new test cases

Sequence Diagram

sequenceDiagram
    participant API as nvte_group_quantize_dbias
    participant GQ as group_quantize
    participant Kernel as group_quantize_mxfp8_kernel
    participant Workspace as Workspace (float32)
    participant Reduce as grouped_reduce_dbias
    participant DBias as GroupedTensor dbias

    API->>GQ: input, output, dbias, workspace
    GQ->>GQ: Validate dbias shape [num_tensors, K]
    GQ->>GQ: Allocate workspace [M/128, K] if needed
    GQ->>Kernel: Launch quantization kernel
    Kernel->>Workspace: Write partial reductions (per 128-row chunks)
    Kernel-->>GQ: Return
    GQ->>Reduce: group_reduce_dbias_kernel
    Note over Reduce: For each tensor_id (blockIdx.y)
    Reduce->>Reduce: Compute workspace offset
    Reduce->>Workspace: Read partial sums for tensor
    Reduce->>Reduce: Sum across rows
    Reduce->>DBias: Write [tensor_id, :] result
    Reduce-->>GQ: Return
    GQ-->>API: Complete

Loading

_{Last reviewed commit: 2eeb836}

[greptile-apps]

_{9 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[Oleg-Goncharov]

/te-ci

[greptile-apps]

_{9 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Verify that tensor_rows is always divisible by chunk_dim_Y (128), otherwise this division silently truncates and skips tail row reduction.

[ptrendx]

Let's see the performance impact of having this here.

[Oleg-Goncharov]

On B300, the difference is within measurement noise. Over 3 runs, nsys shows ~59.69 µs with the check vs. ~59.62 µs without.

[greptile-apps]

_{9 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

For VARYING_FIRST_DIM, the offset computation offsets_ptr[tensor_id] / cols / chunk_dim_Y assumes the data offset is divisible by cols * chunk_dim_Y. However, when tensors have varying first dimensions, the cumulative offset offsets_ptr[tensor_id] equals the sum of M_i * K for all previous tensors. If any M_i % chunk_dim_Y != 0, this division will truncate and compute an incorrect workspace offset, causing data corruption.

The kernel in group_quantize_mxfp8.cuh:109-111 validates each tensor's first dimension is divisible by 128, which ensures M_i % chunk_dim_Y == 0, but the workspace offset depends on the sum of all previous tensor sizes being correctly aligned. Verify this is always satisfied for VARYING_FIRST_DIM case.

[greptile-apps]

_{10 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{10 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[Oleg-Goncharov]

/te-ci