IFU 2.17 upstream feature enablement - #692
Draft
AllenFarcas wants to merge 3 commits into
Draft
Conversation
The grouped MXFP8 scale swizzle refused any group whose members differ in shape, which is the MoE case where each expert receives a different token count. Two things blocked it, both in the device-side block scheduler; the tile kernels it dispatches to were already portable. The workload measurement used a wave32 butterfly over a 32-bit shuffle mask. HIP requires a 64-bit mask and full-wave participation, and the reduction ran under a half-populated wave. Replace it with a serial scan by one thread: the kernel already performs the same O(num_tensors) scan on every persistent-grid iteration, so measuring it once costs strictly less than work already being done, and it leaves no assumption about how many lanes execute in lockstep. The dynamic shared memory request was the tile staging area plus one int for the block count. The staging area alone is exactly 64 KiB, which is also the per-workgroup LDS ceiling on gfx942-class parts, so the extra word made the launch unsatisfiable there. Alias the count into the first word of the staging buffer instead; every thread copies it into a register before the barrier that hands the buffer over. Two defects in the path being enabled, both reachable on either platform: The tile implementation leaves no trailing barrier, so consecutive iterations of the persistent loop overlapped in the staging area. Add the barrier at the top of the loop body rather than at the end of the implementation, which is shared with the non-persistent kernels where it would be pure cost. An expert routed no tokens yields zero tiles, and the load-width expression then evaluated to zero and divided by it. Return an empty tiling instead. The per-tensor geometry was also duplicated verbatim between the counting and resolution passes; drift between the two copies would have selected the wrong tensor's scales rather than failing. Give it one definition. The test builds its scale blocks directly rather than gathering them from test::Tensor, whose MXFP8 scale allocation is unpadded on ROCm and so carries a per-tensor stride this entry point does not expect. That also makes a zero-token expert expressible, and lets one test body serve both platforms. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Blocks ran four to six lines where this file's convention is one or two. The reasoning they carried belongs in review notes rather than the source. Each why is kept -- the gfx942 LDS ceiling, the missing trailing barrier, the unpadded ROCm scale stride, the divide-by-zero -- and the justification around it is not. No code change. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commits changed code CUDA compiles: the block-count reduction, the shared memory request, and the variable-shape test body. None of that was needed to enable the feature on ROCm, and it put the fork ahead of upstream on a path upstream owns and tests. The kernel is now split. CUDA keeps upstream's implementation verbatim, wave32 shuffle and all, including its persistent-loop barrier omission and its zero-tile divide -- those are upstream defects to report, not to patch here. ROCm gets its own arm with the serial scan, the 64 KiB shared memory request, and both fixes. The test follows the same split: upstream's suite is restored unchanged for CUDA, and the ROCm layout builder becomes a separate suite. test::Tensor emits unpadded MXFP8 scales on ROCm, so the two cannot share a body. Every deletion against upstream is now guard scaffolding or the ROCm NVTE_ERROR. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implemented IFU 2.17 upstream features:
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: