Skip to content

Norm Samples updates for B,S,H style tensor inputs #432

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rmhaskarnvidia wants to merge 4 commits into
base: develop
Choose a base branch
from
Open

Norm Samples updates for B,S,H style tensor inputs #432

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bc08ccd
Fix inferencing logic for rms norm bwd and update the samples to b,s,…
rmhaskarnvidia Jul 23, 2026
3b3cb1e
B,S,H updates for layernorm samples
rmhaskarnvidia Jul 23, 2026
6817826
Fix wrong skip error messages for norm samples
rmhaskarnvidia Jul 23, 2026
77f4157
Python samples and notebook changes for B,S,H input tensors
rmhaskarnvidia Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 11 additions & 24 deletions include/cudnn_frontend/node/dln.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,54 +29,41 @@ class DLNNode : public NodeCRTP<DLNNode> {

attributes.fill_from_context(context);

// TODO: Only inferencing from X works today.
auto X = attributes.inputs[Layernorm_backward_attributes::input_names::X];
auto const x_tensor_dim = X->get_dim();
auto X = attributes.inputs[Layernorm_backward_attributes::input_names::X];
auto const x_tensor_dim = X->get_dim();
auto const x_tensor_stride = X->get_stride();

auto DY = attributes.inputs[Layernorm_backward_attributes::input_names::DY];
auto dy_tensor_dim = DY->get_dim();

// Only infer dims and strides if user did not set them
if (dy_tensor_dim.empty()) {
dy_tensor_dim.resize(x_tensor_dim.size());
DY->set_dim(x_tensor_dim);
}
if (DY->get_stride().empty()) {
auto const& DY_dim = DY->get_dim();
// Default to NHWC
auto const& stride_order = detail::generate_NHWC_stride_order(DY_dim.size());
DY->set_stride(detail::generate_stride(DY_dim, stride_order));
DY->set_stride(x_tensor_stride);
}

auto DX = attributes.outputs[Layernorm_backward_attributes::output_names::DX];
auto dx_tensor_dim = DX->get_dim();
// Only infer dims and strides if user did not set them
if (dx_tensor_dim.empty()) {
dx_tensor_dim.resize(x_tensor_dim.size());
DX->set_dim(x_tensor_dim);
}
if (DX->get_stride().empty()) {
auto const& DX_dim = DX->get_dim();
// Default to NHWC
auto const& stride_order = detail::generate_NHWC_stride_order(DX_dim.size());
DX->set_stride(detail::generate_stride(DX_dim, stride_order));
DX->set_stride(x_tensor_stride);
}

auto scale_bias_dim = X->get_dim();
scale_bias_dim[0] = 1;
auto scale = attributes.inputs[Layernorm_backward_attributes::input_names::SCALE];

// Set channel length tensors
auto infer_scale_bias_tensors = [&scale_bias_dim](std::shared_ptr<Tensor_attributes>& T) {
auto tensor_dim = T->get_dim();
// Infer dscale/dbias from scale
auto infer_scale_bias_tensors = [&scale](std::shared_ptr<Tensor_attributes>& T) {
// Only infer dims and strides if user did not set them
if (tensor_dim.empty()) {
T->set_dim(scale_bias_dim);
if (T->get_dim().empty()) {
T->set_dim(scale->get_dim());
}
if (T->get_stride().empty()) {
auto const& T_dim = T->get_dim();
// Default to NHWC
auto const& stride_order = detail::generate_NHWC_stride_order(T_dim.size());
T->set_stride(detail::generate_stride(T_dim, stride_order));
T->set_stride(scale->get_stride());
}
};

Expand Down
35 changes: 11 additions & 24 deletions include/cudnn_frontend/node/rmsnorm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,54 +226,41 @@ class DRMSNormNode : public NodeCRTP<DRMSNormNode> {

attributes.fill_from_context(context);

// TODO: Only inferencing from X works today.
auto X = attributes.inputs[Rmsnorm_backward_attributes::input_names::X];
auto const x_tensor_dim = X->get_dim();
auto X = attributes.inputs[Rmsnorm_backward_attributes::input_names::X];
auto const x_tensor_dim = X->get_dim();
auto const x_tensor_stride = X->get_stride();

auto DY = attributes.inputs[Rmsnorm_backward_attributes::input_names::DY];
auto dy_tensor_dim = DY->get_dim();

// Only infer dims and strides if user did not set them
if (dy_tensor_dim.empty()) {
dy_tensor_dim.resize(x_tensor_dim.size());
DY->set_dim(x_tensor_dim);
}
if (DY->get_stride().empty()) {
auto const& DY_dim = DY->get_dim();
// Default to NHWC
auto const& stride_order = detail::generate_NHWC_stride_order(DY_dim.size());
DY->set_stride(detail::generate_stride(DY_dim, stride_order));
DY->set_stride(x_tensor_stride);
}

auto DX = attributes.outputs[Rmsnorm_backward_attributes::output_names::DX];
auto dx_tensor_dim = DX->get_dim();
// Only infer dims and strides if user did not set them
if (dx_tensor_dim.empty()) {
dx_tensor_dim.resize(x_tensor_dim.size());
DX->set_dim(x_tensor_dim);
}
if (DX->get_stride().empty()) {
auto const& DX_dim = DX->get_dim();
// Default to NHWC
auto const& stride_order = detail::generate_NHWC_stride_order(DX_dim.size());
DX->set_stride(detail::generate_stride(DX_dim, stride_order));
DX->set_stride(x_tensor_stride);
}

auto scale_bias_dim = X->get_dim();
scale_bias_dim[0] = 1;
auto scale = attributes.inputs[Rmsnorm_backward_attributes::input_names::SCALE];

// Set channel length tensors
auto infer_scale_bias_tensors = [&scale_bias_dim](std::shared_ptr<Tensor_attributes>& T) {
auto tensor_dim = T->get_dim();
// Infer dscale/dbias from scale
auto infer_scale_bias_tensors = [&scale](std::shared_ptr<Tensor_attributes>& T) {
// Only infer dims and strides if user did not set them
if (tensor_dim.empty()) {
T->set_dim(scale_bias_dim);
if (T->get_dim().empty()) {
T->set_dim(scale->get_dim());
}
if (T->get_stride().empty()) {
auto const& T_dim = T->get_dim();
// Default to NHWC
auto const& stride_order = detail::generate_NHWC_stride_order(T_dim.size());
T->set_stride(detail::generate_stride(T_dim, stride_order));
T->set_stride(scale->get_stride());
}
};

Expand Down
2 changes: 1 addition & 1 deletion samples/cpp/norm/batchnorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ TEST_CASE("SGBN Add Relu Graph", "[batchnorm][graph]") {
SKIP("single GPU BN is not supported in cudnn versions prior to 8.7");
#endif
if (check_device_arch_newer_than("ampere") == false) {
SKIP("ConvBNFprop requires Ampere and up");
SKIP("SGBN Add Relu requires Ampere and up");
}
// Create a unique_ptr for the cuDNN handle
auto handle_ptr = create_cudnn_handle();
Expand Down
46 changes: 23 additions & 23 deletions samples/cpp/norm/layernorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,17 +181,17 @@ TEST_CASE("LayerNorm Training", "[layernorm][graph]") {

auto X = graph.tensor(fe::graph::Tensor_attributes()
.set_name("X")
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));
auto scale = graph.tensor(fe::graph::Tensor_attributes()
.set_name("scale")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto bias = graph.tensor(fe::graph::Tensor_attributes()
.set_name("bias")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));

float epsilon_cpu = 1e-05f;
Expand Down Expand Up @@ -260,17 +260,17 @@ TEST_CASE("LayerNorm Inference", "[layernorm][graph]") {

auto X = graph.tensor(fe::graph::Tensor_attributes()
.set_name("X")
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));
auto scale = graph.tensor(fe::graph::Tensor_attributes()
.set_name("scale")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto bias = graph.tensor(fe::graph::Tensor_attributes()
.set_name("bias")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));

float epsilon_cpu = 1e-05f;
Expand Down Expand Up @@ -332,27 +332,27 @@ TEST_CASE("LayerNorm Backward", "[layernorm][graph]") {

auto X = graph.tensor(fe::graph::Tensor_attributes()
.set_name("X")
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));
auto DY = graph.tensor(fe::graph::Tensor_attributes()
.set_name("DY")
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));

auto scale = graph.tensor(fe::graph::Tensor_attributes()
.set_name("scale")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto mean = graph.tensor(fe::graph::Tensor_attributes()
.set_name("mean")
.set_dim({batch_size * seq_length, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_dim({batch_size, seq_length, 1})
.set_stride({seq_length, 1, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto inv_variance = graph.tensor(fe::graph::Tensor_attributes()
.set_name("inv_variance")
.set_dim({batch_size * seq_length, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_dim({batch_size, seq_length, 1})
.set_stride({seq_length, 1, 1})
.set_data_type(fe::DataType_t::FLOAT));

auto DLN_options = fe::graph::Layernorm_backward_attributes().set_saved_mean_and_inv_variance(mean, inv_variance);
Expand All @@ -362,7 +362,7 @@ TEST_CASE("LayerNorm Backward", "[layernorm][graph]") {
dbias->set_output(true).set_data_type(fe::DataType_t::FLOAT);

#if (CUDNN_VERSION < 8905)
SKIP("single GPU BN is not supported in cudnn versions prior to 8.7");
SKIP("LayerNorm is not supported in cudnn versions prior to 8.9.5");
#endif
if (check_device_arch_newer_than("ampere") == false) {
SKIP("LayerNorm Backward requires Ampere and up");
Expand Down
45 changes: 23 additions & 22 deletions samples/cpp/norm/rmsnorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,12 @@ TEST_CASE("RmsNorm Training", "[rmsnorm][graph]") {
auto X = graph.tensor(fe::graph::Tensor_attributes()
.set_name("X")
.set_data_type(fe::DataType_t::FLOAT)
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));
auto scale = graph.tensor(fe::graph::Tensor_attributes()
.set_name("scale")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));

float epsilon_cpu = 1e-05f;
Expand Down Expand Up @@ -77,7 +77,7 @@ TEST_CASE("RmsNorm Training", "[rmsnorm][graph]") {
REQUIRE(graph.warmup(handle).is_good());

Surface<float> X_tensor(batch_size * seq_length * hidden_size);
Surface<float> Var_tensor(batch_size * seq_length);
Surface<float> Inv_variance_tensor(batch_size * seq_length);
Surface<float> Scale_tensor(hidden_size);
Surface<float> Y_tensor(batch_size * seq_length * hidden_size);

Expand All @@ -86,7 +86,10 @@ TEST_CASE("RmsNorm Training", "[rmsnorm][graph]") {
Surface<int8_t> workspace(workspace_size);

std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> variant_pack = {
{X, X_tensor.devPtr}, {inv_variance, Var_tensor.devPtr}, {scale, Scale_tensor.devPtr}, {Y, Y_tensor.devPtr}};
{X, X_tensor.devPtr},
{inv_variance, Inv_variance_tensor.devPtr},
{scale, Scale_tensor.devPtr},
{Y, Y_tensor.devPtr}};

REQUIRE(graph.execute(handle, variant_pack, workspace.devPtr).is_good());
}
Expand All @@ -103,17 +106,17 @@ TEST_CASE("RmsNorm Inference", "[rmsnorm][graph]") {
auto X = graph.tensor(fe::graph::Tensor_attributes()
.set_name("X")
.set_data_type(fe::DataType_t::FLOAT)
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));
auto scale = graph.tensor(fe::graph::Tensor_attributes()
.set_name("scale")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto bias = graph.tensor(fe::graph::Tensor_attributes()
.set_name("bias")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));

float epsilon_cpu = 1e-05f;
Expand Down Expand Up @@ -174,23 +177,23 @@ TEST_CASE("RmsNorm Backward", "[rmsnorm][graph]") {
auto X = graph.tensor(fe::graph::Tensor_attributes()
.set_name("X")
.set_data_type(fe::DataType_t::FLOAT)
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));
auto DY = graph.tensor(fe::graph::Tensor_attributes()
.set_name("DY")
.set_data_type(fe::DataType_t::FLOAT)
.set_dim({batch_size * seq_length, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size}));
.set_dim({batch_size, seq_length, hidden_size})
.set_stride({seq_length * hidden_size, hidden_size, 1}));

auto scale = graph.tensor(fe::graph::Tensor_attributes()
.set_name("scale")
.set_dim({1, hidden_size, 1, 1})
.set_stride({hidden_size, 1, hidden_size, hidden_size})
.set_dim({1, 1, hidden_size})
.set_stride({hidden_size, hidden_size, 1})
.set_data_type(fe::DataType_t::FLOAT));
auto inv_variance = graph.tensor(fe::graph::Tensor_attributes()
.set_name("inv_variance")
.set_dim({batch_size * seq_length, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_dim({batch_size, seq_length, 1})
.set_stride({seq_length, 1, 1})
.set_data_type(fe::DataType_t::FLOAT));

auto DRMS_options = fe::graph::Rmsnorm_backward_attributes().has_dbias(false);
Expand Down Expand Up @@ -221,11 +224,9 @@ TEST_CASE("RmsNorm Backward", "[rmsnorm][graph]") {

Surface<float> X_tensor(batch_size * seq_length * hidden_size);
Surface<float> DY_tensor(batch_size * seq_length * hidden_size);
Surface<float> Mean_tensor(batch_size * seq_length);
Surface<float> Inv_variance_tensor(batch_size * seq_length);
Surface<float> Scale_tensor(hidden_size);
Surface<float> Dscale_tensor(hidden_size);
Surface<float> Dbias_tensor(hidden_size);
Surface<float> DX_tensor(batch_size * seq_length * hidden_size);

int64_t workspace_size = 0;
Expand Down
Loading