flashdreams.accelerated API v0 - #486
Conversation
Greptile SummaryThe PR introduces reusable accelerated quantization and multi-head-attention primitives, integrates selectable optimized attention backends into OmniDreams, and adds benchmark tooling.
Confidence Score: 5/5The PR appears safe to merge. The previously reported stale derived-weight issue is resolved because optimized attention now registers a post-load hook that rebuilds fused and quantized projections from the loaded canonical parameters, and no blocking failure remains. Important Files Changed
Sequence DiagramsequenceDiagram
participant Network as OmniDreams Network
participant Attention as Optimized Attention
participant Loader as load_state_dict
participant Derived as Derived Projections
participant Forward as Inference
Network->>Attention: Construct canonical Q/K/V projections
Attention->>Derived: Build initial non-persistent projections
Loader->>Attention: Load checkpoint parameters
Attention->>Derived: Post-load hook refreshes fused/quantized weights
Network->>Attention: Move or cast module
Attention->>Derived: _apply refreshes weights on final device/dtype
Forward->>Derived: Execute accelerated projections
Reviews (7): Last reviewed commit: "Add benchmark and plot scripts for flash..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
3077d4d to
b2685be
Compare
ArielG-NV
left a comment
ArielG-NV
left a comment
There was a problem hiding this comment.
Question on inclusion of file
fc4b03f to
f63a088
Compare
4ad51b5 to
97dfa44
Compare
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); |
There was a problem hiding this comment.
I guess we should restrict the license header stuff to just the first 2 lines, no need for the rest.
|
|
||
|
|
||
| class _TorchMultiHeadAttention(TorchMultiHeadAttention): | ||
| """Canonical Torch attention implementation used by benchmarks.""" |
There was a problem hiding this comment.
What does canonical mean in this context? Is there a better way to describe this?
Same for all the other usages.
| ) | ||
|
|
||
|
|
||
| class _OptimizedHultiHeadAttention(OptimizedHultiHeadAttention): |
There was a problem hiding this comment.
Should this be OptimizedMultiHeadAttention?
| # Static K/V projection is part of this end-to-end cross-attention | ||
| # measurement because fusion changes that stage, not query-only forward. | ||
| def synchronized_cross_forward() -> Tensor: | ||
| cache = attention.compute_kv(context, context_rope) |
There was a problem hiding this comment.
It looks like this calls compute_kv() on every measured round, so the result represents cache setup plus one query rather than steady-state cross-attention. Both numbers seem useful, but combining them may obscure which optimization is helping. Do they need to be reported separately?
| use_tma=use_tma, | ||
| quantization=QuantizationOption( | ||
| projection=torch.float8_e4m3fn, | ||
| quantized_sdpa=True, |
There was a problem hiding this comment.
It looks like every quantized-SDPA configuration also enables FP8 projections, so the results cannot distinguish projection-quantization performance from FP8 SDPA/KV-cache performance. Is that what is happening here?
Does there need to also be a quantized_sdpa=True, projection=None case? Or a projection-only case for FA2?
| OptimizedImplConfig( | ||
| sdpa_backend=sdpa_backend, | ||
| qkv_fusion_option=qkv_fusion_option, | ||
| use_tma=use_tma, |
There was a problem hiding this comment.
Could use_tma be varied only when sdpa_backend is FA2 or does that not matter?
This cross-product also creates cuDNN tma and no-tma rows even though TMA does not apply to that backend, is this a problem? Does it lead to duplicated benchmark work? The quantized generator at 91 has the same pattern.
| cache.before_update(chunk_idx) | ||
| attention(inputs[chunk_idx], cache, rope_freqs[chunk_idx]) | ||
| cache.after_update(chunk_idx) | ||
| cache.before_update(_WINDOW_CHUNKS) |
There was a problem hiding this comment.
Would it be useful to add a separate lifecycle benchmark that includes before_update() and after_update()? The current result intentionally measures forward after the full-cache roll, but it excludes cache movement and I am wondering if it may miss part of the benefit of lower-precision KV-cache storage.
| bias_label, | ||
| ) | ||
| ) | ||
| benchmark.extra_info.update( |
There was a problem hiding this comment.
Could we include query_dim, n_heads, head_dim, device capability, driver version, Triton version, and effective TMA selection in extra_info? The JSON records the requested policy and chunk/window sizes, but does it contain enough information to reconstruct the full geometry or confirm the path that actually ran?
| return f"optimized-{backend}-{fusion}-{tma}{projection}{quantized_sdpa}" | ||
|
|
||
|
|
||
| _SHARED_CONFIGS = tuple( |
There was a problem hiding this comment.
Would it make sense to run a smaller representative policy matrix by default and make the exhaustive sweep opt-in? The current combinations expand to 1,008 cases, each with five warmups and 50 measured rounds, and some implementation combinations are redundant. A smaller default might be easier to run and interpret.
| torch.float8_e4m3fn if quantized_dtype is torch.float8_e5m2 else quantized_dtype | ||
| ) | ||
| scaled_output_dtype = ( | ||
| torch.bfloat16 |
There was a problem hiding this comment.
Maybe we should label this case with its effective output precision rather than grouping it as FP32.
With slice scaling and an FP32 requested output, _scaled_mm produces BF16 and the end-to-end path only casts that result to FP32 afterward. It's not really precision-equivalent to the FP32 baseline and could make the comparison misleading.
|
|
||
| assert dequantized.shape == _SHAPE | ||
| assert dequantized.dtype is torch.float16 | ||
| assert torch.isfinite(dequantized).all() |
There was a problem hiding this comment.
I am not sure that shape, dtype, and finiteness would catch an incorrect but fast kernel.
Do we need to compare the dequantized result against original, and similarly compare the quantization outputs against the Torch reference outside the timed region? It may also be useful to save the numerical error in benchmark.extra_info but I'll leave that up to you.
| ) | ||
| amax_o_desc.set_output(False).set_dim(list(amax_o.shape)).set_stride( | ||
| list(amax_o.stride()) | ||
| ) |
There was a problem hiding this comment.
Should amax_s_desc and amax_o_desc be marked with set_output(True) here? I think cuDNN documents these as required FP8 SDPA outputs. With False, we still allocate and bind buffers for them but discard the calibration information needed to select meaningful scales.
| list(amax_o.stride()) | ||
| ) | ||
| graph.build([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) | ||
| scale = torch.ones_like(amax_s) |
There was a problem hiding this comment.
Is it correct to set these to 1? I am wondering if this leaves Q, K, V, the softmax probabilities, and the output uncalibrated in E4M3. I think cuDNN recommends deriving FP8 scaling from observed AMax values. Do we need to provide calibrated scales or long-context quality evidence before using this as an optimization?
| }, | ||
| workspace, | ||
| ) | ||
| return output |
There was a problem hiding this comment.
Could each call return independent output storage, or could this aliasing contract be made explicit and enforced? output is allocated once inside the cached closure, so repeated calls with the same cache key return the same tensor and overwrite all previously returned results. The current MHA caller immediately converts it to BF16/FP16, but native_cudnn_fp8_sdpa() is exported and direct callers could retain a result that is silently mutated by the next invocation. I am wondering if this will cause issues.
| bhld_strides = (x.stride(0), x.stride(2), x.stride(1), x.stride(3)) | ||
| return bhld_strides[-1] == 1 and all( | ||
| stride > 0 and stride * element_size % 16 == 0 for stride in bhld_strides[:-1] | ||
| ) |
There was a problem hiding this comment.
I think TMA tensor descriptors require the base pointer itself to be 16-byte aligned, but this predicate only checks the strides. Do we also need to require x.data_ptr() % 16 == 0, and add a misaligned-view test?
| raise TypeError(f"use_tma must be a bool; got {self.use_tma!r}") | ||
|
|
||
|
|
||
| class OptimizedHultiHeadAttention(MultiHeadAttention[BlockKVCache]): |
There was a problem hiding this comment.
I'm guessing this should be OptimizedMultiHeadAttention.
| .contiguous() | ||
| ) | ||
| fused_bias = None | ||
| if self.query_projection.bias is not None: |
There was a problem hiding this comment.
I think whether the fused projection has a bias is determined solely from the query bias, is that right? If Q has no bias but K or V does, those biases would be silently dropped. If Q has a bias but K/V does not, this hits an assertion.
I think the fused-KV path has the same issue at line 382. Do we need to preserve mixed biases by filling absent components with zeros or maybe explicitly validate that the relevant bias configurations match?
|
|
||
| write_end = kv_cache.write_end | ||
| write_start = write_end - kv_cache.chunk_size | ||
| return rope_freqs[write_start:write_end], rope_freqs[: kv_cache.size] |
There was a problem hiding this comment.
For a static cross-attention cache, chunk_size and write_end both correspond to the context length S, so this returns S query frequencies. Would _apply_rope() then reject a query whose length L differs from S, even though the public cross-attention contract allows different query and context lengths?
The test avoids this by setting L == S. Should we support independent query/key frequency slices, or explicitly reject and document after-cache RoPE for static cross-attention?
|
|
||
| write_end = kv_cache.write_end | ||
| write_start = write_end - kv_cache.chunk_size | ||
| return rope_freqs[write_start:write_end], rope_freqs[: kv_cache.size] |
There was a problem hiding this comment.
For a static cache created by compute_kv(), chunk_size equals the context length S, so this produces S query frequencies. _apply_rope() then rejects a query with length L != S, but I think the public cross-attention contract permits different query and context lengths. Should we support separate query/key frequency slices or explicitly reject and document after-cache RoPE for static cross-attention?
| self.dtype = dtype | ||
| self.register_buffer( | ||
| "weight_scale", weight_scale.contiguous(), persistent=False | ||
| ) |
There was a problem hiding this comment.
For FP8, both self.weight and weight_scale are floating-point buffers, so something like layer.to(dtype=torch.bfloat16) converts the FP8 weight and FP32 scale to BF16 while self.dtype still advertises FP8.
The optimized MHA rebuilds these derived buffers in its own _apply(), does this class need equivalent protection? Do we need to preserve or rebuild these buffers, or explicitly reject dtype casts, and add a direct regression test?
| return output | ||
|
|
||
| current_dtype = scales[0].dtype | ||
| for index, scale in enumerate(scales): |
There was a problem hiding this comment.
Could multiple scales be applied in one kernel, or could we benchmark whether that is worthwhile?
This loop allocates an intermediate tensor and launches a full-tensor multiply kernel for every scale. The INT8 linear path always supplies activation and weight scales, so it currently performs two dequantization passes after the GEMM. It would be useful to know whether fusing those scales materially improves the end-to-end result.
| use_tma=True, | ||
| quantization=QuantizationOption(projection=torch.float8_e4m3fn), | ||
| ), | ||
| ), |
There was a problem hiding this comment.
Should this case also set minimum_compute_capability=(9, 0)? OptimizedMultiHeadAttention explicitly rejects CUDA devices below compute capability 9.0, but this case currently passes skip_unsupported_device() on SM8.x. The benchmark would then fail during execution instead of skipping cleanly, contrary to the README’s stated behavior. The FA2 optimized case below already declares this requirement.
|
/ok to test 97dfa44 |
| or config.quantization.projection is not None | ||
| ) | ||
| for config in (self_config, cross_config) | ||
| ) |
There was a problem hiding this comment.
Should every non-None optimized configuration require compute capability 9.0 here? OptimizedMultiHeadAttention._validate_device() rejects devices below 9.0 regardless of the selected SDPA backend or projection dtype. As written, a plain optimized cuDNN configuration without FP8 projection gets no minimum capability, passes skip_unsupported_device() on SM8.x, and then fails during execution instead of skipping cleanly.
| _block_case(self_config, cross_config) | ||
| for self_config in _MODULE_SELF_ATTENTION_CONFIGS | ||
| for cross_config in _MODULE_CROSS_ATTENTION_CONFIGS | ||
| ) |
There was a problem hiding this comment.
Do we need all of this for the full-block benchmark? There are 37 self-attention and 25 cross-attention choices, so this creates 925 block benchmark cases, each with five warmups and 50 measured rounds, before the 62 isolated attention cases run. Some configurations are also effectively redundant, such as varying TMA under cuDNN. Could the full-block benchmark use representative self/cross pairs while keeping the exhaustive policy sweep isolated or opt-in?
| patch_volume = network_config.patch_temporal * network_config.patch_spatial**2 | ||
| tokens_per_frame = patch_h * patch_w | ||
| chunk_tokens = patch_t * tokens_per_frame | ||
| _WINDOW_SIZE_T * tokens_per_frame |
There was a problem hiding this comment.
Not sure how this line got in there.
|
|
||
| def teardown_generate() -> None: | ||
| nonlocal next_chunk_index | ||
| pipeline.finalize( |
There was a problem hiding this comment.
Is finalize() intended to be included in the reported full-pipeline latency? Because it is invoked through the benchmark teardown, its execution is excluded from the measured target. Finalization performs additional transformer work, so this may understate recurring per-chunk latency. Could we include it in the timed operation, or report generation and finalization separately? It will affect the total fps after all.
| self_attn_optimized_impl_config=OptimizedImplConfig( | ||
| qkv_fusion_option=QKVFusionOption.FULL, | ||
| sdpa_backend=SDPABackend.FA2, | ||
| ), |
There was a problem hiding this comment.
Is the quantization policy missing from this preset? The description says it uses FP8 projections, but both OptimizedImplConfig instances retain the default quantization, where projection=None and quantized_sdpa=False. As written, the projections remain at the model’s native precision rather than FP8. Should we configure the intended quantization explicitly for both attention branches, or update the description if native precision is intentional?
| pipeline=derive_config( | ||
| SV_2STEPS_CHUNK2_LOC6_LIGHTVAE_LIGHTTAE_PERF, | ||
| diffusion_model=dict(transformer=dict(skip_finalize_kv_cache=True)), | ||
| ), |
There was a problem hiding this comment.
Is skipping the final KV-cache refresh intentional for all public accelerated runners? Finalization is not merely cleanup: it re-noises the generated latent and runs finalize_kv_cache() to prepare the context for the next autoregressive chunk. This file otherwise treats skipping that operation as an experimental ablation, so enabling it here, and repeatedly in the runners below, changes model behavior and may affect rollout quality. Should we retain normal finalization, or document and justify this tradeoff with quality results?
| assert network_config.cross_attention_backend is AttentionBackend.OPTIMIZED | ||
| assert ( | ||
| network_config.self_attn_optimized_impl_config.sdpa_backend is SDPABackend.FA2 | ||
| ) |
There was a problem hiding this comment.
Should this test also assert the preset’s advertised quantization policy? It currently verifies only the FA2 backend and fusion options, so it passes while both OptimizedImplConfig instances default to projection=None and quantized_sdpa=False, despite the preset being documented as using FP8 projections. Should we also assert the intended QuantizationOption for both branches, or compare these policies with the selected RTX PRO 6000 benchmark case?
| matrix, | ||
| aspect="auto", | ||
| cmap="RdYlGn_r", | ||
| norm=CenteredNorm(vcenter=1.0), |
There was a problem hiding this comment.
Should the runtime ratios be transformed to log space before applying the diverging color scale?
CenteredNorm(vcenter=1.0) treats ratios additively, so multiplicatively equivalent changes such as 0.5x and 2x receive different color intensities. A very slow outlier can also determine the range and wash out all useful speedups. Using log2(runtime / reference) centered at zero would represent equivalent speedups and slowdowns symmetrically.
| if reference is None: | ||
| return f"{value:.2f} ms\nreference unavailable" | ||
| if value < reference: | ||
| return f"{value:.2f} ms\n{(1 - value / reference) * 100:.0f}% faster" |
There was a problem hiding this comment.
This formula calculates the percentage reduction in latency, rather than percentage "faster." For example, half the runtime is displayed as 50% faster, although it represents a 2X speedup. Could these labels say % lower/higher latency, or report multiplicative speedup instead, so the benchmark result is unambiguous?
| ) | ||
| axes.set_xlabel("GEMM configuration") | ||
| axes.set_ylabel("Source dtype") | ||
| axes.set_title(f"Quantized GEMM {_SCOPE_LABELS[timing_scope]}\n{subtitle}") |
There was a problem hiding this comment.
The gemm-only plot ranks these configurations as though they produce equivalent results, but INT8 measures raw torch._int_mm output as INT32 without dequantization, while FP8 still applies its scales and full precision returns the source dtype. Should we make these different output semantics explicit, or avoid presenting a single “fastest quantized config” for this scope? Otherwise the comparison could be interpreted as an apples-to-apples speedup.
| matrix, | ||
| aspect="auto", | ||
| cmap="RdYlGn_r", | ||
| norm=CenteredNorm(vcenter=1.0), |
There was a problem hiding this comment.
Same as the comment I left in the other file.
| if reference is None: | ||
| return f"{value:.2f} ms\nreference unavailable" | ||
| if value < reference: | ||
| return f"{value:.2f} ms\n{(1 - value / reference) * 100:.0f}% faster" |
There was a problem hiding this comment.
Same faster comment as the one I left for the other file.
| if reference is None: | ||
| return f"{value:.2f} ms\nreference unavailable" | ||
| if value < reference: | ||
| return f"{value:.2f} ms\n{(1 - value / reference) * 100:.0f}% faster" |
There was a problem hiding this comment.
Same comment as the faster one from the other files.
| matrix, | ||
| aspect="auto", | ||
| cmap="RdYlGn_r", | ||
| norm=CenteredNorm(vcenter=1.0), |
There was a problem hiding this comment.
Same as the comment from the other files.
| np.asarray(matrix), | ||
| aspect="auto", | ||
| cmap="RdYlGn_r", | ||
| norm=CenteredNorm(vcenter=1.0), |
| r"(?:optimized|triton)-(cudnn|fa2)-(bf16)-self-(none|full|fuse-kv)-" | ||
| r"(no-tma|tma)-cross-(none|fuse-kv)-(no-tma|tma)", | ||
| implementation, | ||
| ) |
There was a problem hiding this comment.
Do these label parsers accept the identifiers produced by the current benchmarks?
Module configurations ending in -quantized-sdpa miss this pattern, while the selected end-to-end cases use optimized-cudnn-fp8-... and optimized-fa2-quantized-sdpa-..., neither of which matches the patterns below.
Do these configurations silently fall back to raw parameter text? Should the parsing be updated so all current benchmark cases receive accurate structured labels?
There was a problem hiding this comment.
I guess this also applies to line 196.
|
Could we add a document summarizing the performance findings from this work so far? The benchmark scripts explain how to generate measurements, but it would be useful to have a durable summary of what was learned: the hardware and workloads tested, baseline and optimized results, which optimizations were beneficial at each layer, and any important tradeoffs or limitations. It does not need to be exhaustive, but it would provide context for the configurations selected in the OmniDreams integration and give future changes a reference point. Also, the PR will need to be rebased, and the CI failures will need to be fixed. Other than that, this is looking pretty good. |
3 participants
Add FlashDreams accelerated primitives and integrate them with Omnidreams
Summary
Introduce reusable quantization and optimized multi-head attention primitives under
flashdreams.accelerated, integrate them with Omnidreams, and add reproducible performance benchmarks and plotting tools.1.
flashdreams.accelerated1.1 Quantization toolkit
state_dict, preserving source checkpoint compatibility.1.2 Optimized MHA
BlockKVCacheand preserve caller-managed cache lifecycles.2. Integration: Omnidreams
omnidreams-triton-fa2omnidreams-cuda-cudnnomnidreams-cuda-spargeomnidreams-cuda-sage3fp812.0a3. Benchmark and plot scripts
pytest-benchmarkconfiguration and dependencies.Testing
The change includes CPU, GPU, and manual benchmark coverage. GPU validation requires a supported NVIDIA GPU; the optimized MHA path requires compute capability 9.0 or newer.
Generated benchmark JSON and figures are intentionally not committed.