Skip to content

Support zero-element graphs (e.g. SDPA with batch size 0) #404

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

Draft
Anerudhan wants to merge 1 commit into
base: develop
Choose a base branch
from
Draft

Support zero-element graphs (e.g. SDPA with batch size 0) #404

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
63 changes: 63 additions & 0 deletions docs/zero-element-graphs.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
# Zero-element (no-op) graphs

Graphs whose tensors contain a dimension of size 0 — for example SDPA with a
batch size of 0 — are supported by the frontend as no-ops
(see [issue #101](https://github.com/NVIDIA/cudnn-frontend/issues/101)).

The cuDNN backend does not accept tensor descriptors with 0-sized dimensions.
Instead of lowering such graphs, the frontend detects them during `validate()`:
if the graph references at least one zero-element tensor and **every**
non-virtual output tensor is zero-element, no output byte would ever be
written, so the graph is flagged as a zero-element no-op. For such graphs:

- `validate()`, `build_operation_graph()`, `create_execution_plans()`,
`check_support()`, `build_plans()`, and `build()` all succeed without
touching the backend.
- `get_workspace_size()` reports 0.
- `execute()` returns success without launching any work. Pointers in the
variant pack are ignored (zero-element tensors may have null pointers).
- `populate_cuda_graph()` leaves the provided CUDA graph empty (an empty CUDA
graph is valid and instantiable), and `update_cuda_graph()` is a no-op.
- `Graph::is_zero_element_graph()` (C++) / `pygraph.is_zero_element_graph()`
(Python) report whether the graph was flagged.

## Example

```cpp
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);

int64_t b = 0, h = 4, s_q = 64, s_kv = 64, d = 64; // batch size 0
auto Q = graph.tensor(fe::graph::Tensor_attributes()
.set_dim({b, h, s_q, d})
.set_stride({h * s_q * d, s_q * d, d, 1}));
// ... K, V ...

auto [O, Stats] = graph.sdpa(Q, K, V, sdpa_options);
O->set_output(true).set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1});

// The full pipeline succeeds; execute() is a no-op.
auto status = graph.build(handle, {fe::HeurMode_t::A});
```

## Unsupported mixes

A graph that mixes zero-element tensors with **non-zero-element** output
tensors is rejected at `validate()` with `GRAPH_NOT_SUPPORTED` and a clear
message. For example, a matmul contracting over a dimension of size 0
(`[1, m, 0] x [1, 0, n] -> [1, m, n]`) would require zero-filling the output,
which cuDNN does not do.

## Limitations

- Plan serialization (`serialize(std::vector<uint8_t>&)` /
`deserialize(handle, data)`) is not supported for zero-element no-op graphs,
as there is no execution plan to serialize. JSON graph (structure)
serialization is unaffected: a graph deserialized from JSON is re-detected as
a no-op during its `validate()`.
- Runtime shape overrides (`override_uids` / `override_shapes` /
`override_strides`) cannot be applied to a graph built as a zero-element
no-op; build the graph with non-zero shapes instead.
167 changes: 167 additions & 0 deletions include/cudnn_frontend/graph_interface.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,12 @@ class Graph : public ICudnn, public INode {
int64_t fe_workspace_size = 0;
uint64_t graph_uid;

// Set during validate() when every non-virtual output tensor is zero-element
// (has a dimension of size 0, e.g. SDPA with batch size 0). Executing such a
// graph writes nothing, so backend lowering is skipped and all build/execute
// stages become no-ops. The backend does not accept 0-sized dimensions.
bool is_zero_element_graph_ = false;

std::unordered_set<std::shared_ptr<Tensor_attributes>> deserialized_tensor_properties;
std::unordered_map<uid_t, pass_by_values_t> deserialized_pass_by_value;
std::unordered_map<uid_t, std::tuple<int64_t, int64_t, std::vector<float>>> deserialized_workspace_modifications;
Expand Down Expand Up @@ -93,6 +99,54 @@ class Graph : public ICudnn, public INode {
return {error_code_t::OK, ""};
}

// Detects zero-element (no-op) graphs. Called from validate(), after shape
// inference has filled in all tensor dims.
// A graph is a no-op iff it references at least one zero-element tensor and
// every non-virtual output tensor is zero-element: no output byte would ever
// be written, so all build/execute stages can safely be skipped.
// A graph that mixes zero-element tensors with non-zero-element outputs
// (e.g. a matmul whose contracted dimension is 0, whose output would need
// to be zero-filled) is rejected here with a clear message instead of
// failing later in the backend with a generic CUDNN_STATUS_BAD_PARAM.
error_t
detect_zero_element_graph() {
is_zero_element_graph_ = false;

bool has_zero_element_tensor = false;
bool all_outputs_zero_element = true;
for (auto const &input : full_graph_inputs) {
if (input && input->get_volume() == 0) {
has_zero_element_tensor = true;
}
}
for (auto const &output : full_graph_outputs) {
if (output == nullptr || output->get_is_virtual()) {
continue;
}
if (output->get_volume() == 0) {
has_zero_element_tensor = true;
} else {
all_outputs_zero_element = false;
}
}

if (has_zero_element_tensor == false) {
return {error_code_t::OK, ""};
}

RETURN_CUDNN_FRONTEND_ERROR_IF(
all_outputs_zero_element == false,
error_code_t::GRAPH_NOT_SUPPORTED,
"Graph mixes zero-element tensors (a dimension of size 0) with non-zero-element output tensors. "
"This is not supported: cuDNN cannot compute non-empty outputs from zero-element inputs.");

is_zero_element_graph_ = true;
CUDNN_FE_LOG_LABEL_ENDL(
"INFO: All output tensors are zero-element. Graph is treated as a no-op: build stages are skipped, "
"workspace size is 0, and execute() launches no work.");
return {error_code_t::OK, ""};
}

error_t
log_tensors_to_dump_(cudnnHandle_t handle,
std::unordered_map<int64_t, void *> const &tensor_uid_to_pointer_map) const {
Expand Down Expand Up @@ -646,6 +700,12 @@ class Graph : public ICudnn, public INode {
std::unordered_map<Tensor_attributes::uid_t, void *> &uid_to_device_ptrs,
void *workspace,
cudaGraph_t cudnn_cuda_graph) {
// Zero-element no-op graph: populate_cuda_graph left the cuda graph empty; nothing to update.
if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping cuda graph update of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

// Initializes this cudnn graph
RETURN_CUDNN_FRONTEND_ERROR_IF(
cudnn_cuda_graph == nullptr, error_code_t::INVALID_VALUE, "cudnn_cuda_graph should not be a nullptr");
Expand Down Expand Up @@ -788,6 +848,13 @@ class Graph : public ICudnn, public INode {
std::unordered_map<Tensor_attributes::uid_t, void *> &uid_to_device_ptrs,
void *workspace,
cudaGraph_t cudnn_cuda_graph) {
// Zero-element no-op graph: leave the provided cuda graph empty (an empty
// cuda graph is valid and instantiable).
if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping cuda graph population of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

// Check if the cuda graph is empty
size_t numNodes = 0;
_CUDNN_CHECK_CUDA_ERROR(detail::cuda_graph_get_nodes(cudnn_cuda_graph, nullptr, &numNodes));
Expand Down Expand Up @@ -936,6 +1003,9 @@ class Graph : public ICudnn, public INode {
CHECK_CUDNN_FRONTEND_ERROR(output->validate());
}

// Now that all shapes are known, detect zero-element (no-op) graphs.
CHECK_CUDNN_FRONTEND_ERROR(detect_zero_element_graph());

// Get all the pre assigned uids
CHECK_CUDNN_FRONTEND_ERROR(get_pre_assigned_uids(used_uids));
// Clear state
Expand All @@ -951,6 +1021,11 @@ class Graph : public ICudnn, public INode {
build_operation_graph() {
CUDNN_FE_LOG_BANNER(" BUILD OP GRAPH WITHOUT HANDLE ");

if (is_zero_element_graph_) {
CUDNN_FE_LOG_BANNER(" SKIPPED BUILD OP GRAPH (ZERO-ELEMENT NO-OP GRAPH) ");
return {error_code_t::OK, ""};
}

if (device_properties == nullptr) {
return {error_code_t::ATTRIBUTE_NOT_SET, "Device properties are not set."};
}
Expand All @@ -962,6 +1037,11 @@ class Graph : public ICudnn, public INode {
build_operation_graph(cudnnHandle_t handle) {
CUDNN_FE_LOG_BANNER(" BUILD OP GRAPH ");

if (is_zero_element_graph_) {
CUDNN_FE_LOG_BANNER(" SKIPPED BUILD OP GRAPH (ZERO-ELEMENT NO-OP GRAPH) ");
return {error_code_t::OK, ""};
}

CUDNN_FE_LOG_BANNER(" 1/4 INFER PROPERTIES OF NODES ");

// expand composite nodes
Expand Down Expand Up @@ -1061,6 +1141,13 @@ class Graph : public ICudnn, public INode {

error_t
get_workspace_size_plan_at_index(int64_t plan_index, int64_t &cudnn_workspace_size) const {
// Zero-element no-op graphs have no plans and need no workspace.
if (is_zero_element_graph_) {
cudnn_workspace_size = 0;
CUDNN_FE_LOG_LABEL_ENDL("INFO: get_workspace_size() is 0 (zero-element no-op graph)");
return {error_code_t::OK, ""};
}

// OSS SDPA engine workspace: 16 bytes for tile_id_counter
if (plan_index == graph::Execution_plan_list::OSS_SDPA_ENGINE_CANDIDATE) {
cudnn_workspace_size = fe_workspace_size + experimental::Sm90SdpaPrefillEngine::get_workspace_size();
Expand Down Expand Up @@ -1093,6 +1180,17 @@ class Graph : public ICudnn, public INode {
std::vector<int64_t> const &override_uids,
std::vector<std::vector<int64_t>> const &override_shapes,
std::vector<std::vector<int64_t>> const &override_strides) const {
// Zero-element no-op graphs have no plans and need no workspace.
if (is_zero_element_graph_) {
RETURN_CUDNN_FRONTEND_ERROR_IF(!override_uids.empty(),
error_code_t::GRAPH_NOT_SUPPORTED,
"Graph was built as a zero-element no-op; it has no execution plan to "
"apply shape overrides to. Build the graph with non-zero shapes instead.");
cudnn_workspace_size = 0;
CUDNN_FE_LOG_LABEL_ENDL("INFO: get_workspace_size() is 0 (zero-element no-op graph)");
return {error_code_t::OK, ""};
}

RETURN_CUDNN_FRONTEND_ERROR_IF(override_uids.size() != override_shapes.size(),
error_code_t::INVALID_VALUE,
"override_uids and override_shapes must have the same size.");
Expand Down Expand Up @@ -1218,6 +1316,11 @@ class Graph : public ICudnn, public INode {
void *user_impl = nullptr) {
(void)user_impl; // reserved for future use

if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping autotune of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

const int maxIterCount = 100;
const float threshold = 0.95f;

Expand Down Expand Up @@ -1373,6 +1476,16 @@ class Graph : public ICudnn, public INode {
std::vector<int64_t> const &override_uids = {},
std::vector<std::vector<int64_t>> const &override_shapes = {},
std::vector<std::vector<int64_t>> const &override_strides = {}) const {
// Zero-element no-op graph: all outputs are zero-element, nothing to compute.
if (is_zero_element_graph_) {
RETURN_CUDNN_FRONTEND_ERROR_IF(!override_uids.empty(),
error_code_t::GRAPH_NOT_SUPPORTED,
"Graph was built as a zero-element no-op; it has no execution plan to "
"apply shape overrides to. Build the graph with non-zero shapes instead.");
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping execution of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

if (!varpack_prep_state->prepared.load(std::memory_order_acquire)) {
CHECK_CUDNN_FRONTEND_ERROR(const_cast<Graph *>(this)->prepare_variant_pack_template());
}
Expand Down Expand Up @@ -1404,6 +1517,16 @@ class Graph : public ICudnn, public INode {
std::vector<int64_t> const &override_uids = {},
std::vector<std::vector<int64_t>> const &override_shapes = {},
std::vector<std::vector<int64_t>> const &override_strides = {}) const {
// Zero-element no-op graph: all outputs are zero-element, nothing to compute.
if (is_zero_element_graph_) {
RETURN_CUDNN_FRONTEND_ERROR_IF(!override_uids.empty(),
error_code_t::GRAPH_NOT_SUPPORTED,
"Graph was built as a zero-element no-op; it has no execution plan to "
"apply shape overrides to. Build the graph with non-zero shapes instead.");
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping execution of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

// Lazy init: prepare template if not done (e.g. deserialized graphs, build_plan_at_index)
if (!varpack_prep_state->prepared.load(std::memory_order_acquire)) {
CHECK_CUDNN_FRONTEND_ERROR(const_cast<Graph *>(this)->prepare_variant_pack_template());
Expand Down Expand Up @@ -1502,6 +1625,11 @@ class Graph : public ICudnn, public INode {

error_t
warmup(cudnnHandle_t handle) {
if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping warmup of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

cudaStream_t fake_stream;

cudaStream_t original_stream;
Expand Down Expand Up @@ -1610,6 +1738,9 @@ class Graph : public ICudnn, public INode {
error_t
serialize(std::vector<uint8_t> &data, bool serialize_structure = true) const {
CUDNN_FE_LOG_BANNER(" SERIALIZE PLAN ");
RETURN_CUDNN_FRONTEND_ERROR_IF(is_zero_element_graph_,
error_code_t::GRAPH_NOT_SUPPORTED,
"Zero-element no-op graphs have no execution plan to serialize.");
#ifndef CUDNN_FRONTEND_SKIP_JSON_LIB
json j;
// Optionally serialize the graph structure (nodes/tensors).
Expand Down Expand Up @@ -2119,9 +2250,24 @@ class Graph : public ICudnn, public INode {
return check_support();
}

// Returns true if validate() determined this graph to be a zero-element no-op:
// it references zero-element tensors (a dimension of size 0, e.g. batch size 0)
// and every non-virtual output tensor is zero-element. Such graphs skip backend
// lowering, report a workspace size of 0, and execute() launches no work.
bool
is_zero_element_graph() const {
return is_zero_element_graph_;
}

// overload for deviceless AoT compilation
error_t
check_support() {
// Zero-element no-op graphs are always supported: no backend plan is needed.
if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: check_support() OK for zero-element no-op graph.");
return {error_code_t::OK, ""};
}

// Check OSS engine first if registered

CHECK_CUDNN_FRONTEND_ERROR(context.populate_sm_version_from_device());
Expand Down Expand Up @@ -2826,6 +2972,11 @@ inline error_t
Graph::create_execution_plans(std::vector<HeurMode_t> const &mode) {
CUDNN_FE_LOG_BANNER(" CREATE EXECUTION PLANS (HEURISTICS QUERY) ");

if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping execution plan creation of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

// CHECK IF NEED TO OVERRIDE HEURISTICS QUERY
for (auto &sub_node : sub_nodes) {
if (auto [engine_id, user_knobs] = sub_node->override_heuristics_query(); engine_id != -1) {
Expand Down Expand Up @@ -2896,6 +3047,12 @@ Graph::create_execution_plan(int64_t const engine_id, std::unordered_map<KnobTyp
// first create the engine
// this just uses the global engine id and operation graph
CUDNN_FE_LOG_BANNER(" CREATE EXECUTION PLAN for engine id " << engine_id << " ");

if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping execution plan creation of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

detail::backend_descriptor engine(CUDNN_BACKEND_ENGINE_DESCRIPTOR);
RETURN_CUDNN_FRONTEND_ERROR_IF(engine.get_status() != CUDNN_STATUS_SUCCESS,
error_code_t::CUDNN_BACKEND_API_FAILED,
Expand All @@ -2919,6 +3076,11 @@ Graph::create_execution_plan(int64_t const engine_id, std::unordered_map<KnobTyp

inline error_t
Graph::build_plan_at_index(int64_t plan_index) {
if (is_zero_element_graph_) {
CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping plan build of zero-element no-op graph.");
return {error_code_t::OK, ""};
}

CHECK_CUDNN_FRONTEND_ERROR(plans.build_plan_at_index(plan_index));
return {error_code_t::OK, ""};
}
Expand All @@ -2931,6 +3093,11 @@ Graph::build_plans(BuildPlanPolicy_t const policy, bool const do_multithreaded_b
CUDNN_FE_LOG_BANNER(" BUILD PLANS for policy " << static_cast<int>(policy) << " ");
#endif

if (is_zero_element_graph_) {
CUDNN_FE_LOG_BANNER(" SKIPPED BUILD PLANS (ZERO-ELEMENT NO-OP GRAPH) ");
return {error_code_t::OK, ""};
}

// Build OSS SDPA engine if it passed check_support
if (plans.has_oss_sdpa_engine()) {
auto oss_status = plans.build_oss_sdpa_engine();
Expand Down
1 change: 1 addition & 0 deletions llms.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel
- [Dynamic kernel cache](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/dynamic-kernel-cache.md)
- [Custom execution plans](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/custom-execution-plan.md)
- [Compile-time constants](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/CompileTimeConstants.md)
- [Zero-element (no-op) graphs](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/zero-element-graphs.md)
- [Adding PyTorch custom ops](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/adding_torch_custom_ops.md)
- [Python graph and execution backends](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/python_graph_and_execution_backends.md)

Expand Down
Loading