Skip to content

Commit 3181072

Browse files
Caleb Duclaude
Caleb Du
and
claude
committed
[Common] Rebuild per-group col-quant tensors per tile in grouped RHT kernels
Keep only the group index and global scales live across k-tiles in the column-quant epilogue and rebuild the per-group output/scale tensors per tile. This lowers register pressure and removes the remaining register spills in the grouped (graph-safe and non-graph-safe) RHT cast-fusion kernels. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caleb Du <cadu@nvidia.com>
1 parent cae55a2 commit 3181072

2 files changed

Lines changed: 59 additions & 100 deletions

File tree

‎transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu‎

Lines changed: 32 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -742,35 +742,6 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g
742742
shape_rep, num_tensors, scheduler.tile_n_base() * size<1>(epilogue_tiler),
743743
/*cached_id=*/0, packed_N, M, shared_storage.group_offsets, tensor_range_cache);
744744

745-
// Determine quantization scale factor layouts/output splits for this group
746-
TSFDLayout sfd_layout;
747-
int cur_N = shared_storage.group_first_dims[group_idx];
748-
if constexpr (kEnableSwizzleSFOutput) {
749-
sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});
750-
} else {
751-
sfd_layout = make_layout(make_shape(M, make_shape(Int<SFVecSize>{}, cur_N / SFVecSize)),
752-
make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));
753-
}
754-
// Build output tensors for columns and their quant scales
755-
// TODO(zhongbo): double check the logic here
756-
int64_t group_elem_offset = shared_storage.group_offsets[group_idx];
757-
Tensor mD = make_tensor(cute::subbyte_iterator<TD>(reinterpret_cast<TD *>(
758-
reinterpret_cast<char *>(QA_COLWISE) + group_elem_offset / 2)),
759-
make_shape(M, cur_N), DStride{}); // (M,packed_N)
760-
Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),
761-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
762-
763-
// for every tensor [x, y] row major, x y both a multiple of 128
764-
// both of its rowwise and colwise scaling factors will have exactly x * y / 16 elements in FP8 E4M3
765-
Tensor mSFD = make_tensor(
766-
make_gmem_ptr<TSFD>(reinterpret_cast<TSFD *>(reinterpret_cast<char *>(SFA_COLWISE) +
767-
group_elem_offset / kNVFP4BlockSize)),
768-
sfd_layout);
769-
Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),
770-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
771-
772-
Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));
773-
774745
// Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors
775746
auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));
776747
auto tiled_r2g =
@@ -819,32 +790,39 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g
819790
c_global_amax_val);
820791
global_decode_scale = 1.0f / global_encode_scale;
821792
global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;
822-
// TODO(zhongbo): double check the logic here
823-
cur_N = shared_storage.group_first_dims[group_idx];
824-
if constexpr (kEnableSwizzleSFOutput) {
825-
sfd_layout =
826-
tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});
827-
} else {
828-
sfd_layout =
829-
make_layout(make_shape(M, make_shape(Int<SFVecSize>{}, cur_N / SFVecSize)),
830-
make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));
831-
}
832-
// update tensor
833-
group_elem_offset = shared_storage.group_offsets[group_idx];
834-
mD = make_tensor(cute::subbyte_iterator<TD>(reinterpret_cast<TD *>(
835-
reinterpret_cast<char *>(QA_COLWISE) + group_elem_offset / 2)),
836-
make_shape(M, cur_N), DStride{});
837-
gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),
838-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
839-
mSFD = make_tensor(make_gmem_ptr<TSFD>(reinterpret_cast<TSFD *>(
840-
reinterpret_cast<char *>(SFA_COLWISE) +
841-
group_elem_offset / kNVFP4BlockSize)),
842-
sfd_layout);
843-
gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),
844-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
845-
846-
gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));
847793
}
794+
// Determine quantization scale factor layouts/output splits for this group
795+
TSFDLayout sfd_layout;
796+
const int cur_N = shared_storage.group_first_dims[group_idx];
797+
if constexpr (kEnableSwizzleSFOutput) {
798+
sfd_layout =
799+
tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});
800+
} else {
801+
sfd_layout =
802+
make_layout(make_shape(M, make_shape(Int<SFVecSize>{}, cur_N / SFVecSize)),
803+
make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));
804+
}
805+
// Build output tensors for columns and their quant scales
806+
// TODO(zhongbo): double check the logic here
807+
const int64_t group_elem_offset = shared_storage.group_offsets[group_idx];
808+
Tensor mD = make_tensor(
809+
cute::subbyte_iterator<TD>(reinterpret_cast<TD *>(
810+
reinterpret_cast<char *>(QA_COLWISE) + group_elem_offset / 2)),
811+
make_shape(M, cur_N), DStride{}); // (M,packed_N)
812+
Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),
813+
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
814+
815+
// for every tensor [x, y] row major, x y both a multiple of 128
816+
// both of its rowwise and colwise scaling factors will have exactly x * y / 16 elements in FP8 E4M3
817+
Tensor mSFD = make_tensor(
818+
make_gmem_ptr<TSFD>(reinterpret_cast<TSFD *>(
819+
reinterpret_cast<char *>(SFA_COLWISE) + group_elem_offset / kNVFP4BlockSize)),
820+
sfd_layout);
821+
Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),
822+
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
823+
824+
Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));
825+
848826
int local_tile_n_idx =
849827
(global_tile_n_offset - tensor_range_cache.begin) / size<1>(epilogue_tiler);
850828
Tensor tDgD_mn = gD_mn_view(_, _, _, tile_m, local_tile_n_idx);

‎transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu‎

Lines changed: 27 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -701,30 +701,6 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device(
701701
}
702702
int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler));
703703

704-
// Determine quantization scale factor layouts/output splits for this group
705-
TSFDLayout sfd_layout;
706-
int cur_N = args.split_sections[group_idx];
707-
if constexpr (kEnableSwizzleSFOutput) {
708-
sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});
709-
} else {
710-
sfd_layout = make_layout(make_shape(M, make_shape(Int<SFVecSize>{}, cur_N / SFVecSize)),
711-
make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));
712-
}
713-
// Build output tensors for columns and their quant scales
714-
Tensor mD = make_tensor(
715-
cute::subbyte_iterator<TD>(reinterpret_cast<TD *>(args.output_colwise_list[group_idx])),
716-
make_shape(M, cur_N), DStride{}); // (M,packed_N)
717-
Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),
718-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
719-
720-
Tensor mSFD = make_tensor(make_gmem_ptr<TSFD>(reinterpret_cast<TSFD *>(
721-
args.output_colwise_scale_inv_list[group_idx])),
722-
sfd_layout);
723-
Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),
724-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
725-
726-
Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));
727-
728704
// Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors
729705
auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{}));
730706
auto tiled_r2g =
@@ -770,29 +746,34 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device(
770746
c_global_amax_val);
771747
global_decode_scale = 1.0f / global_encode_scale;
772748
global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;
773-
cur_N = args.split_sections[group_idx];
774-
if constexpr (kEnableSwizzleSFOutput) {
775-
sfd_layout =
776-
tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});
777-
} else {
778-
sfd_layout =
779-
make_layout(make_shape(M, make_shape(Int<SFVecSize>{}, cur_N / SFVecSize)),
780-
make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));
781-
}
782-
// update tensor
783-
mD = make_tensor(cute::subbyte_iterator<TD>(
784-
reinterpret_cast<TD *>(args.output_colwise_list[group_idx])),
785-
make_shape(M, cur_N), DStride{});
786-
gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),
787-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
788-
mSFD = make_tensor(make_gmem_ptr<TSFD>(reinterpret_cast<TSFD *>(
789-
args.output_colwise_scale_inv_list[group_idx])),
790-
sfd_layout);
791-
gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),
792-
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
793-
794-
gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));
795749
}
750+
// Determine quantization scale factor layouts/output splits for this group
751+
TSFDLayout sfd_layout;
752+
const int cur_N = args.split_sections[group_idx];
753+
if constexpr (kEnableSwizzleSFOutput) {
754+
sfd_layout =
755+
tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{});
756+
} else {
757+
sfd_layout =
758+
make_layout(make_shape(M, make_shape(Int<SFVecSize>{}, cur_N / SFVecSize)),
759+
make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{})));
760+
}
761+
// Build output tensors for columns and their quant scales
762+
Tensor mD = make_tensor(
763+
cute::subbyte_iterator<TD>(
764+
reinterpret_cast<TD *>(args.output_colwise_list[group_idx])),
765+
make_shape(M, cur_N), DStride{}); // (M,packed_N)
766+
Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _),
767+
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
768+
769+
Tensor mSFD = make_tensor(make_gmem_ptr<TSFD>(reinterpret_cast<TSFD *>(
770+
args.output_colwise_scale_inv_list[group_idx])),
771+
sfd_layout);
772+
Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _),
773+
Step<_1, _1, X>{}); // (BLK_M,BLK_N)
774+
775+
Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler));
776+
796777
int group_start_offset = args.split_sections_range[group_idx];
797778
int local_tile_n_idx =
798779
(global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler);

0 commit comments

Comments
 (0)