diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index b1f165f491..d35eccd5c5 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -394,6 +394,21 @@ def _identity_score_mod(_graph, score, _tensors): return score +# Soft-cap / relative-bias score_mods are idiomatically written as module-level lambdas, which +# all share __qualname__ == ""; keep them in a tuple so the cache key must tell them apart. +_MODULE_LEVEL_LAMBDA_SCORE_MODS = ( + lambda graph, score, tensors: graph.mul(a=score, b=tensors["scale"]), + lambda graph, score, tensors: graph.add(a=score, b=tensors["scale"]), +) + +# Same body, differing only in a default argument. The constant lives in __defaults__, not in +# the code object, so keying on __code__ alone is not enough to tell these apart. +_DEFAULT_ARG_LAMBDA_SCORE_MODS = ( + lambda graph, score, tensors, cap=10.0: graph.mul(a=score, b=cap), + lambda graph, score, tensors, cap=50.0: graph.mul(a=score, b=cap), +) + + def _install_fake_flax_fused_attn(monkeypatch, *, kernel_available=True): captured = {} @@ -730,6 +745,46 @@ def forward(self, _graph, score, _tensors): assert tex_attention._graph_cache_key("fwd", config_1, ()) is None +def test_fused_attn_score_mod_config_separates_module_level_lambdas(): + """Distinct module-level lambdas must not share one compiled cuDNN graph.""" + first, second = _MODULE_LEVEL_LAMBDA_SCORE_MODS + assert first.__qualname__ == second.__qualname__ == "" + + key_1 = tex_attention._score_mod_callback_cache_key(first) + key_2 = tex_attention._score_mod_callback_cache_key(second) + assert key_1 != key_2 + assert key_1 == tex_attention._score_mod_callback_cache_key(first) + + # Named module-level functions stay cacheable and stable. + named_key = tex_attention._score_mod_callback_cache_key(_identity_score_mod) + assert not tex_attention._score_mod_key_is_uncacheable(named_key) + assert named_key == tex_attention._score_mod_callback_cache_key(_identity_score_mod) + + config_1, _, _ = make_fused_attn_score_mod_config( + first, None, {"scale": 1.0}, None, _CONFIG_TEST_SCALING_FACTOR, True + ) + config_2, _, _ = make_fused_attn_score_mod_config( + second, None, {"scale": 1.0}, None, _CONFIG_TEST_SCALING_FACTOR, True + ) + assert config_1 != config_2 + assert tex_attention._graph_cache_key("fwd", config_1, ()) != tex_attention._graph_cache_key( + "fwd", config_2, () + ) + + +def test_fused_attn_score_mod_config_separates_default_arg_lambdas(): + """Lambdas differing only in a default argument must not share one compiled graph.""" + first, second = _DEFAULT_ARG_LAMBDA_SCORE_MODS + # Same code object: the differing constant is in __defaults__, not co_consts. + assert first.__code__.co_consts == second.__code__.co_consts + assert first.__defaults__ != second.__defaults__ + + key_1 = tex_attention._score_mod_callback_cache_key(first) + key_2 = tex_attention._score_mod_callback_cache_key(second) + assert key_1 != key_2 + assert not tex_attention._score_mod_key_is_uncacheable(key_1) + + @pytest.mark.skipif(not _has_cudnn_frontend_python(), reason="cuDNN Python frontend is required") def test_fused_attn_score_mod_post_scale_bias_optional_bprop(): """Post-scale-bias score_mod matches the JAX reference without explicit bprop.""" diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fc21771297..97e993b752 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -502,23 +502,25 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( "Please upgrade your cuDNN version if possible." << std::endl; } - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && sm_arch_ == 120) { + // cudnn-frontend gates SDPA on the compute-capability major version, so sm_121 shares these. + const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130); + if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && is_sm12x) { if (cudnn_runtime_version < 91801) { backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of sm_arch_ == 120 and cudnn_runtime_version < " + std::cout << "Warning: Given combination of sm_arch_ == 12x and cudnn_runtime_version < " "91801 is not supported. " << " Please upgrade your cuDNN version if possible." << std::endl; } else if (deterministic && is_training) { backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Deterministic fused attention on SM120 is not supported." + std::cout << "Warning: Deterministic fused attention on SM12x is not supported." << std::endl; } else { - // Known missing support for T3HD/TH3D layouts on SM120 + // Known missing support for T3HD/TH3D layouts on SM12x const bool is_t3hd_or_th3d = (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD || qkv_layout == NVTE_QKV_Layout::NVTE_TH3D); if (is_t3hd_or_th3d) { backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM120 is not supported. " + std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM12x is not supported. " << " Please consider using other THD layouts if possible." << std::endl; } } diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index bf34758a35..5e11ab0d9c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -86,7 +86,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const auto cudnn_runtime_version = cudnnGetVersion(); const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; + // cuDNN gates these behaviours on the compute capability family, so sm_121 etc. behave + // like sm_120 and must take the same workarounds. + const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130); + bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && !is_sm12x; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); @@ -112,10 +115,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] + // On SM 12x, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] // as interleaved and rejects them. Use BHSD-like dimensions/strides with max_seqlen at plan build // so the check passes; ragged offset still provides variable-length boundaries. - if (sm_arch_ != 120) { + if (!is_sm12x) { // replace batch size and maximum sequence lengths with maximum token counts // for query and key/value so the graph is static within each quantization bucket. // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size: @@ -179,6 +182,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( cudnn_frontend::DataType_t::NOT_SET, cudnn_frontend::DataType_t::NOT_SET, return_max_logit, + device_id, }; namespace fe = cudnn_frontend; @@ -678,7 +682,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( const auto cudnn_runtime_version = cudnnGetVersion(); const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; + const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130); + bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && !is_sm12x; NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); @@ -690,8 +695,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd). - if (sm_arch_ != 120) { + // On SM 12x, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd). + if (!is_sm12x) { // replace batch size and maximum sequence lengths with maximum token counts // for query and key/value so the graph is static within each quantization bucket b = max_b; @@ -743,6 +748,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( cudnn_frontend::DataType_t::NOT_SET, cudnn_frontend::DataType_t::NOT_SET, false, + device_id, }; namespace fe = cudnn_frontend; @@ -893,7 +899,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( if (use_ragged_stats) { sdpa_backward_options.set_max_total_seq_len_q(s_q); } - if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + if (is_ragged_kv && cudnn_runtime_version >= 90600 && !is_sm12x) { sdpa_backward_options.set_max_total_seq_len_kv(s_kv); } @@ -1204,6 +1210,7 @@ void fused_attn_arbitrary_seqlen_fwd( const int device_id = cuda::current_device(); const int sm_arch_ = cuda::sm_arch(device_id); + const bool is_sm12x = (sm_arch_ >= 120 && sm_arch_ < 130); void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; @@ -1231,10 +1238,9 @@ void fused_attn_arbitrary_seqlen_fwd( Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; - // sm120 does not use ragged stats: the graph declares a dense + // sm12x does not use ragged stats: the graph declares a dense // [b, h, s_q, 1] stats tensor, so allocate to match (same as Max below). - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && !is_sm12x) { output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1244,8 +1250,7 @@ void fused_attn_arbitrary_seqlen_fwd( if (return_max_logit) { Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Max->data.dptr = nullptr; - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && !is_sm12x) { output_Max->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 000af41aee..80168144d3 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -6,6 +6,7 @@ #include "../common.h" #include "../cudnn_utils.h" +#include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -116,7 +117,8 @@ void fused_attn_fp8_fwd_impl( o_tensor_type, cudnn_frontend::DataType_t::NOT_SET, cudnn_frontend::DataType_t::NOT_SET, - false}; + false, + cuda::current_device()}; namespace fe = cudnn_frontend; using graph_and_tensors = @@ -584,7 +586,8 @@ void fused_attn_fp8_bwd_impl( o_tensor_type, do_tensor_type, dqkv_tensor_type, - false}; + false, + cuda::current_device()}; namespace fe = cudnn_frontend; using graph_and_tensors = diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 1864f9417d..80cf5cad8f 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -288,6 +288,9 @@ struct FADescriptor_v1 { cudnn_frontend::DataType_t do_tensor_type; cudnn_frontend::DataType_t dqkv_tensor_type; bool return_max_logit; + // The plan caches are keyed by this descriptor but the cuDNN handle is per-device, so a + // plan built on one device must never be replayed on another. + int device_id; bool operator<(const FADescriptor_v1 &rhs) const { return std::tie(b, h, hg, s_q, s_kv, d_qk, d_v, num_pages_k, num_pages_v, page_size_k, @@ -296,7 +299,7 @@ struct FADescriptor_v1 { do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, bias_type, qkv_tensor_type, o_tensor_type, do_tensor_type, - dqkv_tensor_type, return_max_logit) < + dqkv_tensor_type, return_max_logit, device_id) < std::tie(rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d_qk, rhs.d_v, rhs.num_pages_k, rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k, rhs.max_pages_per_seq_v, rhs.bias_b, rhs.bias_h, rhs.bias_sq, rhs.bias_skv, @@ -305,7 +308,7 @@ struct FADescriptor_v1 { rhs.do_scale_inv_format, rhs.mask_type, rhs.softmax_type, rhs.window_size_left, rhs.window_size_right, rhs.bottom_right_diagonal, rhs.deterministic, rhs.bias_type, rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, - rhs.dqkv_tensor_type, rhs.return_max_logit); + rhs.dqkv_tensor_type, rhs.return_max_logit, rhs.device_id); } }; diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index ff9f5fb146..7ef8ad7c5c 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -202,7 +202,27 @@ def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: and callback.__closure__ is None and "" not in callback.__qualname__ ): - return ("function", callback.__module__, callback.__qualname__) + # __qualname__ is "" for every module-level lambda, so it does not identify + # the body; include the code object so distinct score_mods get distinct graphs. + # Defaults live outside the code object and must be keyed separately: baking a + # scalar in via a default argument is the usual way to write a soft-cap score_mod. + code = callback.__code__ + kwdefaults = callback.__kwdefaults__ + key = ( + "function", + callback.__module__, + callback.__qualname__, + code.co_code, + code.co_consts, + code.co_names, + callback.__defaults__, + tuple(sorted(kwdefaults.items())) if kwdefaults else None, + ) + try: + hash(key) + except TypeError: + return _UncacheableScoreModKey() + return key return _UncacheableScoreModKey() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 995dd2e90a..7508a15760 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1616,11 +1616,11 @@ def forward( softmax_lse_in_packed_format = False if qkv_format == "thd": if use_fused_attention: - softmax_lse_in_packed_format = get_cudnn_version() >= ( - 9, - 6, - 0, - ) and get_device_compute_capability() != (12, 0) + # sm12x (sm_120, sm_121, ...) keeps the dense BHS1 stats layout; cuDNN + # classifies by compute-capability family, so match on the major version. + softmax_lse_in_packed_format = ( + get_cudnn_version() >= (9, 6, 0) and get_device_compute_capability()[0] != 12 + ) else: softmax_lse_in_packed_format = fa_utils.v2_6_0_plus or use_flash_attn_3 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8a4c6ce286..8187e7a099 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -459,6 +459,9 @@ def get_attention_backend( if not logger.hasHandlers(): logger.addHandler(AttentionLogging._stream_handler) device_compute_capability = get_device_compute_capability() + # cuDNN classifies by compute-capability family, so sm_121 (GB10) hits the same cuDNN + # limitations as sm_120 and must take the same workarounds. + is_sm12x = device_compute_capability[0] == 12 cudnn_version = get_cudnn_version() run_config = { "transformer_engine_version": te.__version__, @@ -708,16 +711,16 @@ def _disable_all_flash_attention() -> None: logger.debug("Disabling FusedAttention for %s", fp8_recipe.__class__.__name__) use_fused_attention = False - if device_compute_capability == (12, 0): + if is_sm12x: if use_flash_attention: logger.debug( "Disabling FlashAttention as FP8 is not supported" - " for compute capability = sm120" + " for compute capability = sm12x" ) if use_fused_attention: logger.debug( "Disabling FusedAttention as FP8 is not supported" - " for compute capability = sm120" + " for compute capability = sm12x" ) use_flash_attention = False use_fused_attention = False @@ -826,14 +829,14 @@ def _disable_all_flash_attention() -> None: # Flash v4 | FP16/BF16 | TODO | sm80+ | bshd,sbhd,thd | TODO # Unfused | FP32/FP16/BF16 | non-paged/paged | all | bshd,sbhd,thd | >= 1 if inference_params is not None: - # Temporarily disabling fused attention for kv caching for sm89/sm120 irrespective of + # Temporarily disabling fused attention for kv caching for sm89/sm12x irrespective of # cuDNN version until the cuDNN bug is resolved. - if device_compute_capability in ((8, 9), (12, 0)): - logger.debug("Disabling FusedAttention for KV caching for sm89/sm120") + if device_compute_capability == (8, 9) or is_sm12x: + logger.debug("Disabling FusedAttention for KV caching for sm89/sm12x") use_fused_attention = False - # Temporarily disable FlashAttention for KV caching on sm120 - if device_compute_capability == (12, 0): - logger.debug("Disabling FlashAttention for KV caching for sm120") + # Temporarily disable FlashAttention for KV caching on sm12x + if is_sm12x: + logger.debug("Disabling FlashAttention for KV caching for sm12x") use_flash_attention = False if context_parallel: logger.debug("Disabling all backends for KV caching with context parallelism") @@ -895,15 +898,11 @@ def _disable_all_flash_attention() -> None: qkv_layout, ) use_fused_attention = False - if ( - device_compute_capability == (12, 0) - and (head_dim_qk > 128 or head_dim_qk % 8 != 0) - and is_training - ): + if is_sm12x and (head_dim_qk > 128 or head_dim_qk % 8 != 0) and is_training: if use_fused_attention: logger.debug( "Disabling FusedAttention as MLA for backward pass is not supported for compute" - " capability = sm120 for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:" + " capability = sm12x for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:" " head_dim_qk = %s", head_dim_qk, ) @@ -1057,19 +1056,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_unfused_attention: logger.debug("Disabling UnfusedDotProductAttention for pad_between_seqs = True") use_unfused_attention = False - if device_compute_capability == (12, 0): + if is_sm12x: if cudnn_version < (9, 18, 1): if use_fused_attention: logger.debug( "Disabling FusedAttention as qkv_format = thd is" - " not supported for compute capability = sm120 and cuDNN version < 9.18.1" + " not supported for compute capability = sm12x and cuDNN version < 9.18.1" ) use_fused_attention = False elif qkv_layout in {"t3hd", "th3d"}: if use_fused_attention: logger.debug( "Disabling FusedAttention as qkv_layout = %s is not supported for" - " compute capability = sm120", + " compute capability = sm12x", qkv_layout, ) use_fused_attention = False