From da86d099cd0a5999342c8f65f6f6ff566761f33c Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Sun, 16 Aug 2026 20:43:52 +0000 Subject: [PATCH 1/5] [None][feat] Wire SkipSoftmax sparse attention into the CuTeDSL backend SkipSoftmax was previously TRTLLM-only. This lowers the shared SkipSoftmaxAttentionConfig/SkipSoftmaxParams path into CuTeDSLAttention too, so the CuTeDSL dense FMHA kernel gets the same timestep-aware sparsity scheduling. - AttentionConfig: skip_softmax now composes with backend=CUTEDSL (and quant_attention_config), not just TRTLLM. VSA remains CUTEDSL-only and still mutually exclusive with quant_attention_config. - cute_dsl/fmha.py: cute_dsl_fmha_fwd and CuTeDSLAttention accept sparse_params (SkipSoftmaxParams); the runtime threshold is resolved per-timestep via the shared scheduler instead of a static value. - attention.py: sparse_params is now lowered for backend in (TRTLLM, CUTEDSL), not just TRTLLM. - sparse_attention.py: factor disabled_until_timestep resolution into SkipSoftmaxAttentionConfig.resolve_disabled_until_timestep(), reused by both to_sparse_params() and BaseDiffusionModel's CUDA-graph phase tracking (modeling.py), which previously only read the raw field. - test_skip_softmax.py: new coverage for CuTeDSL param lowering, timestep-tracked threshold resolution, and CUDA-graph phase cutoff under the checkpoint-provided disabled_until_timestep. Signed-off-by: Kanghwan Jang Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/fmha.py | 38 +++++- .../visual_gen/attention_backend/utils.py | 4 +- .../_torch/visual_gen/models/modeling.py | 4 +- .../_torch/visual_gen/modules/attention.py | 8 +- tensorrt_llm/visual_gen/args.py | 21 ++-- tensorrt_llm/visual_gen/sparse_attention.py | 25 ++-- .../sparse_attention/test_skip_softmax.py | 108 +++++++++++++++++- 7 files changed, 182 insertions(+), 26 deletions(-) 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..eadcc837b855 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,20 @@ # ============================================================================ +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: + 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 +254,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 +269,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: @@ -584,6 +608,7 @@ def __init__( dtype: torch.dtype | None = None, quant_attention_config: QuantAttentionConfig | None = None, skip_softmax_threshold_scale: float | None = None, + sparse_params: SkipSoftmaxParams | None = None, **kwargs, ): self.layer_idx = layer_idx @@ -592,7 +617,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 +755,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..9e9669521200 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,20 @@ def to_sparse_params(self, **kwargs): return None return SkipSoftmaxParams(scheduler=scheduler) + def resolve_disabled_until_timestep(self, **kwargs) -> Optional[float]: + """Resolve the user override or checkpoint-provided timestep cutoff.""" + 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(kwargs) + 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/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py index 3ecaabacc85e..3c8e908082f6 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,6 +5,7 @@ import json import math +from types import SimpleNamespace from typing import Optional import pytest @@ -15,7 +16,14 @@ SkipSoftmaxParams, SkipSoftmaxScheduler, ) -from tensorrt_llm.visual_gen.args import AttentionConfig, VisualGenArgs +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 +109,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 +409,91 @@ 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 + + class TestVisualGenSkipSoftmaxPipelineConfig: """Pipeline config: multi-transformer checkpoints keep metadata separated.""" From 1206f15abd08594e7e56383657206b309b47b3a6 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:56:50 +0000 Subject: [PATCH 2/5] [None][test] Register test_skip_softmax.py in the l0_cpu test-db Addresses CodeRabbit review comment on #17781: the file was never registered in test-db (pre-existing gap, not introduced by the prior commit), so its new CuTeDSL SkipSoftmax tests were not selected by CI. Registered as a whole-file entry alongside its visual_gen siblings, matching the existing convention in this list. The other review comment (regenerate the LLM-args golden manifest) is a false positive: `python3 scripts/generate_llm_args_golden_manifest.py --check` exits 0 against this branch. The prior commit changed only AttentionConfig validator *logic* (which backends skip_softmax may compose with), not the schema shape the manifest captures, so regeneration produces no diff. Signed-off-by: Kanghwan Jang Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_cpu.yml | 1 + 1 file changed, 1 insertion(+) 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 From a6ed96fc6c9b0702f237011ccf02dd96faf5c1a3 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Mon, 17 Aug 2026 04:56:29 +0000 Subject: [PATCH 3/5] [None][fix] Address review: explicit kwargs, dead-param comment, timestep test - sparse_attention.py: SkipSoftmaxAttentionConfig.resolve_disabled_until_timestep takes explicit keyword-only checkpoint_config/pretrained_config instead of bare **kwargs, so a mistyped keyword at a call site raises TypeError instead of silently resolving to None (both existing call sites -- to_sparse_params's checkpoint_config= and BaseDiffusionModel's pretrained_config= -- already used the right names, so behavior is unchanged; this only tightens the contract). - cute_dsl/fmha.py: document CuTeDSLAttention's skip_softmax_threshold_scale as a legacy/debug-only knob superseded by sparse_params for every in-tree construction path (create_attention never forwards it). - test_skip_softmax.py: add test_forward_threads_timestep_and_sparse_params_to_kernel_call, which monkeypatches cute_dsl_fmha_fwd and asserts CuTeDSLAttention.forward's timestep kwarg reaches the kernel call unchanged -- the one link in the timestep-gating chain the existing CPU tests didn't cover. Signed-off-by: Kanghwan Jang Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/fmha.py | 5 +++ tensorrt_llm/visual_gen/sparse_attention.py | 19 ++++++-- .../sparse_attention/test_skip_softmax.py | 45 +++++++++++++++++++ 3 files changed, 66 insertions(+), 3 deletions(-) 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 eadcc837b855..510723fb456d 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 @@ -607,6 +607,11 @@ 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, diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 9e9669521200..d261c770ca9e 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -99,8 +99,19 @@ def to_sparse_params(self, **kwargs): return None return SkipSoftmaxParams(scheduler=scheduler) - def resolve_disabled_until_timestep(self, **kwargs) -> Optional[float]: - """Resolve the user override or checkpoint-provided timestep cutoff.""" + 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 @@ -108,7 +119,9 @@ def resolve_disabled_until_timestep(self, **kwargs) -> Optional[float]: skip_softmax_disabled_until_timestep_from_ckpt_sparse_attention_config, ) - ckpt_sparse_attention_config = self._ckpt_sparse_attention_config_from_kwargs(kwargs) + 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 ) 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 3c8e908082f6..f05fa1191223 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 @@ -9,6 +9,7 @@ from typing import Optional import pytest +import torch import yaml from pydantic import ValidationError @@ -493,6 +494,50 @@ def register_extra_key_fn(self, name, fn): assert phase_fn(timestep=0.6) == 0 assert phase_fn(timestep=0.59) == 1 + 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.""" From 2b5f9f7790b26a03a9cc63f06b84ad28d9cd72a6 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Mon, 17 Aug 2026 07:22:05 +0000 Subject: [PATCH 4/5] [None][doc] Address review: document CUTEDSL SkipSoftmax support, add block-scaled numerical coverage - sparse-attention.md previously said SkipSoftmax only works with TRTLLM; update it to cover the CUTEDSL backend added in this PR, including that it composes with quant_attention_config block-scaled Q/K recipes. - Add test_cute_dsl_fmha_blockscaled_forward_skip_softmax: numerical coverage for SkipSoftmax combined with the CUTEDSL MXFP8/NVFP4 block-scaled Q/K path, checked against dense SDPA via cosine similarity (mirrors test_attention_trtllm_sage.py's skip_softmax check). Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../visual-gen/features/sparse-attention.md | 24 ++++- .../visual_gen/test_attention_cute_dsl.py | 94 +++++++++++++++++++ 2 files changed, 117 insertions(+), 1 deletion(-) 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/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py index e122a84fc104..c7ad02114912 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,97 @@ 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"), + ], +) +@pytest.mark.parametrize( + "skip_softmax_threshold_scale_factor", + [ + pytest.param(0.0, id="dense"), + pytest.param(0.3, id="skip_softmax"), + ], +) +def test_cute_dsl_fmha_blockscaled_forward_skip_softmax( + qk_sf_vec: int, + qk_cutlass_dtype_name: str, + skip_softmax_threshold_scale_factor: float, +) -> None: + """SkipSoftmax composed with the block-scaled (MXFP8 / NVFP4) Q@K path. + + SkipSoftmax approximates the dense softmax, so unlike + ``test_cute_dsl_fmha_blockscaled_forward`` this asserts cosine similarity + against the dense SDPA reference rather than a tight elementwise + tolerance, mirroring ``test_attention_trtllm_sage.py``'s skip_softmax + check for the TRTLLM backend. + """ + _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 + + 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) + + 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=skip_softmax_threshold_scale_factor, + ) + torch.cuda.synchronize() + + out_ref = _sdpa_ref(q_bf16, k_bf16, v_bf16, is_causal=False, sm_scale=sm_scale) + assert torch.isfinite(out).all(), "SkipSoftmax block-scaled FMHA produced NaN / Inf" + + cos_sim = F.cosine_similarity( + out.reshape(-1).float(), out_ref.reshape(-1).float(), dim=0 + ).item() + assert cos_sim > 0.990, f"Cosine similarity {cos_sim:.6f} below threshold" From 8cac68a020ed5a404d4750a13ab8d6f4a92cec89 Mon Sep 17 00:00:00 2001 From: Kanghwan Jang <861393+karljang@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:31:48 +0000 Subject: [PATCH 5/5] [None][fix] Address review: CPU-lane skip guard, calibrated skip_softmax test, fail-open warning - test_skip_softmax.py: test_forward_threads_timestep_and_sparse_params_to_kernel_call called CuTeDSLAttention.forward(), which raises ImportError in _prepare_inputs before ever reaching the monkeypatched kernel launcher. This file is registered in l0_cpu.yml, so it errored (not skipped) on a CPU-only image without the CuTe DSL runtime. Add a skipif guard. - test_attention_cute_dsl.py: the new block-scaled skip_softmax test used skip_softmax_threshold_scale_factor=0.3, which at seq_len_kv=512 is far below the kernel skip threshold and never actually skips a block, so the test exercised the dense path under a different name. Calibrated (via a threshold sweep on this seed/shape/dtype) to 700.0, past the first block-skip transition with margin, and added an assertion that the skip-softmax output measurably diverges from a same-inputs dense run to confirm skipping actually happened. Loosened the SDPA-reference cosine bound from 0.99 to 0.95 to match: SkipSoftmax discards real information once a block is skipped, so it cannot both truly skip and hit the same bound as a threshold picked specifically not to skip. Verified on a B200 (sm_100a). - fmha.py: _resolve_skip_softmax_threshold_scale_factor fails open when sparse_params has disabled_until_timestep configured but timestep is None (runs unthrottled instead of respecting the cutoff). Add a logger.warning_once so a caller that forgets to thread timestep through gets a log line instead of a silent quality regression. Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com> --- .../attention_backend/cute_dsl/fmha.py | 14 +++ .../sparse_attention/test_skip_softmax.py | 5 + .../visual_gen/test_attention_cute_dsl.py | 99 +++++++++++-------- 3 files changed, 79 insertions(+), 39 deletions(-) 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 510723fb456d..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 @@ -72,6 +72,20 @@ def _resolve_skip_softmax_threshold_scale_factor( ) -> 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: 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 f05fa1191223..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 @@ -17,6 +17,7 @@ SkipSoftmaxParams, SkipSoftmaxScheduler, ) +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, @@ -494,6 +495,10 @@ def register_extra_key_fn(self, name, fn): 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 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 c7ad02114912..0410d4fe418d 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl.py @@ -337,25 +337,34 @@ def test_cute_dsl_fmha_blockscaled_forward( pytest.param(16, "Float4E2M1FN", id="nvfp4"), ], ) -@pytest.mark.parametrize( - "skip_softmax_threshold_scale_factor", - [ - pytest.param(0.0, id="dense"), - pytest.param(0.3, id="skip_softmax"), - ], -) def test_cute_dsl_fmha_blockscaled_forward_skip_softmax( qk_sf_vec: int, qk_cutlass_dtype_name: str, - skip_softmax_threshold_scale_factor: float, ) -> None: """SkipSoftmax composed with the block-scaled (MXFP8 / NVFP4) Q@K path. - SkipSoftmax approximates the dense softmax, so unlike - ``test_cute_dsl_fmha_blockscaled_forward`` this asserts cosine similarity - against the dense SDPA reference rather than a tight elementwise - tolerance, mirroring ``test_attention_trtllm_sage.py``'s skip_softmax - check for the TRTLLM backend. + 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() @@ -364,6 +373,7 @@ def test_cute_dsl_fmha_blockscaled_forward_skip_softmax( 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) @@ -390,35 +400,46 @@ def test_cute_dsl_fmha_blockscaled_forward_skip_softmax( 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) - 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) + 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 - 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=skip_softmax_threshold_scale_factor, + 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." ) - torch.cuda.synchronize() out_ref = _sdpa_ref(q_bf16, k_bf16, v_bf16, is_causal=False, sm_scale=sm_scale) - assert torch.isfinite(out).all(), "SkipSoftmax block-scaled FMHA produced NaN / Inf" - cos_sim = F.cosine_similarity( - out.reshape(-1).float(), out_ref.reshape(-1).float(), dim=0 + out_skip.reshape(-1).float(), out_ref.reshape(-1).float(), dim=0 ).item() - assert cos_sim > 0.990, f"Cosine similarity {cos_sim:.6f} below threshold" + assert cos_sim > 0.95, f"Cosine similarity {cos_sim:.6f} below threshold"