Skip to content

Support ragged offset multiplier in SDPA #290

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

Merged
Anerudhan merged 1 commit into from
Jun 10, 2026
Merged

Support ragged offset multiplier in SDPA #290

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 3 additions & 0 deletions include/cudnn_frontend/cudnn_interface.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,9 @@ create_cudnn_tensor(
CHECK_CUDNN_FRONTEND_ERROR(create_cudnn_tensor(ragged_offset_props, tensors, potential_uid, used_uids));
tensor_builder.setRaggedOffset(tensors.at(ragged_offset_props->get_uid()));
}
if (props->has_ragged_offset_multiplier()) {
tensor_builder.setRaggedOffsetMultiplier(props->get_ragged_offset_multiplier());
}

#ifdef NV_CUDNN_DISABLE_EXCEPTION
// disable exception macro is defined. Calling build will not throw.
Expand Down
32 changes: 29 additions & 3 deletions include/cudnn_frontend/graph_properties.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,15 @@ class Tensor_attributes {
error_code_t::ATTRIBUTE_NOT_SET,
"Tensor '" + name + "' can't have both compile-time constant and runtime pass_by_value.");

RETURN_CUDNN_FRONTEND_ERROR_IF(
has_ragged_offset_multiplier() && !ragged_offset,
error_code_t::ATTRIBUTE_NOT_SET,
"Tensor '" + name + "' with ragged offset multiplier must also have a ragged offset tensor.");

RETURN_CUDNN_FRONTEND_ERROR_IF(ragged_offset_multiplier <= 0,
error_code_t::INVALID_VALUE,
"Tensor '" + name + "' ragged offset multiplier must be positive.");

return {error_code_t::OK, ""};
}

Expand All @@ -118,9 +127,10 @@ class Tensor_attributes {
bool uid_assigned = false;

std::shared_ptr<Tensor_attributes> ragged_offset;
int64_t alignment = 16; // Default to 16 bytes
int64_t vector_count = 1; // Default to 1 (no vectorization)
int64_t vector_dimension = -1; // Default to -1 (not set)
int64_t ragged_offset_multiplier = 1;
int64_t alignment = 16; // Default to 16 bytes
int64_t vector_count = 1; // Default to 1 (no vectorization)
int64_t vector_dimension = -1; // Default to -1 (not set)

auto
fill_from_context(detail::Context const& context) -> Tensor_attributes& {
Expand Down Expand Up @@ -323,6 +333,16 @@ class Tensor_attributes {
return ragged_offset;
}

int64_t
get_ragged_offset_multiplier() const {
return ragged_offset_multiplier;
}

bool
has_ragged_offset_multiplier() const {
return ragged_offset_multiplier != 1;
}

auto
set_is_virtual(bool const value) -> Tensor_attributes& {
is_virtual = value;
Expand Down Expand Up @@ -447,6 +467,12 @@ class Tensor_attributes {
ragged_offset = value;
return *this;
}

auto
set_ragged_offset_multiplier(int64_t value) -> Tensor_attributes& {
ragged_offset_multiplier = value;
return *this;
}
};

class Batchnorm_attributes;
Expand Down
12 changes: 12 additions & 0 deletions include/cudnn_frontend/node/sdpa_support_surface.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -404,6 +404,18 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co
error_code_t::GRAPH_NOT_SUPPORTED,
"Composite SDPA node doesn't support CU_SEQ_LEN_Q / CU_SEQ_LEN_KV inputs");
}
// The ragged offset multiplier is only supported by the unified forward engine.
// Reject it here so auto-select routes such graphs to the unified implementation.
for (const auto& [key, value] : inputs) {
RETURN_CUDNN_FRONTEND_ERROR_IF(value != nullptr && value->has_ragged_offset_multiplier(),
error_code_t::GRAPH_NOT_SUPPORTED,
"Composite SDPA node doesn't support a ragged offset multiplier");
}
for (const auto& [key, value] : outputs) {
RETURN_CUDNN_FRONTEND_ERROR_IF(value != nullptr && value->has_ragged_offset_multiplier(),
error_code_t::GRAPH_NOT_SUPPORTED,
"Composite SDPA node doesn't support a ragged offset multiplier");
}
break;
case AttentionImplementation_t::UNIFIED: {
auto effective_cudnn_ver = std::min(detail::get_backend_version(), detail::get_compiled_version());
Expand Down
15 changes: 15 additions & 0 deletions include/cudnn_frontend/utils/serialize.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -606,6 +606,9 @@ to_json(nlohmann::json& j, const Tensor_attributes& ta) {
j["ragged_offset_uid"] = ta.ragged_offset->get_uid();
j["ragged_offset_name"] = ta.ragged_offset->get_name();
}
if (ta.has_ragged_offset_multiplier()) {
j["ragged_offset_multiplier"] = ta.ragged_offset_multiplier;
}
}

inline void
Expand All @@ -623,6 +626,18 @@ from_json(const nlohmann::json& j, Tensor_attributes& ta) {
if (ta.is_pass_by_value && !j["pass_by_value"].is_null()) {
ta.pass_by_value = j.at("pass_by_value");
}
if (j.contains("ragged_offset_uid")) {
auto ragged_offset = std::make_shared<Tensor_attributes>();
ragged_offset->uid = j.at("ragged_offset_uid").get<Tensor_attributes::uid_t>();
ragged_offset->uid_assigned = true;
if (j.contains("ragged_offset_name")) {
ragged_offset->name = j.at("ragged_offset_name").get<std::string>();
}
ta.ragged_offset = ragged_offset;
}
if (j.contains("ragged_offset_multiplier")) {
Comment thread
egilliam-nv marked this conversation as resolved.
ta.ragged_offset_multiplier = j.at("ragged_offset_multiplier").get<int64_t>();
}
}

NLOHMANN_JSON_SERIALIZE_ENUM(KnobType_t,
Expand Down
61 changes: 61 additions & 0 deletions include/cudnn_frontend_Tensor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,9 @@ class Tensor_v8 : public BackendDescriptor {
if (raggedOffset != nullptr) {
ss << "\n raggedOffset: Enabled UID: " << raggedOffset->getId();
}
if (hasRaggedOffsetMultiplier()) {
ss << "\n raggedOffsetMultiplier: " << raggedOffsetMultiplier;
}
return ss.str();
}

Expand Down Expand Up @@ -146,6 +149,16 @@ class Tensor_v8 : public BackendDescriptor {
return alignment;
}

int64_t
getRaggedOffsetMultiplier() const {
return raggedOffsetMultiplier;
}

bool
hasRaggedOffsetMultiplier() const {
return raggedOffsetMultiplier != 1;
}

bool
isVirtualTensor() const {
return isVirtual;
Expand Down Expand Up @@ -187,6 +200,7 @@ class Tensor_v8 : public BackendDescriptor {
cudnn_frontend::TensorReordering_t reorder_type =
cudnn_frontend::TensorReordering_t::NONE; //! Type of reordering in the tensor
std::shared_ptr<Tensor_v8> raggedOffset; //! Ragged offsets for ragged tensors
int64_t raggedOffsetMultiplier = 1; //! Unit size of ragged offsets in tensor elements
};

///
Expand Down Expand Up @@ -293,6 +307,12 @@ class TensorBuilder_v8 {
return *this;
}

auto
setRaggedOffsetMultiplier(int64_t const multiplier) -> TensorBuilder_v8 & {
m_tensor.raggedOffsetMultiplier = multiplier;
return *this;
}

// Clone parameters of another tensor. Make sure to still set the UID since UID of two tensors shouldn't be the
// same.
auto
Expand Down Expand Up @@ -499,6 +519,47 @@ class TensorBuilder_v8 {
}
#endif

// Set ragged offset multiplier
if (m_tensor.hasRaggedOffsetMultiplier()) {
auto ragged_offset_multiplier_cudnn_ver_error =
"CUDNN_BACKEND_TENSOR_DESCRIPTOR: SetAttribute "
"CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER requires cudnn version 9.24.0";
#if (CUDNN_VERSION >= 92400)
NV_CUDNN_FE_DYNAMIC_CHECK_BACKEND_DESCRIPTOR(92400, m_tensor, ragged_offset_multiplier_cudnn_ver_error);
if (m_tensor.raggedOffset == nullptr) {
set_error_and_throw_exception(
&m_tensor,
CUDNN_STATUS_BAD_PARAM,
"CUDNN_BACKEND_TENSOR_DESCRIPTOR: CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER requires "
"CUDNN_ATTR_TENSOR_RAGGED_OFFSET_DESC");
return std::move(m_tensor);
}
if (m_tensor.raggedOffsetMultiplier <= 0) {
set_error_and_throw_exception(
&m_tensor,
CUDNN_STATUS_BAD_PARAM,
"CUDNN_BACKEND_TENSOR_DESCRIPTOR: CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER must be positive");
return std::move(m_tensor);
}
status = detail::set_attribute(m_tensor.pointer->get_backend_descriptor(),
CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER,
CUDNN_TYPE_INT64,
1,
&m_tensor.raggedOffsetMultiplier);
if (status != CUDNN_STATUS_SUCCESS) {
set_error_and_throw_exception(&m_tensor,
status,
"CUDNN_BACKEND_TENSOR_DESCRIPTOR: SetAttribute "
"CUDNN_ATTR_TENSOR_RAGGED_OFFSET_MULTIPLIER Failed");
return std::move(m_tensor);
}
#else
set_error_and_throw_exception(
&m_tensor, CUDNN_STATUS_INVALID_VALUE, ragged_offset_multiplier_cudnn_ver_error);
return std::move(m_tensor);
#endif // CUDNN_VERSION >= 92400
}

// Set the reorder_type
if (m_tensor.reorder_type != cudnn_frontend::TensorReordering_t::NONE) {
cudnnBackendTensorReordering_t cudnn_reordering_type;
Expand Down
3 changes: 3 additions & 0 deletions python/cudnn/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ def _tensor(
reordering_type=tensor_reordering.NONE,
name="",
uid=-1,
ragged_offset_multiplier=1,
):
"""
Create a tensor.
Expand All @@ -80,6 +81,7 @@ def _tensor(
ragged_offset (cudnn_tensor): The ragged offset tensor.
reordering_type (cudnn.tensor_reordering): The reordering type of the tensor.
name (str): The name of the tensor.
ragged_offset_multiplier (int): Unit size of ragged offsets in tensor elements. A value of 1 means no multiplier.

Returns:
cudnn_tensor: The created tensor.
Expand All @@ -94,6 +96,7 @@ def _tensor(
reordering_type=reordering_type,
name=name,
uid=uid,
ragged_offset_multiplier=ragged_offset_multiplier,
)


Expand Down
4 changes: 4 additions & 0 deletions python/properties.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,10 @@ init_properties(py::module_& m) {
&cudnn_frontend::graph::Tensor_attributes::set_vector_count_and_dimension,
py::return_value_policy::reference)
.def("set_ragged_offset", &cudnn_frontend::graph::Tensor_attributes::set_ragged_offset)
.def("get_ragged_offset_multiplier", &cudnn_frontend::graph::Tensor_attributes::get_ragged_offset_multiplier)
.def("set_ragged_offset_multiplier",
&cudnn_frontend::graph::Tensor_attributes::set_ragged_offset_multiplier,
py::return_value_policy::reference)
.def("__repr__", [](cudnn_frontend::graph::Tensor_attributes const& props) {
std::ostringstream out;
out << json{props};
Expand Down
9 changes: 6 additions & 3 deletions python/pygraph/pygraph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,8 @@ PyGraph::tensor(std::vector<int64_t> const& dim,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> const& ragged_offset,
cudnn_frontend::TensorReordering_t const reordering_type,
std::string const& name,
int64_t const& uid) {
int64_t const& uid,
int64_t const& ragged_offset_multiplier) {
auto props = cudnn_frontend::graph::Tensor_attributes()
.set_data_type(data_type)
.set_is_virtual(is_virtual)
Expand All @@ -96,7 +97,8 @@ PyGraph::tensor(std::vector<int64_t> const& dim,
.set_stride(stride)
.set_ragged_offset(ragged_offset)
.set_reordering_type(reordering_type)
.set_name(name);
.set_name(name)
.set_ragged_offset_multiplier(ragged_offset_multiplier);

if (uid != -1) {
props.set_uid(uid);
Expand Down Expand Up @@ -867,7 +869,8 @@ init_pygraph_submodule(py::module_& m) {
py::arg_v{"ragged_offset", nullptr},
py::arg_v{"reordering_type", cudnn_frontend::TensorReordering_t::NONE},
py::arg_v("name", ""),
py::arg_v("uid", -1))
py::arg_v("uid", -1),
py::arg_v("ragged_offset_multiplier", int64_t{1}))
.def("genstats",
&PyGraph::genstats,
py::arg("input"),
Expand Down
3 changes: 2 additions & 1 deletion python/pygraph/pygraph.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,8 @@ class PyGraph {
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> const& ragged_offset,
cudnn_frontend::TensorReordering_t const reordering_type,
std::string const& name,
int64_t const& uid);
int64_t const& uid,
int64_t const& ragged_offset_multiplier);

std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>
tensor_like(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> const& pyobj, std::string const&);
Expand Down
30 changes: 30 additions & 0 deletions test/cpp/tensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,36 @@ TEST_CASE("tensor query checks", "[query_tensor_attributes_of_uid]") {
REQUIRE(t.get_name() == name);
}

TEST_CASE("tensor ragged offset multiplier attributes", "[tensor_ragged_offset_multiplier]") {
namespace fe = cudnn_frontend;

fe::graph::Graph graph;
graph.set_io_data_type(fe::DataType_t::HALF)
.set_intermediate_data_type(fe::DataType_t::FLOAT)
.set_compute_data_type(fe::DataType_t::FLOAT);

auto ragged_offset = graph.tensor(fe::graph::Tensor_attributes()
.set_name("offsets")
.set_dim({3, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_data_type(fe::DataType_t::INT64)
.set_uid(2));

auto X = graph.tensor(fe::graph::Tensor_attributes()
.set_name("ragged")
.set_dim({2, 4, 8, 16})
.set_stride({512, 128, 16, 1})
.set_data_type(fe::DataType_t::HALF)
.set_uid(3)
.set_ragged_offset(ragged_offset)
.set_ragged_offset_multiplier(4));

fe::graph::Tensor_attributes t;
REQUIRE(graph.query_tensor_attributes_of_uid(X->get_uid(), t).is_good());
REQUIRE(t.get_ragged_offset_multiplier() == 4);
REQUIRE(t.validate().is_good());
}

TEST_CASE("Block_scale_dequantize graph creation with negative scales", "[block_scale_dequantize_graph]") {
namespace fe = cudnn_frontend;

Expand Down
22 changes: 18 additions & 4 deletions test/python/sdpa/fp16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,10 +170,18 @@ def allocate_tensors(cfg, rng_data_gen, perf=False):
allocs[TensorUid.seq_len_kv] = (seq_len_kv_gpu, None, None)

if cfg.is_ragged:
allocs[TensorUid.q_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * cfg.d_qk).to(torch.int64), None, None)
allocs[TensorUid.k_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.h_k * cfg.d_qk).to(torch.int64), None, None)
allocs[TensorUid.v_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.h_v * cfg.d_v).to(torch.int64), None, None)
allocs[TensorUid.o_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * cfg.d_v).to(torch.int64), None, None)
# When testing the ragged offset multiplier, store the offsets in coarser
# units by dividing out a per-tensor multiplier (M_q=M_k=d_qk, M_v=M_o=d_v);
# the engine recovers the true element offset by multiplying back. stats keeps
# the default multiplier of 1. The chosen multipliers always divide evenly.
q_off_mult = cfg.d_qk if cfg.with_ragged_offset_multiplier else 1
k_off_mult = cfg.d_qk if cfg.with_ragged_offset_multiplier else 1
v_off_mult = cfg.d_v if cfg.with_ragged_offset_multiplier else 1
o_off_mult = cfg.d_v if cfg.with_ragged_offset_multiplier else 1
allocs[TensorUid.q_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * cfg.d_qk // q_off_mult).to(torch.int64), None, None)
allocs[TensorUid.k_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.h_k * cfg.d_qk // k_off_mult).to(torch.int64), None, None)
allocs[TensorUid.v_ragged_offset] = ((prefix_sum(seq_len_kv_gpu) * cfg.h_v * cfg.d_v // v_off_mult).to(torch.int64), None, None)
allocs[TensorUid.o_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * cfg.d_v // o_off_mult).to(torch.int64), None, None)
allocs[TensorUid.stats_ragged_offset] = ((prefix_sum(seq_len_q_gpu) * cfg.h_q * 1).to(torch.int64), None, None)

if cfg.is_bias:
Expand Down Expand Up @@ -282,6 +290,10 @@ def create_forward_graph(cfg, tensors, cudnn_handle):
q.set_ragged_offset(q_ragged_offset)
k.set_ragged_offset(k_ragged_offset)
v.set_ragged_offset(v_ragged_offset)
if cfg.with_ragged_offset_multiplier:
q.set_ragged_offset_multiplier(cfg.d_qk)
k.set_ragged_offset_multiplier(cfg.d_qk)
v.set_ragged_offset_multiplier(cfg.d_v)

# Create graph tensors for score_max and score_sum_exp
score_max = score_sum_exp = sink_token = None
Expand Down Expand Up @@ -340,6 +352,8 @@ def create_forward_graph(cfg, tensors, cudnn_handle):
o.set_uid(int(TensorUid.o)).set_output(True).set_dim(cfg.shape_o).set_stride(cfg.stride_o)
if cfg.is_ragged:
o.set_ragged_offset(o_ragged_offset)
if cfg.with_ragged_offset_multiplier:
o.set_ragged_offset_multiplier(cfg.d_v)

if cfg.with_score_max:
score_max.set_ragged_offset(stats_ragged_offset)
Expand Down
8 changes: 5 additions & 3 deletions test/python/sdpa/random_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ class ExecConfig:
with_sink_token: bool = False
with_unfuse_fma: bool = False
with_rope: bool = False
with_ragged_offset_multiplier: bool = False
rescale_threshold: float = None

diag_align: cudnn.diagonal_alignment = None
Expand Down Expand Up @@ -256,9 +257,10 @@ def __call__(self, rng, rng_data_seed, rng_geom_seed=None):
randoms_.d_qk, randoms_.d_v = randoms["d_qk_d_v"]
randoms_.h_q, randoms_.h_k, randoms_.h_v = randoms["head_count"]

randoms_.is_ragged = randoms["is_ragged_or_padded_or_full"] in ("ragged", "cu_ragged")
randoms_.is_padding = randoms["is_ragged_or_padded_or_full"] in ("padded", "ragged", "cu_padded", "cu_ragged")
randoms_.is_cu_seq_len = randoms["is_ragged_or_padded_or_full"] in ("cu_padded", "cu_ragged")
randoms_.is_ragged = randoms["is_ragged_or_padded_or_full"] in ("ragged", "cu_ragged", "ragged_mult", "cu_ragged_mult")
randoms_.is_padding = randoms["is_ragged_or_padded_or_full"] in ("padded", "ragged", "cu_padded", "cu_ragged", "ragged_mult", "cu_ragged_mult")
randoms_.is_cu_seq_len = randoms["is_ragged_or_padded_or_full"] in ("cu_padded", "cu_ragged", "cu_ragged_mult")
randoms_.with_ragged_offset_multiplier = randoms["is_ragged_or_padded_or_full"] in ("ragged_mult", "cu_ragged_mult")

if randoms["is_ragged_or_padded_or_full"] != "full":
# ~10% chance of 0-length sequence for each batch
Expand Down
Loading