[PyTorch][common] NVFP4: enable row-scaled transpose quantization for backward#3206
[PyTorch][common] NVFP4: enable row-scaled transpose quantization for backward#3206cael-ling wants to merge 2 commits into
Conversation
Greptile SummaryThis PR extends the row-scaled NVFP4 quantization path (originally forward-only from #2931) to emit a columnwise (transposed) output in the same kernel pass, enabling the wgrad GEMM in
Confidence Score: 5/5Safe to merge. The core quantization logic, dispatch guards, bilateral GEMM post-scaling, and amax buffer sizing are all correct. The only finding is a memory access inefficiency in the new pre-pass kernel. The two-pass amax-then-quantize design is sound, the ROW_SCALED_NVFP4 template path correctly overrides S_enc_colwise_block before use, the bilateral outer-product scaling in gemm.py reproduces the #2931 fprop result bit-for-bit for the scalar case and extends it correctly to the vector case, and the quantizer.cpp amax buffer resizing handles both the scalar-to-vector and vector-to-scalar transitions. A column-strided access pattern in the new pre-pass amax kernel is the only thing worth addressing before the next performance-sensitive deployment. transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh — the new compute_columnwise_amax_kernel should be reviewed for memory access efficiency before scaling to very large matrices. Important Files Changed
Sequence DiagramsequenceDiagram
participant FWD as Forward Pass
participant QFwd as quantize_fwd_helper
participant ColAmax as compute_columnwise_amax
participant RowAmax as compute_rowwise_amax
participant Kernel as quantize_transpose_tuned_1D ROW_SCALED=true
participant GEMM as general_gemm bilateral post-scale
participant BWD as Backward Pass
participant QBwd as quantize_bwd_helper
FWD->>QFwd: "input X row_scaled=true columnwise=true"
QFwd->>RowAmax: pre-pass amax_rowwise M
QFwd->>ColAmax: pre-pass amax_columnwise K
QFwd->>Kernel: quantize both directions
Kernel-->>FWD: rowwise FP4 fprop + columnwise FP4 wgrad
FWD->>GEMM: "fprop X_fp4 row @ W_fp4 T TN layout"
GEMM-->>FWD: out times B.amax_row times A.amax_scalar
BWD->>QBwd: "grad dY row_scaled=true if applicable"
QBwd->>RowAmax: pre-pass
QBwd->>ColAmax: pre-pass if columnwise
QBwd->>Kernel: quantize
BWD->>GEMM: "wgrad X_fp4 col T @ dY_fp4 bilateral scales"
GEMM-->>BWD: "dW = out times output_row_scales outer output_col_scales"
Reviews (5): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| static void* persistent_buffer(size_t bytes, cudaStream_t stream, int which) { | ||
| static void* bufs[2] = {nullptr, nullptr}; | ||
| static size_t caps[2] = {0, 0}; | ||
| if (bytes > caps[which]) { | ||
| if (bufs[which] != nullptr) { | ||
| NVTE_CHECK_CUDA(cudaFreeAsync(bufs[which], stream)); | ||
| } | ||
| const size_t newcap = bytes + bytes / 2; // slack to avoid frequent regrows | ||
| NVTE_CHECK_CUDA(cudaMallocAsync(&bufs[which], newcap, stream)); | ||
| caps[which] = newcap; | ||
| } | ||
| return bufs[which]; |
There was a problem hiding this comment.
When two grouped NVFP4 GEMMs run on different streams, both calls reuse the same static metadata and workspace buffers. One launch can overwrite or free the scratch while another stream's CUTLASS kernel is still reading it, which can produce wrong grouped GEMM outputs or allocator failures.
| cutlass::KernelHardwareInfo hw_info; | ||
| hw_info.device_id = 0; | ||
| hw_info.sm_count = cached_sm_count(); |
There was a problem hiding this comment.
This launch always reports device 0 to CUTLASS, but the API runs on the caller's current CUDA device. On a multi-GPU job where the active device is not 0, the grouped scheduler can use the wrong SM count/device id and size persistent work incorrectly, causing poor scheduling or launch failures.
| cutlass::KernelHardwareInfo hw_info; | |
| hw_info.device_id = 0; | |
| hw_info.sm_count = cached_sm_count(); | |
| int device_id = 0; | |
| NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); | |
| cutlass::KernelHardwareInfo hw_info; | |
| hw_info.device_id = device_id; | |
| hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(device_id); |
| NVTE_CHECK(a_t->data.dtype == DType::kFloat4E2M1 && b_t->data.dtype == DType::kFloat4E2M1, | ||
| "group ", g, ": A/B must be FP4 e2m1"); | ||
| NVTE_CHECK((d_t->data.dtype == DType::kFloat32) == d_is_fp32, "group ", g, | ||
| ": D dtype must be uniform across groups"); |
There was a problem hiding this comment.
The grouped API validates A/B/D and alpha tensors but never checks a_sf[g] or b_sf[g] before passing their bytes to CUTLASS as NVFP4 scale factors. A caller that passes an FP32 or otherwise wrong scale tensor gets silently corrupted GEMM results instead of the same clear error the dense APIs return.
| NVTE_CHECK(a_t->data.dtype == DType::kFloat4E2M1 && b_t->data.dtype == DType::kFloat4E2M1, | |
| "group ", g, ": A/B must be FP4 e2m1"); | |
| NVTE_CHECK((d_t->data.dtype == DType::kFloat32) == d_is_fp32, "group ", g, | |
| ": D dtype must be uniform across groups"); | |
| NVTE_CHECK(a_t->data.dtype == DType::kFloat4E2M1 && b_t->data.dtype == DType::kFloat4E2M1, | |
| "group ", g, ": A/B must be FP4 e2m1"); | |
| NVTE_CHECK(sa_t->data.dtype == DType::kFloat8E4M3 || sa_t->data.dtype == DType::kByte, | |
| "group ", g, ": A scale must be FP8 e4m3 (or raw uint8 byte)"); | |
| NVTE_CHECK(sb_t->data.dtype == DType::kFloat8E4M3 || sb_t->data.dtype == DType::kByte, | |
| "group ", g, ": B scale must be FP8 e4m3 (or raw uint8 byte)"); | |
| NVTE_CHECK((d_t->data.dtype == DType::kFloat32) == d_is_fp32, "group ", g, | |
| ": D dtype must be uniform across groups"); |
github-actions
Bot
added
the
community-contribution
e5f6f9b to
bec279d
Compare
|
The kernels are written based on the assumption that moe parameters are treated as separate tensors, and inputs to the kernel are tensors that are pointers list, which is then not graph safe, instead of a grouped tensor, which is then graph safe. For recipes that target RL, you can still assume this design - non-grouped tensor, used by the grouped_linear module, where we have the most flexibility in terms of adding new recipes. For recipes targeting pretraining, it needs to justify itself by showing better performance than mxfp8 path with grouepd quantization + grouped gemm with activation & FC2 quant fusion + being safe with cuda graph. |
bec279d to
ffda05f
Compare
|
Reworked per review feedback: dropped the standalone CUTLASS operators, folded the backward capability into the #2931 row-scaled path as a minimal dense extension. Grouped/MoE support will follow in a separate PR. |
cael-ling
changed the title
… backward The NVIDIA#2931 row-scaled NVFP4 path only produced the rowwise forward activation; its columnwise/transpose output was rejected. That made the per-token activation unusable in the backward wgrad GEMM, so row-scaled training had to fall back to a dequantized/high-precision backward. This change extends the existing row-scaled path to also emit the columnwise (transpose) direction, so a training Linear with row_scaled_activation=True now quantizes the forward activation row-scaled in both directions and the wgrad GEMM consumes the row-scaled transpose directly. It is a minimal extension of NVIDIA#2931 (no new CUTLASS kernels, no grouped path, no RHT/4over6 transpose fusion). Signed-off-by: Cael Ling <caell@nvidia.com>
44c7da3 to
9d42b33
Compare
for more information, see https://pre-commit.ci
|
Hi @zhongbozhu, thanks a lot for the earlier feedback. I've reworked this PR, accordingly: I dropped the standalone CUTLASS operators and folded the backward capability into the #2931 row-scaled path as a minimal dense extension. To be clear on scope, this targets the RL path only — dense |
| const IType *__restrict__ input, | ||
| float *__restrict__ output_columnwise_amax, | ||
| const float *__restrict__ noop) { | ||
| #if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) |
There was a problem hiding this comment.
No reason to limit this to Blackwell, this is actually a very general kernel. We could also optimize by vectorizing reads so that each warp loads a full 128 byte cache line, but we can defer until we actually observe that this is a perf bottleneck.
|
/te-ci |
3 participants
Description
Extends the #2931 row-scaled NVFP4 path to emit the column-wise direction, enabling backward NVFP4 computation for dense
Linearlayers withrow_scaled_activation=True. Previously row-scaled path produced only the row-wise forward activation; its column-wiseoutput was rejected, so the per-token activation could not feed the backward wgrad GEMM and row-scaled training had to fall back to a dequantized/high-precision backward.With this change the forward activation is quantized row-scaled in both directions and the wgrad GEMM consumes the row-scaled transpose directly.
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: