From 00b8d306c9f101303de5c11a11c86c705811d629 Mon Sep 17 00:00:00 2001 From: Cael Ling Date: Wed, 8 Jul 2026 21:00:17 -0700 Subject: [PATCH 1/2] [PyTorch] NVFP4: emit GEMM-swizzled scales in the non-RHT 2D quantize kernel Fixes a CUDA graph hang/corruption with NVFP4 weight caching. Scope: the non-RHT 2D quantize path on 128-aligned shapes, which is the default cached-weight case. Other paths (1D weight scaling, non-128-aligned, etc) keep the compact + post-quantize swizzle fallback and are left for follow-ups. - swizzle.cuh: NVFP4 128x4 GEMM-swizzled scale-index helper (byte-compatible with nvte_swizzle_scaling_factors). - quantize_transpose_nvfp4_2D_kernel: add WITH_GEMM_SWIZZLED_SCALES; write rowwise and columnwise scales at swizzled offsets; relax the compact-only guard to allow the 2D swizzled path on 128-aligned dims. - dispatch: thread with_gemm_swizzled_scales through the blockwise fallback (which already implements kSwizzledScale) instead of hardcoding false. - quantizer: is_eligible_for_2d_swizzle_fusion + shared gate so the non-RHT 2D path enables swizzled SF on 128-aligned shapes; the post-quantize swizzle fallback auto-skips when the flag is set. - tests: byte-equal SF vs nvte_swizzle_scaling_factors (rowwise/columnwise/ both, multiple shapes) + shape gate; cached-weight scale-pointer stability and CUDA graph capture/replay regression. Signed-off-by: Cael Ling --- .../test_nvfp4_2d_quantize_swizzle_fusion.py | 241 ++++++++++++++++++ .../nvfp4/test_nvfp4_weight_swizzle_cache.py | 220 ++++++++++++++++ .../common/cast/dispatch/quantize.cuh | 4 +- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 73 ++++-- .../common/cast/nvfp4/swizzle.cuh | 58 +++++ transformer_engine/pytorch/csrc/common.h | 10 + transformer_engine/pytorch/csrc/quantizer.cpp | 42 ++- 7 files changed, 622 insertions(+), 26 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_2d_quantize_swizzle_fusion.py create mode 100644 tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py create mode 100644 transformer_engine/common/cast/nvfp4/swizzle.cuh diff --git a/tests/pytorch/nvfp4/test_nvfp4_2d_quantize_swizzle_fusion.py b/tests/pytorch/nvfp4/test_nvfp4_2d_quantize_swizzle_fusion.py new file mode 100644 index 0000000000..cfc279f886 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_2d_quantize_swizzle_fusion.py @@ -0,0 +1,241 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fidelity tests for NVFP4 (non-RHT) 2D-quantize swizzled SF output. + +Mirrors ``test_nvfp4_rht_quantize_swizzle_fusion.py`` but for the plain +2D-block-scaling quantize kernel used by cached weights: when the quantizer +has ``optimize_for_gemm=True``, ``with_2d_quantization=True``, ``with_rht=False`` +and the shape is 128-aligned (so no SF padding is needed), the quantize kernel +should emit scale factors directly in the GEMM-swizzled layout ``cuBLAS LT`` +consumes, eliminating the otherwise-required ``nvte_swizzle_scaling_factors`` +pass between quantize and GEMM. + +The fidelity contract is: ``quantizer_swizzle_fusion(x)`` produces SF that are +byte-equal to ``swizzle_nvfp4_scale(quantizer(x).sx)``. The ``_rowwise_data`` / +``_columnwise_data`` quantized FP4 buffers and amaxes must also be byte-equal +(the swizzle optimization changes only the SF layout, not the FP4 data itself). +""" + +from typing import Tuple + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex # noqa: F401 +from transformer_engine.pytorch import NVFP4Quantizer +from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import ( + NVFP4TensorStorage, +) + +from nvfp4_utils import ( + swizzle_nvfp4_scale, + get_nvfp4_scale_shape_no_padding, +) + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def _unpack_quantized_tensor( + quantized_tensor: NVFP4TensorStorage, +) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Extract ``(qx, sx, amax_row, qx_t, sx_t, amax_col)`` for byte comparison.""" + qx, sx, amax_row = None, None, None + qx_t, sx_t, amax_col = None, None, None + if quantized_tensor._rowwise_data is not None: + qx = quantized_tensor._rowwise_data.view(dtype=torch.uint8) + if quantized_tensor._rowwise_scale_inv is not None: + sx = quantized_tensor._rowwise_scale_inv + if quantized_tensor._amax_rowwise is not None: + amax_row = quantized_tensor._amax_rowwise + if quantized_tensor._columnwise_data is not None: + qx_t = quantized_tensor._columnwise_data.view(dtype=torch.uint8) + if quantized_tensor._columnwise_scale_inv is not None: + sx_t = quantized_tensor._columnwise_scale_inv + if quantized_tensor._amax_columnwise is not None: + amax_col = quantized_tensor._amax_columnwise + return qx, sx, amax_row, qx_t, sx_t, amax_col + + +def _make_quantizer(return_rowwise: bool, return_transpose: bool) -> NVFP4Quantizer: + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=return_rowwise, + columnwise=return_transpose, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=True, + ) + + +def _check_nvfp4_2d_quantize_swizzle_fusion( + x_dtype: torch.dtype, + M: int, + N: int, + return_rowwise: bool, + return_transpose: bool, +) -> None: + device = "cuda" + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + x = torch.randn((M, N), dtype=x_dtype, device=device) + + # Reference quantizer: compact SF (default, no optimize_for_gemm). + quantizer = _make_quantizer(return_rowwise, return_transpose) + + # SUT: same quantizer but with swizzled-SF emission enabled. The kernel must + # bake the swizzled layout in directly (no standalone swizzle pass). + quantizer_swizzle_fusion = quantizer.copy() + quantizer_swizzle_fusion.optimize_for_gemm = True + + sut = quantizer_swizzle_fusion(x) + assert sut._with_gemm_swizzled_scales, ( + f"expected in-kernel swizzled SF for aligned shape ({M}, {N}); " + "the 2D quantize kernel should have baked the swizzled layout in" + ) + + qx_swf, sx_swf, amax_row_swf, qx_t_swf, sx_t_swf, amax_col_swf = _unpack_quantized_tensor(sut) + qx_ref, sx_ref, amax_row_ref, qx_t_ref, sx_t_ref, amax_col_ref = _unpack_quantized_tensor( + quantizer(x) + ) + + if return_rowwise: + # FP4 data buffer and amax must be byte-equal (swizzle only changes SF + # layout, not the quantized data). + torch.testing.assert_close(qx_swf, qx_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(amax_row_swf, amax_row_ref, atol=0.0, rtol=0.0) + + # SF tensor must match the swizzle of the reference compact SF. + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x.shape, False) + assert valid_scale_shape == sx_swf.shape, ( + "rowwise SF shape mismatch; this test assumes the input shape needs no " + f"SF padding (got valid={valid_scale_shape}, got_swf={sx_swf.shape})." + ) + sx_ref_swizzled = swizzle_nvfp4_scale(M, N, sx_ref, columnwise=False) + torch.testing.assert_close(sx_swf, sx_ref_swizzled, atol=0.0, rtol=0.0) + + if return_transpose: + torch.testing.assert_close(qx_t_swf, qx_t_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(amax_col_swf, amax_col_ref, atol=0.0, rtol=0.0) + + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x.shape, True) + assert valid_scale_shape == sx_t_swf.shape, ( + "columnwise SF shape mismatch; this test assumes the input shape needs no " + f"SF padding (got valid={valid_scale_shape}, got_swf={sx_t_swf.shape})." + ) + sx_t_ref_swizzled = swizzle_nvfp4_scale(M, N, sx_t_ref, columnwise=True) + torch.testing.assert_close(sx_t_swf, sx_t_ref_swizzled, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # 128-aligned shapes (no SF padding needed) -> in-kernel swizzle eligible. + (128, 128), + (256, 256), + (1024, 1024), + (2048, 2048), + (8192, 1024), + (8192, 5120), + (8192, 10240), + (16384, 8192), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) +def test_nvfp4_2d_quantize_swizzle_fusion( + x_dtype: torch.dtype, + M: int, + N: int, + quantize_mode: str, +) -> None: + if quantize_mode == "rowwise_only": + return_rowwise, return_transpose = True, False + elif quantize_mode == "both_directions": + return_rowwise, return_transpose = True, True + elif quantize_mode == "columnwise_only": + return_rowwise, return_transpose = False, True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + + _check_nvfp4_2d_quantize_swizzle_fusion( + x_dtype=x_dtype, + M=M, + N=N, + return_rowwise=return_rowwise, + return_transpose=return_transpose, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N, expected_swizzled", + [ + # Eligible: both dims % 128 == 0 -> 2D quantize kernel bakes swizzled SF. + (128, 128, True), + (256, 512, True), + # Ineligible (cols % 128 != 0) -> kernel emits compact SF; create_tensor + # reports swizzled=False (the fused kernel did not bake it in). + (128, 192, False), + # Ineligible (rows % 128 != 0). + (64, 128, False), + ], +) +def test_nvfp4_2d_swizzle_fusion_shape_gate(M: int, N: int, expected_swizzled: bool) -> None: + """``create_tensor`` must report in-kernel swizzle only for 128-aligned shapes. + + ``make_empty`` runs ``create_tensor`` only (no quantize kernel, no + post-quantize ``inplace_swizzle_scale_for_gemm`` fallback), so its + ``_with_gemm_swizzled_scales`` observes the in-kernel gate in isolation. + """ + quantizer = _make_quantizer(return_rowwise=True, return_transpose=True) + quantizer.optimize_for_gemm = True + out = quantizer.make_empty([M, N], dtype=torch.bfloat16) + assert out._with_gemm_swizzled_scales is expected_swizzled, ( + f"2D-kernel shape gate expected _with_gemm_swizzled_scales={expected_swizzled} " + f"for shape ({M}, {N}), got {out._with_gemm_swizzled_scales}" + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + (128, 128), + (256, 512), + # Ineligible shapes still end up swizzled via the post-quantize fallback. + (128, 192), + (64, 128), + ], +) +def test_nvfp4_2d_swizzle_fusion_end_to_end_swizzled(M: int, N: int) -> None: + """End-to-end quantize must never crash and must always yield swizzled SF. + + Eligible (128-aligned) shapes get swizzled SF directly from the 2D quantize + kernel; ineligible shapes are gated off the in-kernel path and get swizzled + by the post-quantize ``inplace_swizzle_scale_for_gemm`` fallback. Either way + the final ``_with_gemm_swizzled_scales`` must be True. + """ + quantizer = _make_quantizer(return_rowwise=True, return_transpose=True) + quantizer.optimize_for_gemm = True + x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda") + + result = quantizer(x) + assert result._with_gemm_swizzled_scales is True, ( + "End-to-end quantize expected _with_gemm_swizzled_scales=True for shape " + f"({M}, {N}) with optimize_for_gemm=True + with_2d_quantization=True, " + f"got {result._with_gemm_swizzled_scales}" + ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py b/tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py new file mode 100644 index 0000000000..b4cb2ee1ce --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py @@ -0,0 +1,220 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Tests for cached NVFP4 weight scale swizzle (in-kernel swizzled SF output). + +The NVFP4 2D quantize kernel writes GEMM-swizzled scale factors directly into +the cached weight's scale buffer, so there is no separate out-of-place swizzle +pass that could reallocate and rebind the scale pointer on a cache hit. These +tests guard that (a) the cached scale buffers keep a stable address across +weight updates, (b) weight caching stays numerically correct, and (c) a CUDA +graph that captured a GEMM against the cached weight survives an eager +re-quantize (the original hang/corruption).""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import NVFP4BlockScaling + + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + +pytestmark = pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) + + +def _make_module(kind, in_features, out_features, device, num_gemms=1): + common = dict(bias=True, params_dtype=torch.bfloat16) + if kind == "Linear": + return te.Linear(in_features, out_features, **common).to(device) + if kind == "LayerNormLinear": + return te.LayerNormLinear(in_features, out_features, **common).to(device) + if kind == "LayerNormMLP": + return te.LayerNormMLP(in_features, out_features, **common).to(device) + if kind == "GroupedLinear": + return te.GroupedLinear(num_gemms, in_features, out_features, **common).to(device) + raise ValueError(f"unknown module kind {kind}") + + +def _clone_params(src, dst): + with torch.no_grad(): + dst_params = dict(dst.named_parameters()) + for name, param in src.named_parameters(): + dst_params[name].copy_(param) + + +def _step(module, x, is_first, recipe, m_splits=None): + x = x.detach().clone().requires_grad_(True) + module.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + if m_splits is None: + out = module(x, is_first_microbatch=is_first) + else: + out = module(x, m_splits, is_first_microbatch=is_first) + out.sum().backward() + return out.detach().float(), x.grad.detach().float() + + +_MODULE_KINDS = ["Linear", "LayerNormLinear", "LayerNormMLP", "GroupedLinear"] + + +def _grouped_m_splits(kind, batch, num_gemms): + if kind != "GroupedLinear": + return None + return [batch // num_gemms] * num_gemms + + +@pytest.mark.parametrize("kind", ["Linear", "LayerNormMLP"]) +def test_cached_weight_scale_pointers_stable(kind): + """Cached weight scale buffers must not be reallocated across weight updates.""" + torch.manual_seed(0) + device = "cuda" + batch = 512 + recipe = NVFP4BlockScaling(disable_stochastic_rounding=True) + module = _make_module(kind, 1024, 1024, device) + m_splits = _grouped_m_splits(kind, batch, 1) + + x0 = torch.randn(batch, 1024, dtype=torch.bfloat16, device=device) + x1 = torch.randn(batch, 1024, dtype=torch.bfloat16, device=device) + + _step(module, x0, True, recipe, m_splits) + workspaces = module._fp8_workspaces + assert workspaces, "expected cached weight workspaces after first microbatch" + + ptrs_before = {} + for name, ws in workspaces.items(): + assert getattr(ws, "_with_gemm_swizzled_scales", False) + ptrs_before[name] = ( + ws._rowwise_scale_inv.data_ptr(), + ws._columnwise_scale_inv.data_ptr(), + ) + + with torch.no_grad(): + for param in module.parameters(): + if param.ndim >= 2: + param.add_(torch.randn_like(param) * 1e-3) + + _step(module, x1, True, recipe, m_splits) + + for name, ws in workspaces.items(): + ptrs_after = ( + ws._rowwise_scale_inv.data_ptr(), + ws._columnwise_scale_inv.data_ptr(), + ) + assert ptrs_after == ptrs_before[name], ( + f"cached weight workspace {name!r} reallocated scale buffers on weight update" + ) + assert getattr(ws, "_with_gemm_swizzled_scales", False) + + +@pytest.mark.parametrize("kind", _MODULE_KINDS) +@pytest.mark.parametrize("microbatches", [1, 4]) +def test_weight_swizzle_cache_numerics(kind, microbatches): + torch.manual_seed(1234) + device = "cuda" + in_features, out_features = 1024, 1024 + batch = 512 + num_gemms = 2 if kind == "GroupedLinear" else 1 + m_splits = _grouped_m_splits(kind, batch, num_gemms) + recipe = NVFP4BlockScaling(disable_stochastic_rounding=True) + + ref = _make_module(kind, in_features, out_features, device, num_gemms) + opt = _make_module(kind, in_features, out_features, device, num_gemms) + _clone_params(ref, opt) + + inputs = [ + torch.randn(batch, in_features, dtype=torch.bfloat16, device=device) + for _ in range(microbatches) + ] + + atol, rtol = 1e-3, 1e-3 + for mb in range(microbatches): + ref_out, ref_dgrad = _step(ref, inputs[mb], None, recipe, m_splits) + opt_out, opt_dgrad = _step(opt, inputs[mb], mb == 0, recipe, m_splits) + torch.testing.assert_close(opt_out, ref_out, atol=atol, rtol=rtol) + torch.testing.assert_close(opt_dgrad, ref_dgrad, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize("kind", ["Linear", "LayerNormLinear"]) +def test_cached_weight_swizzle_scale_cuda_graph_replay(kind): + """A captured CUDA graph must stay valid after an eager cached-weight update. + + Reproduces the reviewer's hang/corruption: a CUDA graph captures a forward + that consumes the cached weight (``is_first_microbatch=False`` -> no + re-quantize inside the graph), so the captured GEMM hard-codes the swizzled + scale address. The next batch's first microbatch re-quantizes the same cached + tensor eagerly (outside the graph). On the buggy path this reallocates the + swizzled scale and frees the captured address, so graph replay reads stale + memory. With persistent scale buffers the address is fixed and replay is + unaffected. + """ + torch.manual_seed(0) + device = "cuda" + batch = 512 + recipe = NVFP4BlockScaling(disable_stochastic_rounding=True) + module = _make_module(kind, 1024, 1024, device) + + # Populate the weight cache (first microbatch). + x_warm = torch.randn(batch, 1024, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=recipe): + module(x_warm, is_first_microbatch=True) + + ws = module._fp8_workspaces["weight"] + addr_before = ws._rowwise_scale_inv.data_ptr() + + # Static input the captured graph reads from. + x_static = torch.randn(batch, 1024, dtype=torch.bfloat16, device=device) + + # Warmup on a side stream (required before graph capture), then capture a + # forward that consumes the cached weight without re-quantizing it. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + with te.autocast(enabled=True, recipe=recipe): + module(x_static, is_first_microbatch=False) + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + with te.autocast(enabled=True, recipe=recipe): + y_static = module(x_static, is_first_microbatch=False) + + graph.replay() + torch.cuda.synchronize() + + # Next batch's first microbatch: eager weight update + re-quantize of the + # cached weight. With persistent buffers this rewrites the SAME scale buffer + # in place; the buggy path reallocates it and frees the captured address. + with torch.no_grad(): + module.weight.add_(torch.randn_like(module.weight) * 0.5) + with te.autocast(enabled=True, recipe=recipe): + module( + torch.randn(batch, 1024, dtype=torch.bfloat16, device=device), + is_first_microbatch=True, + ) + addr_after = ws._rowwise_scale_inv.data_ptr() + + # Ground truth for the updated weights: an eager forward that consumes the + # (now updated) cached weight, i.e. what graph replay must reproduce. + with te.autocast(enabled=True, recipe=recipe): + expected = module(x_static, is_first_microbatch=False).detach().float().clone() + + # Churn the allocator so a freed scale buffer would be reused, turning a + # dangling captured pointer into observable corruption on replay. + _churn = [torch.randn(4096, 1024, dtype=torch.bfloat16, device=device) for _ in range(8)] + torch.cuda.synchronize() + + graph.replay() + torch.cuda.synchronize() + replayed = y_static.detach().float().clone() + + assert addr_after == addr_before, ( + "cached weight swizzled scale buffer moved on eager re-quantize; a captured " + "CUDA graph would replay against a stale address" + ) + # Replay reads the persistent scale buffer in place, so it must match an + # eager forward with the updated weights. On the buggy path it instead reads + # freed/reused memory and diverges. + torch.testing.assert_close(replayed, expected, atol=1e-2, rtol=1e-2) + del graph diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 031122d966..033d464bcf 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -150,7 +150,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, /*output=*/output_tensor->data, /*output_t=*/output_tensor->columnwise_data, /*epsilon=*/0.0f, /*return_identity=*/output_tensor->has_data(), /*return_transpose=*/output_tensor->has_columnwise_data(), /*pow2_scale=*/false, - /*swizzled_scale=*/false, + /*swizzled_scale=*/output_tensor->with_gemm_swizzled_scales, /*use_stochastic_rounding=*/quant_config_cpp.stochastic_rounding, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, @@ -310,7 +310,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens /*output=*/output_tensor->data, /*output_t=*/output_tensor->columnwise_data, /*epsilon=*/0.0f, /*return_identity=*/output_tensor->has_data(), /*return_transpose=*/output_tensor->has_columnwise_data(), /*pow2_scale=*/false, - /*swizzled_scale=*/false, + /*swizzled_scale=*/output_tensor->with_gemm_swizzled_scales, /*use_stochastic_rounding=*/quant_config_cpp.stochastic_rounding, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a4979f8b81..243c1e93f5 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -24,6 +24,7 @@ #include "../../utils.cuh" #include "core_nvfp4.cuh" #include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh" +#include "swizzle.cuh" namespace transformer_engine { namespace dispatch { @@ -778,7 +779,8 @@ __global__ void __launch_bounds__(THREADS_NUM) } template + typename IType, bool USE_STOCHASTIC_ROUNDING, bool RETURN_ROWWISE, bool RETURN_TRANSPOSE, + bool WITH_GEMM_SWIZZLED_SCALES = false> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_2D_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, @@ -1206,7 +1208,16 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t scales_offset_Y = scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; const size_t scales_offset_X = scales_offset_X_rowwise; - const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + size_t scale_idx_global; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // Write the scale directly into the cuBLAS GEMM-swizzled layout so no + // separate swizzle pass is needed. SFs_per_row (= cols / SCALE_DIM) is + // the number of compact scale columns. + scale_idx_global = + swizzle::gemm_swizzled_scale_idx(scales_offset_Y, scales_offset_X, SFs_per_row); + } else { + scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + } const bool rowwise_scale_is_within_bounds_Y = (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; @@ -1289,20 +1300,35 @@ __global__ void __launch_bounds__(THREADS_NUM) // Vectorized store scaling factors through SHMEM if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { - using ScalesVec = Vec; const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; - ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); - const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; const size_t count = // number of scales in Y dimension of this chunk (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); - nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; - constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); - if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { - // Fast path: vectorized store when destination is properly aligned - scales_vec.store_to(dst); + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + // The swizzled layout scatters the contiguous columnwise scales, so the + // vectorized store cannot be used. Emit each scale at its swizzled offset. + // The transposed scale matrix has `rows / SCALE_DIM` (= M/16) column tiles + // (exact because the swizzled path requires 128-aligned dims). Read + // SCALE_DIM by value; passing the namespace-scope constexpr by reference + // (e.g. via DIVUP) would ODR-use it and fail to compile in device code. + const size_t col_length_t = rows / SCALE_DIM; + for (size_t k = 0; k < count; ++k) { + const size_t off = + swizzle::gemm_swizzled_scale_idx(scales_offset_Y_t, scales_offset_X_t + k, col_length_t); + scales_t_ptr[off] = out_colwise_scales_sh[scale_idx_sh + k]; + } } else { - // Safe path: element-wise store for tails or unaligned destinations - scales_vec.store_to_elts(dst, 0, count); + using ScalesVec = Vec; + ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); + const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; + nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; + constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); + if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { + // Fast path: vectorized store when destination is properly aligned + scales_vec.store_to(dst); + } else { + // Safe path: element-wise store for tails or unaligned destinations + scales_vec.store_to_elts(dst, 0, count); + } } } @@ -1358,7 +1384,12 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, "Row-scaled NVFP4 quantization requires rowwise amax."); NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), "Row-scaled NVFP4 quantization does not produce columnwise output."); - NVTE_CHECK(!output->with_gemm_swizzled_scales, "Output must have scales in compact format."); + // In-kernel GEMM-swizzled scale output is only implemented on the 2D quantization + // path; every other configuration must emit compact scales. + const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; + NVTE_CHECK(!with_gemm_swizzled_scales || use_2d_quantization, + "In-kernel GEMM-swizzled NVFP4 scales are only supported on the 2D quantization path; " + "other paths must emit scales in compact format."); if (return_transpose) { NVTE_CHECK(output->has_columnwise_data(), "NVFP4 transposed output tensor must be allocated."); NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), @@ -1373,6 +1404,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA NVTE_CHECK(cols % 32 == 0, "Number of tensor cols must be a multiple of 32"); // 16B alignment for TMA + // The swizzled scale-factor layout is tiled 128x4; only 128-aligned dims avoid + // needing padded (zeroed) scale tiles, so restrict in-kernel swizzle to them. + NVTE_CHECK(!with_gemm_swizzled_scales || (rows % 128 == 0 && cols % 128 == 0), + "In-kernel GEMM-swizzled NVFP4 scales require rows and cols to be multiples of 128."); const size_t blocks_Y = DIVUP(rows, CHUNK_DIM_Y); const size_t blocks_X = DIVUP(cols, CHUNK_DIM_X); @@ -1451,9 +1486,15 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, ROW_SCALED_NVFP4>; if constexpr (use_2d_quantization) { - kernel = quantize_transpose_nvfp4_2D_kernel; + if (with_gemm_swizzled_scales) { + kernel = quantize_transpose_nvfp4_2D_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, + RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; + } else { + kernel = quantize_transpose_nvfp4_2D_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, + RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; + } } cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); diff --git a/transformer_engine/common/cast/nvfp4/swizzle.cuh b/transformer_engine/common/cast/nvfp4/swizzle.cuh new file mode 100644 index 0000000000..6e82974b60 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/swizzle.cuh @@ -0,0 +1,58 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file swizzle.cuh + * \brief Helper for emitting GEMM-swizzled NVFP4 scale factors in-kernel. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_NVFP4_SWIZZLE_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_NVFP4_SWIZZLE_CUH_ + +#include + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { +namespace swizzle { + +/*! \brief Map compact NVFP4 scale-factor indices to the GEMM-swizzled offset. + * + * cuBLAS block-scaled GEMM consumes scale factors in a "swizzled" 512B-base-block + * layout (128 rows x 4 cols per base block, itself split into 4 column blocks of + * 32 rows x 4 cols): + * https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout + * + * This converts an index (row_idx, col_idx) in the compact scale matrix -- the + * layout that matches the FP4 data -- into the flat offset of the swizzled buffer. + * \p col_length is the number of scale columns of the compact matrix (i.e. the + * last-dim tile count, K/16 for the rowwise operand or M/16 for the columnwise + * operand). This is byte-compatible with nvte_swizzle_scaling_factors and with + * the fallback kernel's scale_factor_swizzled_offset. + */ +__device__ __forceinline__ size_t gemm_swizzled_scale_idx(size_t row_idx, size_t col_idx, + size_t col_length) { + constexpr size_t kRowsPerBaseBlock = 128; + constexpr size_t kRowsPerBaseBlockCol = 32; + constexpr size_t kColsPerBaseBlockCol = 4; + + const size_t rb = row_idx / kRowsPerBaseBlock; + const size_t rem = row_idx % kRowsPerBaseBlock; + const size_t d4 = rem / kRowsPerBaseBlockCol; + const size_t d3 = rem % kRowsPerBaseBlockCol; + const size_t cbg = col_idx / kColsPerBaseBlockCol; + const size_t d5 = col_idx % kColsPerBaseBlockCol; + + const size_t cbg_cnt = (col_length + kColsPerBaseBlockCol - 1) / kColsPerBaseBlockCol; + // Row-major offset in the logical (rb_cnt, cbg_cnt, 32, 4, 4) shape. + return ((rb * cbg_cnt + cbg) * kRowsPerBaseBlockCol + d3) * 16 + d4 * kColsPerBaseBlockCol + d5; +} + +} // namespace swizzle +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_NVFP4_SWIZZLE_CUH_ diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index e4aba728b6..aa0e0c87fe 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -406,6 +406,16 @@ class NVFP4Quantizer : public Quantizer { static bool is_eligible_for_rht_cast_fusion(const std::vector& shape, bool for_grouped_kernel = false); + /*! @brief Whether a tensor of the given shape can have GEMM-swizzled scale + * factors written directly by the (non-RHT) 2D NVFP4 quantize kernel. + * + * The swizzled scale layout is tiled 128x4; requiring both flattened dims to + * be multiples of 128 guarantees no padded scale tiles are needed for either + * the rowwise or columnwise operand, so the standalone swizzle pass can be + * skipped entirely. + */ + static bool is_eligible_for_2d_swizzle_fusion(const std::vector& shape); + private: void quantize_with_rht_unfused_helper(const TensorWrapper& input, TensorWrapper& out, TensorWrapper& rht_output_t_cpp, diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index d2a3be888f..4308464c52 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1922,6 +1922,32 @@ bool NVFP4Quantizer::is_eligible_for_rht_cast_fusion(const std::vector& transformer_engine::cuda::sm_arch() <= 110; } +bool NVFP4Quantizer::is_eligible_for_2d_swizzle_fusion(const std::vector& shape) { + const auto [rows, cols] = get_2d_dims(shape); + return rows % 128 == 0 && cols % 128 == 0 && transformer_engine::cuda::sm_arch() >= 100 && + transformer_engine::cuda::sm_arch() <= 110; +} + +namespace { + +// Whether this quantizer should emit GEMM-swizzled scale factors directly from the +// quantize kernel (RHT cast-fusion path or the plain 2D quantize path), avoiding the +// standalone post-quantize swizzle pass. +bool nvfp4_emits_gemm_swizzled_scales(const NVFP4Quantizer& q, const std::vector& shape) { + if (!q.optimize_for_gemm) { + return false; + } + if (q.with_rht) { + return NVFP4Quantizer::is_eligible_for_rht_cast_fusion(shape); + } + // Plain (non-RHT) 2D quantize kernel can bake the swizzled layout for aligned shapes. + return q.with_2d_quantization && !q.row_scaled_nvfp4 && + q.nvfp4_4over6_mode == kNVTENVFP44Over6Disabled && + NVFP4Quantizer::is_eligible_for_2d_swizzle_fusion(shape); +} + +} // namespace + std::pair NVFP4Quantizer::create_tensor( const std::vector& shape, DType dtype, std::optional device_opt, bool pin_memory) const { @@ -1932,10 +1958,10 @@ std::pair NVFP4Quantizer::create_tensor( const std::vector shape_int64(shape.begin(), shape.end()); const auto [flat_first_dim, flat_last_dim] = get_2d_dims(shape); - // Swizzled SF is only valid when the RHT cast-fusion path runs; - // other quantize paths reject it. - const bool with_gemm_swizzled_scales = this->optimize_for_gemm && this->with_rht && - NVFP4Quantizer::is_eligible_for_rht_cast_fusion(shape); + // Swizzled SF is emitted directly by the RHT cast-fusion kernel or by the plain + // 2D quantize kernel (for aligned shapes); other quantize paths reject it and go + // through the standalone post-quantize swizzle pass instead. + const bool with_gemm_swizzled_scales = nvfp4_emits_gemm_swizzled_scales(*this, shape); NVTE_CHECK(flat_first_dim % NVFP4_BLOCK_SIZE == 0, "First dim for NVFP4 must be divisible by ", NVFP4_BLOCK_SIZE, " (got shape=", shape, ")"); NVTE_CHECK(flat_last_dim % NVFP4_BLOCK_SIZE == 0, @@ -2263,10 +2289,10 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( const auto [flat_first_dim, flat_last_dim] = get_2d_dims(shape); - // Swizzled SF is only valid when the RHT cast-fusion path runs; - // other quantize paths reject it. - const bool with_gemm_swizzled_scales = this->optimize_for_gemm && this->with_rht && - NVFP4Quantizer::is_eligible_for_rht_cast_fusion(shape); + // Swizzled SF is emitted directly by the RHT cast-fusion kernel or by the plain + // 2D quantize kernel (for aligned shapes); other quantize paths reject it and go + // through the standalone post-quantize swizzle pass instead. + const bool with_gemm_swizzled_scales = nvfp4_emits_gemm_swizzled_scales(*this, shape); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; From 24d088c98ae09def90f18651a3dfd5e5014aa704 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 04:56:29 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py | 6 +++--- .../common/cast/nvfp4/quantize_transpose_nvfp4.cuh | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py b/tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py index b4cb2ee1ce..8b3ccf4797 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py +++ b/tests/pytorch/nvfp4/test_nvfp4_weight_swizzle_cache.py @@ -101,9 +101,9 @@ def test_cached_weight_scale_pointers_stable(kind): ws._rowwise_scale_inv.data_ptr(), ws._columnwise_scale_inv.data_ptr(), ) - assert ptrs_after == ptrs_before[name], ( - f"cached weight workspace {name!r} reallocated scale buffers on weight update" - ) + assert ( + ptrs_after == ptrs_before[name] + ), f"cached weight workspace {name!r} reallocated scale buffers on weight update" assert getattr(ws, "_with_gemm_swizzled_scales", False) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 243c1e93f5..3a34c76de8 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -1312,8 +1312,8 @@ __global__ void __launch_bounds__(THREADS_NUM) // (e.g. via DIVUP) would ODR-use it and fail to compile in device code. const size_t col_length_t = rows / SCALE_DIM; for (size_t k = 0; k < count; ++k) { - const size_t off = - swizzle::gemm_swizzled_scale_idx(scales_offset_Y_t, scales_offset_X_t + k, col_length_t); + const size_t off = swizzle::gemm_swizzled_scale_idx(scales_offset_Y_t, + scales_offset_X_t + k, col_length_t); scales_t_ptr[off] = out_colwise_scales_sh[scale_idx_sh + k]; } } else { @@ -1488,12 +1488,12 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, if constexpr (use_2d_quantization) { if (with_gemm_swizzled_scales) { kernel = quantize_transpose_nvfp4_2D_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, - RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/true>; } else { kernel = quantize_transpose_nvfp4_2D_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, - RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; } }