diff --git a/docs/source/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 3f9e20e1ef11..b4620cdc0618 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/sparse-attention.md @@ -90,7 +90,7 @@ User configuration is supplied through Python or YAML and controls how the check `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers. -Skip Softmax Attention only works with the **TRTLLM** attention backend in VisualGen. Set `attention_config.backend` to `TRTLLM` when enabling it. +Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA is the only CUTEDSL sparse-attention algorithm that is mutually exclusive with quantized attention. #### Python API @@ -130,6 +130,19 @@ args = VisualGenArgs( ) ``` +```python +# CUTEDSL backend: +args = VisualGenArgs( + model="", + attention_config=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=SkipSoftmaxAttentionConfig( + threshold_scale_factor=5000.0, + ), + ), +) +``` + #### YAML ```yaml @@ -151,6 +164,15 @@ attention_config: disabled_until_timestep: 0.6 ``` +```yaml +# CUTEDSL backend: +attention_config: + backend: CUTEDSL + sparse_attention_config: + algorithm: skip_softmax + threshold_scale_factor: 5000.0 +``` + ### CUDA Graphs `disabled_until_timestep` creates two sparse-attention phases when it is set: the high-timestep disabled phase and the enabled phase after the cutoff. VisualGen includes that phase in CUDA graph keys so graph capture does not reuse a graph across different Skip Softmax Attention settings. See [VisualGen CUDA Graphs](cuda-graph.md) for the general capture and replay design. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py index ada69e1acd9e..93aa23fe6809 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py @@ -13,11 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -CuTe DSL (NVIDIA kernels) Dense FMHA Backend for Visual Generation Models +CuTe DSL (NVIDIA kernels) FMHA Backend for Visual Generation Models -JIT-compiles the dense FMHA kernel and caches the compiled artifact for each kernel configuration. -Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16 inputs. The VSA sparse path uses -VSAAttention from vsa.py instead. +JIT-compiles dense or SkipSoftmax FMHA and caches the compiled artifact for each kernel +configuration. Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16 inputs. The VSA +sparse path uses VSAAttention from vsa.py instead. """ import math @@ -25,6 +25,7 @@ import torch +from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import SkipSoftmaxParams from tensorrt_llm.logger import logger from tensorrt_llm.visual_gen.args import QuantAttentionConfig @@ -64,6 +65,34 @@ # ============================================================================ +def _resolve_skip_softmax_threshold_scale_factor( + threshold_scale_factor: float | None, + sparse_params: SkipSoftmaxParams | None, + timestep: Any, +) -> float | None: + """Resolve the active CuTeDSL threshold for the current denoising phase.""" + if sparse_params is not None: + if timestep is None and sparse_params.scheduler.disabled_until_timestep is not None: + # Fail-open: a missing timestep resolves to the full (unthrottled) + # threshold, i.e. skip-softmax runs during the high-noise steps + # `disabled_until_timestep` exists to protect. This is silent + # elsewhere (a quality regression, not an error), so surface it + # once per process instead of only in this function's return value. + logger.warning_once( + "SkipSoftmax scheduler has disabled_until_timestep=" + f"{sparse_params.scheduler.disabled_until_timestep} configured, but no " + "`timestep` was passed to the CuTeDSL attention forward call. Skip-softmax " + "will run unthrottled (as if past the cutoff) until `timestep` is threaded " + "through.", + key="cute_dsl_skip_softmax_missing_timestep", + ) + runtime_params = sparse_params.scheduler.get_runtime_params(timestep=timestep) + threshold_scale_factor = runtime_params.threshold_scale_factor_prefill + if threshold_scale_factor is None or threshold_scale_factor <= 0.0: + return None + return threshold_scale_factor + + def _check_cute_runtime_available() -> None: if _cute_dsl_import_error is None: return @@ -239,6 +268,8 @@ def cute_dsl_fmha_fwd( scale_o: float | torch.Tensor = 1.0, is_persistent: bool = True, skip_softmax_threshold_scale_factor: float | None = None, + sparse_params: SkipSoftmaxParams | None = None, + timestep: Any = None, qk_sf_vec: int = 0, q_sf: torch.Tensor | None = None, k_sf: torch.Tensor | None = None, @@ -252,7 +283,14 @@ def cute_dsl_fmha_fwd( When `qk_sf_vec` is non-zero, dispatches to the block-scaled kernel class: 32 selects MXFP8 (Q/K stored as FP8 e4m3, SFs as Float8E8M0FNU uint8 storage); 16 selects NVFP4 (Q/K stored as packed FP4 in torch.uint8, SFs as Float8E4M3FN in uint8 storage). + When `sparse_params` is set, its timestep-aware scheduler overrides the direct threshold. """ + skip_softmax_threshold_scale_factor = _resolve_skip_softmax_threshold_scale_factor( + skip_softmax_threshold_scale_factor, + sparse_params, + timestep, + ) + _check_cute_runtime_available() _validate_inputs(q, k, v, o) if qk_sf_vec != 0: @@ -583,7 +621,13 @@ def __init__( num_kv_heads: int | None = None, dtype: torch.dtype | None = None, quant_attention_config: QuantAttentionConfig | None = None, + # Legacy static threshold, superseded by `sparse_params`'s timestep-aware + # scheduler for every in-tree construction path (`create_attention` never + # forwards this). Kept only for direct-construction debug/testing use + # (e.g. unit tests that want a fixed threshold without a scheduler); + # mutually exclusive with `sparse_params` below. skip_softmax_threshold_scale: float | None = None, + sparse_params: SkipSoftmaxParams | None = None, **kwargs, ): self.layer_idx = layer_idx @@ -592,7 +636,10 @@ def __init__( self.num_kv_heads = num_kv_heads or num_heads self.dtype = dtype self.quant_attention_config = quant_attention_config + if skip_softmax_threshold_scale is not None and sparse_params is not None: + raise ValueError("Set either skip_softmax_threshold_scale or sparse_params, not both.") self.skip_softmax_threshold_scale = skip_softmax_threshold_scale + self.sparse_params = sparse_params self.scale = 1.0 / math.sqrt(head_dim) # CuTe DSL expects [B, S, H, D] format @@ -727,6 +774,8 @@ def _fwd( scale_v_channels=scale_v_channels, scale_o=kwargs.get("scale_o", 1.0), skip_softmax_threshold_scale_factor=skip_softmax_threshold_scale, + sparse_params=self.sparse_params, + timestep=kwargs.get("timestep"), qk_sf_vec=qk_sf_vec, q_sf=q_sf, k_sf=k_sf, diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index ce012e16aea3..12108ad84ece 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -49,8 +49,8 @@ def get_visual_gen_attention_backend( Better performance but requires fused QKV - "FA4": Flash Attention 4; provides higher speedup on Blackwell GPUs (sm100) Requires flash-attn package with cute interface - - "CUTEDSL": CuTe DSL kernels. create_attention selects dense FMHA or VSA from - AttentionConfig.sparse_attention_config. + - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA + from AttentionConfig.sparse_attention_config. """ # Lazy imports to avoid circular dependency from .cute_dsl import CuTeDSLAttention diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 7244a724649b..d549cfce3d35 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -77,7 +77,9 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: if not isinstance(sparse_config, SkipSoftmaxAttentionConfig): return - disabled_until_timestep = sparse_config.disabled_until_timestep + disabled_until_timestep = sparse_config.resolve_disabled_until_timestep( + pretrained_config=self.model_config.pretrained_config, + ) if disabled_until_timestep is None: return diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 96256ea9a82a..8c19484229c7 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -220,10 +220,14 @@ def __init__( backend_num_heads = self.local_num_attention_heads backend_num_kv_heads = self.local_num_key_value_heads - # Resolve sparse attention params for TRTLLM backend + # Lower the shared SkipSoftmax user/checkpoint config for each backend + # whose kernel consumes SkipSoftmaxParams. sparse_params = None ss_cfg = config.attention.sparse_attention_config - if isinstance(ss_cfg, SkipSoftmaxAttentionConfig) and backend_name == "TRTLLM": + if isinstance(ss_cfg, SkipSoftmaxAttentionConfig) and backend_name in ( + "TRTLLM", + "CUTEDSL", + ): sparse_params = ss_cfg.to_sparse_params( module_name=self.module_name, pretrained_config=config.pretrained_config, diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index bd742ac16922..d314c92b24ae 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -116,7 +116,7 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM backend) or VSA (CUTEDSL backend)." + "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)." ), ) @@ -177,29 +177,34 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": return self algo = self.sparse_attention_config.algorithm - required_backend = {"skip_softmax": "TRTLLM", "vsa": "CUTEDSL"}.get(algo) - if required_backend is None: + supported_backends = { + "skip_softmax": ("TRTLLM", "CUTEDSL"), + "vsa": ("CUTEDSL",), + }.get(algo) + if supported_backends is None: return self - if self.backend != required_backend: + if self.backend not in supported_backends: raise ValueError( f"sparse_attention_config with algorithm='{algo}' requires " - f"backend='{required_backend}', got backend='{self.backend}'. " - f"Either set backend='{required_backend}' or remove " + f"backend in {supported_backends}, got backend='{self.backend}'. " + f"Either select a supported backend or remove " f"sparse_attention_config." ) return self @model_validator(mode="after") def _validate_cutedsl_quant_sparse_mutex(self) -> "AttentionConfig": - # quant_attention_config and sparse_attention_config are mutually exclusive. + # VSA replaces the dense CuTeDSL path and cannot compose with quantized + # attention. SkipSoftmax is part of that dense path and can compose. if ( self.backend == "CUTEDSL" and self.quant_attention_config is not None and self.sparse_attention_config is not None + and self.sparse_attention_config.algorithm == "vsa" ): raise ValueError( - "CUTEDSL backend: quant_attention_config and " + "CUTEDSL backend: quant_attention_config and VSA " "sparse_attention_config are mutually exclusive (the " "CuTeDSLAttention dispatcher selects either the dense path " "or the sparse VSA path, not both)." diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 09d7e7e10d21..d261c770ca9e 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -72,7 +72,6 @@ def to_sparse_params(self, **kwargs): from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import ( SkipSoftmaxParams, SkipSoftmaxScheduler, - skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config, ) module_name = kwargs.get("module_name", None) @@ -81,13 +80,9 @@ def to_sparse_params(self, **kwargs): if module_name is not None and self._is_disabled(module_name, ckpt_sparse_attention_config): return None - disabled_until_timestep = self.disabled_until_timestep - if disabled_until_timestep is None: - disabled_until_timestep = ( - skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config( - ckpt_sparse_attention_config - ) - ) + disabled_until_timestep = self.resolve_disabled_until_timestep( + checkpoint_config=ckpt_sparse_attention_config, + ) threshold_scale_factor = self.resolve_threshold_scale_factor(ckpt_sparse_attention_config) if threshold_scale_factor is None: @@ -104,6 +99,33 @@ def to_sparse_params(self, **kwargs): return None return SkipSoftmaxParams(scheduler=scheduler) + def resolve_disabled_until_timestep( + self, + *, + checkpoint_config: Optional[Dict[str, Any]] = None, + pretrained_config: Any = None, + ) -> Optional[float]: + """Resolve the user override or checkpoint-provided timestep cutoff. + + Exactly one of ``checkpoint_config`` (a raw checkpoint config dict) or + ``pretrained_config`` (the model's pretrained config object/dict) is + expected per call site; both default to ``None`` so a mistyped keyword + raises ``TypeError`` here instead of silently resolving to ``None``. + """ + if self.disabled_until_timestep is not None: + return self.disabled_until_timestep + + from tensorrt_llm._torch.attention_backend.sparse.skip_softmax import ( + skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config, + ) + + ckpt_sparse_attention_config = self._ckpt_sparse_attention_config_from_kwargs( + {"checkpoint_config": checkpoint_config, "pretrained_config": pretrained_config} + ) + return skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config( + ckpt_sparse_attention_config + ) + def _is_disabled( self, module_name: str, diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 2b46f28c6873..726cf5aac9c4 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -43,6 +43,7 @@ l0_cpu: - unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py + - unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_cache_dit.py - unittest/_torch/visual_gen/test_flux_infer.py diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py index 3ecaabacc85e..02118d078d1f 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py @@ -5,9 +5,11 @@ import json import math +from types import SimpleNamespace from typing import Optional import pytest +import torch import yaml from pydantic import ValidationError @@ -15,7 +17,15 @@ SkipSoftmaxParams, SkipSoftmaxScheduler, ) -from tensorrt_llm.visual_gen.args import AttentionConfig, VisualGenArgs +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import fmha as cute_dsl_fmha +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import ( + CuTeDSLAttention, + _resolve_skip_softmax_threshold_scale_factor, +) +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel +from tensorrt_llm._torch.visual_gen.modules.attention import Attention +from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig, VisualGenArgs from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig pytestmark = pytest.mark.cpu_only @@ -101,6 +111,19 @@ def test_python_api_parses_skip_softmax_config(self): assert sparse_config.disabled_until_timestep == 0.6 assert AttentionConfig(**config.model_dump()).model_dump() == config.model_dump() + def test_cutedsl_api_accepts_skip_softmax_with_quantized_attention(self): + config = AttentionConfig( + backend="CUTEDSL", + quant_attention_config=QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8"), + sparse_attention_config=SkipSoftmaxAttentionConfig( + threshold_scale_factor=5000.0, + disabled_until_timestep=0.6, + ), + ) + + assert isinstance(config.sparse_attention_config, SkipSoftmaxAttentionConfig) + assert config.quant_attention_config is not None + def test_yaml_api_parses_skip_softmax_config(self): # YAML config should deserialize to the same public config object as # the Python API. @@ -388,6 +411,139 @@ def test_graph_phase_tracks_disabled_until_timestep_boundary( ) +class TestVisualGenSkipSoftmaxCuTeDSL: + """CuTeDSL lowering and runtime scheduling use the shared SkipSoftmax params.""" + + def test_attention_lowers_skip_softmax_params_for_cutedsl(self): + sparse_config = SkipSoftmaxAttentionConfig( + threshold_scale_factor=5000.0, + disabled_until_timestep=0.6, + ) + quant_config = QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8") + model_config = DiffusionModelConfig( + component_name="transformer", + pretrained_config=SimpleNamespace(), + attention=AttentionConfig( + backend="CUTEDSL", + quant_attention_config=quant_config, + sparse_attention_config=sparse_config, + ), + skip_create_weights_in_init=True, + ) + + attention = Attention( + hidden_size=16, + num_attention_heads=2, + head_dim=8, + qk_norm=False, + config=model_config, + module_name="blocks.0.attn1", + ) + + assert isinstance(attention.attn, CuTeDSLAttention) + assert isinstance(attention.sparse_params, SkipSoftmaxParams) + assert attention.attn.sparse_params is attention.sparse_params + assert attention.attn.quant_attention_config is quant_config + + @pytest.mark.parametrize( + ("timestep", "expected"), + [ + (1.0, None), + (0.6, None), + (0.59, 5000.0), + (None, 5000.0), + ], + ) + def test_runtime_threshold_tracks_timestep(self, timestep, expected): + sparse_params = SkipSoftmaxAttentionConfig( + threshold_scale_factor=5000.0, + disabled_until_timestep=0.6, + ).to_sparse_params() + + threshold = _resolve_skip_softmax_threshold_scale_factor( + None, + sparse_params, + timestep, + ) + + assert threshold == expected + + def test_cuda_graph_phase_uses_checkpoint_timestep_cutoff(self): + sparse_config = SkipSoftmaxAttentionConfig(threshold_scale_factor=5000.0) + model_config = DiffusionModelConfig( + pretrained_config=SimpleNamespace( + sparse_attention_config=_ckpt_sparse_attention_config(disabled_until_timestep=0.6) + ), + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=sparse_config, + ), + ) + model = BaseDiffusionModel(model_config) + + class _Runner: + def __init__(self): + self.extra_key_fns = {} + + def register_extra_key_fn(self, name, fn): + self.extra_key_fns[name] = fn + + runner = _Runner() + model.register_cuda_graph_extra_key_fns(runner) + + phase_fn = runner.extra_key_fns["skip_softmax_phase"] + assert phase_fn(timestep=0.6) == 0 + assert phase_fn(timestep=0.59) == 1 + + @pytest.mark.skipif( + cute_dsl_fmha._cute_dsl_import_error is not None, + reason="CuTe DSL runtime unavailable", + ) + def test_forward_threads_timestep_and_sparse_params_to_kernel_call(self, monkeypatch): + """The timestep-gating feature hinges on `_fwd`'s `kwargs.get("timestep")` + reaching `cute_dsl_fmha_fwd` unchanged; nothing else in this chain would + raise if that silently dropped to `None` (the scheduler would just apply + the full threshold during early, high-noise steps -- a quality + regression, not an error). This exercises `CuTeDSLAttention.forward` + (the boundary `_attn_impl`/`Attention.forward` call into, and the one + `Attention._attn_impl` also threads `timestep` through unchanged to) + directly, monkeypatching the actual kernel launcher so no CUDA/cutlass + runtime is required. + """ + sparse_params = SkipSoftmaxAttentionConfig( + threshold_scale_factor=5000.0, + disabled_until_timestep=0.6, + ).to_sparse_params() + + captured_kwargs = {} + + def _fake_cute_dsl_fmha_fwd(q, k, v, o, **kwargs): + captured_kwargs.update(kwargs) + o.zero_() + + monkeypatch.setattr( + "tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha.cute_dsl_fmha_fwd", + _fake_cute_dsl_fmha_fwd, + ) + + attn = CuTeDSLAttention( + layer_idx=0, + num_heads=2, + head_dim=8, + sparse_params=sparse_params, + ) + + batch, seq_len, num_heads, head_dim = 1, 4, 2, 8 + q = torch.zeros(batch, seq_len, num_heads, head_dim, dtype=torch.bfloat16) + k = torch.zeros(batch, seq_len, num_heads, head_dim, dtype=torch.bfloat16) + v = torch.zeros(batch, seq_len, num_heads, head_dim, dtype=torch.bfloat16) + + attn.forward(q, k, v, timestep=0.59) + + assert captured_kwargs.get("timestep") == 0.59 + assert captured_kwargs.get("sparse_params") is sparse_params + + class TestVisualGenSkipSoftmaxPipelineConfig: """Pipeline config: multi-transformer checkpoints keep metadata separated.""" diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py index e122a84fc104..0410d4fe418d 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py @@ -328,3 +328,118 @@ def test_cute_dsl_fmha_blockscaled_forward( out_ref = _sdpa_ref(q_bf16, k_bf16, v_bf16, is_causal, sm_scale) assert torch.isfinite(out).all(), "Block-scaled FMHA produced NaN / Inf" torch.testing.assert_close(out, out_ref, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize( + ("qk_sf_vec", "qk_cutlass_dtype_name"), + [ + pytest.param(32, "Float8E4M3FN", id="mxfp8"), + pytest.param(16, "Float4E2M1FN", id="nvfp4"), + ], +) +def test_cute_dsl_fmha_blockscaled_forward_skip_softmax( + qk_sf_vec: int, + qk_cutlass_dtype_name: str, +) -> None: + """SkipSoftmax composed with the block-scaled (MXFP8 / NVFP4) Q@K path. + + The kernel's skip threshold in log2-space is + ``log2(skip_softmax_threshold_scale_factor / seq_len_kv)``, and at this + ``seq_len_kv`` the skip decision is block-granular (only a handful of + K-blocks), so which blocks get skipped is a step function of the + threshold rather than a smooth one: values calibrated empirically (via a + threshold sweep on this exact seed/shape/dtype) below ~650 skip nothing + (dense-equivalent, just quantization noise) and above ~750 over-skip + into a near-degenerate output (cosine similarity vs. the dense SDPA + reference collapsing well under 0.5). 700 sits past the first skip + transition with margin from the next one, so it reliably skips at least + one block without over-skipping, for both MXFP8 and NVFP4. To confirm + skipping actually happened (not just that the kernel didn't crash), the + skip-softmax output is asserted to measurably differ from the dense + (``threshold=0``) output computed from the *same* quantized inputs, in + addition to a cosine-similarity check against the dense SDPA reference. + Unlike ``test_cute_dsl_fmha_blockscaled_forward``, this doesn't use a + tight elementwise tolerance or the 0.99 bound + ``test_attention_trtllm_sage.py`` uses for the TRTLLM backend's + (differently-thresholded) skip_softmax check: SkipSoftmax intentionally + discards information once a block is actually skipped, so 0.95 is used + instead -- still a strong correlation, but consistent with a threshold + picked to guarantee real skipping over reproducing TRTLLM's bound. + """ + _require_supported_gpu_arch() + + device = torch.device("cuda:0") + batch_size, seq_len_q, seq_len_kv = 1, 512, 512 + num_heads, num_heads_kv = 4, 2 + head_dim = 128 # kernel-imposed for block-scaled MXFP8 / NVFP4 + sm_scale = head_dim**-0.5 + skip_softmax_threshold_scale_factor = 700.0 + + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + + q_bf16 = ( + torch.randn(batch_size, seq_len_q, num_heads, head_dim, dtype=torch.bfloat16, device=device) + * 0.5 + ) + k_bf16 = ( + torch.randn( + batch_size, seq_len_kv, num_heads_kv, head_dim, dtype=torch.bfloat16, device=device + ) + * 0.5 + ) + v_bf16 = ( + torch.randn( + batch_size, seq_len_kv, num_heads_kv, head_dim, dtype=torch.bfloat16, device=device + ) + * 0.5 + ) + + q_q, q_sf, scale_q = _quantize_blockscaled_one(q_bf16, qk_sf_vec) + k_q, k_sf, scale_k = _quantize_blockscaled_one(k_bf16, qk_sf_vec) + v_q, scale_v, scale_v_channels = _quantize_fp8_v(v_bf16, per_head_channel=False) + qk_cutlass_dtype = getattr(cutlass, qk_cutlass_dtype_name) + + def _run(threshold: float) -> torch.Tensor: + out = torch.empty( + batch_size, seq_len_q, num_heads, head_dim, dtype=torch.bfloat16, device=device + ) + lse = torch.empty(batch_size, seq_len_q, num_heads, dtype=torch.float32, device=device) + cute_dsl_fmha_fwd( + q_q, + k_q, + v_q, + out, + is_causal=False, + scale_v=scale_v, + scale_v_channels=scale_v_channels, + sm_scale=sm_scale, + lse=lse, + scale_q=scale_q, + scale_k=scale_k, + qk_sf_vec=qk_sf_vec, + q_sf=q_sf, + k_sf=k_sf, + qk_cutlass_dtype=qk_cutlass_dtype, + skip_softmax_threshold_scale_factor=threshold, + ) + torch.cuda.synchronize() + return out + + out_dense = _run(0.0) + out_skip = _run(skip_softmax_threshold_scale_factor) + + assert torch.isfinite(out_dense).all(), "Dense block-scaled FMHA produced NaN / Inf" + assert torch.isfinite(out_skip).all(), "SkipSoftmax block-scaled FMHA produced NaN / Inf" + + max_abs_diff = (out_skip.float() - out_dense.float()).abs().max().item() + assert max_abs_diff > 1e-3, ( + f"SkipSoftmax output matches the dense output (max abs diff {max_abs_diff:.6f}); " + "skip_softmax_threshold_scale_factor did not actually skip any blocks." + ) + + out_ref = _sdpa_ref(q_bf16, k_bf16, v_bf16, is_causal=False, sm_scale=sm_scale) + cos_sim = F.cosine_similarity( + out_skip.reshape(-1).float(), out_ref.reshape(-1).float(), dim=0 + ).item() + assert cos_sim > 0.95, f"Cosine similarity {cos_sim:.6f} below threshold"