Variable Grouped Swizzle by int-smart · Pull Request #2914 · NVIDIA/TransformerEngine

int-smart · 2026-04-22T07:41:28Z

Description

Grouped Swizzle with variable shape. Not sure if this is needed but if not can be closed.

Fixes #2451

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

… update C++ operator interface Signed-off-by: Abhishek <abhi.dtu11@gmail.com>

Signed-off-by: Abhishek <abhi.dtu11@gmail.com>

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

greptile-apps · 2026-04-22T07:46:05Z

Greptile Summary

This PR extends swizzle_grouped_scaling_factors to handle grouped tensors with non-uniform (M, K) shapes per tensor. It adds a new persistent-grid CUDA kernel that uses an intra-block warp reduction to compute a total-block count and then dispatches each logical block to its owning tensor via an inline prefix-sum scan, while retaining the existing uniform-shape fast path unchanged. The test helper build_grouped_tensor is also corrected to produce num_tensors + 1 offsets, aligning it with the production-side contract already enforced in transformer_engine.cpp.

Confidence Score: 5/5

Safe to merge — only P2 findings (unchecked sync return value in test, O(num_tensors) kernel scan); no correctness bugs identified.

All findings are P2 (style/performance). The core kernel logic, shared-memory usage, and persistent-grid dispatch are correct. The offsets fix in the test helper resolves a latent mismatch with the production validator.

transformer_engine/common/swizzle/swizzle.cu — variable-shape kernel inner scan loop; tests/cpp/operator/test_swizzle.cu — unchecked cudaDeviceSynchronize return value.

Important Files Changed

Filename	Overview
transformer_engine/common/swizzle/swizzle.cu	Adds `grouped_swizzle_scaling_variable_shape_kernel` — a persistent-grid kernel that handles per-tensor variable M/K for grouped MXFP8 swizzle. Core logic is correct but the per-block O(num_tensors) scan is a scalability concern.
tests/cpp/test_common.cu	Grows the offsets vector from `num_tensors` to `num_tensors + 1` elements — a correctness fix aligning the test helper with the production-side contract enforced in `transformer_engine.cpp`.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp	Removes the guard that rejected non-uniform `first_dims`/`last_dims`, enabling the variable-shape code path. No logic changes beyond unlocking the feature.
tests/cpp/operator/test_swizzle.cu	Adds `performTestGroupedSwizzleMXFP8Variable` covering six shape permutations. Minor: `cudaDeviceSynchronize()` return value is unchecked.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[swizzle_grouped_scaling_factors] --> B{is_variable_shape?}
    B -- No --> C[Uniform-shape path]
    B -- Yes --> D[Variable-shape path]
    D --> E[Compute common_m / common_k]
    E --> F[Size persistent grid via cudaOccupancy]
    F --> G[grouped_swizzle_scaling_variable_shape_kernel]
    G --> H[Warp 0 computes s_total_blocks via warp reduction]
    H --> I[__syncthreads]
    I --> J[Persistent loop: linear_block_id in 0..total_blocks]
    J --> K[O-n scan to find tensor_id and current_scale_base]
    K --> L{rowwise?}
    L -- Yes --> M[swizzle_row_scaling_kernel_impl]
    L -- No --> N[swizzle_col_scaling_kernel_impl]

_{Reviews (3): Last reviewed commit: "Merge branch 'main' into feat/variable_s..." | Re-trigger Greptile}

greptile-apps · 2026-04-22T07:46:09Z

+      int device_id;
+      cudaGetDevice(&device_id);
+      int num_SMs;
+      cudaDeviceGetAttribute(&num_SMs, cudaDevAttrMultiProcessorCount, device_id);
+      // Find out how many blocks of this specific kernel can fit on one SM
+      int max_active_blocks_per_sm;
+      cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+          &max_active_blocks_per_sm,
+          grouped_swizzle_scaling_variable_shape_kernel<SF_TILE_DIM_M, SF_TILE_DIM_K>,
+          TB_DIM * TB_DIM,  // block size
+          max_slm_size      // dynamic shared memory
+      );
+      int persistent_blocks = num_SMs * max_active_blocks_per_sm;


Unchecked CUDA API calls can silently produce zero-block launches

cudaGetDevice, cudaDeviceGetAttribute, and cudaOccupancyMaxActiveBlocksPerMultiprocessor are all called without NVTE_CHECK_CUDA. If any of these fail, max_active_blocks_per_sm is left with an indeterminate (or zero) value, making persistent_blocks = 0. Launching the persistent kernel with 0 blocks is legal in CUDA — it silently does nothing — so the output buffer stays uninitialized with no error raised.

Suggested change

int device_id;

cudaGetDevice(&device_id);

int num_SMs;

cudaDeviceGetAttribute(&num_SMs, cudaDevAttrMultiProcessorCount, device_id);

// Find out how many blocks of this specific kernel can fit on one SM

int max_active_blocks_per_sm;

cudaOccupancyMaxActiveBlocksPerMultiprocessor(

&max_active_blocks_per_sm,

grouped_swizzle_scaling_variable_shape_kernel<SF_TILE_DIM_M, SF_TILE_DIM_K>,

TB_DIM * TB_DIM, // block size

max_slm_size // dynamic shared memory

);

int persistent_blocks = num_SMs * max_active_blocks_per_sm;

int device_id;

NVTE_CHECK_CUDA(cudaGetDevice(&device_id));

int num_SMs;

NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&num_SMs, cudaDevAttrMultiProcessorCount, device_id));

// Find out how many blocks of this specific kernel can fit on one SM

int max_active_blocks_per_sm;

NVTE_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor(

&max_active_blocks_per_sm,

grouped_swizzle_scaling_variable_shape_kernel<SF_TILE_DIM_M, SF_TILE_DIM_K>,

TB_DIM * TB_DIM, // block size

max_slm_size // dynamic shared memory

));

NVTE_CHECK(max_active_blocks_per_sm > 0, "Occupancy query returned 0 blocks per SM.");

int persistent_blocks = num_SMs * max_active_blocks_per_sm;

greptile-apps · 2026-04-22T07:46:10Z

+  if (!is_variable_shape) {
+    // Fallback to uniform shape implementation
+    NVTE_CHECK(input->all_same_shape(), "Grouped swizzle requires uniform tensor shapes.");
+    NVTE_CHECK(input->all_same_last_dim() && input->all_same_first_dim(),
+               "Grouped swizzle requires uniform tensor shapes.");


Dead code: redundant assertions inside !is_variable_shape branch

is_variable_shape is defined as !input->all_same_shape(), so inside if (!is_variable_shape) the two NVTE_CHECK calls are tautologies — they can never fire. They add noise and could mislead future readers into thinking the branch can handle non-uniform shapes. Consider removing them or converting them to a comment.

Suggested change

if (!is_variable_shape) {

// Fallback to uniform shape implementation

NVTE_CHECK(input->all_same_shape(), "Grouped swizzle requires uniform tensor shapes.");

NVTE_CHECK(input->all_same_last_dim() && input->all_same_first_dim(),

"Grouped swizzle requires uniform tensor shapes.");

if (!is_variable_shape) {

// All tensors share the same shape; use the optimised uniform-shape path.

greptile-apps · 2026-04-22T07:46:10Z

+    if (int_stride % 2 != 0) int_stride++;
+    int* d_block_offsets = reinterpret_cast<int*>(workspace);
+    int* d_global_counter = d_block_offsets + num_tensors + 1;
+    int* d_total_blocks = d_global_counter + 1;


d_total_blocks is written but never consumed

d_total_blocks is populated by compute_grouped_swizzle_setup (as *total_blocks = current_block_offset) but is never read by the persistent kernel or any host code afterward. The persistent grid terminates via the tensor_id == -1 sentinel, not via a stored total. If this field was intended as a diagnostic or future guard, a comment would clarify its purpose; otherwise it can be removed to avoid confusing future maintainers and wasting a device-side write.

greptile-apps · 2026-04-22T07:46:11Z

+  size_t num_tensors = input.num_tensors();
+  size_t num_int_elems = num_tensors + 3;       // n+1 block_offsets + gc + tb
+  if (num_int_elems % 2 != 0) num_int_elems++;  // pad to even for size_t alignment
+  size_t workspace_size = num_int_elems * sizeof(int) + (num_tensors + 1) * sizeof(size_t);
+  workspace_size = roundup(workspace_size, 256);
+  auto workspace =
+      allocateSpace(std::vector<size_t>{workspace_size}, transformer_engine::DType::kByte, false);
+
  NVTE_SCOPED_GIL_RELEASE({
    nvte_swizzle_grouped_scaling_factors(swizzle_input.data(), swizzle_output.data(),
-                                         at::cuda::getCurrentCUDAStream());
+                                         getDataPtr(workspace), at::cuda::getCurrentCUDAStream());


Workspace allocated unconditionally even for uniform-shape inputs

The workspace is only consumed by the variable-shape code path in swizzle_grouped_scaling_factors. For uniform shapes the pointer is accepted but immediately ignored. Gating the allocation on whether variable shapes are present (e.g., first_dims.data_ptr != nullptr || last_dims.data_ptr != nullptr) would avoid a small but unnecessary device allocation on every invocation with uniform tensors. This is a performance suggestion, not a correctness issue.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

vthumbe1503 · 2026-04-22T22:43:30Z

+          cudaFuncAttributeMaxDynamicSharedMemorySize, max_slm_size));
+
+      int device_id;
+      cudaGetDevice(&device_id);


Caching the attrbues like number of sms and max active blocks per device would be ideal to reduce CPU overheads on each call.

We already have a function in transformer_engine/common/util/cuda_runtime.cpp called "sm_count". Could you please use that here?

vthumbe1503 · 2026-04-22T22:46:15Z

+    // Fallback to uniform shape implementation
+    NVTE_CHECK(input->all_same_shape(), "Grouped swizzle requires uniform tensor shapes.");
+    NVTE_CHECK(input->all_same_last_dim() && input->all_same_first_dim(),
+               "Grouped swizzle requires uniform tensor shapes.");


These checks might not be needed. Given we used input->all_same_shape() to reach this stage

vthumbe1503

I think that workspace allocation + small kernel for computing offsets + persistent kernel might be an overkill for swizzling. @int-smart Do you have some performance numbers by any chance for the swizzling kernel on Blackwell?

How about we follow a SM filling grid pattern like in grouped_bias_add kernel in this PR?
https://github.com/NVIDIA/TransformerEngine/pull/2885/changes/BASE..b64559af9b89d816b8d7ffba4f5273e556d90c8e#diff-fa75cbeb11caf588f79b811be355c8f00b0cf5d4b807c259b94f2a40ffc8db6f

With this pattern thread block id is dynamically decided based on sum(first_dims) and at the same time we divide the rows of grouped_tensor uniformly among the SMs. However it only handles variable first_dims(Need to extend the idea for other cases like all dims being variable)

int-smart · 2026-04-22T23:38:34Z

@vthumbe1503 Will check the PR and get back

int-smart · 2026-04-22T23:41:37Z

With regards to Blackwell I dont have the numbers tbh. I can generate it for RTX 40 series

Signed-off-by: Abhishek <abhi.dtu11@gmail.com>

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

int-smart · 2026-04-25T23:48:14Z

@vthumbe1503 Consolidated to one kernel, removed shared memory allocation and tried to stick to the PR you mentioned. If this works let me know. Seems to perform better on rtx 4070 than my last approach. There are still some optimizations can be done but that would need more shared memory alloc.

vthumbe1503 · 2026-04-27T18:38:31Z

+      int num_SMs;
+      cudaDeviceGetAttribute(&num_SMs, cudaDevAttrMultiProcessorCount, device_id);
+      int max_active_blocks_per_sm;
+      cudaOccupancyMaxActiveBlocksPerMultiprocessor(


Can we add NVTE_CHECK_CUDA around these APIs.

vthumbe1503 · 2026-04-27T19:48:13Z

Thanks for the PR @int-smart. Overall LGTM. Have a few minor comments w.r.t CPU overheads and error handling.

vthumbe1503 · 2026-04-27T19:48:19Z

/te-ci

int-smart and others added 3 commits April 21, 2026 23:54

feat: add support for grouped GEMM swizzling with variable shapes and…

e362860

… update C++ operator interface Signed-off-by: Abhishek <abhi.dtu11@gmail.com>

Added confirmation with uniformity in one of the dimensions

ab16ec7

Signed-off-by: Abhishek <abhi.dtu11@gmail.com>

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

16735fd

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

greptile-apps Bot reviewed Apr 22, 2026

View reviewed changes

vthumbe1503 self-requested a review April 22, 2026 22:33

vthumbe1503 reviewed Apr 22, 2026

View reviewed changes

vthumbe1503 self-assigned this Apr 22, 2026

int-smart and others added 2 commits April 25, 2026 15:55

Using single kernel for variable m and k

749692c

Signed-off-by: Abhishek <abhi.dtu11@gmail.com>

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

5836d2c

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

vthumbe1503 reviewed Apr 27, 2026

View reviewed changes

Merge branch 'main' into feat/variable_swizzle

56d23c0

Conversation

int-smart commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

greptile-apps Bot Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Apr 22, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 left a comment

Choose a reason for hiding this comment

Uh oh!

int-smart commented Apr 22, 2026

Uh oh!

int-smart commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

int-smart commented Apr 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vthumbe1503 Apr 27, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 commented Apr 27, 2026

Uh oh!

vthumbe1503 commented Apr 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

int-smart commented Apr 22, 2026 •

edited

Loading

greptile-apps Bot commented Apr 22, 2026 •

edited

Loading

int-smart commented Apr 22, 2026 •

edited

Loading

int-smart commented Apr 25, 2026 •

edited

Loading