Skip to content

[PyTorch][Core] Fix CUBLAS GGEMM when weight dims are not divisible by 128 - #2954

Merged
vthumbe1503 merged 7 commits into
NVIDIA:mainfrom
vthumbe1503:fix_cublas_grouped_gemm_gptoss_sizes
May 4, 2026
Merged

vthumbe1503 merged 7 commits into
NVIDIA:mainfrom
vthumbe1503:fix_cublas_grouped_gemm_gptoss_sizes

Conversation

@vthumbe1503

Copy link
Copy Markdown
Collaborator

Description

Please include a brief summary of the changes, relevant motivation and context.

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

Please list the changes introduced in this PR:

  • Change A
  • Change B

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

vthumbe1503 and others added 2 commits May 1, 2026 21:00
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@greptile-apps

greptile-apps Bot commented May 1, 2026 •

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR fixes incorrect MXFP8 scale-inverse pointer offsets in the cuBLAS grouped GEMM kernel when weight dimensions are not divisible by 128. The old code computed scale offsets as data_byte_offset / 32, which is only correct when dimensions are exact multiples of 32 and there is no tile-padding; the new code uses padded_mxfp8_scale_inv_bytes() that correctly accounts for the 128×4 swizzle-tile padding boundary. A new test shape (2, 256, 2880, 2880) (2880 = 90×32, not divisible by 128) directly exercises the fixed path.

Confidence Score: 4/5

PR is safe to merge; the fix is narrowly scoped to the MXFP8 scale-offset calculation and is well-validated by a new non-128-divisible test shape.

The core logic in padded_mxfp8_scale_inv_bytes and compute_grouped_tensor_mxfp8_scale_inv_offset is consistent with the existing compute_grouped_tensor_offset pattern and with the swizzle tile constants. rowwise is correctly propagated in all three entry-point code paths. No P0 issues found. Score capped at 4 because the fix touches GPU kernel pointer arithmetic that is hard to validate statically.

transformer_engine/common/gemm/cublaslt_grouped_gemm.cu — specifically padded_mxfp8_scale_inv_bytes and its consumers; reviewer should confirm that on-the-fly swizzling produces a buffer whose stride exactly matches the padded formula for the grouped (non-discrete) weight path.

Important Files Changed

Filename Overview
transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Core fix: replaces data_offset/32 with compute_grouped_tensor_mxfp8_scale_inv_offset() for padded MXFP8 scale layout; adds rowwise field to GroupedOperandSelection and threads it through setup_grouped_gemm_kernel; logic is correct across all three grouped-GEMM entry-points.
transformer_engine/common/cast/mxfp8/swizzle.cuh Promotes TILE_DIM_X=4 and TILE_DIM_Y=128 to namespace-level constexpr so they can be reused in the grouped-GEMM kernel; no functional change to gemm_swizzled_scale_idx.
tests/pytorch/test_numerics.py Refactors _make_grouped_tensor_quantized_mxfp8 to use explicit rowwise/columnwise/is_weight args, adds _per_tensor_quantize_mxfp8 helper, and adds the non-128-divisible shape (2880) as a parametrize case.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["nvte_grouped_gemm / nvte_grouped_gemm_with_discrete_out\n(GroupedTensor A & B)"] --> B["select_grouped_operand(A/B)\nsets sel.rowwise = true/false\nbased on MXFP8 + trans flags"]
    C["nvte_grouped_gemm_with_discrete_inputA\n(discrete A_list, GroupedTensor B)"] --> D["choose_grouped_operand_storage\nA_sel.rowwise = choice.use_rowwise\n(new in this PR)"]
    C --> E["select_grouped_operand(B)\nsets B_sel.rowwise"]
    B --> F["launch_grouped_gemm_setup\na_rowwise = A_sel.rowwise\nb_rowwise = B_sel.rowwise"]
    D --> F
    E --> F
    F --> G["setup_grouped_gemm_kernel<<<...>>>"]
    G --> H{"scaling_mode ==\nNVTE_MXFP8_1D_SCALING?"}
    H -- "Yes (old)" --> I["scale_ptr = base + data_offset/32\n❌ wrong when dims not divisible by 128"]
    H -- "Yes (new)" --> J["scale_offset = compute_grouped_tensor_mxfp8_scale_inv_offset\n(uses padded_mxfp8_scale_inv_bytes with 128x4 tile padding)\n✅ correct for all shapes"]
    H -- "No (tensor scaling)" --> K["scale_ptr = base + tensor_idx\n(unchanged)"]
    J --> L["a_scale_inv_ptrs[idx] / b_scale_inv_ptrs[idx]\npointed at correct per-expert scale block"]
Loading

Reviews (4): Last reviewed commit: "Fix test case dimensions in test_numeric..." | Re-trigger Greptile

Comment thread transformer_engine/common/gemm/cublaslt_grouped_gemm.cu
Comment thread transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Outdated
Comment thread tests/pytorch/test_numerics.py Outdated
Comment thread tests/pytorch/test_numerics.py Outdated

@zhongbozhu zhongbozhu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@zhongbozhu

Copy link
Copy Markdown
Collaborator

For future reference, this fix PR should be applied to NVFP4 recipe as well.

@vthumbe1503

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

@vthumbe1503
vthumbe1503 requested a review from timmoon10 May 1, 2026 22:22

@timmoon10 timmoon10 left a comment •

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall this looks good. My only serious suggestion is to change the name of GroupedOperandSelection.scale_rowwise to rowwise, since that actually changes the intent of the kernel.

Comment thread transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Outdated
Comment thread transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Outdated
Comment thread tests/pytorch/test_numerics.py Outdated
timmoon10
timmoon10 previously approved these changes May 2, 2026
@vthumbe1503

Copy link
Copy Markdown
Collaborator Author

/te-ci

@vthumbe1503 vthumbe1503 changed the title Fix CUBLAS GGEMM when weight dims are not divisible by 128 [PyTorch][Core] Fix CUBLAS GGEMM when weight dims are not divisible by 128 May 2, 2026
Comment on lines +3163 to +3175
quantizer.optimize_for_gemm = not is_weight
grouped_input = torch.cat(tensors, dim=0)
first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device)
if is_weight:
first_dims = None
else:
first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device)
return tex.group_quantize(grouped_input, quantizer, len(tensors), first_dims)


def _per_tensor_quantize_mxfp8(
tensors: List[torch.Tensor],
*,
rowwise: bool,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 _per_tensor_quantize_mxfp8 produces non-swizzled scales that break discrete_in case

MXFP8Quantizer.optimize_for_gemm defaults to False (set in Quantizer.__init__), so every tensor returned by this helper has with_gemm_swizzled_scales=False. When case == "discrete_in", the test passes A_fp8 directly to general_grouped_gemm_for_grouped_tensor, which routes to nvte_grouped_gemm_with_discrete_inputA. That function contains a hard NVTE_CHECK(A_list_info.with_gemm_swizzled_scales, "MXFP8 grouped GEMM: A scales must be swizzled for GEMM."), so the test will throw for every MXFP8 shape when case == "discrete_in".

The old code called grouped_A.split_into_quantized_tensors() on a grouped tensor built with optimize_for_gemm=True, so the split tensors inherited with_gemm_swizzled_scales=True.

Suggested change
quantizer.optimize_for_gemm = not is_weight
grouped_input = torch.cat(tensors, dim=0)
first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device)
if is_weight:
first_dims = None
else:
first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device)
return tex.group_quantize(grouped_input, quantizer, len(tensors), first_dims)
def _per_tensor_quantize_mxfp8(
tensors: List[torch.Tensor],
*,
rowwise: bool,
def _per_tensor_quantize_mxfp8(
tensors: List[torch.Tensor],
*,
rowwise: bool,
columnwise: bool,
) -> List:
"""Quantize each tensor individually with MXFP8.
Used to build reference discrete inputs for grouped GEMM.
"""
quantizer = MXFP8Quantizer(
fp8_dtype=tex.DType.kFloat8E4M3,
rowwise=rowwise,
columnwise=columnwise,
)
quantizer.optimize_for_gemm = True
return [quantizer(t) for t in tensors]

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If inputs are not swizzled already, then TE ensures to do that before GEMM

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So not needed for a functionality test

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503

Copy link
Copy Markdown
Collaborator Author

/te-ci

timmoon10
timmoon10 previously approved these changes May 2, 2026
Total dim should be divisible by 128

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
@vthumbe1503

Copy link
Copy Markdown
Collaborator Author

/te-ci pytorch

@vthumbe1503
vthumbe1503 merged commit ad4b3fd into NVIDIA:main May 4, 2026
21 of 24 checks passed
pggPL added a commit to pggPL/TransformerEngine that referenced this pull request May 11, 2026
Apply the same fix as upstream PR NVIDIA#2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL added a commit to pggPL/TransformerEngine that referenced this pull request May 11, 2026
Apply the same fix as upstream PR NVIDIA#2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL added a commit to pggPL/TransformerEngine that referenced this pull request May 11, 2026
Apply the same fix as upstream PR NVIDIA#2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
faradawn pushed a commit to faradawn/TransformerEngine that referenced this pull request May 14, 2026
…y 128 (NVIDIA#2954)

* fix CUBLAS for GPT oss sizes

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* ci fix

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* Fix test case dimensions in test_numerics.py

Total dim should be divisible by 128

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

---------

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
vthumbe1503 added a commit that referenced this pull request May 27, 2026
…hopper (#2971)

* code init

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* code drop

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Remove redundant nvte_set/get_grouped_tensor_swizzled_scales

Use existing nvte_set_grouped_tensor_param with kNVTEGroupedWithGEMMSwizzledScales
instead of the dedicated set/get functions.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Add Hopper support for grouped GEMM and refactor cuBLAS version checks

- Add CUBLAS_NVFP4_GROUPED_GEMM_VERSION and CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION macros (13.4+)
- Update check_grouped_gemm_requirements to allow SM90 with cuBLAS 13.4+
- Refactor execute_grouped_gemm to use GroupedGemmConfig struct
- Add divisibility-by-128 validation for FP8 block scaling in setup kernel and quantizer
- Support scalar alpha/beta for Hopper (no per-group alpha/beta)
- Expose get_grouped_gemm_setup_workspace_size to PyTorch via pybind
- Update PyTorch tests to run grouped GEMM on Hopper with cuBLAS 13.4+

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add NVFP4 support for discrete-input grouped GEMM and skip FP8 tensor scaling tests on Hopper

Extend nvte_grouped_gemm_with_discrete_inputA to handle NVFP4 (Float4E2M1)
inputs: accept kFloat4E2M1 dtype, propagate scale_inv pointers, collect
contiguous amax from discrete tensors, and enforce swizzled-scales checks
for NVFP4 alongside MXFP8. Also add GTEST_SKIP for FP8 tensor scaling
grouped GEMM on Hopper since cuBLAS does not support it there.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add alignment assertions for MXFP8/NVFP4 scale offsets in grouped GEMM tests

The setup kernel computes per-tensor scale pointers as data_offset /
block_size, which assumes no padding in the scale buffer. This is only
correct when first_dim % 128 == 0 and last_dim % 128 == 0 (MXFP8) or
last_dim % 64 == 0 (NVFP4). Add explicit assertions in
build_grouped_tensor to catch any future test shapes that violate this.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Fix grouped GEMM: NVFP4 columnwise transa=N + relax MXFP8 alignment for swizzle tests

cublaslt_grouped_gemm.cu:
- Fix incorrect handling of NVFP4/MXFP8 columnwise data in
  build_grouped_gemm_multi_inputA_args by adding a swap_dims flag
  consistent with choose_grouped_operand_storage. Use A_sel.trans
  (post-flip) for gemm_config.avg_k so K is selected from the
  correct dim with discrete A_list.

tests/cpp/test_common.{h,cu}:
- Add enforce_grouped_gemm_alignment parameter (default true) to
  build_grouped_tensor; the MXFP8/NVFP4 first/last_dim 128/64
  alignment asserts are only relevant for the grouped GEMM setup
  kernel, so callers that bypass it (swizzle/unswizzle) opt out.

tests/cpp/operator/test_swizzle.cu:
- Pass enforce_grouped_gemm_alignment=false to build_grouped_tensor
  in MXFP8 swizzle/unswizzle/roundtrip tests, which intentionally
  exercise non-padded shapes.

tests/cpp/operator/test_grouped_gemm.cu:
- Sync GPU/cuBLAS skip rules across all 3 sub-tests, add
  cudaDeviceSynchronize() after nvte_multi_tensor_gemm reference for
  defensive sync, and skip NVFP4 + AllDifferent in all 3 sub-tests
  due to a known flaky bug in the nvte_multi_tensor_gemm reference.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Clarify swap_dims comment in build_grouped_gemm_multi_inputA_args

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix grouped GEMM scale_inv offsets for NVFP4 and FP8 block scaling

Apply the same fix as upstream PR #2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Relax NVFP4 amax contiguity; consolidate scale_inv offset helpers; test cleanup

Production:
- nvte_grouped_gemm_with_discrete_inputA no longer requires per-expert amax
  buffers to be contiguous. Add `amax_ptrs[kMaxGroups]` to MultiTensorGroupGemmInputArgs
  and read each tensor's amax via indirection in setup_grouped_gemm_kernel
  (mirrors the existing scale_inv_ptrs pattern). The launcher enables the
  NVFP4 alpha computation when amax is available from either source.
- Consolidate four near-identical
  compute_grouped_tensor_{mxfp8,nvfp4,block_1d,block_2d}_scale_inv_offset
  into a single template `compute_grouped_scale_inv_offset<PaddedFn>` and
  collapse the A/B recipe-switch in setup_grouped_gemm_kernel into a local
  `fill_scale_ptr` lambda.

Tests:
- Drop the per-test amax staging workaround in run_grouped_gemm_discrete_in_case
  (no longer needed after the contiguity relax).
- Fix amax management in make_nvfp4_operand: copy values into result's own
  amax buffers instead of aliasing pointers (prevents double-free).
- Extract the three duplicated cuBLAS-version/compute-capability skip blocks
  into a shared `grouped_gemm_skip_reason` helper.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused float_size in GroupedGemmSetupWorkspace::from_buffers

Silences -Wunused-variable (#177-D in nvcc).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix Hopper grouped GEMM alpha beta handling

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address code review: NVFP4 amax check, swap_dims default, test refactor

- nvte_grouped_gemm and nvte_grouped_gemm_with_discrete_out now validate
  per-operand amax for NVFP4 (previously silently dropped the global-scale
  factor when amax was missing). discrete_inputA path also checks B's amax.
- Remove unused ShapeCase::kUnalignedAllSameNVFP4 enum and its comment.
- OperandStorageChoice::swap_dims now defaults to false; rowwise returns
  no longer pass spurious swap_dims=true.
- Unify GroupedGemmSetupWorkspace layout: from_buffers(nullptr, n) returns
  the total byte count, and required_setup_size derives its result from it
  so the layout cannot drift between the two.
- test_common.cu: consolidate the three gather_*_scales lambdas into a
  single gather_scale_inv(bytes_per_elem, get_shape, get_cpu_ptr) helper.
- test_grouped_gemm.cu: extract make_grouped_gemm_ref / make_alpha_beta /
  compare_grouped_d_to_multi helpers; the three run_* variants drop from
  ~1029 to 774 lines with no behavior change.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address code review feedback (#2971)

- discrete-A path: force non_tn_fp8_ok=false for FP8 block scaling to match
  select_grouped_operand logic for B.
- Setup workspace: explicit 16-byte base alignment NVTE_CHECK before
  GroupedGemmSetupWorkspace::from_buffers; matches contract every standard
  allocator (cudaMalloc / PyTorch / XLA) already satisfies.
- Tests: remove compile-time cuBLAS gating, run-time check via
  cuda::cublas_version() suffices since the test doesn't call cuBLAS directly.
- Tests: replace InputCase enum with std::optional<NVTEScalingMode>;
  nullopt = BF16, otherwise scaling mode drives dispatch.
- Tests: NVFP4 operand uses one shared BF16 input transposed via
  nvte_transpose for the columnwise direction (no more duplicate fillUniform
  with different seeds across row/col).
- Tests: parametrize output dtype (BF16 default, plus FP16 and FP32 cases on
  BF16/FP8/NVFP4 recipes); implementation already accepts all three.
- Tests: add NVFP4 2D-quantization coverage (verifies VEC16 scale layout
  fed to cuBLAS is unchanged vs 1D).
- Tests: tighten NVFP4 alignment check from %16 to %32 (TMA requirement of
  the optimized BF16 quantize path), fix misleading kUnalignedAllSame
  comment ({160,288,416} are multiples of 32, common to MXFP8 and NVFP4).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Simplify make_*_operand signatures to take use_rowwise directly

cuBLAS scaled-GEMM kernels all run in TN, so each operand only needs the
single direction matching its transpose flag. Mapping (is_A, transposed) ->
use_rowwise is uniform across MXFP8 / NVFP4 / FP8 block scaling; move it to
make_grouped_gemm_ref once instead of duplicating the if/else in each helper.

For NVFP4 specifically: drop nvte_transpose + the two-step intermediate
tensor assembly. nvte_quantize_v2 with a single-direction NVFP4 output
(rowwise XOR columnwise) directly produces what cuBLAS needs — columnwise
output goes via fallback quantize kernel (rowwise-only NVFP4 quantize kernel
hard-fails when output has no rowwise data, see quantize.cuh:111).

Matches the production pattern used by PyTorch and JAX bindings: never
allocate both NVFP4 directions on a single tensor (swizzle hard-fails when
both scale_invs are set, see swizzle.cu:985).

Net change: -112 lines.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Test infrastructure fixes for NVFP4 columnwise + skip unsupported combos

- test_common.cu: move NVFP4_1D_SCALING to the DELAYED/BLOCK_SCALING switch
  arm that allocates the columnwise data buffer as transpose(passed_shape).
  Required by TE's NVFP4 convention (column-wise data is transposed-then-
  quantized; see Tensor::shape() in common/common.h). Was incorrectly grouped
  with MXFP8 (whose columnwise data shape matches the logical shape).
  Without this, allocating columnwise-only NVFP4 in a test wires up a
  scale_inv shape that disagrees with what TE's CheckScaleTensorShape
  expects.

- libtransformer_engine.version: export
  transformer_engine::cuda::cublas_version so the test can read the run-time
  cuBLAS version (analogous to the already-exported sm_arch, sm_count,
  current_device).

- test_grouped_gemm.cu: simplify make_*_operand signatures to take
  use_rowwise directly, with the (is_A, transposed) -> use_rowwise mapping
  centralized in make_grouped_gemm_ref. NVFP4 now uses single-direction
  allocation (no nvte_transpose, no two-step assembly) — same pattern as
  MXFP8 / FP8 block scaling, working thanks to the test_common.cu fix.

- test_grouped_gemm.cu: tighten skip logic via grouped_gemm_skip_reason
  (now takes TestParams):
    * Skip NVFP4 + FP16 output  (cuBLAS hard-errors in cublaslt_gemm.cu)
    * Skip NVFP4 2D quantization in non-TN layouts (fallback quantize path
      doesn't support 2D; production weight quantization always allocates
      both directions and hits the optimized path).
  Drop the BF16 + FP16 output test case (cuBLAS grouped GEMM has no
  algorithm for that combination, even in TN).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Trim verbose comments in grouped GEMM test

Drop multi-line rationale comments that were conversation residue (linking
back to specific other files/lines, restating GEMM TN convention, listing
PR numbers) and keep just the short ones that name what the code does.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fix FP8 block scaling NT/NN failures on Hopper

Two related bugs in the grouped GEMM path were causing FP8 block scaling
NT and NN cases to fail on Hopper for non-Mul128 dims (Mul32 tests).

1. build_grouped_gemm_multi_inputA_args (discrete-A path) used
   t->shape() (LOGICAL) for cuBLAS rows/cols, while the symmetric
   grouped path (build_grouped_tensor / select_grouped_operand) uses
   tensors[i]->{rowwise,columnwise}_shape() (PHYSICAL). For FP8 block
   columnwise the two differ (physical is transposed of logical), so
   args.rows became N instead of K. Switch the discrete-A path to use
   data.shape directly — this aligns both paths and makes swap_dims
   redundant for this function.

2. padded_block_{1d,2d}_scale_inv_floats had a columnwise branch that
   swapped first/last in the formula, but the quantizer
   (test_common.cu:get_scales) always uses logical dims regardless of
   direction. Combined with grouped meta passing transposed dims for
   columnwise data, the formula produced wrong per-expert scale stride
   — only visible for dims not divisible by 128. The unified formula
   ceil(last/128) * roundup(first, 4) is correct for both directions
   because the meta swap and quantizer swap cancel out.

Also remove now-redundant check_fp4_output_compat duplication and trim
stale comments.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use logical grouped GEMM shapes

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix mixed FP8 block grouped GEMM scaling

Handle per-operand FP8 block scaling modes when computing grouped scale offsets, and tighten NVFP4 validation to reject unsupported mixed or non-per-group-alpha paths.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim unsupported FP8 block grouped GEMM tests

Avoid test cases that require direct columnwise 2D block quantization, which is not a supported test setup path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address review comment

Added validation for FP8 block scaling support in grouped GEMM.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tests to relax alignment requirements for swizzling tests

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Fix alignment calculation in required_setup_size

Adjust required_setup_size to account for alignment.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use nvfp4 alpha only for nvfp4

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Clean up grouped GEMM tests and document alignment

- Clarify scale_inv padding comments in test_grouped_gemm.cu so it is
  obvious that kAllSameMul32 exercises per-expert scale_inv offsets only
  for recipes that pad grouped scale_inv storage.
- Centralize the 32 MiB cuBLAS workspace size into a single
  kCublasWorkspaceBytes constant in test_grouped_gemm.cu.
- Drop the test-only enforce_grouped_gemm_alignment flag from
  build_grouped_tensor and remove all its call-site overrides; alignment
  is already validated inside the grouped GEMM implementation.
- Add a short comment in cublaslt_grouped_gemm.cu explaining why the int
  and float arrays in GroupedGemmSetupWorkspace do not require an
  explicit align_ptr() call (cuBLAS only mandates 16-byte alignment for
  the pointer arrays; 4-byte natural alignment is sufficient for the
  int/float arrays and is preserved by the layout).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Post-merge cleanup in grouped GEMM helpers and swizzle tests

- test_swizzle.cu: drop the enforce_grouped_gemm_alignment=false argument
  introduced upstream after the flag was removed from build_grouped_tensor.
- cublaslt_grouped_gemm.cu (_with_discrete_inputA): remove redundant B
  swizzled-scales check; the same condition is already enforced inside
  validate_grouped_gemm_inputs({inputB}, ...) called just above.
- build_grouped_gemm_multi_inputA_args: drop the requires_scale_inv
  parameter and derive it per-tensor from t->dtype() / t->scaling_mode.
  validate_grouped_gemm_multi_inputA_list already enforces uniform
  scaling mode, so per-iteration evaluation is safe and removes the
  duplicated computation at the call site.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
KshitijLakhani pushed a commit that referenced this pull request May 27, 2026
…hopper (#2971)

* code init

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* code drop

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Remove redundant nvte_set/get_grouped_tensor_swizzled_scales

Use existing nvte_set_grouped_tensor_param with kNVTEGroupedWithGEMMSwizzledScales
instead of the dedicated set/get functions.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Add Hopper support for grouped GEMM and refactor cuBLAS version checks

- Add CUBLAS_NVFP4_GROUPED_GEMM_VERSION and CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION macros (13.4+)
- Update check_grouped_gemm_requirements to allow SM90 with cuBLAS 13.4+
- Refactor execute_grouped_gemm to use GroupedGemmConfig struct
- Add divisibility-by-128 validation for FP8 block scaling in setup kernel and quantizer
- Support scalar alpha/beta for Hopper (no per-group alpha/beta)
- Expose get_grouped_gemm_setup_workspace_size to PyTorch via pybind
- Update PyTorch tests to run grouped GEMM on Hopper with cuBLAS 13.4+

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add NVFP4 support for discrete-input grouped GEMM and skip FP8 tensor scaling tests on Hopper

Extend nvte_grouped_gemm_with_discrete_inputA to handle NVFP4 (Float4E2M1)
inputs: accept kFloat4E2M1 dtype, propagate scale_inv pointers, collect
contiguous amax from discrete tensors, and enforce swizzled-scales checks
for NVFP4 alongside MXFP8. Also add GTEST_SKIP for FP8 tensor scaling
grouped GEMM on Hopper since cuBLAS does not support it there.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add alignment assertions for MXFP8/NVFP4 scale offsets in grouped GEMM tests

The setup kernel computes per-tensor scale pointers as data_offset /
block_size, which assumes no padding in the scale buffer. This is only
correct when first_dim % 128 == 0 and last_dim % 128 == 0 (MXFP8) or
last_dim % 64 == 0 (NVFP4). Add explicit assertions in
build_grouped_tensor to catch any future test shapes that violate this.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Fix grouped GEMM: NVFP4 columnwise transa=N + relax MXFP8 alignment for swizzle tests

cublaslt_grouped_gemm.cu:
- Fix incorrect handling of NVFP4/MXFP8 columnwise data in
  build_grouped_gemm_multi_inputA_args by adding a swap_dims flag
  consistent with choose_grouped_operand_storage. Use A_sel.trans
  (post-flip) for gemm_config.avg_k so K is selected from the
  correct dim with discrete A_list.

tests/cpp/test_common.{h,cu}:
- Add enforce_grouped_gemm_alignment parameter (default true) to
  build_grouped_tensor; the MXFP8/NVFP4 first/last_dim 128/64
  alignment asserts are only relevant for the grouped GEMM setup
  kernel, so callers that bypass it (swizzle/unswizzle) opt out.

tests/cpp/operator/test_swizzle.cu:
- Pass enforce_grouped_gemm_alignment=false to build_grouped_tensor
  in MXFP8 swizzle/unswizzle/roundtrip tests, which intentionally
  exercise non-padded shapes.

tests/cpp/operator/test_grouped_gemm.cu:
- Sync GPU/cuBLAS skip rules across all 3 sub-tests, add
  cudaDeviceSynchronize() after nvte_multi_tensor_gemm reference for
  defensive sync, and skip NVFP4 + AllDifferent in all 3 sub-tests
  due to a known flaky bug in the nvte_multi_tensor_gemm reference.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Clarify swap_dims comment in build_grouped_gemm_multi_inputA_args

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix grouped GEMM scale_inv offsets for NVFP4 and FP8 block scaling

Apply the same fix as upstream PR #2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Relax NVFP4 amax contiguity; consolidate scale_inv offset helpers; test cleanup

Production:
- nvte_grouped_gemm_with_discrete_inputA no longer requires per-expert amax
  buffers to be contiguous. Add `amax_ptrs[kMaxGroups]` to MultiTensorGroupGemmInputArgs
  and read each tensor's amax via indirection in setup_grouped_gemm_kernel
  (mirrors the existing scale_inv_ptrs pattern). The launcher enables the
  NVFP4 alpha computation when amax is available from either source.
- Consolidate four near-identical
  compute_grouped_tensor_{mxfp8,nvfp4,block_1d,block_2d}_scale_inv_offset
  into a single template `compute_grouped_scale_inv_offset<PaddedFn>` and
  collapse the A/B recipe-switch in setup_grouped_gemm_kernel into a local
  `fill_scale_ptr` lambda.

Tests:
- Drop the per-test amax staging workaround in run_grouped_gemm_discrete_in_case
  (no longer needed after the contiguity relax).
- Fix amax management in make_nvfp4_operand: copy values into result's own
  amax buffers instead of aliasing pointers (prevents double-free).
- Extract the three duplicated cuBLAS-version/compute-capability skip blocks
  into a shared `grouped_gemm_skip_reason` helper.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused float_size in GroupedGemmSetupWorkspace::from_buffers

Silences -Wunused-variable (#177-D in nvcc).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix Hopper grouped GEMM alpha beta handling

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address code review: NVFP4 amax check, swap_dims default, test refactor

- nvte_grouped_gemm and nvte_grouped_gemm_with_discrete_out now validate
  per-operand amax for NVFP4 (previously silently dropped the global-scale
  factor when amax was missing). discrete_inputA path also checks B's amax.
- Remove unused ShapeCase::kUnalignedAllSameNVFP4 enum and its comment.
- OperandStorageChoice::swap_dims now defaults to false; rowwise returns
  no longer pass spurious swap_dims=true.
- Unify GroupedGemmSetupWorkspace layout: from_buffers(nullptr, n) returns
  the total byte count, and required_setup_size derives its result from it
  so the layout cannot drift between the two.
- test_common.cu: consolidate the three gather_*_scales lambdas into a
  single gather_scale_inv(bytes_per_elem, get_shape, get_cpu_ptr) helper.
- test_grouped_gemm.cu: extract make_grouped_gemm_ref / make_alpha_beta /
  compare_grouped_d_to_multi helpers; the three run_* variants drop from
  ~1029 to 774 lines with no behavior change.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address code review feedback (#2971)

- discrete-A path: force non_tn_fp8_ok=false for FP8 block scaling to match
  select_grouped_operand logic for B.
- Setup workspace: explicit 16-byte base alignment NVTE_CHECK before
  GroupedGemmSetupWorkspace::from_buffers; matches contract every standard
  allocator (cudaMalloc / PyTorch / XLA) already satisfies.
- Tests: remove compile-time cuBLAS gating, run-time check via
  cuda::cublas_version() suffices since the test doesn't call cuBLAS directly.
- Tests: replace InputCase enum with std::optional<NVTEScalingMode>;
  nullopt = BF16, otherwise scaling mode drives dispatch.
- Tests: NVFP4 operand uses one shared BF16 input transposed via
  nvte_transpose for the columnwise direction (no more duplicate fillUniform
  with different seeds across row/col).
- Tests: parametrize output dtype (BF16 default, plus FP16 and FP32 cases on
  BF16/FP8/NVFP4 recipes); implementation already accepts all three.
- Tests: add NVFP4 2D-quantization coverage (verifies VEC16 scale layout
  fed to cuBLAS is unchanged vs 1D).
- Tests: tighten NVFP4 alignment check from %16 to %32 (TMA requirement of
  the optimized BF16 quantize path), fix misleading kUnalignedAllSame
  comment ({160,288,416} are multiples of 32, common to MXFP8 and NVFP4).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Simplify make_*_operand signatures to take use_rowwise directly

cuBLAS scaled-GEMM kernels all run in TN, so each operand only needs the
single direction matching its transpose flag. Mapping (is_A, transposed) ->
use_rowwise is uniform across MXFP8 / NVFP4 / FP8 block scaling; move it to
make_grouped_gemm_ref once instead of duplicating the if/else in each helper.

For NVFP4 specifically: drop nvte_transpose + the two-step intermediate
tensor assembly. nvte_quantize_v2 with a single-direction NVFP4 output
(rowwise XOR columnwise) directly produces what cuBLAS needs — columnwise
output goes via fallback quantize kernel (rowwise-only NVFP4 quantize kernel
hard-fails when output has no rowwise data, see quantize.cuh:111).

Matches the production pattern used by PyTorch and JAX bindings: never
allocate both NVFP4 directions on a single tensor (swizzle hard-fails when
both scale_invs are set, see swizzle.cu:985).

Net change: -112 lines.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Test infrastructure fixes for NVFP4 columnwise + skip unsupported combos

- test_common.cu: move NVFP4_1D_SCALING to the DELAYED/BLOCK_SCALING switch
  arm that allocates the columnwise data buffer as transpose(passed_shape).
  Required by TE's NVFP4 convention (column-wise data is transposed-then-
  quantized; see Tensor::shape() in common/common.h). Was incorrectly grouped
  with MXFP8 (whose columnwise data shape matches the logical shape).
  Without this, allocating columnwise-only NVFP4 in a test wires up a
  scale_inv shape that disagrees with what TE's CheckScaleTensorShape
  expects.

- libtransformer_engine.version: export
  transformer_engine::cuda::cublas_version so the test can read the run-time
  cuBLAS version (analogous to the already-exported sm_arch, sm_count,
  current_device).

- test_grouped_gemm.cu: simplify make_*_operand signatures to take
  use_rowwise directly, with the (is_A, transposed) -> use_rowwise mapping
  centralized in make_grouped_gemm_ref. NVFP4 now uses single-direction
  allocation (no nvte_transpose, no two-step assembly) — same pattern as
  MXFP8 / FP8 block scaling, working thanks to the test_common.cu fix.

- test_grouped_gemm.cu: tighten skip logic via grouped_gemm_skip_reason
  (now takes TestParams):
    * Skip NVFP4 + FP16 output  (cuBLAS hard-errors in cublaslt_gemm.cu)
    * Skip NVFP4 2D quantization in non-TN layouts (fallback quantize path
      doesn't support 2D; production weight quantization always allocates
      both directions and hits the optimized path).
  Drop the BF16 + FP16 output test case (cuBLAS grouped GEMM has no
  algorithm for that combination, even in TN).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Trim verbose comments in grouped GEMM test

Drop multi-line rationale comments that were conversation residue (linking
back to specific other files/lines, restating GEMM TN convention, listing
PR numbers) and keep just the short ones that name what the code does.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fix FP8 block scaling NT/NN failures on Hopper

Two related bugs in the grouped GEMM path were causing FP8 block scaling
NT and NN cases to fail on Hopper for non-Mul128 dims (Mul32 tests).

1. build_grouped_gemm_multi_inputA_args (discrete-A path) used
   t->shape() (LOGICAL) for cuBLAS rows/cols, while the symmetric
   grouped path (build_grouped_tensor / select_grouped_operand) uses
   tensors[i]->{rowwise,columnwise}_shape() (PHYSICAL). For FP8 block
   columnwise the two differ (physical is transposed of logical), so
   args.rows became N instead of K. Switch the discrete-A path to use
   data.shape directly — this aligns both paths and makes swap_dims
   redundant for this function.

2. padded_block_{1d,2d}_scale_inv_floats had a columnwise branch that
   swapped first/last in the formula, but the quantizer
   (test_common.cu:get_scales) always uses logical dims regardless of
   direction. Combined with grouped meta passing transposed dims for
   columnwise data, the formula produced wrong per-expert scale stride
   — only visible for dims not divisible by 128. The unified formula
   ceil(last/128) * roundup(first, 4) is correct for both directions
   because the meta swap and quantizer swap cancel out.

Also remove now-redundant check_fp4_output_compat duplication and trim
stale comments.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use logical grouped GEMM shapes

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix mixed FP8 block grouped GEMM scaling

Handle per-operand FP8 block scaling modes when computing grouped scale offsets, and tighten NVFP4 validation to reject unsupported mixed or non-per-group-alpha paths.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim unsupported FP8 block grouped GEMM tests

Avoid test cases that require direct columnwise 2D block quantization, which is not a supported test setup path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address review comment

Added validation for FP8 block scaling support in grouped GEMM.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tests to relax alignment requirements for swizzling tests

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Fix alignment calculation in required_setup_size

Adjust required_setup_size to account for alignment.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use nvfp4 alpha only for nvfp4

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Clean up grouped GEMM tests and document alignment

- Clarify scale_inv padding comments in test_grouped_gemm.cu so it is
  obvious that kAllSameMul32 exercises per-expert scale_inv offsets only
  for recipes that pad grouped scale_inv storage.
- Centralize the 32 MiB cuBLAS workspace size into a single
  kCublasWorkspaceBytes constant in test_grouped_gemm.cu.
- Drop the test-only enforce_grouped_gemm_alignment flag from
  build_grouped_tensor and remove all its call-site overrides; alignment
  is already validated inside the grouped GEMM implementation.
- Add a short comment in cublaslt_grouped_gemm.cu explaining why the int
  and float arrays in GroupedGemmSetupWorkspace do not require an
  explicit align_ptr() call (cuBLAS only mandates 16-byte alignment for
  the pointer arrays; 4-byte natural alignment is sufficient for the
  int/float arrays and is preserved by the layout).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Post-merge cleanup in grouped GEMM helpers and swizzle tests

- test_swizzle.cu: drop the enforce_grouped_gemm_alignment=false argument
  introduced upstream after the flag was removed from build_grouped_tensor.
- cublaslt_grouped_gemm.cu (_with_discrete_inputA): remove redundant B
  swizzled-scales check; the same condition is already enforced inside
  validate_grouped_gemm_inputs({inputB}, ...) called just above.
- build_grouped_gemm_multi_inputA_args: drop the requires_scale_inv
  parameter and derive it per-tensor from t->dtype() / t->scaling_mode.
  validate_grouped_gemm_multi_inputA_list already enforces uniform
  scaling mode, so per-iteration evaluation is safe and removes the
  duplicated computation at the call site.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Baibaifan pushed a commit to Baibaifan/TransformerEngine that referenced this pull request Jun 1, 2026
…hopper (NVIDIA#2971)

* code init

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* code drop

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Remove redundant nvte_set/get_grouped_tensor_swizzled_scales

Use existing nvte_set_grouped_tensor_param with kNVTEGroupedWithGEMMSwizzledScales
instead of the dedicated set/get functions.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Add Hopper support for grouped GEMM and refactor cuBLAS version checks

- Add CUBLAS_NVFP4_GROUPED_GEMM_VERSION and CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION macros (13.4+)
- Update check_grouped_gemm_requirements to allow SM90 with cuBLAS 13.4+
- Refactor execute_grouped_gemm to use GroupedGemmConfig struct
- Add divisibility-by-128 validation for FP8 block scaling in setup kernel and quantizer
- Support scalar alpha/beta for Hopper (no per-group alpha/beta)
- Expose get_grouped_gemm_setup_workspace_size to PyTorch via pybind
- Update PyTorch tests to run grouped GEMM on Hopper with cuBLAS 13.4+

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add NVFP4 support for discrete-input grouped GEMM and skip FP8 tensor scaling tests on Hopper

Extend nvte_grouped_gemm_with_discrete_inputA to handle NVFP4 (Float4E2M1)
inputs: accept kFloat4E2M1 dtype, propagate scale_inv pointers, collect
contiguous amax from discrete tensors, and enforce swizzled-scales checks
for NVFP4 alongside MXFP8. Also add GTEST_SKIP for FP8 tensor scaling
grouped GEMM on Hopper since cuBLAS does not support it there.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add alignment assertions for MXFP8/NVFP4 scale offsets in grouped GEMM tests

The setup kernel computes per-tensor scale pointers as data_offset /
block_size, which assumes no padding in the scale buffer. This is only
correct when first_dim % 128 == 0 and last_dim % 128 == 0 (MXFP8) or
last_dim % 64 == 0 (NVFP4). Add explicit assertions in
build_grouped_tensor to catch any future test shapes that violate this.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Fix grouped GEMM: NVFP4 columnwise transa=N + relax MXFP8 alignment for swizzle tests

cublaslt_grouped_gemm.cu:
- Fix incorrect handling of NVFP4/MXFP8 columnwise data in
  build_grouped_gemm_multi_inputA_args by adding a swap_dims flag
  consistent with choose_grouped_operand_storage. Use A_sel.trans
  (post-flip) for gemm_config.avg_k so K is selected from the
  correct dim with discrete A_list.

tests/cpp/test_common.{h,cu}:
- Add enforce_grouped_gemm_alignment parameter (default true) to
  build_grouped_tensor; the MXFP8/NVFP4 first/last_dim 128/64
  alignment asserts are only relevant for the grouped GEMM setup
  kernel, so callers that bypass it (swizzle/unswizzle) opt out.

tests/cpp/operator/test_swizzle.cu:
- Pass enforce_grouped_gemm_alignment=false to build_grouped_tensor
  in MXFP8 swizzle/unswizzle/roundtrip tests, which intentionally
  exercise non-padded shapes.

tests/cpp/operator/test_grouped_gemm.cu:
- Sync GPU/cuBLAS skip rules across all 3 sub-tests, add
  cudaDeviceSynchronize() after nvte_multi_tensor_gemm reference for
  defensive sync, and skip NVFP4 + AllDifferent in all 3 sub-tests
  due to a known flaky bug in the nvte_multi_tensor_gemm reference.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Clarify swap_dims comment in build_grouped_gemm_multi_inputA_args

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix grouped GEMM scale_inv offsets for NVFP4 and FP8 block scaling

Apply the same fix as upstream PR NVIDIA#2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Relax NVFP4 amax contiguity; consolidate scale_inv offset helpers; test cleanup

Production:
- nvte_grouped_gemm_with_discrete_inputA no longer requires per-expert amax
  buffers to be contiguous. Add `amax_ptrs[kMaxGroups]` to MultiTensorGroupGemmInputArgs
  and read each tensor's amax via indirection in setup_grouped_gemm_kernel
  (mirrors the existing scale_inv_ptrs pattern). The launcher enables the
  NVFP4 alpha computation when amax is available from either source.
- Consolidate four near-identical
  compute_grouped_tensor_{mxfp8,nvfp4,block_1d,block_2d}_scale_inv_offset
  into a single template `compute_grouped_scale_inv_offset<PaddedFn>` and
  collapse the A/B recipe-switch in setup_grouped_gemm_kernel into a local
  `fill_scale_ptr` lambda.

Tests:
- Drop the per-test amax staging workaround in run_grouped_gemm_discrete_in_case
  (no longer needed after the contiguity relax).
- Fix amax management in make_nvfp4_operand: copy values into result's own
  amax buffers instead of aliasing pointers (prevents double-free).
- Extract the three duplicated cuBLAS-version/compute-capability skip blocks
  into a shared `grouped_gemm_skip_reason` helper.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused float_size in GroupedGemmSetupWorkspace::from_buffers

Silences -Wunused-variable (NVIDIA#177-D in nvcc).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix Hopper grouped GEMM alpha beta handling

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address code review: NVFP4 amax check, swap_dims default, test refactor

- nvte_grouped_gemm and nvte_grouped_gemm_with_discrete_out now validate
  per-operand amax for NVFP4 (previously silently dropped the global-scale
  factor when amax was missing). discrete_inputA path also checks B's amax.
- Remove unused ShapeCase::kUnalignedAllSameNVFP4 enum and its comment.
- OperandStorageChoice::swap_dims now defaults to false; rowwise returns
  no longer pass spurious swap_dims=true.
- Unify GroupedGemmSetupWorkspace layout: from_buffers(nullptr, n) returns
  the total byte count, and required_setup_size derives its result from it
  so the layout cannot drift between the two.
- test_common.cu: consolidate the three gather_*_scales lambdas into a
  single gather_scale_inv(bytes_per_elem, get_shape, get_cpu_ptr) helper.
- test_grouped_gemm.cu: extract make_grouped_gemm_ref / make_alpha_beta /
  compare_grouped_d_to_multi helpers; the three run_* variants drop from
  ~1029 to 774 lines with no behavior change.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address code review feedback (NVIDIA#2971)

- discrete-A path: force non_tn_fp8_ok=false for FP8 block scaling to match
  select_grouped_operand logic for B.
- Setup workspace: explicit 16-byte base alignment NVTE_CHECK before
  GroupedGemmSetupWorkspace::from_buffers; matches contract every standard
  allocator (cudaMalloc / PyTorch / XLA) already satisfies.
- Tests: remove compile-time cuBLAS gating, run-time check via
  cuda::cublas_version() suffices since the test doesn't call cuBLAS directly.
- Tests: replace InputCase enum with std::optional<NVTEScalingMode>;
  nullopt = BF16, otherwise scaling mode drives dispatch.
- Tests: NVFP4 operand uses one shared BF16 input transposed via
  nvte_transpose for the columnwise direction (no more duplicate fillUniform
  with different seeds across row/col).
- Tests: parametrize output dtype (BF16 default, plus FP16 and FP32 cases on
  BF16/FP8/NVFP4 recipes); implementation already accepts all three.
- Tests: add NVFP4 2D-quantization coverage (verifies VEC16 scale layout
  fed to cuBLAS is unchanged vs 1D).
- Tests: tighten NVFP4 alignment check from %16 to %32 (TMA requirement of
  the optimized BF16 quantize path), fix misleading kUnalignedAllSame
  comment ({160,288,416} are multiples of 32, common to MXFP8 and NVFP4).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Simplify make_*_operand signatures to take use_rowwise directly

cuBLAS scaled-GEMM kernels all run in TN, so each operand only needs the
single direction matching its transpose flag. Mapping (is_A, transposed) ->
use_rowwise is uniform across MXFP8 / NVFP4 / FP8 block scaling; move it to
make_grouped_gemm_ref once instead of duplicating the if/else in each helper.

For NVFP4 specifically: drop nvte_transpose + the two-step intermediate
tensor assembly. nvte_quantize_v2 with a single-direction NVFP4 output
(rowwise XOR columnwise) directly produces what cuBLAS needs — columnwise
output goes via fallback quantize kernel (rowwise-only NVFP4 quantize kernel
hard-fails when output has no rowwise data, see quantize.cuh:111).

Matches the production pattern used by PyTorch and JAX bindings: never
allocate both NVFP4 directions on a single tensor (swizzle hard-fails when
both scale_invs are set, see swizzle.cu:985).

Net change: -112 lines.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Test infrastructure fixes for NVFP4 columnwise + skip unsupported combos

- test_common.cu: move NVFP4_1D_SCALING to the DELAYED/BLOCK_SCALING switch
  arm that allocates the columnwise data buffer as transpose(passed_shape).
  Required by TE's NVFP4 convention (column-wise data is transposed-then-
  quantized; see Tensor::shape() in common/common.h). Was incorrectly grouped
  with MXFP8 (whose columnwise data shape matches the logical shape).
  Without this, allocating columnwise-only NVFP4 in a test wires up a
  scale_inv shape that disagrees with what TE's CheckScaleTensorShape
  expects.

- libtransformer_engine.version: export
  transformer_engine::cuda::cublas_version so the test can read the run-time
  cuBLAS version (analogous to the already-exported sm_arch, sm_count,
  current_device).

- test_grouped_gemm.cu: simplify make_*_operand signatures to take
  use_rowwise directly, with the (is_A, transposed) -> use_rowwise mapping
  centralized in make_grouped_gemm_ref. NVFP4 now uses single-direction
  allocation (no nvte_transpose, no two-step assembly) — same pattern as
  MXFP8 / FP8 block scaling, working thanks to the test_common.cu fix.

- test_grouped_gemm.cu: tighten skip logic via grouped_gemm_skip_reason
  (now takes TestParams):
    * Skip NVFP4 + FP16 output  (cuBLAS hard-errors in cublaslt_gemm.cu)
    * Skip NVFP4 2D quantization in non-TN layouts (fallback quantize path
      doesn't support 2D; production weight quantization always allocates
      both directions and hits the optimized path).
  Drop the BF16 + FP16 output test case (cuBLAS grouped GEMM has no
  algorithm for that combination, even in TN).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Trim verbose comments in grouped GEMM test

Drop multi-line rationale comments that were conversation residue (linking
back to specific other files/lines, restating GEMM TN convention, listing
PR numbers) and keep just the short ones that name what the code does.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fix FP8 block scaling NT/NN failures on Hopper

Two related bugs in the grouped GEMM path were causing FP8 block scaling
NT and NN cases to fail on Hopper for non-Mul128 dims (Mul32 tests).

1. build_grouped_gemm_multi_inputA_args (discrete-A path) used
   t->shape() (LOGICAL) for cuBLAS rows/cols, while the symmetric
   grouped path (build_grouped_tensor / select_grouped_operand) uses
   tensors[i]->{rowwise,columnwise}_shape() (PHYSICAL). For FP8 block
   columnwise the two differ (physical is transposed of logical), so
   args.rows became N instead of K. Switch the discrete-A path to use
   data.shape directly — this aligns both paths and makes swap_dims
   redundant for this function.

2. padded_block_{1d,2d}_scale_inv_floats had a columnwise branch that
   swapped first/last in the formula, but the quantizer
   (test_common.cu:get_scales) always uses logical dims regardless of
   direction. Combined with grouped meta passing transposed dims for
   columnwise data, the formula produced wrong per-expert scale stride
   — only visible for dims not divisible by 128. The unified formula
   ceil(last/128) * roundup(first, 4) is correct for both directions
   because the meta swap and quantizer swap cancel out.

Also remove now-redundant check_fp4_output_compat duplication and trim
stale comments.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use logical grouped GEMM shapes

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix mixed FP8 block grouped GEMM scaling

Handle per-operand FP8 block scaling modes when computing grouped scale offsets, and tighten NVFP4 validation to reject unsupported mixed or non-per-group-alpha paths.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim unsupported FP8 block grouped GEMM tests

Avoid test cases that require direct columnwise 2D block quantization, which is not a supported test setup path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address review comment

Added validation for FP8 block scaling support in grouped GEMM.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tests to relax alignment requirements for swizzling tests

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Fix alignment calculation in required_setup_size

Adjust required_setup_size to account for alignment.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use nvfp4 alpha only for nvfp4

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Clean up grouped GEMM tests and document alignment

- Clarify scale_inv padding comments in test_grouped_gemm.cu so it is
  obvious that kAllSameMul32 exercises per-expert scale_inv offsets only
  for recipes that pad grouped scale_inv storage.
- Centralize the 32 MiB cuBLAS workspace size into a single
  kCublasWorkspaceBytes constant in test_grouped_gemm.cu.
- Drop the test-only enforce_grouped_gemm_alignment flag from
  build_grouped_tensor and remove all its call-site overrides; alignment
  is already validated inside the grouped GEMM implementation.
- Add a short comment in cublaslt_grouped_gemm.cu explaining why the int
  and float arrays in GroupedGemmSetupWorkspace do not require an
  explicit align_ptr() call (cuBLAS only mandates 16-byte alignment for
  the pointer arrays; 4-byte natural alignment is sufficient for the
  int/float arrays and is preserved by the layout).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Post-merge cleanup in grouped GEMM helpers and swizzle tests

- test_swizzle.cu: drop the enforce_grouped_gemm_alignment=false argument
  introduced upstream after the flag was removed from build_grouped_tensor.
- cublaslt_grouped_gemm.cu (_with_discrete_inputA): remove redundant B
  swizzled-scales check; the same condition is already enforced inside
  validate_grouped_gemm_inputs({inputB}, ...) called just above.
- build_grouped_gemm_multi_inputA_args: drop the requires_scale_inv
  parameter and derive it per-tensor from t->dtype() / t->scaling_mode.
  validate_grouped_gemm_multi_inputA_list already enforces uniform
  scaling mode, so per-iteration evaluation is safe and removes the
  duplicated computation at the call site.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: yangfan.bai <yangfan.bai@shopee.com>
Baibaifan pushed a commit to Baibaifan/TransformerEngine that referenced this pull request Jun 1, 2026
…hopper (NVIDIA#2971)

* code init

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* code drop

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Remove redundant nvte_set/get_grouped_tensor_swizzled_scales

Use existing nvte_set_grouped_tensor_param with kNVTEGroupedWithGEMMSwizzledScales
instead of the dedicated set/get functions.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Add Hopper support for grouped GEMM and refactor cuBLAS version checks

- Add CUBLAS_NVFP4_GROUPED_GEMM_VERSION and CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION macros (13.4+)
- Update check_grouped_gemm_requirements to allow SM90 with cuBLAS 13.4+
- Refactor execute_grouped_gemm to use GroupedGemmConfig struct
- Add divisibility-by-128 validation for FP8 block scaling in setup kernel and quantizer
- Support scalar alpha/beta for Hopper (no per-group alpha/beta)
- Expose get_grouped_gemm_setup_workspace_size to PyTorch via pybind
- Update PyTorch tests to run grouped GEMM on Hopper with cuBLAS 13.4+

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add NVFP4 support for discrete-input grouped GEMM and skip FP8 tensor scaling tests on Hopper

Extend nvte_grouped_gemm_with_discrete_inputA to handle NVFP4 (Float4E2M1)
inputs: accept kFloat4E2M1 dtype, propagate scale_inv pointers, collect
contiguous amax from discrete tensors, and enforce swizzled-scales checks
for NVFP4 alongside MXFP8. Also add GTEST_SKIP for FP8 tensor scaling
grouped GEMM on Hopper since cuBLAS does not support it there.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add alignment assertions for MXFP8/NVFP4 scale offsets in grouped GEMM tests

The setup kernel computes per-tensor scale pointers as data_offset /
block_size, which assumes no padding in the scale buffer. This is only
correct when first_dim % 128 == 0 and last_dim % 128 == 0 (MXFP8) or
last_dim % 64 == 0 (NVFP4). Add explicit assertions in
build_grouped_tensor to catch any future test shapes that violate this.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Fix grouped GEMM: NVFP4 columnwise transa=N + relax MXFP8 alignment for swizzle tests

cublaslt_grouped_gemm.cu:
- Fix incorrect handling of NVFP4/MXFP8 columnwise data in
  build_grouped_gemm_multi_inputA_args by adding a swap_dims flag
  consistent with choose_grouped_operand_storage. Use A_sel.trans
  (post-flip) for gemm_config.avg_k so K is selected from the
  correct dim with discrete A_list.

tests/cpp/test_common.{h,cu}:
- Add enforce_grouped_gemm_alignment parameter (default true) to
  build_grouped_tensor; the MXFP8/NVFP4 first/last_dim 128/64
  alignment asserts are only relevant for the grouped GEMM setup
  kernel, so callers that bypass it (swizzle/unswizzle) opt out.

tests/cpp/operator/test_swizzle.cu:
- Pass enforce_grouped_gemm_alignment=false to build_grouped_tensor
  in MXFP8 swizzle/unswizzle/roundtrip tests, which intentionally
  exercise non-padded shapes.

tests/cpp/operator/test_grouped_gemm.cu:
- Sync GPU/cuBLAS skip rules across all 3 sub-tests, add
  cudaDeviceSynchronize() after nvte_multi_tensor_gemm reference for
  defensive sync, and skip NVFP4 + AllDifferent in all 3 sub-tests
  due to a known flaky bug in the nvte_multi_tensor_gemm reference.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Clarify swap_dims comment in build_grouped_gemm_multi_inputA_args

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix grouped GEMM scale_inv offsets for NVFP4 and FP8 block scaling

Apply the same fix as upstream PR NVIDIA#2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Relax NVFP4 amax contiguity; consolidate scale_inv offset helpers; test cleanup

Production:
- nvte_grouped_gemm_with_discrete_inputA no longer requires per-expert amax
  buffers to be contiguous. Add `amax_ptrs[kMaxGroups]` to MultiTensorGroupGemmInputArgs
  and read each tensor's amax via indirection in setup_grouped_gemm_kernel
  (mirrors the existing scale_inv_ptrs pattern). The launcher enables the
  NVFP4 alpha computation when amax is available from either source.
- Consolidate four near-identical
  compute_grouped_tensor_{mxfp8,nvfp4,block_1d,block_2d}_scale_inv_offset
  into a single template `compute_grouped_scale_inv_offset<PaddedFn>` and
  collapse the A/B recipe-switch in setup_grouped_gemm_kernel into a local
  `fill_scale_ptr` lambda.

Tests:
- Drop the per-test amax staging workaround in run_grouped_gemm_discrete_in_case
  (no longer needed after the contiguity relax).
- Fix amax management in make_nvfp4_operand: copy values into result's own
  amax buffers instead of aliasing pointers (prevents double-free).
- Extract the three duplicated cuBLAS-version/compute-capability skip blocks
  into a shared `grouped_gemm_skip_reason` helper.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused float_size in GroupedGemmSetupWorkspace::from_buffers

Silences -Wunused-variable (NVIDIA#177-D in nvcc).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix Hopper grouped GEMM alpha beta handling

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address code review: NVFP4 amax check, swap_dims default, test refactor

- nvte_grouped_gemm and nvte_grouped_gemm_with_discrete_out now validate
  per-operand amax for NVFP4 (previously silently dropped the global-scale
  factor when amax was missing). discrete_inputA path also checks B's amax.
- Remove unused ShapeCase::kUnalignedAllSameNVFP4 enum and its comment.
- OperandStorageChoice::swap_dims now defaults to false; rowwise returns
  no longer pass spurious swap_dims=true.
- Unify GroupedGemmSetupWorkspace layout: from_buffers(nullptr, n) returns
  the total byte count, and required_setup_size derives its result from it
  so the layout cannot drift between the two.
- test_common.cu: consolidate the three gather_*_scales lambdas into a
  single gather_scale_inv(bytes_per_elem, get_shape, get_cpu_ptr) helper.
- test_grouped_gemm.cu: extract make_grouped_gemm_ref / make_alpha_beta /
  compare_grouped_d_to_multi helpers; the three run_* variants drop from
  ~1029 to 774 lines with no behavior change.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address code review feedback (NVIDIA#2971)

- discrete-A path: force non_tn_fp8_ok=false for FP8 block scaling to match
  select_grouped_operand logic for B.
- Setup workspace: explicit 16-byte base alignment NVTE_CHECK before
  GroupedGemmSetupWorkspace::from_buffers; matches contract every standard
  allocator (cudaMalloc / PyTorch / XLA) already satisfies.
- Tests: remove compile-time cuBLAS gating, run-time check via
  cuda::cublas_version() suffices since the test doesn't call cuBLAS directly.
- Tests: replace InputCase enum with std::optional<NVTEScalingMode>;
  nullopt = BF16, otherwise scaling mode drives dispatch.
- Tests: NVFP4 operand uses one shared BF16 input transposed via
  nvte_transpose for the columnwise direction (no more duplicate fillUniform
  with different seeds across row/col).
- Tests: parametrize output dtype (BF16 default, plus FP16 and FP32 cases on
  BF16/FP8/NVFP4 recipes); implementation already accepts all three.
- Tests: add NVFP4 2D-quantization coverage (verifies VEC16 scale layout
  fed to cuBLAS is unchanged vs 1D).
- Tests: tighten NVFP4 alignment check from %16 to %32 (TMA requirement of
  the optimized BF16 quantize path), fix misleading kUnalignedAllSame
  comment ({160,288,416} are multiples of 32, common to MXFP8 and NVFP4).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Simplify make_*_operand signatures to take use_rowwise directly

cuBLAS scaled-GEMM kernels all run in TN, so each operand only needs the
single direction matching its transpose flag. Mapping (is_A, transposed) ->
use_rowwise is uniform across MXFP8 / NVFP4 / FP8 block scaling; move it to
make_grouped_gemm_ref once instead of duplicating the if/else in each helper.

For NVFP4 specifically: drop nvte_transpose + the two-step intermediate
tensor assembly. nvte_quantize_v2 with a single-direction NVFP4 output
(rowwise XOR columnwise) directly produces what cuBLAS needs — columnwise
output goes via fallback quantize kernel (rowwise-only NVFP4 quantize kernel
hard-fails when output has no rowwise data, see quantize.cuh:111).

Matches the production pattern used by PyTorch and JAX bindings: never
allocate both NVFP4 directions on a single tensor (swizzle hard-fails when
both scale_invs are set, see swizzle.cu:985).

Net change: -112 lines.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Test infrastructure fixes for NVFP4 columnwise + skip unsupported combos

- test_common.cu: move NVFP4_1D_SCALING to the DELAYED/BLOCK_SCALING switch
  arm that allocates the columnwise data buffer as transpose(passed_shape).
  Required by TE's NVFP4 convention (column-wise data is transposed-then-
  quantized; see Tensor::shape() in common/common.h). Was incorrectly grouped
  with MXFP8 (whose columnwise data shape matches the logical shape).
  Without this, allocating columnwise-only NVFP4 in a test wires up a
  scale_inv shape that disagrees with what TE's CheckScaleTensorShape
  expects.

- libtransformer_engine.version: export
  transformer_engine::cuda::cublas_version so the test can read the run-time
  cuBLAS version (analogous to the already-exported sm_arch, sm_count,
  current_device).

- test_grouped_gemm.cu: simplify make_*_operand signatures to take
  use_rowwise directly, with the (is_A, transposed) -> use_rowwise mapping
  centralized in make_grouped_gemm_ref. NVFP4 now uses single-direction
  allocation (no nvte_transpose, no two-step assembly) — same pattern as
  MXFP8 / FP8 block scaling, working thanks to the test_common.cu fix.

- test_grouped_gemm.cu: tighten skip logic via grouped_gemm_skip_reason
  (now takes TestParams):
    * Skip NVFP4 + FP16 output  (cuBLAS hard-errors in cublaslt_gemm.cu)
    * Skip NVFP4 2D quantization in non-TN layouts (fallback quantize path
      doesn't support 2D; production weight quantization always allocates
      both directions and hits the optimized path).
  Drop the BF16 + FP16 output test case (cuBLAS grouped GEMM has no
  algorithm for that combination, even in TN).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Trim verbose comments in grouped GEMM test

Drop multi-line rationale comments that were conversation residue (linking
back to specific other files/lines, restating GEMM TN convention, listing
PR numbers) and keep just the short ones that name what the code does.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fix FP8 block scaling NT/NN failures on Hopper

Two related bugs in the grouped GEMM path were causing FP8 block scaling
NT and NN cases to fail on Hopper for non-Mul128 dims (Mul32 tests).

1. build_grouped_gemm_multi_inputA_args (discrete-A path) used
   t->shape() (LOGICAL) for cuBLAS rows/cols, while the symmetric
   grouped path (build_grouped_tensor / select_grouped_operand) uses
   tensors[i]->{rowwise,columnwise}_shape() (PHYSICAL). For FP8 block
   columnwise the two differ (physical is transposed of logical), so
   args.rows became N instead of K. Switch the discrete-A path to use
   data.shape directly — this aligns both paths and makes swap_dims
   redundant for this function.

2. padded_block_{1d,2d}_scale_inv_floats had a columnwise branch that
   swapped first/last in the formula, but the quantizer
   (test_common.cu:get_scales) always uses logical dims regardless of
   direction. Combined with grouped meta passing transposed dims for
   columnwise data, the formula produced wrong per-expert scale stride
   — only visible for dims not divisible by 128. The unified formula
   ceil(last/128) * roundup(first, 4) is correct for both directions
   because the meta swap and quantizer swap cancel out.

Also remove now-redundant check_fp4_output_compat duplication and trim
stale comments.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use logical grouped GEMM shapes

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix mixed FP8 block grouped GEMM scaling

Handle per-operand FP8 block scaling modes when computing grouped scale offsets, and tighten NVFP4 validation to reject unsupported mixed or non-per-group-alpha paths.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim unsupported FP8 block grouped GEMM tests

Avoid test cases that require direct columnwise 2D block quantization, which is not a supported test setup path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address review comment

Added validation for FP8 block scaling support in grouped GEMM.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tests to relax alignment requirements for swizzling tests

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Fix alignment calculation in required_setup_size

Adjust required_setup_size to account for alignment.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use nvfp4 alpha only for nvfp4

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Clean up grouped GEMM tests and document alignment

- Clarify scale_inv padding comments in test_grouped_gemm.cu so it is
  obvious that kAllSameMul32 exercises per-expert scale_inv offsets only
  for recipes that pad grouped scale_inv storage.
- Centralize the 32 MiB cuBLAS workspace size into a single
  kCublasWorkspaceBytes constant in test_grouped_gemm.cu.
- Drop the test-only enforce_grouped_gemm_alignment flag from
  build_grouped_tensor and remove all its call-site overrides; alignment
  is already validated inside the grouped GEMM implementation.
- Add a short comment in cublaslt_grouped_gemm.cu explaining why the int
  and float arrays in GroupedGemmSetupWorkspace do not require an
  explicit align_ptr() call (cuBLAS only mandates 16-byte alignment for
  the pointer arrays; 4-byte natural alignment is sufficient for the
  int/float arrays and is preserved by the layout).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Post-merge cleanup in grouped GEMM helpers and swizzle tests

- test_swizzle.cu: drop the enforce_grouped_gemm_alignment=false argument
  introduced upstream after the flag was removed from build_grouped_tensor.
- cublaslt_grouped_gemm.cu (_with_discrete_inputA): remove redundant B
  swizzled-scales check; the same condition is already enforced inside
  validate_grouped_gemm_inputs({inputB}, ...) called just above.
- build_grouped_gemm_multi_inputA_args: drop the requires_scale_inv
  parameter and derive it per-tensor from t->dtype() / t->scaling_mode.
  validate_grouped_gemm_multi_inputA_list already enforces uniform
  scaling mode, so per-iteration evaluation is safe and removes the
  duplicated computation at the call site.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: yangfan.bai <yangfan.bai@shopee.com>
shangxiaokang pushed a commit to shangxiaokang/TransformerEngine that referenced this pull request Jul 10, 2026
* fix CUBLAS for GPT oss sizes

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* ci fix

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* Fix test case dimensions in test_numerics.py

Total dim should be divisible by 128

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

---------

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
shangxiaokang pushed a commit to shangxiaokang/TransformerEngine that referenced this pull request Jul 10, 2026
…VIDIA#2971)

* code init

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* code drop

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Remove redundant nvte_set/get_grouped_tensor_swizzled_scales

Use existing nvte_set_grouped_tensor_param with kNVTEGroupedWithGEMMSwizzledScales
instead of the dedicated set/get functions.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Add Hopper support for grouped GEMM and refactor cuBLAS version checks

- Add CUBLAS_NVFP4_GROUPED_GEMM_VERSION and CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION macros (13.4+)
- Update check_grouped_gemm_requirements to allow SM90 with cuBLAS 13.4+
- Refactor execute_grouped_gemm to use GroupedGemmConfig struct
- Add divisibility-by-128 validation for FP8 block scaling in setup kernel and quantizer
- Support scalar alpha/beta for Hopper (no per-group alpha/beta)
- Expose get_grouped_gemm_setup_workspace_size to PyTorch via pybind
- Update PyTorch tests to run grouped GEMM on Hopper with cuBLAS 13.4+

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add NVFP4 support for discrete-input grouped GEMM and skip FP8 tensor scaling tests on Hopper

Extend nvte_grouped_gemm_with_discrete_inputA to handle NVFP4 (Float4E2M1)
inputs: accept kFloat4E2M1 dtype, propagate scale_inv pointers, collect
contiguous amax from discrete tensors, and enforce swizzled-scales checks
for NVFP4 alongside MXFP8. Also add GTEST_SKIP for FP8 tensor scaling
grouped GEMM on Hopper since cuBLAS does not support it there.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Add alignment assertions for MXFP8/NVFP4 scale offsets in grouped GEMM tests

The setup kernel computes per-tensor scale pointers as data_offset /
block_size, which assumes no padding in the scale buffer. This is only
correct when first_dim % 128 == 0 and last_dim % 128 == 0 (MXFP8) or
last_dim % 64 == 0 (NVFP4). Add explicit assertions in
build_grouped_tensor to catch any future test shapes that violate this.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Fix grouped GEMM: NVFP4 columnwise transa=N + relax MXFP8 alignment for swizzle tests

cublaslt_grouped_gemm.cu:
- Fix incorrect handling of NVFP4/MXFP8 columnwise data in
  build_grouped_gemm_multi_inputA_args by adding a swap_dims flag
  consistent with choose_grouped_operand_storage. Use A_sel.trans
  (post-flip) for gemm_config.avg_k so K is selected from the
  correct dim with discrete A_list.

tests/cpp/test_common.{h,cu}:
- Add enforce_grouped_gemm_alignment parameter (default true) to
  build_grouped_tensor; the MXFP8/NVFP4 first/last_dim 128/64
  alignment asserts are only relevant for the grouped GEMM setup
  kernel, so callers that bypass it (swizzle/unswizzle) opt out.

tests/cpp/operator/test_swizzle.cu:
- Pass enforce_grouped_gemm_alignment=false to build_grouped_tensor
  in MXFP8 swizzle/unswizzle/roundtrip tests, which intentionally
  exercise non-padded shapes.

tests/cpp/operator/test_grouped_gemm.cu:
- Sync GPU/cuBLAS skip rules across all 3 sub-tests, add
  cudaDeviceSynchronize() after nvte_multi_tensor_gemm reference for
  defensive sync, and skip NVFP4 + AllDifferent in all 3 sub-tests
  due to a known flaky bug in the nvte_multi_tensor_gemm reference.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* Clarify swap_dims comment in build_grouped_gemm_multi_inputA_args

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix grouped GEMM scale_inv offsets for NVFP4 and FP8 block scaling

Apply the same fix as upstream PR NVIDIA#2954 (MXFP8 unaligned dims) to the
analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel.

Background: cuBLAS grouped GEMM expects each expert's scale_inv to live
at a specific offset in the contiguous grouped buffer. The quantizer
allocates each per-expert scale_inv tensor padded to the layout cuBLAS
needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for
block scaling). The setup kernel was computing these offsets as
data_offset / block_size for everything except MXFP8 — silently correct
when dims align to 128, but pointing at the middle of the previous
expert's scale tile when they do not. In MoE forward this is reachable
through variable per-expert token counts.

Add three device helpers mirroring compute_grouped_tensor_mxfp8_-
scale_inv_offset:
- compute_grouped_tensor_nvfp4_scale_inv_offset
- compute_grouped_tensor_block_1d_scale_inv_offset
- compute_grouped_tensor_block_2d_scale_inv_offset
Each sums the same padded per-tensor sizes the quantizer uses at alloc
time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_-
shape).

NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so
sel.shape is already pre-transposed for that recipe — the rowwise
formula on (first, last) recovers the colwise alloc. For block scaling
the formula depends on the canonical orientation, so propagate a new
swap_dims field on GroupedOperandSelection and pass effective_rowwise
(sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant
under this change because swap_dims is always false there and its
helper's byte count is invariant under the rowwise flag anyway.

Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416)
— all multiples of 32/16 (per-recipe block size) but none multiples of
128, so each expert's scale tile is padded. Exercise it across MXFP8 /
NVFP4 / FP8 block scaling and the three transpose configs that match
the existing parameter grid. Relax build_grouped_tensor's defensive
%128 / %64 alignment assertions to %32 / %16 (block-size only), which
is the actual quantizer requirement now that the offset arithmetic no
longer assumes zero padding.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Relax NVFP4 amax contiguity; consolidate scale_inv offset helpers; test cleanup

Production:
- nvte_grouped_gemm_with_discrete_inputA no longer requires per-expert amax
  buffers to be contiguous. Add `amax_ptrs[kMaxGroups]` to MultiTensorGroupGemmInputArgs
  and read each tensor's amax via indirection in setup_grouped_gemm_kernel
  (mirrors the existing scale_inv_ptrs pattern). The launcher enables the
  NVFP4 alpha computation when amax is available from either source.
- Consolidate four near-identical
  compute_grouped_tensor_{mxfp8,nvfp4,block_1d,block_2d}_scale_inv_offset
  into a single template `compute_grouped_scale_inv_offset<PaddedFn>` and
  collapse the A/B recipe-switch in setup_grouped_gemm_kernel into a local
  `fill_scale_ptr` lambda.

Tests:
- Drop the per-test amax staging workaround in run_grouped_gemm_discrete_in_case
  (no longer needed after the contiguity relax).
- Fix amax management in make_nvfp4_operand: copy values into result's own
  amax buffers instead of aliasing pointers (prevents double-free).
- Extract the three duplicated cuBLAS-version/compute-capability skip blocks
  into a shared `grouped_gemm_skip_reason` helper.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused float_size in GroupedGemmSetupWorkspace::from_buffers

Silences -Wunused-variable (NVIDIA#177-D in nvcc).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix Hopper grouped GEMM alpha beta handling

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* fix

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address code review: NVFP4 amax check, swap_dims default, test refactor

- nvte_grouped_gemm and nvte_grouped_gemm_with_discrete_out now validate
  per-operand amax for NVFP4 (previously silently dropped the global-scale
  factor when amax was missing). discrete_inputA path also checks B's amax.
- Remove unused ShapeCase::kUnalignedAllSameNVFP4 enum and its comment.
- OperandStorageChoice::swap_dims now defaults to false; rowwise returns
  no longer pass spurious swap_dims=true.
- Unify GroupedGemmSetupWorkspace layout: from_buffers(nullptr, n) returns
  the total byte count, and required_setup_size derives its result from it
  so the layout cannot drift between the two.
- test_common.cu: consolidate the three gather_*_scales lambdas into a
  single gather_scale_inv(bytes_per_elem, get_shape, get_cpu_ptr) helper.
- test_grouped_gemm.cu: extract make_grouped_gemm_ref / make_alpha_beta /
  compare_grouped_d_to_multi helpers; the three run_* variants drop from
  ~1029 to 774 lines with no behavior change.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address code review feedback (NVIDIA#2971)

- discrete-A path: force non_tn_fp8_ok=false for FP8 block scaling to match
  select_grouped_operand logic for B.
- Setup workspace: explicit 16-byte base alignment NVTE_CHECK before
  GroupedGemmSetupWorkspace::from_buffers; matches contract every standard
  allocator (cudaMalloc / PyTorch / XLA) already satisfies.
- Tests: remove compile-time cuBLAS gating, run-time check via
  cuda::cublas_version() suffices since the test doesn't call cuBLAS directly.
- Tests: replace InputCase enum with std::optional<NVTEScalingMode>;
  nullopt = BF16, otherwise scaling mode drives dispatch.
- Tests: NVFP4 operand uses one shared BF16 input transposed via
  nvte_transpose for the columnwise direction (no more duplicate fillUniform
  with different seeds across row/col).
- Tests: parametrize output dtype (BF16 default, plus FP16 and FP32 cases on
  BF16/FP8/NVFP4 recipes); implementation already accepts all three.
- Tests: add NVFP4 2D-quantization coverage (verifies VEC16 scale layout
  fed to cuBLAS is unchanged vs 1D).
- Tests: tighten NVFP4 alignment check from %16 to %32 (TMA requirement of
  the optimized BF16 quantize path), fix misleading kUnalignedAllSame
  comment ({160,288,416} are multiples of 32, common to MXFP8 and NVFP4).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Simplify make_*_operand signatures to take use_rowwise directly

cuBLAS scaled-GEMM kernels all run in TN, so each operand only needs the
single direction matching its transpose flag. Mapping (is_A, transposed) ->
use_rowwise is uniform across MXFP8 / NVFP4 / FP8 block scaling; move it to
make_grouped_gemm_ref once instead of duplicating the if/else in each helper.

For NVFP4 specifically: drop nvte_transpose + the two-step intermediate
tensor assembly. nvte_quantize_v2 with a single-direction NVFP4 output
(rowwise XOR columnwise) directly produces what cuBLAS needs — columnwise
output goes via fallback quantize kernel (rowwise-only NVFP4 quantize kernel
hard-fails when output has no rowwise data, see quantize.cuh:111).

Matches the production pattern used by PyTorch and JAX bindings: never
allocate both NVFP4 directions on a single tensor (swizzle hard-fails when
both scale_invs are set, see swizzle.cu:985).

Net change: -112 lines.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Test infrastructure fixes for NVFP4 columnwise + skip unsupported combos

- test_common.cu: move NVFP4_1D_SCALING to the DELAYED/BLOCK_SCALING switch
  arm that allocates the columnwise data buffer as transpose(passed_shape).
  Required by TE's NVFP4 convention (column-wise data is transposed-then-
  quantized; see Tensor::shape() in common/common.h). Was incorrectly grouped
  with MXFP8 (whose columnwise data shape matches the logical shape).
  Without this, allocating columnwise-only NVFP4 in a test wires up a
  scale_inv shape that disagrees with what TE's CheckScaleTensorShape
  expects.

- libtransformer_engine.version: export
  transformer_engine::cuda::cublas_version so the test can read the run-time
  cuBLAS version (analogous to the already-exported sm_arch, sm_count,
  current_device).

- test_grouped_gemm.cu: simplify make_*_operand signatures to take
  use_rowwise directly, with the (is_A, transposed) -> use_rowwise mapping
  centralized in make_grouped_gemm_ref. NVFP4 now uses single-direction
  allocation (no nvte_transpose, no two-step assembly) — same pattern as
  MXFP8 / FP8 block scaling, working thanks to the test_common.cu fix.

- test_grouped_gemm.cu: tighten skip logic via grouped_gemm_skip_reason
  (now takes TestParams):
    * Skip NVFP4 + FP16 output  (cuBLAS hard-errors in cublaslt_gemm.cu)
    * Skip NVFP4 2D quantization in non-TN layouts (fallback quantize path
      doesn't support 2D; production weight quantization always allocates
      both directions and hits the optimized path).
  Drop the BF16 + FP16 output test case (cuBLAS grouped GEMM has no
  algorithm for that combination, even in TN).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Trim verbose comments in grouped GEMM test

Drop multi-line rationale comments that were conversation residue (linking
back to specific other files/lines, restating GEMM TN convention, listing
PR numbers) and keep just the short ones that name what the code does.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Fix FP8 block scaling NT/NN failures on Hopper

Two related bugs in the grouped GEMM path were causing FP8 block scaling
NT and NN cases to fail on Hopper for non-Mul128 dims (Mul32 tests).

1. build_grouped_gemm_multi_inputA_args (discrete-A path) used
   t->shape() (LOGICAL) for cuBLAS rows/cols, while the symmetric
   grouped path (build_grouped_tensor / select_grouped_operand) uses
   tensors[i]->{rowwise,columnwise}_shape() (PHYSICAL). For FP8 block
   columnwise the two differ (physical is transposed of logical), so
   args.rows became N instead of K. Switch the discrete-A path to use
   data.shape directly — this aligns both paths and makes swap_dims
   redundant for this function.

2. padded_block_{1d,2d}_scale_inv_floats had a columnwise branch that
   swapped first/last in the formula, but the quantizer
   (test_common.cu:get_scales) always uses logical dims regardless of
   direction. Combined with grouped meta passing transposed dims for
   columnwise data, the formula produced wrong per-expert scale stride
   — only visible for dims not divisible by 128. The unified formula
   ceil(last/128) * roundup(first, 4) is correct for both directions
   because the meta swap and quantizer swap cancel out.

Also remove now-redundant check_fp4_output_compat duplication and trim
stale comments.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use logical grouped GEMM shapes

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix mixed FP8 block grouped GEMM scaling

Handle per-operand FP8 block scaling modes when computing grouped scale offsets, and tighten NVFP4 validation to reject unsupported mixed or non-per-group-alpha paths.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim unsupported FP8 block grouped GEMM tests

Avoid test cases that require direct columnwise 2D block quantization, which is not a supported test setup path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Address review comment

Added validation for FP8 block scaling support in grouped GEMM.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tests to relax alignment requirements for swizzling tests

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Fix alignment calculation in required_setup_size

Adjust required_setup_size to account for alignment.

Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use nvfp4 alpha only for nvfp4

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

* Clean up grouped GEMM tests and document alignment

- Clarify scale_inv padding comments in test_grouped_gemm.cu so it is
  obvious that kAllSameMul32 exercises per-expert scale_inv offsets only
  for recipes that pad grouped scale_inv storage.
- Centralize the 32 MiB cuBLAS workspace size into a single
  kCublasWorkspaceBytes constant in test_grouped_gemm.cu.
- Drop the test-only enforce_grouped_gemm_alignment flag from
  build_grouped_tensor and remove all its call-site overrides; alignment
  is already validated inside the grouped GEMM implementation.
- Add a short comment in cublaslt_grouped_gemm.cu explaining why the int
  and float arrays in GroupedGemmSetupWorkspace do not require an
  explicit align_ptr() call (cuBLAS only mandates 16-byte alignment for
  the pointer arrays; 4-byte natural alignment is sufficient for the
  int/float arrays and is preserved by the layout).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Post-merge cleanup in grouped GEMM helpers and swizzle tests

- test_swizzle.cu: drop the enforce_grouped_gemm_alignment=false argument
  introduced upstream after the flag was removed from build_grouped_tensor.
- cublaslt_grouped_gemm.cu (_with_discrete_inputA): remove redundant B
  swizzled-scales check; the same condition is already enforced inside
  validate_grouped_gemm_inputs({inputB}, ...) called just above.
- build_grouped_gemm_multi_inputA_args: drop the requires_scale_inv
  parameter and derive it per-tensor from t->dtype() / t->scaling_mode.
  validate_grouped_gemm_multi_inputA_list already enforces uniform
  scaling mode, so per-iteration evaluation is safe and removes the
  duplicated computation at the call site.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants