[WanZzzzzz]

Description

Skip fetching the next unused tile in the for-loop logic. This optimization give 5% additional perf improvement for mxfp8 cast-only + swizzle kernel.

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

[greptile-apps]

Greptile Summary

This PR optimizes the MXFP8 cast-only and swizzle kernel by avoiding a speculative fetch after the final iteration-space tile.

• Adds an iteration-bound check to the pipelined prefetch condition.
• Preserves fetching for all four tiles consumed by the reachable two-stage pipeline.
• Avoids unnecessary wait and commit operations when there is no next tile.

Confidence Score: 5/5

The PR appears safe to merge because the added guard skips only an unused out-of-range prefetch while preserving every fetch required by the consumer pipeline.

The changed kernel has one reachable configuration with four iteration tiles and two pipeline stages. Tiles 0–3 are fetched before consumption, while the new condition becomes false only for next tile index 4 after tile 3 has already been staged.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh	Adds a correct iteration-space bound to prevent fetching a nonexistent tile after the last pipelined MXFP8 quantization iteration.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  P["Prologue: fetch tile 0"] --> I0["Iteration 0: fetch tile 1; consume tile 0"]
  I0 --> I1["Iteration 1: fetch tile 2; consume tile 1"]
  I1 --> I2["Iteration 2: fetch tile 3; consume tile 2"]
  I2 --> I3["Iteration 3: skip next=4 fetch; consume tile 3"]
  I3 --> O["Commit output"]

Loading

_{Reviews (1): Last reviewed commit: "mxfp8 quant: skip fetching next tile" | Re-trigger Greptile}

[tdophung]

LGTM

[Oleg-Goncharov]

/te-ci

[Oleg-Goncharov]

LGTM. We can merge it once the CI tests pass.