Repository navigation
Conversation
… MLP forward When the FC1 GEMM + activation wrapper that will be launched advertises supports_canonical_layouts and the NVFP4 input scales are swizzled, pass A as (sum(m), k), prob as (sum(m),) and the swizzled scale buffer as-is instead of the 3-D and 6-D kernel-facing views. The decision follows the kernel choice, so the fused GEMM + activation + RHT + quant kernel, which does not advertise canonical layouts, keeps the legacy layouts, as do unswizzled scales, FC2 and the backward. Signed-off-by: Hwanseo Choi <hwanseoc@nvidia.com>
 [Medium impact] The PR appears safe to merge; the remaining test suggestion is non-blocking. Findings
SummaryThe PR lets supported cuDNN wrappers receive NVFP4 operands in their natural layouts.
hwanseoc acknowledges the known flat-scale check cost in cuDNN FE 1.32.1 and keeps this path because the measured overall speedup remains positive. Diagram%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Select FC1 or dGLU wrapper] --> B{Wrapper supports canonical layouts?}
B -->|No| L[Use legacy views]
B -->|Yes| C{NVFP4 scales are swizzled?}
C -->|No| L
C -->|Yes| N[Pass natural operand layouts]
W[Weight-gradient wrapper] --> S{Supports canonical layouts?}
S -->|Yes| F[Pass scale buffers as-is]
S -->|No| V[Build two-dimensional scale views]
Reviews (5) · Last reviewed commit: "Merge branch 'main' into hwanseoc/canoni..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
hwanseoc
changed the title
50d94c2 to
f16bfe8
Compare
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
|
Tested locally at f16bfe8 with
TE CI's container has cuDNN FE 1.29, so CI only exercises the legacy layout path. The canonical paths need FE 1.31 (GLU) or 1.32 (Hadamard GLU and dGLU). |
… MLP backward When the dGLU wrapper advertises supports_canonical_layouts and the NVFP4 grad output scales are swizzled, pass the grad output as (sum(m), k), the probabilities as (sum(m),) and the swizzled scale buffer as-is. NVFP4 runs the FC1 dgrad on cuBLAS, so it does not need the quant wrapper's support. A single grouped NVFP4 weight keeps its legacy layout. The NVFP4 wgrad now shares the MXFP8 scale handling and passes flat scale buffers when the wgrad wrapper supports them. With cuDNN frontend 1.32.x, checking flat FP4 scale buffers costs about 15 us per wgrad call; NVIDIA/cudnn-frontend#2289 makes that check cheap. Signed-off-by: Hwanseo Choi <hwanseoc@nvidia.com>
f16bfe8 to
4ee7c8f
Compare
|
Update: the head is now 4ee7c8f, which also passes flat wgrad scales for NVFP4. |
|
/te-ci pytorch |
| and getattr(dactivation_kernel, "supports_canonical_layouts", False) | ||
| and getattr(self.grouped_gemm_quant_kernel(), "supports_canonical_layouts", False) | ||
| # NVFP4 runs the FC1 dgrad on cuBLAS, so only the dGLU kernel needs | ||
| # canonical support, with swizzled scales. |
There was a problem hiding this comment.
Nit: This is invalidating the previous "Pack data tensors" comment. Previously that one comment was enough to describe this entire section, but if we need this additional comment then that's a sign that we should restructure. I'd suggest having a separate section to figure out whether or not to use a canonical layout, and leave this code section dedicated to the non-canonical layout reshapes.
As a general pattern, I observe that agents prefer minimal line changes even if it makes the code quality worse.
| # NVFP4 runs the FC1 dgrad on cuBLAS, so only the dGLU kernel needs | ||
| # canonical support, with swizzled scales. |
There was a problem hiding this comment.
Nit: This comment is awkward if you're not aware of the edit history. Compare "NVFP4 doesn't use the cuDNN kernel, so we don't need to check" vs "MXFP8 uses the cuDNN kernel, so we need to check". Agents will overemphasize the last thing they worked on and don't consider how it affects the final code.
Description
Reduce host overhead in the NVFP4 fused grouped MLP by passing operands to the cuDNN grouped GEMM wrappers in their natural layouts: the forward FC1 GEMM, and the backward dGLU and wgrad.
#3641 did this for MXFP8 but kept NVFP4 on the legacy layouts. The cuDNN frontend GLU wrapper accepts canonical layouts since 1.31; the Hadamard GLU (NVIDIA/cudnn-frontend#1601), dGLU and wgrad wrappers since 1.32.
Type of change
Changes
supports_canonical_layoutsand the NVFP4 input scales are swizzled, TE passes A as(sum(m), k), the probabilities as(sum(m),)and the swizzled scale buffer as is, instead of 3-D and 6-D kernel-facing views. The check reads the flag of the kernel that is actually launched; MXFP8 behaves as before.single_grouped_weight=True, which NVFP4 does not support yet.Which path NVFP4 takes, by cuDNN frontend version:
Measurements
NVFP4 (
NVFP4BlockScaling()),ScaledSwiGLU(glu_interleave_size=32), per-expert weights. Host p50 in µs,main(720b252) → this PR, median of three interleaved rounds per cuDNN frontend:With FE 1.32.1, most of the backward saving goes to the flat wgrad scale check (about 15 µs per FP4 call); NVIDIA/cudnn-frontend#2289 removes it. GPU kernel time is unchanged.
Setup: SM100, cuDNN 9.25.1, CUDA 13.4, torch 2.14.0a0, CuTeDSL 4.7.1, process pinned to an idle CPU core, 300 synchronized samples per process.
Testing
tests/pytorch/test_grouped_mlp.pyat this PR's head:TE CI's container has FE 1.29, so CI exercises only the legacy path.
NVFP4 RHT, 4 experts, canonical vs legacy layouts on the same inputs, FE 1.32.1: with canonical layouts on, all four cuDNN calls (FC1, dGLU and both wgrads) receive canonical operands. The forward output, input and probability gradients, and every weight gradient are bitwise identical for SwiGLU (Hadamard GLU) and GeGLU (GLU). Bias gradients use atomics and differ between any two runs, legacy included.
Checklist:
@coderabbitai ignore