Skip to content

Support mixed-form sequence lengths in SDPA forward (cuDNN 9.26+) #430

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
egilliam-nv wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Support mixed-form sequence lengths in SDPA forward (cuDNN 9.26+) #430

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
10 changes: 6 additions & 4 deletions include/cudnn_frontend/node/diagonal_band_mask.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,12 @@ class DiagonalBandMaskNodeBase : public NodeCRTP<DerivedT> {
RETURN_CUDNN_FRONTEND_ERROR_IF(has_left_bound() && has_shift_right_bound(),
error_code_t::INVALID_VALUE,
"DiagonalBandMaskNode cannot have both left_bound and shift_right_bound");
RETURN_CUDNN_FRONTEND_ERROR_IF(
(has_seq_len_q() || has_seq_len_kv()) && (has_cu_seq_len_q() || has_cu_seq_len_kv()),
error_code_t::INVALID_VALUE,
"SEQ_LEN_Q / SEQ_LEN_KV and CU_SEQ_LEN_Q / CU_SEQ_LEN_KV are mutually exclusive");
RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_q() && has_cu_seq_len_q(),
error_code_t::INVALID_VALUE,
"SEQ_LEN_Q and CU_SEQ_LEN_Q are mutually exclusive");
RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_kv() && has_cu_seq_len_kv(),
error_code_t::INVALID_VALUE,
"SEQ_LEN_KV and CU_SEQ_LEN_KV are mutually exclusive");

return pre_validate_node_extra();
}
Expand Down
69 changes: 43 additions & 26 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 @@ -84,26 +84,32 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
error_code_t::GRAPH_NOT_SUPPORTED,
"Bias mask data type cannot be boolean");

// validate options for padding mask: padding requires per-sequence length tensors,
// either as (SEQ_LEN_Q + SEQ_LEN_KV) or as (CU_SEQ_LEN_Q + CU_SEQ_LEN_KV).
RETURN_CUDNN_FRONTEND_ERROR_IF(
padding_mask && (!has_seq_len_q || !has_seq_len_kv) && (!has_cu_seq_len_q || !has_cu_seq_len_kv),
error_code_t::ATTRIBUTE_NOT_SET,
"Padding mask requires seq_len_q/seq_len_kv (or cu_seq_len_q/cu_seq_len_kv) to be set.");
// validate options for padding mask: padding requires a per-sequence length
// tensor on each side. Each side independently uses exactly one
// representation -- per-batch (seq_len_*) or cumulative (cu_seq_len_*) --
// and the two sides may use different forms, e.g. cumulative Q lengths
// with per-batch KV lengths.
bool const has_any_seq_len_q = has_seq_len_q || has_cu_seq_len_q;
bool const has_any_seq_len_kv = has_seq_len_kv || has_cu_seq_len_kv;
RETURN_CUDNN_FRONTEND_ERROR_IF(padding_mask && !(has_any_seq_len_q && has_any_seq_len_kv),
error_code_t::ATTRIBUTE_NOT_SET,
"Padding mask requires seq_len_q/seq_len_kv (or cu_seq_len_q/cu_seq_len_kv) "
"to be set.");
RETURN_CUDNN_FRONTEND_ERROR_IF(
(!padding_mask && !attention_score_modifier) &&
(has_seq_len_q || has_seq_len_kv || has_cu_seq_len_q || has_cu_seq_len_kv),
(!padding_mask && !attention_score_modifier) && (has_any_seq_len_q || has_any_seq_len_kv),
error_code_t::ATTRIBUTE_NOT_SET,
"seq_len_q/seq_len_kv (or cu_seq_len_q/cu_seq_len_kv) needs to be set only if padding mask is enabled.");

// Cumulative sequence length tensors must be set together, and are mutually
// exclusive with the (per-batch) seq_len tensors.
RETURN_CUDNN_FRONTEND_ERROR_IF(has_cu_seq_len_q != has_cu_seq_len_kv,
error_code_t::ATTRIBUTE_NOT_SET,
"cu_seq_len_q and cu_seq_len_kv must both be set or both unset.");
RETURN_CUDNN_FRONTEND_ERROR_IF((has_cu_seq_len_q || has_cu_seq_len_kv) && (has_seq_len_q || has_seq_len_kv),
RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_q && has_cu_seq_len_q,
error_code_t::INVALID_VALUE,
"seq_len_q and cu_seq_len_q cannot both be set.");
RETURN_CUDNN_FRONTEND_ERROR_IF(has_seq_len_kv && has_cu_seq_len_kv,
error_code_t::INVALID_VALUE,
"Cannot specify both seq_len tensors and cu_seq_len tensors.");
"seq_len_kv and cu_seq_len_kv cannot both be set.");
RETURN_CUDNN_FRONTEND_ERROR_IF(has_any_seq_len_q != has_any_seq_len_kv,
error_code_t::ATTRIBUTE_NOT_SET,
"A Q-side and a KV-side sequence length tensor must be provided together "
"(each side may independently use seq_len_* or cu_seq_len_*).");

RETURN_CUDNN_FRONTEND_ERROR_IF(is_ragged && ((padding_mask == false) && (attention_score_modifier == nullptr)),
error_code_t::GRAPH_NOT_SUPPORTED,
Expand Down Expand Up @@ -301,10 +307,10 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
"Paged caches are not supported in combination with ragged offsets.");

RETURN_CUDNN_FRONTEND_ERROR_IF(
is_paged && !((has_seq_len_q && has_seq_len_kv) || (has_cu_seq_len_q && has_cu_seq_len_kv)),
is_paged && !(has_any_seq_len_q && has_any_seq_len_kv),
error_code_t::GRAPH_NOT_SUPPORTED,
"Paged caches can only be used in combination with padding mask and variable sequence lengths "
"for both Q and KV (via seq_len_q/seq_len_kv or cu_seq_len_q/cu_seq_len_kv).");
"for both Q and KV (each side independently via seq_len_* or cu_seq_len_*).");

RETURN_CUDNN_FRONTEND_ERROR_IF(
!is_paged && max_seq_kv_explicit,
Expand Down Expand Up @@ -389,21 +395,24 @@ SDPA_attributes::validate_sdpa_support_surface(const detail::Context& context,
inline error_t
SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Context& context,
AttentionImplementation_t impl) const {
auto const has_input = [this](input_names name) {
auto const it = inputs.find(name);
return it != inputs.end() && it->second != nullptr;
};

switch (impl) {
case AttentionImplementation_t::AUTO:
// This function should not be called with AUTO.
return {error_code_t::INVALID_VALUE,
"Can't call verify_sdpa_support_surface_for_implementation with impl=AUTO"};
case AttentionImplementation_t::COMPOSITE:
for (const auto& [key, value] : inputs) {
RETURN_CUDNN_FRONTEND_ERROR_IF(key == input_names::Block_mask && value != nullptr,
error_code_t::GRAPH_NOT_SUPPORTED,
"Composite SDPA node doesn't support Block_mask input");
RETURN_CUDNN_FRONTEND_ERROR_IF(
(key == input_names::CU_SEQ_LEN_Q || key == input_names::CU_SEQ_LEN_KV) && value != nullptr,
error_code_t::GRAPH_NOT_SUPPORTED,
"Composite SDPA node doesn't support CU_SEQ_LEN_Q / CU_SEQ_LEN_KV inputs");
}
RETURN_CUDNN_FRONTEND_ERROR_IF(has_input(input_names::Block_mask),
error_code_t::GRAPH_NOT_SUPPORTED,
"Composite SDPA node doesn't support Block_mask input");
RETURN_CUDNN_FRONTEND_ERROR_IF(
has_input(input_names::CU_SEQ_LEN_Q) || has_input(input_names::CU_SEQ_LEN_KV),
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) {
Expand Down Expand Up @@ -459,6 +468,14 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co
allowed_input_msg += ", CU_SEQ_LEN_Q, CU_SEQ_LEN_KV";
}

if (effective_cudnn_ver < 92600) {
RETURN_CUDNN_FRONTEND_ERROR_IF(
has_input(input_names::CU_SEQ_LEN_Q) != has_input(input_names::CU_SEQ_LEN_KV),
error_code_t::GRAPH_NOT_SUPPORTED,
"Mixed-form sequence lengths (cu_seq_len on one side with seq_len on the other) "
"require cuDNN 9.26.0 or above");
}

if (effective_cudnn_ver >= 92500) {
allowed_input_names.insert({input_names::Descale_Q,
input_names::Descale_K,
Expand Down
26 changes: 17 additions & 9 deletions test/python/sdpa/fp16.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ def validate_config(cfg):
if cfg.is_cu_seq_len:
assert cfg.is_padding == True, "is_cu_seq_len=True requires is_padding=True"
assert cfg.is_train == False, "is_cu_seq_len=True is forward-only (cu_seq_len is not plumbed for backward)"
assert cfg.cu_seq_len_sides in ("both", "q", "kv"), f"invalid cu_seq_len_sides={cfg.cu_seq_len_sides}"

assert isinstance(cfg.seq_len_q, (list, tuple)), "input 'seq_len_q' must be list or tuple"
if cfg.is_padding:
Expand Down Expand Up @@ -119,6 +120,10 @@ def validate_config(cfg):
print("@@@@ Overall result: WAIVED, cu_seq_len_q/cu_seq_len_kv require cuDNN 9.24.0 or higher.")
pytest.skip("cu_seq_len_q/cu_seq_len_kv require cuDNN 9.24.0 or higher")

if cudnn_version < "9.26.0" and cfg.is_cu_seq_len and cfg.cu_seq_len_sides != "both":
print("@@@@ Overall result: WAIVED, mixed-form sequence lengths require cuDNN 9.26.0 or higher.")
pytest.skip("mixed-form sequence lengths (cumulative on one side only) require cuDNN 9.26.0 or higher")


def allocate_tensors(cfg, rng_data_gen, perf=False):
allocs = {}
Expand Down Expand Up @@ -159,14 +164,17 @@ def allocate_tensors(cfg, rng_data_gen, perf=False):
seq_len_q_gpu = torch.tensor(cfg.seq_len_q, dtype=torch.int32, device="cuda").view(-1) if len(cfg.seq_len_q) > 0 else None
seq_len_kv_gpu = torch.tensor(cfg.seq_len_kv, dtype=torch.int32, device="cuda").view(-1) if len(cfg.seq_len_kv) > 0 else None

if cfg.is_cu_seq_len:
# When using cu_seq_len, the seq_len_q/seq_len_kv tensors are not part of the
# graph; instead, supply 1-D (b+1,) int32 prefix-sums of the per-batch seq_lens
# (the frontend promotes these to the 4-D form the backend requires).
# A side using the cumulative form gets a 1-D (b+1,) int32 prefix-sum of its
# per-batch seq_lens (the frontend promotes these to the 4-D form the backend
# requires) instead of the regular per-batch tensor. The two sides choose
# independently (cu_seq_len_sides), so mixed forms are exercised too.
if cfg.is_cu_seq_len_q():
allocs[TensorUid.cu_seq_len_q] = (prefix_sum(seq_len_q_gpu).to(torch.int32).view(-1), None, None)
allocs[TensorUid.cu_seq_len_kv] = (prefix_sum(seq_len_kv_gpu).to(torch.int32).view(-1), None, None)
else:
allocs[TensorUid.seq_len_q] = (seq_len_q_gpu, None, None)
if cfg.is_cu_seq_len_kv():
allocs[TensorUid.cu_seq_len_kv] = (prefix_sum(seq_len_kv_gpu).to(torch.int32).view(-1), None, None)
else:
allocs[TensorUid.seq_len_kv] = (seq_len_kv_gpu, None, None)

if cfg.is_ragged:
Expand Down Expand Up @@ -267,11 +275,11 @@ def create_forward_graph(cfg, tensors, cudnn_handle):
block_mask_dim = (cfg.batches, cfg.h_q, (cfg.s_q + TILE_M - 1) // TILE_M, ((cfg.s_kv + TILE_N - 1) // TILE_N + 7) // 8)
block_mask = graph.tensor(uid=int(TensorUid.block_mask), dim=block_mask_dim, stride=(block_mask_dim[1]*block_mask_dim[2]*block_mask_dim[3], block_mask_dim[2]*block_mask_dim[3], block_mask_dim[3], 1), data_type=cudnn.data_type.UINT8) if cfg.is_block_mask else None

seq_len_q = graph.tensor(uid=int(TensorUid.seq_len_q), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len) else None
seq_len_kv = graph.tensor(uid=int(TensorUid.seq_len_kv), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len) else None
seq_len_q = graph.tensor(uid=int(TensorUid.seq_len_q), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len_q()) else None
seq_len_kv = graph.tensor(uid=int(TensorUid.seq_len_kv), dim=(cfg.batches,), stride=(1,), data_type=cudnn.data_type.INT32) if (cfg.is_padding and not cfg.is_cu_seq_len_kv()) else None

cu_seq_len_q = graph.tensor(uid=int(TensorUid.cu_seq_len_q), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len else None
cu_seq_len_kv = graph.tensor(uid=int(TensorUid.cu_seq_len_kv), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len else None
cu_seq_len_q = graph.tensor(uid=int(TensorUid.cu_seq_len_q), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len_q() else None
cu_seq_len_kv = graph.tensor(uid=int(TensorUid.cu_seq_len_kv), dim=(cfg.batches + 1,), stride=(1,), data_type=cudnn.data_type.INT32) if cfg.is_cu_seq_len_kv() else None

seed = offset = dropout_tuple = rng_dump = None
if cfg.is_dropout:
Expand Down
10 changes: 10 additions & 0 deletions test/python/sdpa/random_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,10 @@ class ExecConfig:
# (cumulative sequence-length tensors of shape (b+1, 1, 1, 1)) instead of
# the regular per-batch seq_len_q/seq_len_kv tensors. Implies is_padding=True.
is_cu_seq_len: bool = None
# Which sides use the cumulative form when is_cu_seq_len is True: "both"
# (default), "q" (cu_seq_len_q with seq_len_kv), or "kv" (seq_len_q with
# cu_seq_len_kv). The mixed forms require cuDNN 9.26+.
cu_seq_len_sides: str = "both"
is_ragged: bool = None
is_dropout: bool = None
is_determin: bool = None
Expand Down Expand Up @@ -144,6 +148,12 @@ class ExecConfig:
def is_train(self):
return not self.is_infer

def is_cu_seq_len_q(self):
return bool(self.is_cu_seq_len) and self.cu_seq_len_sides in ("both", "q")

def is_cu_seq_len_kv(self):
return bool(self.is_cu_seq_len) and self.cu_seq_len_sides in ("both", "kv")

def fill_derived_fields(self):
"""
Fill in derived fields (shapes, strides) from basic dims.
Expand Down
59 changes: 59 additions & 0 deletions test/python/test_mhas_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1024,6 +1024,65 @@ def test_sdpa_mxfp8_bwd_L0(env_info, test_no, request, cudnn_handle):
# # Single repro test
# # ===================

MIXED_SEQ_LEN_FORM_CASES = [
("q", cudnn.diagonal_alignment.TOP_LEFT, None),
("kv", cudnn.diagonal_alignment.TOP_LEFT, None),
("q", cudnn.diagonal_alignment.BOTTOM_RIGHT, 0),
("kv", cudnn.diagonal_alignment.BOTTOM_RIGHT, 0),
]


@pytest.mark.parametrize(
"cu_sides,diag_align,right_bound",
MIXED_SEQ_LEN_FORM_CASES,
ids=["cu_q", "cu_kv", "cu_q_brcm", "cu_kv_brcm"],
)
@pytest.mark.L0
def test_sdpa_mixed_seq_len_forms_L0(env_info, cu_sides, diag_align, right_bound, request, cudnn_handle):
"""Mixed-form sequence lengths: cumulative on one side, per-batch on the other.

Deterministic configs with non-uniform per-batch lengths, so misreading one
side's form cannot produce a passing result. The bottom-right causal cases
guard the DiagonalBandMask alignment derivation, which must treat the two
sides' forms independently. Requires cuDNN 9.26+ (skips below via exec_sdpa).
"""
test = SDPATestConfig(**env_info, implementation=cudnn.attention_implementation.AUTO)
test.cfg = ExecConfig(
data_type=torch.bfloat16,
rng_data_seed=1234,
rng_geom_seed=5678,
is_alibi=False,
is_infer=True,
is_paged=False,
is_bias=False,
is_block_mask=False,
is_padding=True,
is_cu_seq_len=True,
cu_seq_len_sides=cu_sides,
is_ragged=False,
is_dropout=False,
is_determin=False,
batches=4,
d_qk=64,
d_v=64,
s_q=256,
s_kv=512,
h_q=3,
h_k=3,
h_v=3,
diag_align=diag_align,
left_bound=None,
right_bound=right_bound,
seq_len_q=[128, 100, 256, 37],
seq_len_kv=[96, 64, 512, 200],
implementation=cudnn.attention_implementation.AUTO,
)
test.cfg.fill_derived_fields()
test.showConfig((request.node.name, len(MIXED_SEQ_LEN_FORM_CASES)), request)

exec_sdpa(test.cfg, request, cudnn_handle)


@pytest.mark.skipif("not config.getoption('--repro')", reason="used with '--repro' only")
@pytest.mark.L0
@pytest.mark.L1
Expand Down