[perf] FastH3 two-Spark: gather packed FP4 inputs before head projections

Opt-in h3_vsa_fp4.sp_input_gather (default false; requires sp_projection_overlap) gathers the
already-packed FP4 inputs and rotary tables, then computes each rank's assigned Q/K/V heads locally;
the gate projection overlaps the smaller transfer. Supported on the measured GB10/cuDNN shape only;
other shapes keep the existing exchange. Enabled in the Spark-pair example.
T2AV on 2 Sparks: 19.35 -> 18.76 s/clip (merge gate check, same SP reduction order rules).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

.agents/lessons/2026-10-23_asynchronous-gemm-parity.md

@@ -0,0 +1,27 @@
+---
+date: 2026-10-23
+category: evaluation
+severity: important
+---
+
+# Synchronized GEMM parity does not establish request determinism
+
+A fixed cuDNN FP4 GEMM plan matched isolated captured outputs, a complete step,
+and three scored requests on one GB10. Repeated requests on another GB10 later
+produced different latents. Restoring the prior row path reproduced the known
+latents on three fresh requests. The replacement was withdrawn.
+
+Comparing every GEMM output with `torch.equal` made all 400 calls and the complete
+request match. That comparison synchronizes execution; it can conceal an
+ordering or lifetime problem. Repeated component tests also passed. The exact
+cause remains unresolved, so these observations do not establish a vendor bug.
+
+A separate larger workspace happened to make two requests match, but cuDNN
+reported zero required workspace for the selected plan. Treat that as an
+allocation/timing perturbation, not evidence of a workspace fix.
+
+After component parity, repeat complete requests without additional per-kernel
+synchronization, including another device when the change is intended to be
+portable. Compare latent and decoded-output hashes. Preserve failing evidence
+and restore the stable path when the proposed gain is small and the cause is
+unresolved; passing a finite scored seed set does not establish determinism.

.agents/lessons/2026-10-23_compiled-forward-deepcopy.md

@@ -0,0 +1,38 @@
+---
+date: 2026-10-23
+category: evaluation
+severity: important
+---
+
+# Copying an instance-compiled module can benchmark the original module
+
+## What happened
+
+A decoder quantization experiment replaced every intended linear in a copied
+module, yet measured exactly the original output and almost the original time.
+The replacement layers were present in `named_modules()`, which was insufficient
+evidence that the new computation actually ran.
+
+## Root cause
+
+`ComposedPipelineBase._compile_with_conditions` assigns
+`submodule.forward = torch.compile(submodule.forward)`. The returned function
+closes over a bound method of the original module. Python's `copy.deepcopy`
+copies module state but retains function objects and their closures. The copied
+instance therefore still has a compiled `forward` that calls the original
+instance. Unwrapping `_orig_mod` alone does not handle this form of compilation.
+
+## Fix
+
+Prefer constructing a fresh module and loading the intended state. If copying
+this known instance-compiled form for an experiment, remove the copied instance's
+`forward` override before modifying children, then compile its newly bound class
+method. Do not remove arbitrary user-defined forward overrides blindly.
+
+## Prevention
+
+Inspect the callable as well as the module tree. Verify a replacement executes
+with a temporary hook or profiler trace, compare its output against the intended
+numerical oracle, and run final timing in a fresh process. An unchanged output
+from a supposedly lossy quantization is a reason to investigate dispatch before
+accepting the measurement.

.agents/lessons/2026-10-23_triton-affine-gather.md

@@ -0,0 +1,26 @@
+---
+date: 2026-10-23
+category: performance
+severity: important
+---
+
+# Wide Triton gathers can turn an affine slice into expensive selection code
+
+A decoder RMSNorm prototype kept a 2,048-element row in registers, then used
+`tl.gather` with indices `512 * chunk + 4 * lane + component` to extract the
+128 lanes needed by its original reduction order. Although every index was
+an affine expression, this lowered poorly on the Spark's Triton/SM121 toolchain.
+The eight-row variant emitted PTX with 168,062 virtual 32-bit registers and
+spent minutes in ptxas. Earlier variants were numerically exact but took
+392–1,475 ms for work the existing kernel completed in about 2.4 ms.
+
+For this layout, reshaping the row into four 512-element chunks and extracting
+them with `tl.trans` and `tl.split` avoided the selection expansion. The same
+ordered FMAs then compiled promptly and matched the original outputs. Keeping
+the row in registers still gave no meaningful speedup, so the rewrite was dropped.
+
+Inspect generated IR and benchmark one small launch before compiling a whole
+parameter sweep of wide gathers. Prefer static reshape/split operations for
+known affine partitions; reserve gathers for indexing that actually needs them.
+After changing an in-place residual path, also zero invalid padded rows before
+quantization: the old two-pass kernel obtained those zeros from its masked reload.

docs/design/inference_schema_parity_inventory.yaml

@@ -89,6 +89,7 @@ surfaces:
       taeh3_checkpoint: "Optional local TAEH3 safetensors path; model-specific optimization not yet represented in the typed public schema."
       taeh3_chunk_size: "TAEH3 temporal chunk length; model-specific optimization not yet represented in the typed public schema."
       h3_vae_nvfp4_decoder: "Exported NVFP4 MiniMax-H3 video decoder file that replaces the VAE decoder at load; model-specific optimization not yet represented in the typed public schema."
+      h3_vae_nvfp4_quantize_dense: "Quantize the loaded MiniMax-H3 decoder to NVFP4 in memory without changing its depth; carried through generator.pipeline.experimental."
       vae_parallel_encode: "MiniMax-H3 sequence-parallel reference VAE encode opt-in; model-specific optimization not yet represented in the typed public schema."
       vae_parallel_decode_strategy: "Chunk-transport collective for vae_parallel_decode; model-specific optimization not yet represented in the typed public schema."
       attention_backend: "Process-wide default attention-backend request applied per component at load time; kernel-selection knob not yet represented in the typed public schema."

docs/getting_started/installation/spark_pair.md

@@ -149,6 +149,39 @@ fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair.yam
 
 Stop the cluster when you are done: `ray stop` on both nodes.
 
+For eager FastH3 FP4 attention, `generator.pipeline.experimental.h3_vsa_fp4.sp_head_chunks`
+overlaps head-group communication with independent attention work. It defaults to `1`.
+The head count must divide evenly by `sp_size * sp_head_chunks`; on a two-Spark pair FastH3's
+56 heads allow 2, 4, 7, 14 or 28 chunks. The pair example config uses 7 (four heads per group). The grouped path packs directly into each exchange buffer and writes
+received heads into their original output positions. Changing the attention batch shape can change
+floating-point reductions, so compare generated quality as well as wall time for your checkpoint.
+
+The optional `h3_vsa_fp4.sp_projection_overlap: true` starts each Q/K/V exchange as soon as its
+projection is packed, overlapping traffic with the remaining projection GEMMs. It defaults to
+`false` and requires two ranks, head dimension 128, eager BF16 inference and compatible NVFP4
+QKV/gate projections with shared input quantization. Other configurations use the ordinary
+head-group exchange. Packing arithmetic and FP8 transport precision are unchanged.
+
+The optional `h3_vsa_fp4.sp_input_gather: true` gathers the existing packed FP4 inputs and
+rotary tables before computing each rank's assigned Q/K/V heads. The gate projection overlaps
+that smaller input transfer. It defaults to `false`, requires `sp_projection_overlap: true`,
+and currently supports the measured GB10/cuDNN shape: 7430 local rows, 5376 input channels,
+56 heads and bias-free NVFP4 QKV projections. Other shapes use the existing exchange path.
+All projections, normalization, rotary operations and FP8 packet rounding remain present;
+no steps or tokens are omitted. Validate generated output when changing this GEMM shape.
+
+The optional `h3_vsa_fp4.sp_pooled_key_mean: true` reuses tile statistics for key smoothing
+on the grouped tile-64 path with FP8 prefix attention and `smooth_k: true`. It defaults to
+`false`. All key rows still contribute, including zero padding in the normalization length,
+but the changed reduction order can alter generated samples. Validate quality for the checkpoint
+and workload before enabling it.
+
+The optional `h3_vsa_fp4.sp_direct_keys: true` avoids the intermediate BF16 key tensor on
+that pooled-mean path. It defaults to `false`. Key preparation reads the current exchange's
+FP8 packets and applies the same FP32 multiply, BF16 rounding and smoothing as the ordinary
+receive path. All keys and their tile statistics still contribute; other attention configurations
+retain ordinary key stores.
+
 ## Released V2, FastH3-QAD-v2 and Trim NVFP4 stacks
 
 The public eight-forward V2 and Trim stacks include the NVFP4 encoder and

docs/getting_started/installation/spark_performance.md

@@ -390,10 +390,9 @@ off on other GPUs unless a matched benchmark validates the split there.
 `FASTVIDEO_H3_VAE_NVFP4_DECODER=/path/to/decoder.pt` swaps in an NVFP4
 decoder distilled with quantization-aware training and exported by
 `scripts/distill/minimax_h3_nvfp4_decoder/export_deploy.py`. At load the
-INT8 overlay is skipped, the VAE is depth-cut to the decoder's block count
-(the 26-layer light VAE hosts an 8-block decoder; a 36-block decoder needs
-the full `MiniMaxAI/MiniMax-H3` `vae/` in the stack's `vae/` folder), and
-only the packed FP4 weights stay resident (0.31 GiB for 8 blocks instead of
+INT8 overlay is skipped, the VAE uses the exported decoder's block count
+(a complete export can also replace a shallower light decoder when its
+head/width geometry matches), and only the packed FP4 weights stay resident (0.31 GiB for 8 blocks instead of
 2.3 GiB). The decode stage runs bf16 autocast for these decoders, the
 precision they were trained in. The decoder GEMMs use the cuDNN backend on
 GB10; outputs match CUTLASS bit for bit.
@@ -406,6 +405,16 @@ read straight onto the GPU and never materializes the full-precision weights
 of the quantized linears; it decodes bit-identically to packing the bf16 file
 at load.
 
+To quantize the original decoder without replacing it with a distilled export,
+set `generator.pipeline.experimental.h3_vae_nvfp4_quantize_dense: true` and
+`generator.engine.compile.vae_enabled: false`. This experimental option keeps
+every decoder block, quantizes its block linears to NVFP4 with unit activation
+scale during loading, and uses the existing fused BF16 decoder path. It leaves
+the checkpoint files unchanged and skips the INT8 overlay. Clear
+`FASTVIDEO_H3_VAE_NVFP4_DECODER` and `h3_vae_nvfp4_decoder`: exported and runtime
+quantization selections are mutually exclusive. The output is lossy and needs
+quality validation for the selected checkpoint. Blackwell FP4 hardware is required.
+
 V2 at 832x480, 124 frames, seed 1234, one Spark, warm calls:
 
 | Video decoder | End to end | Denoise | Video decode | Peak allocated |

examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml

@@ -1,5 +1,6 @@
 # Start Ray on both Sparks and source spark_pair_env.sh first; see spark_pair.md.
 # FastH3-QAD-v2 eight-forward video+audio on two DGX Sparks over QSFP RoCE.
+# Uses the published fast-8L NVFP4 decoder on both workers (FASTVIDEO_H3_VAE_NVFP4_DECODER) with parallel temporal chunks.
 # Download the complete release checkpoint as described in
 # docs/getting_started/installation/spark_performance.md.
 #
@@ -9,8 +10,8 @@
 # settings below are an inference choice, not its linears-only QAD recipe.
 # The NVFP4 encoder and lightweight VAE ship in the same repository.
 #
-# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cutlass \
-# FASTVIDEO_H3_VAE_TILE_BATCH=1 \
+# FASTVIDEO_H3_VSA_FP4=1 FASTVIDEO_MINIMAX_H3_FUSIONS=all FASTVIDEO_NVFP4_MM_BACKEND=cudnn FASTVIDEO_NVFP4_MM_ROW_SPLIT=table \
+# FASTVIDEO_H3_VAE_TILE_BATCH=12 FASTVIDEO_H3_VAE_INT8_OVERLAY=0 \
 # FASTVIDEO_VSA_TRITON=1 FASTVIDEO_VSA_PTX_BLACKWELL=0 FASTVIDEO_FA4=0 \
 # FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN_H3 FASTVIDEO_STAGE_LOGGING=1 \
 # fastvideo generate --config examples/inference/basic/basic_fasth3_spark_pair_qad_v2_nvfp4.yaml
@@ -36,7 +37,7 @@ generator:
       lazy_module_load: false
     compile:
       enabled: false
-      vae_enabled: false
+      vae_enabled: false  # Published NVFP4 decoder uses its fused eager inference path.
   pipeline:
     workload_type: t2v
     vae_tiling: true
@@ -47,6 +48,11 @@ generator:
         prefix_precision: fp8
         two_level_p: true
         smooth_k: true
+        sp_head_chunks: 7
+        sp_pooled_key_mean: true
+        sp_direct_keys: true
+        sp_projection_overlap: true
+        sp_input_gather: true
       VSA_sparsity: 0.8
       VSA_tile_size: 64
       h3_sequential_load: false

examples/inference/basic/basic_fasth3_spark_qad_v2_nvfp4.yaml

@@ -1,4 +1,5 @@
 # FastH3-QAD-v2 eight-forward video+audio on one DGX Spark.
+# Uses the published fast-8L NVFP4 decoder selected with FASTVIDEO_H3_VAE_NVFP4_DECODER (see spark_performance.md).
 # Download the complete release checkpoint as described in
 # docs/getting_started/installation/spark_performance.md.
 #
@@ -34,7 +35,7 @@ generator:
       lazy_module_load: false
     compile:
       enabled: false
-      vae_enabled: true
+      vae_enabled: false  # Packed NVFP4 decoder has its own fused inference path.
   pipeline:
     workload_type: t2v
     vae_tiling: true

examples/inference/optimizations/spark_pair_env.sh

@@ -19,3 +19,7 @@ export NCCL_NVLS_ENABLE="${NCCL_NVLS_ENABLE:-0}"
 # runaway and SIGTERMs the worker around shard 11/14.
 export RAY_memory_monitor_refresh_ms="${RAY_memory_monitor_refresh_ms:-0}"
 export RAY_memory_usage_threshold="${RAY_memory_usage_threshold:-1.0}"
+# Measured on a two-Spark QSFP pair: eight NCCL channels per peer and high-priority
+# NCCL streams let the sequence-parallel exchanges overlap the projection GEMMs.
+export NCCL_NCHANNELS_PER_NET_PEER="${NCCL_NCHANNELS_PER_NET_PEER:-8}"
+export TORCH_NCCL_HIGH_PRIORITY="${TORCH_NCCL_HIGH_PRIORITY:-1}"

fastvideo-kernel/attn_qat_infer/blackwell/launch.h

@@ -89,7 +89,11 @@ void run_flash_fwd(Flash_fwd_params &params, cudaStream_t stream) {
     static constexpr int ctaSize = Kernel_traits::kNWarps * 32;
     params.m_block_divmod = cutlass::FastDivmod(num_blocks_m);
     params.total_blocks = num_blocks_m * params.h * params.b;
-    dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, 170);
+    // One persistent worker per SM. A fixed 170-worker grid leaves a partial
+    // final wave on smaller GPUs such as GB10, while every worker still walks
+    // its full static share of query tiles.
+    const int num_sms = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+    dim3 grid_dims = Scheduler::get_grid_dim(scheduler_args, num_sms);
     dim3 block_dims(ctaSize);
     dim3 cluster_dims(size<0>(ClusterShape{}), size<1>(ClusterShape{}), size<2>(ClusterShape{}));
     cutlass::ClusterLaunchParams launch_params{grid_dims, block_dims, cluster_dims, smem_size, stream};

fastvideo-kernel/csrc/attention/minimax_h3_prefix.cu

@@ -0,0 +1,84 @@
+// Exact Omni prefix attention: Torch-pinned FlashAttention, larger query tile only.
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cmath>
+#include "flash_fwd_launch_template.h"
+
+using namespace FLASH_NAMESPACE;
+
+template<int M, int W, bool QRegs = false, bool Share = false>
+struct Variant {
+    using Traits = Flash_fwd_kernel_traits<128, M, 64, W, QRegs, Share, cutlass::bfloat16_t>;
+    static auto kernel() {
+        // Keep the original K tile and reduction order. All masking stays enabled.
+        return flash_fwd_kernel<Traits, false, false, false, false, false, true, false, false>;
+    }
+    static void setup() {
+        static const bool initialized = []() {
+            if (Traits::kSmemSize >= 48 * 1024) {
+                C10_CUDA_CHECK(cudaFuncSetAttribute(kernel(), cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                                   Traits::kSmemSize));
+            }
+            return true;
+        }();
+        (void) initialized;
+    }
+    static void launch(Flash_fwd_params &params, cudaStream_t stream) {
+        setup();
+        dim3 grid((params.seqlen_q + M - 1) / M, params.b, params.h);
+        auto function = kernel();
+        function<<<grid, W * 32, Traits::kSmemSize, stream>>>(params);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+    }
+    static std::vector<int64_t> metadata() {
+        setup();
+        cudaFuncAttributes attributes;
+        C10_CUDA_CHECK(cudaFuncGetAttributes(&attributes, kernel()));
+        int blocks = 0;
+        C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+            &blocks, kernel(), W * 32, Traits::kSmemSize));
+        return {M, W, Traits::kSmemSize, attributes.numRegs,
+                static_cast<int64_t>(attributes.localSizeBytes), blocks};
+    }
+};
+
+torch::Tensor prefix_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v) {
+    TORCH_CHECK(q.is_cuda() && k.device() == q.device() && v.device() == q.device(), "CUDA operands required");
+    TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.sizes() == k.sizes(), "BHSD operands required");
+    TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 &&
+                v.scalar_type() == at::kBFloat16, "BF16 operands required");
+    TORCH_CHECK(q.size(3) == 128 && k.size(3) == 128 && q.stride(3) == 1 &&
+                k.stride(3) == 1 && v.stride(3) == 1, "Contiguous 128-wide heads required");
+    TORCH_CHECK(q.size(0) == k.size(0) && q.size(1) == k.size(1), "Batch/head mismatch");
+    TORCH_CHECK(q.size(0) > 0 && q.size(1) > 0 && q.size(2) > 0 && k.size(2) > 0,
+                "Nonempty attention required");
+    c10::cuda::CUDAGuard guard(q.device());
+    auto out = torch::empty({q.size(0), q.size(2), q.size(1), 128}, q.options()).transpose(1, 2);
+    auto lse = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options().dtype(at::kFloat));
+    Flash_fwd_params params{};
+    params.q_ptr = q.data_ptr(); params.k_ptr = k.data_ptr(); params.v_ptr = v.data_ptr();
+    params.o_ptr = out.data_ptr(); params.softmax_lse_ptr = lse.data_ptr();
+    params.q_batch_stride = q.stride(0); params.q_head_stride = q.stride(1); params.q_row_stride = q.stride(2);
+    params.k_batch_stride = k.stride(0); params.k_head_stride = k.stride(1); params.k_row_stride = k.stride(2);
+    params.v_batch_stride = v.stride(0); params.v_head_stride = v.stride(1); params.v_row_stride = v.stride(2);
+    params.o_batch_stride = out.stride(0); params.o_head_stride = out.stride(1); params.o_row_stride = out.stride(2);
+    params.b = q.size(0); params.h = q.size(1); params.h_k = k.size(1); params.h_h_k_ratio = 1;
+    params.seqlen_q = q.size(2); params.seqlen_k = k.size(2); params.total_q = params.b * params.seqlen_q;
+    params.seqlen_q_rounded = (params.seqlen_q + 127) / 128 * 128;
+    params.seqlen_k_rounded = (params.seqlen_k + 127) / 128 * 128;
+    params.d = 128; params.d_rounded = 128;
+    params.scale_softmax = static_cast<float>(1.0 / std::sqrt(128.0));
+    params.scale_softmax_log2 = params.scale_softmax * M_LOG2E;
+    params.p_dropout = 1.0f; params.p_dropout_in_uint8_t = 255; params.rp_dropout = 1.0f;
+    params.scale_softmax_rp_dropout = params.scale_softmax;
+    params.window_size_left = -1; params.window_size_right = -1;
+    params.is_bf16 = true; params.is_seqlens_k_cumulative = true; params.num_splits = 1;
+    auto stream = at::cuda::getCurrentCUDAStream();
+    Variant<256, 8>::launch(params, stream);
+    return out;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, binding) {
+    binding.def("forward", &prefix_forward);
+}

fastvideo-kernel/csrc/quantization/minimax_h3_swiglu_gemm.cu

@@ -0,0 +1,203 @@
+// SPDX-License-Identifier: Apache-2.0
+// Experimental full-K SM120 GEMM with value/gate-paired columns and FP4 output.
+#include "minimax_h3_swiglu_quant.cuh"
+#include <cutlass/cutlass.h>
+#include <cutlass/epilogue/collective/collective_builder.hpp>
+#include <cutlass/gemm/collective/collective_builder.hpp>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/tile_scheduler.hpp>
+#include <cutlass/util/packed_stride.hpp>
+
+namespace h3_epilogue {
+using namespace cute;
+using namespace cutlass::epilogue::fusion;
+struct Operation : LinearCombination<cutlass::bfloat16_t, float, void> {
+  using ElementAux = cutlass::bfloat16_t;
+  using GmemLayoutTagAux = cutlass::layout::RowMajor;
+  static constexpr int AlignmentAux = 8;
+};
+
+template<class CtaTile, class EpiTile>
+struct Store {
+  static constexpr int EM = size<0>(EpiTile{});
+  static constexpr int EN = size<1>(EpiTile{});
+  static_assert(EN % 32 == 0);
+  struct SharedStorage { cutlass::array_aligned<__nv_bfloat16, EM * EN / 2> values; };
+  struct Arguments {
+    const float* alpha_ptr = nullptr;
+    const float* scale_ptr = nullptr;
+    uint8_t* q = nullptr;
+    uint8_t* scales = nullptr;
+    float* amax = nullptr;
+    __nv_bfloat16* y = nullptr;
+  };
+  using Params = Arguments;
+  Params const* params;
+  __nv_bfloat16* shared;
+  CUTLASS_HOST_DEVICE Store() = default;
+  CUTLASS_HOST_DEVICE Store(Params const& p, SharedStorage const& s)
+      : params(&p), shared(const_cast<__nv_bfloat16*>(s.values.data())) {}
+  template<class Shape>
+  static Params to_underlying_arguments(Shape const&, Arguments const& a, void*) { return a; }
+  template<class Shape> static bool can_implement(Shape const& s, Arguments const&) {
+    return get<1>(s) % 128 == 0;
+  }
+  template<class Shape> static size_t get_workspace_size(Shape const&, Arguments const&) { return 0; }
+  template<class Shape>
+  static cutlass::Status initialize_workspace(Shape const&, Arguments const&, void*, cudaStream_t,
+                                               cutlass::CudaHostAdapter* = nullptr) {
+    return cutlass::Status::kSuccess;
+  }
+  CUTLASS_DEVICE bool is_producer_load_needed() const { return false; }
+  CUTLASS_DEVICE bool is_C_load_needed() const { return false; }
+  template<class... A> CUTLASS_DEVICE auto get_producer_load_callbacks(ProducerLoadArgs<A...> const&) {
+    return EmptyProducerLoadCallbacks{};
+  }
+
+  template<class Coords, int Threads>
+  struct Consumer : EmptyConsumerStoreCallbacks {
+    Coords coords;
+    Params p;
+    __nv_bfloat16* shared;
+    int tid, rows, width, row_base, col_base;
+    float alpha, global_sf;
+    CUTLASS_DEVICE Consumer(Coords c, Params p_, __nv_bfloat16* sm, int t, int m, int n, int rb, int cb)
+        : coords(c), p(p_), shared(sm), tid(t), rows(m), width(n/2), row_base(rb), col_base(cb),
+          alpha(*p_.alpha_ptr), global_sf(*p_.scale_ptr) {}
+    template<class Acc, int F>
+    CUTLASS_DEVICE auto visit(cutlass::Array<Acc, F> const& acc, int v, int em, int en) {
+      static_assert(F % 2 == 0);
+      auto c = coords(_,_,_,em,en);
+      CUTLASS_PRAGMA_UNROLL
+      for (int i=0; i<F; i+=2) {
+        auto xy = c(v*F+i);
+        const int row = get<0>(xy), col = get<1>(xy);
+        const float value = __bfloat162float(__float2bfloat16_rn(__fmul_rn(alpha, float(acc[i]))));
+        const float gate = __bfloat162float(__float2bfloat16_rn(__fmul_rn(alpha, float(acc[i+1]))));
+        const float y = row < rows ? swiglu_bf16(value,gate) : 0;
+        shared[(row % EM) * (EN/2) + (col % EN)/2] = __float2bfloat16_rn(y);
+      }
+      return acc;
+    }
+    template<class Buffer, class Sync, class Values>
+    CUTLASS_DEVICE void reduce(Buffer&&, Sync const& sync, int em, int en, bool, Values) {
+      sync();
+      CUTLASS_PRAGMA_UNROLL
+      for (int i=tid; i<EM*EN/2; i+=Threads) {
+        const int row = row_base + em*EM + i/(EN/2);
+        const int col = col_base/2 + en*(EN/2) + i%(EN/2);
+        const float y = __bfloat162float(shared[i]);
+        float amax = fabsf(y);
+        CUTLASS_PRAGMA_UNROLL
+        for (int d=1; d<16; d*=2)
+          amax=fmaxf(amax,__shfl_xor_sync(0xffffffff,amax,d,16));
+        const uint8_t sf=scale_fp8(__fmul_rn(global_sf,__fmul_rn(amax,rcp_approx(6.0f))));
+        const float inv=amax!=0 ? rcp_approx(__fmul_rn(scale_float(sf),rcp_approx(global_sf))) : 0;
+        const float code=__fmul_rn(y,inv);
+        const float hi=__shfl_xor_sync(0xffffffff,code,1);
+        if (row<rows && col<width) {
+          if ((col&1)==0) p.q[int64_t(row)*(width/2)+col/2]=fp4_pair(code,hi);
+          if (p.y) p.y[int64_t(row)*width+col]=__float2bfloat16_rn(y);
+        }
+        if (row < (rows + 127) / 128 * 128 && col<width && (col&15)==0) {
+          const int group=col/16;
+          const int64_t idx=(((((int64_t(row)/128)*(width/64)+group/4)*32+row%32)*4+(row/32)%4)*4+group%4);
+          p.scales[idx]=sf;
+          if (p.amax) p.amax[int64_t(row)*(width/16)+group]=amax;
+        }
+      }
+      sync();
+    }
+  };
+  template<bool ReferenceSrc, class... A>
+  CUTLASS_DEVICE auto get_consumer_store_callbacks(ConsumerStoreArgs<A...> const& args) {
+    auto [M,N,K,L]=args.problem_shape_mnkl;
+    auto [m,n,k,l]=args.tile_coord_mnkl;
+    auto identity=make_identity_tensor(make_shape(M,N));
+    auto tile=local_tile(identity,take<0,2>(args.tile_shape_mnk),make_coord(m,n));
+    auto coords=sm90_partition_for_epilogue<ReferenceSrc>(tile,args.epi_tile,args.tiled_copy,args.thread_idx);
+    constexpr int threads=decltype(size(args.tiled_copy))::value;
+    return Consumer<decltype(coords),threads>(coords,*params,shared,args.thread_idx,M,N,
+                                              m*size<0>(CtaTile{}),n*size<1>(CtaTile{}));
+  }
+};
+} // namespace h3_epilogue
+
+namespace cutlass::epilogue::fusion {
+template<int SC,int SD,int F,bool R,bool D,class CT,class ET,class... Extra>
+struct FusionCallbacks<cutlass::epilogue::Sm90TmaWarpSpecialized<SC,SD,F,R,D>,
+                       h3_epilogue::Operation,CT,ET,Extra...> : h3_epilogue::Store<CT,ET> {
+  using Impl=h3_epilogue::Store<CT,ET>;
+  using Impl::Impl;
+  using Operation=h3_epilogue::Operation;
+};
+}
+
+namespace h3_epilogue {
+using Shape=cute::Shape<int,int,int,int>;
+#ifndef H3_TILE_M
+#define H3_TILE_M 128
+#define H3_TILE_N 128
+#define H3_TILE_K 256
+#endif
+#ifndef H3_PINGPONG
+#define H3_PINGPONG 0
+#endif
+#ifndef H3_STAGES
+#define H3_STAGES 0
+#endif
+using Tile=cute::Shape<cute::Int<H3_TILE_M>,cute::Int<H3_TILE_N>,cute::Int<H3_TILE_K>>;
+using Cluster=cute::Shape<cute::_1,cute::_1,cute::_1>;
+using Epilogue=typename cutlass::epilogue::collective::CollectiveBuilder<
+  cutlass::arch::Sm120,cutlass::arch::OpClassTensorOp,Tile,Cluster,
+  cutlass::epilogue::collective::EpilogueTileAuto,float,float,
+  void,cutlass::layout::RowMajor,8,void,cutlass::layout::RowMajor,8,
+  cutlass::epilogue::TmaWarpSpecialized,Operation>::CollectiveOp;
+using Mainloop=typename cutlass::gemm::collective::CollectiveBuilder<
+  cutlass::arch::Sm120,cutlass::arch::OpClassBlockScaledTensorOp,
+  cutlass::nv_float4_t<cutlass::float_e2m1_t>,cutlass::layout::RowMajor,32,
+  cutlass::nv_float4_t<cutlass::float_e2m1_t>,cutlass::layout::ColumnMajor,32,
+  float,Tile,Cluster,
+  cute::conditional_t<H3_STAGES == 0,
+    cutlass::gemm::collective::StageCountAutoCarveout<sizeof(typename Epilogue::SharedStorage)>,
+    cutlass::gemm::collective::StageCount<H3_STAGES>>, 
+  cute::conditional_t<H3_PINGPONG != 0, cutlass::gemm::KernelTmaWarpSpecializedPingpong,
+                      cutlass::gemm::KernelTmaWarpSpecializedCooperative>>::CollectiveOp;
+using Kernel=cutlass::gemm::kernel::GemmUniversal<Shape,Mainloop,Epilogue,cutlass::gemm::StaticPersistentScheduler>;
+using Gemm=cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+}
+extern "C" int launch_gemm_swiglu(const void* a,const void* b,const void* sa,const void* sb,
+                                 const float* alpha,const float* sf,void* q,void* scales,
+                                 float* amax,void* y,int m,int n,int k,void* workspace,
+                                 uint64_t workspace_bytes,int raster,int swizzle,void* raw_stream) {
+  using namespace h3_epilogue;
+  auto stream=static_cast<cudaStream_t>(raw_stream);
+  Gemm::Arguments args{};
+  args.mode=cutlass::gemm::GemmUniversalMode::kGemm;
+  args.problem_shape=cute::make_shape(m,n,k,1);
+  args.mainloop.ptr_A=static_cast<const cutlass::float_e2m1_t*>(a);
+  args.mainloop.ptr_B=static_cast<const cutlass::float_e2m1_t*>(b);
+  args.mainloop.ptr_SFA=static_cast<const cutlass::float_ue4m3_t*>(sa);
+  args.mainloop.ptr_SFB=static_cast<const cutlass::float_ue4m3_t*>(sb);
+  args.mainloop.dA=cutlass::make_cute_packed_stride(Kernel::StrideA{}, {m,k,1});
+  args.mainloop.dB=cutlass::make_cute_packed_stride(Kernel::StrideB{}, {n,k,1});
+  using Config=Mainloop::Sm1xxBlkScaledConfig;
+  args.mainloop.layout_SFA=Config::tile_atom_to_shape_SFA(args.problem_shape);
+  args.mainloop.layout_SFB=Config::tile_atom_to_shape_SFB(args.problem_shape);
+  args.epilogue.dC=cutlass::make_cute_packed_stride(Kernel::StrideC{}, {m,n,1});
+  args.epilogue.dD=args.epilogue.dC;
+  args.epilogue.thread={alpha,sf,static_cast<uint8_t*>(q),static_cast<uint8_t*>(scales),amax,
+                       static_cast<__nv_bfloat16*>(y)};
+  args.scheduler.max_swizzle_size=swizzle;
+  args.scheduler.raster_order=static_cast<decltype(args.scheduler.raster_order)>(raster);
+  args.hw_info.cluster_shape=dim3(1,1,1);
+  args.hw_info.cluster_shape_fallback=dim3(1,1,1);
+  Gemm gemm;
+  if (gemm.get_workspace_size(args)>workspace_bytes) return -1;
+  auto status=gemm.can_implement(args);
+  if (status!=cutlass::Status::kSuccess) return 100+int(status);
+  status=gemm.initialize(args,workspace,stream);
+  if (status!=cutlass::Status::kSuccess) return 200+int(status);
+  status=gemm.run(stream);
+  return status==cutlass::Status::kSuccess ? int(cudaGetLastError()) : 300+int(status);
+}

fastvideo-kernel/csrc/quantization/minimax_h3_swiglu_quant.cuh

@@ -0,0 +1,88 @@
+// SPDX-License-Identifier: Apache-2.0
+// Private building block for a SwiGLU/NVFP4 GEMM epilogue. No runtime dispatch.
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+__device__ __forceinline__ float rcp_approx(float x) {
+  float y;
+  asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+  return y;
+}
+
+__device__ __forceinline__ float swiglu_bf16(float value, float gate) {
+  float e;
+  const float log2e = -1.4426950408889634f;
+  const float exponent = __fmul_rn(gate, log2e);
+  asm("ex2.approx.f32 %0, %1;" : "=f"(e) : "f"(exponent));
+  float sigmoid;
+  const float denominator = __fadd_rn(1.0f, e);
+  // Match Triton sigmoid, including subnormal reciprocals for large negative gates.
+  asm("div.full.f32 %0, %1, %2;" : "=f"(sigmoid) : "f"(1.0f), "f"(denominator));
+  return __bfloat162float(__float2bfloat16_rn(__fmul_rn(value, __fmul_rn(gate, sigmoid))));
+}
+
+__device__ __forceinline__ uint8_t scale_fp8(float x) {
+  uint32_t packed;
+  asm("{ .reg .b16 s; cvt.rn.satfinite.e4m3x2.f32 s, %1, %1; mov.b32 %0, {s,s}; }"
+      : "=r"(packed) : "f"(x));
+  return static_cast<uint8_t>(packed);
+}
+
+__device__ __forceinline__ float scale_float(uint8_t x) {
+  // E4M3 positive finite values; use an exact bit construction including subnormals.
+  const int exponent = (x >> 3) & 15;
+  const int mantissa = x & 7;
+  return exponent == 0 ? __fmul_rn(float(mantissa), 0.001953125f)
+                       : __uint_as_float(((exponent + 120) << 23) | (mantissa << 20));
+}
+
+__device__ __forceinline__ uint8_t fp4_pair(float lo, float hi) {
+  uint32_t packed;
+  asm("{ .reg .b8 q; cvt.rn.satfinite.e2m1x2.f32 q, %1, %2; mov.b32 %0, {q,q,q,q}; }"
+      : "=r"(packed) : "f"(hi), "f"(lo));
+  return static_cast<uint8_t>(packed);
+}
+
+__global__ void quant_probe(const __nv_bfloat16* x, const float* sf, uint8_t* q,
+                            uint8_t* scales, float* amax_out, __nv_bfloat16* y_out,
+                            int rows, int width, bool paired) {
+  const int row = blockIdx.x;
+  const int col = blockIdx.y * blockDim.x + threadIdx.x;
+  float value = 0, gate = 0;
+  if (row < rows && col < width) {
+    const int64_t base = int64_t(row) * width * 2;
+    value = __bfloat162float(x[base + (paired ? 2 * col : col)]);
+    gate = __bfloat162float(x[base + (paired ? 2 * col + 1 : width + col)]);
+  }
+  const float y = swiglu_bf16(value, gate);
+  float amax = fabsf(y);
+  for (int delta = 1; delta < 16; delta *= 2)
+    amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, delta, 16));
+  const float global_sf = *sf;
+  const uint8_t scale = scale_fp8(__fmul_rn(global_sf, __fmul_rn(amax, rcp_approx(6.0f))));
+  const float inv = amax != 0 ? rcp_approx(__fmul_rn(scale_float(scale), rcp_approx(global_sf))) : 0;
+  const float code = __fmul_rn(y, inv);
+  const float hi = __shfl_xor_sync(0xffffffff, code, 1);
+  if (row < rows && col < width) {
+    if ((col & 1) == 0) q[int64_t(row) * (width / 2) + col / 2] = fp4_pair(code, hi);
+    if (y_out) y_out[int64_t(row) * width + col] = __float2bfloat16_rn(y);
+  }
+  if (col < width && (col & 15) == 0) {
+    const int group = col / 16;
+    const int64_t index = (((((int64_t(row) / 128) * (width / 64) + group / 4) * 32
+                           + row % 32) * 4 + (row / 32) % 4) * 4 + group % 4);
+    scales[index] = scale;
+    if (amax_out) amax_out[int64_t(row) * (width / 16) + group] = amax;
+  }
+}
+
+extern "C" int launch_quant_probe(const void* x, const float* sf, void* q, void* scales,
+                                  float* amax_out, void* y_out, int rows, int width,
+                                  int paired, void* stream) {
+  quant_probe<<<dim3((rows + 127) / 128 * 128, (width + 255) / 256), 256, 0,
+                static_cast<cudaStream_t>(stream)>>>(static_cast<const __nv_bfloat16*>(x), sf,
+      static_cast<uint8_t*>(q), static_cast<uint8_t*>(scales), amax_out,
+      static_cast<__nv_bfloat16*>(y_out), rows, width, paired);
+  return static_cast<int>(cudaGetLastError());
+}

fastvideo-kernel/python/fastvideo_kernel/h3_swiglu_gemm_binding.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Private eager binding for the SM121 H3 SwiGLU/FP4 GEMM epilogue.
+
+The Spark prebuild installs this module beside its fingerprinted native library.
+Inputs are prequantized, with value/gate weight columns paired in memory.
+"""
+import ctypes
+from pathlib import Path
+
+import torch
+
+_library = ctypes.CDLL(str(Path(__file__).with_name('lib_fastvideo_h3_swiglu_gemm.so')))
+_launch = _library.launch_gemm_swiglu
+_launch.argtypes = ([ctypes.c_void_p] * 10 + [ctypes.c_int] * 3
+                    + [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_int, ctypes.c_int, ctypes.c_void_p])
+_launch.restype = ctypes.c_int
+_workspaces: dict[torch.device, torch.Tensor] = {}
+
+
+def run(a, b, a_scale, b_scale, alpha, global_scale, q, scales, amax=None, values=None,
+        raster=0, swizzle=1):
+    """Write packed SwiGLU outputs; b and b_scale have contiguous transposes."""
+    workspace = _workspaces.get(a.device)
+    if workspace is None:
+        workspace = torch.empty(16 * 1024 * 1024, device=a.device, dtype=torch.uint8)
+        _workspaces[a.device] = workspace
+    rc = _launch(
+        a.data_ptr(), b.data_ptr(), a_scale.data_ptr(), b_scale.data_ptr(), alpha.data_ptr(),
+        global_scale.data_ptr(), q.data_ptr(), scales.data_ptr(),
+        0 if amax is None else amax.data_ptr(), 0 if values is None else values.data_ptr(),
+        a.shape[0], b.shape[1], a.shape[1] * 2, workspace.data_ptr(), workspace.numel(),
+        raster, swizzle, torch.cuda.current_stream(a.device).cuda_stream,
+    )
+    if rc:
+        raise RuntimeError(f'H3 SwiGLU GEMM returned CUDA/CUTLASS error {rc}')

fastvideo/distributed/async_collectives.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tensor collectives with explicit input ownership until their consumers wait."""
+
+from dataclasses import dataclass
+import math
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+
+
+@dataclass
+class PendingTensor:
+    output: torch.Tensor
+    source: torch.Tensor
+    work: dist.Work
+    logical_shape: torch.Size | None = None
+
+    def wait(self) -> torch.Tensor:
+        """Order consumers on the current CUDA stream after the collective."""
+        self.work.wait()
+        if self.logical_shape is None:
+            return self.output
+        peer_elements = math.prod(self.logical_shape) // self.output.shape[0]
+        return self.output[:, :peer_elements].contiguous().view(self.logical_shape)
+
+
+def all_to_all_tensor_async(source: torch.Tensor, group: dist.ProcessGroup, *, align_bytes: int = 1) -> PendingTensor:
+    """Exchange equal peer slices, optionally aligning their NCCL buffer offsets.
+
+    Small FP32 scale slices with an odd head count can otherwise start eight
+    bytes off a 16-byte boundary and defeat NCCL's vectorized copies. Padding
+    only the transport buffer retains the caller's logical tensor layout.
+    """
+    logical_shape = None
+    if align_bytes > 1:
+        if align_bytes % source.element_size():
+            raise ValueError("Collective alignment must be a multiple of the element size")
+        world = dist.get_world_size(group)
+        if source.numel() % world:
+            raise ValueError("Equal all-to-all slices require a tensor divisible by the group size")
+        elements = source.numel() // world
+        padding = -elements % (align_bytes // source.element_size())
+        if padding:
+            logical_shape = source.shape
+            source = F.pad(source.reshape(world, elements), (0, padding))
+    output = torch.empty_like(source)
+    work = dist.all_to_all_single(output, source, group=group, async_op=True)
+    assert work is not None
+    return PendingTensor(output, source, work, logical_shape)
+
+
+def all_gather_tensor_async(source: torch.Tensor, group: dist.ProcessGroup, world: int) -> PendingTensor:
+    output = torch.empty((world, *source.shape), device=source.device, dtype=source.dtype)
+    work = dist.all_gather_into_tensor(output, source, group=group, async_op=True)
+    assert work is not None
+    return PendingTensor(output, source, work)
+
+
+@dataclass
+class PeerSlices:
+    """Two-rank receive whose local slice stays in the producer's storage."""
+    local: torch.Tensor
+    remote: torch.Tensor
+    template: torch.Tensor
+    rank: int
+
+    def materialize(self) -> torch.Tensor:
+        """Restore source-rank order for consumers without a two-pointer kernel."""
+        slices = (self.local, self.remote) if self.rank == 0 else (self.remote, self.local)
+        return torch.stack([part.view(torch.uint8) for part in slices]).view(self.template.dtype)
+
+
+@dataclass
+class PendingPeerSlices:
+    source: torch.Tensor
+    remote: torch.Tensor
+    rank: int
+    works: list[dist.Work]
+
+    def wait(self) -> PeerSlices:
+        for work in self.works:
+            work.wait()
+        return PeerSlices(self.source[self.rank], self.remote.view(self.source.dtype), self.source, self.rank)
+
+
+def exchange_peer_tensor_async(source: torch.Tensor, group: dist.ProcessGroup) -> PendingPeerSlices:
+    """Exchange only the remote half of a contiguous two-rank payload.
+
+    Every rank must enqueue calls in the same order. The producer is retained
+    through completion and consumption, including the local slice that bypasses
+    NCCL. Receivers index the two slices directly to avoid a concatenation copy.
+    """
+    if dist.get_world_size(group) != 2 or source.shape[0] != 2 or not source.is_contiguous():
+        raise ValueError("Peer-slice exchange requires two contiguous destination slices")
+    rank = dist.get_rank(group)
+    peer = dist.get_global_rank(group, 1 - rank)
+    send = source[1 - rank].view(torch.uint8)
+    remote = torch.empty_like(send)
+    works = dist.batch_isend_irecv(
+        [dist.P2POp(dist.irecv, remote, peer, group),
+         dist.P2POp(dist.isend, send, peer, group)])
+    return PendingPeerSlices(source, remote, rank, works)

fastvideo/distributed/projection_exchange.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Two-rank QKV traffic that can start before all projections are ready."""
+
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+
+class Completion:
+
+    def __init__(self, device: torch.device) -> None:
+        self.device = device
+        self.works: list[dist.Work] = []
+        self.waited_streams: set[int] = set()
+        self.ready = torch.cuda.Event()
+        self.producer = torch.cuda.current_stream(device).cuda_stream
+
+    def wait(self) -> None:
+        stream = torch.cuda.current_stream(self.device).cuda_stream
+        if stream not in self.waited_streams:
+            if stream != self.producer:
+                torch.cuda.current_stream(self.device).wait_event(self.ready)
+            for work in self.works:
+                work.wait()
+            self.waited_streams.add(stream)
+
+
+@dataclass
+class PendingInput:
+    source: torch.Tensor
+    remote: torch.Tensor
+    rank: int
+    completion: Completion
+
+    def wait(self) -> PeerSlices:
+        self.completion.wait()
+        return PeerSlices(self.source[self.rank], self.remote.view(self.source.dtype), self.source, self.rank)
+
+
+@dataclass
+class PendingScale:
+    source: torch.Tensor
+    output: torch.Tensor
+    completion: Completion
+
+    def wait(self) -> torch.Tensor:
+        self.completion.wait()
+        return self.output
+
+
+class ProjectionExchange:
+    """Exchange Q/K/V as each projection becomes ready, retaining their storage.
+
+    Input is [head groups, destinations, QKV, rows, heads, 128]. Call start
+    in Q/K/V order on the producing CUDA stream, then packets on that stream.
+    Input payload and scales must remain unmodified until all consumers wait.
+    The local payload aliases its producer; only the remote slice travels.
+    """
+
+    def __init__(self, payload: torch.Tensor, scale: torch.Tensor, group: dist.ProcessGroup) -> None:
+        if (dist.get_world_size(group) != 2 or payload.ndim != 6 or payload.shape[1:3] != (2, 3)
+                or payload.shape[-1] != 128 or not payload.is_cuda or not payload.is_contiguous()
+                or payload.dtype != torch.float8_e4m3fn or scale.shape != payload.shape[:-1]
+                or scale.device != payload.device or scale.dtype != torch.float32 or not scale.is_contiguous()):
+            raise ValueError("Projection exchange requires contiguous two-rank FP8 QKV and FP32 scales")
+        self.payload, self.scale, self.group = payload, scale, group
+        self.rank = dist.get_rank(group)
+        self.peer = dist.get_global_rank(group, 1 - self.rank)
+        self.remote = torch.empty((payload.shape[0], *payload.shape[2:]), device=payload.device, dtype=torch.uint8)
+        self.received_scale = torch.empty_like(scale)
+        self.completion = Completion(payload.device)
+        self.next_projection = 0
+
+    def start(self, projection: int) -> None:
+        if projection != self.next_projection or projection not in (0, 1, 2):
+            raise ValueError("Every rank must start projections in Q, K, V order")
+        self.next_projection += 1
+        ops: list[dist.P2POp] = []
+        for chunk in range(self.payload.shape[0]):
+            send = self.payload[chunk, 1 - self.rank, projection].view(torch.uint8)
+            remote = self.remote[chunk, projection]
+            send_scale = self.scale[chunk, 1 - self.rank, projection]
+            remote_scale = self.received_scale[chunk, 1 - self.rank, projection]
+            ops.extend((dist.P2POp(dist.irecv, remote, self.peer,
+                                   self.group), dist.P2POp(dist.isend, send, self.peer, self.group),
+                        dist.P2POp(dist.irecv, remote_scale, self.peer,
+                                   self.group), dist.P2POp(dist.isend, send_scale, self.peer, self.group)))
+        self.completion.works.extend(dist.batch_isend_irecv(ops))
+
+    def packets(self) -> list[tuple[PendingInput, PendingScale]]:
+        """Return consumers after all projections have been queued on this stream."""
+        if self.next_projection != 3:
+            raise RuntimeError("All three projection sends must start before consumption")
+        self.received_scale[:, self.rank].copy_(self.scale[:, self.rank])
+        self.completion.ready.record()
+        return [(PendingInput(self.payload[c], self.remote[c], self.rank,
+                              self.completion), PendingScale(self.scale[c], self.received_scale[c], self.completion))
+                for c in range(self.payload.shape[0])]
+
+
+def completed_projection_packets(payload: torch.Tensor, scale: torch.Tensor,
+                                 rank: int) -> list[tuple[PendingInput, PendingScale]]:
+    """Expose already source-ordered QKV packets through the usual consumers.
+
+    Input is [head groups, source ranks, QKV, rows, heads, 128]. An event orders
+    consumers on other CUDA streams after production; handles retain both
+    source slices and scales throughout their consumption.
+    """
+    if (payload.ndim != 6 or payload.shape[1:3] != (2, 3) or payload.shape[-1] != 128 or rank not in (0, 1)
+            or not payload.is_cuda or not payload.is_contiguous() or payload.dtype != torch.float8_e4m3fn
+            or scale.shape != payload.shape[:-1] or scale.device != payload.device or scale.dtype != torch.float32
+            or not scale.is_contiguous()):
+        raise ValueError("Completed projection packets require two contiguous source-rank slices")
+    completion = Completion(payload.device)
+    completion.ready.record()
+    return [(PendingInput(payload[c], payload[c, 1 - rank].view(torch.uint8), rank,
+                          completion), PendingScale(scale[c], scale[c], completion)) for c in range(payload.shape[0])]

fastvideo/fastvideo_args.py

@@ -234,6 +234,9 @@ class FastVideoArgs:
     # Exported NVFP4 MiniMax-H3 decoder file (decoder-nvfp4.safetensors or decoder.safetensors) that replaces
     # the VAE decoder at load. Overrides FASTVIDEO_H3_VAE_NVFP4_DECODER. Needs Blackwell FP4 hardware.
     h3_vae_nvfp4_decoder: str | None = None
+    # Quantize the loaded dense decoder in memory, preserving its full depth and original weights file.
+    # Mutually exclusive with an exported NVFP4 decoder; skips the INT8 overlay and uses BF16 decode autocast.
+    h3_vae_nvfp4_quantize_dense: bool = False
 
     # Load each heavy component on first use and free it once the last stage
     # that holds it has run, instead of keeping every component resident from
@@ -443,12 +446,19 @@ class FastVideoArgs:
             env_backend = envs.FASTVIDEO_ATTENTION_BACKEND.get()
             if env_backend is not None and backend_name_to_enum(env_backend) is not None:
                 self.attention_backend = env_backend
-        allowed_fp4 = {"enabled", "prefix_precision", "two_level_p", "smooth_k"}
+        allowed_fp4 = {
+            "enabled", "prefix_precision", "two_level_p", "smooth_k", "sp_head_chunks", "sp_pooled_key_mean",
+            "sp_projection_overlap", "sp_direct_keys", "sp_input_gather"
+        }
         if not isinstance(self.h3_vsa_fp4, dict) or self.h3_vsa_fp4.keys() - allowed_fp4:
             raise ValueError(f"h3_vsa_fp4 must be a mapping with keys from {sorted(allowed_fp4)}")
-        for key in ("enabled", "two_level_p", "smooth_k"):
+        for key in ("enabled", "two_level_p", "smooth_k", "sp_pooled_key_mean", "sp_projection_overlap",
+                    "sp_direct_keys", "sp_input_gather"):
             if key in self.h3_vsa_fp4 and not isinstance(self.h3_vsa_fp4[key], bool):
                 raise ValueError(f"h3_vsa_fp4.{key} must be a boolean")
+        head_chunks = self.h3_vsa_fp4.get("sp_head_chunks", 1)
+        if type(head_chunks) is not int or head_chunks < 1:
+            raise ValueError("h3_vsa_fp4.sp_head_chunks must be a positive integer")
         if self.h3_vsa_fp4.get("prefix_precision", "") not in ("", "fp8", "bf16"):
             raise ValueError("h3_vsa_fp4.prefix_precision must be empty, fp8 or bf16")
         if self.h3_vsa_fp4.get("enabled") and (self.inference_torch_compile or self.enable_torch_compile):

fastvideo/layers/quantization/nvfp4_config.py

@@ -418,6 +418,14 @@ def _mm_fp4(
         raise TypeError(f"Unsupported kwargs for _mm_fp4: {sorted(kwargs)}")
     if alpha is not None and not torch.is_tensor(alpha):
         alpha = torch.tensor(alpha, device=a.device, dtype=torch.float32)
+    # Compact H3 FFN-up: direct output slices avoid the concatenation that
+    # otherwise erases cuDNN's gain from these smaller row plans on GB10.
+    if (out is None and backend == "cudnn" and block_size == 16 and use_nvfp4 and not use_8x4_sf_layout
+            and (a.shape[0], a.shape[1], b.shape[1]) == (14859, 2688, 28672) and a.dtype == torch.uint8
+            and a_scale.ndim == 2 and a_scale.shape[0] >= a.shape[0] and a.is_cuda
+            and not torch.compiler.is_compiling() and envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.get() != "off"
+            and _is_dgx_spark(a.device.index or 0)):
+        return _mm_fp4_direct_rows(a, b, a_scale, b_scale, alpha, out_dtype, 8)
     if (out is None and backend == "cudnn" and a.shape[0] >= _SPLIT_MIN_ROWS
             and envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.get() != "off" and not torch.compiler.is_compiling()):
         return _mm_fp4_row_split(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend,
@@ -462,6 +470,28 @@ def _mm_fp4_row_chunks(a: torch.Tensor, a_scale: torch.Tensor, parts: int) -> li
     return chunks
 
 
+def _mm_fp4_direct_rows(a, b, a_scale, b_scale, alpha, out_dtype, parts: int) -> torch.Tensor:
+    """Run aligned cuDNN row chunks into one contiguous destination, including the final partial tile."""
+    _, mm_fp4, _ = _require_flashinfer()
+    out = torch.empty((a.shape[0], b.shape[1]), device=a.device, dtype=out_dtype)
+    offset = 0
+    for chunk, scale in _mm_fp4_row_chunks(a, a_scale, parts):
+        rows = chunk.shape[0]
+        mm_fp4(chunk,
+               b,
+               scale,
+               b_scale,
+               alpha,
+               out_dtype,
+               out[offset:offset + rows],
+               block_size=16,
+               use_8x4_sf_layout=False,
+               backend="cudnn",
+               use_nvfp4=True)
+        offset += rows
+    return out
+
+
 def _mm_fp4_chunked(a, b, a_scale, b_scale, alpha, out_dtype, block_size, use_8x4_sf_layout, backend, use_nvfp4,
                     parts: int) -> torch.Tensor:
     chunks = _mm_fp4_row_chunks(a, a_scale, parts) if parts > 1 else [(a, a_scale)]
@@ -920,7 +950,7 @@ def load_minimax_h3_nvfp4_dit_export(
         for name, module in modules.items() if isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod)
     }
     groups: dict[str, dict[str, str]] = {}
-    with safe_open(path, framework="pt", device="cpu") as reader:
+    with safe_open(path, framework="pt", device=str(device)) as reader:
         for key in reader.keys():  # noqa: SIM118
             if H3_NVFP4_DIT_KEY_SEP not in key:
                 raise ValueError(f"MiniMax-H3 NVFP4 DiT export key {key!r} is missing {H3_NVFP4_DIT_KEY_SEP!r}")
@@ -944,7 +974,7 @@ def load_minimax_h3_nvfp4_dit_export(
             for buffer_name in H3_NVFP4_DIT_BUFFER_NAMES + (H3_NVFP4_DIT_INPUT_SF_NAME, ):
                 if buffer_name not in buffers:
                     continue
-                tensor = reader.get_tensor(buffers[buffer_name]).to(device=device)
+                tensor = reader.get_tensor(buffers[buffer_name])
                 module.register_buffer(buffer_name, tensor, persistent=False)
             module.register_parameter("weight", None)
             loaded_names.add(next(name for name, candidate in modules.items() if candidate is module))

fastvideo/models/dits/minimax_h3.py

@@ -29,6 +29,7 @@ from fastvideo.distributed.parallel_state import (get_sp_group, get_sp_world_siz
 from fastvideo.forward_context import get_forward_context
 from fastvideo.layers.linear import ReplicatedLinear
 from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod
+from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
 from fastvideo.layers.mlp import MLP
 from fastvideo.layers.pdd import PDDReplicatedLinear, fuse_pdd_heads
 from fastvideo.layers.quantization import QuantizationConfig
@@ -163,7 +164,10 @@ class MiniMaxH3FeedForward(nn.Module):
         # (4.5 GiB at 78k tokens), so chunks bound the activation peak on 24-32 GB GPUs.
         self.chunk_tokens = envs.FASTVIDEO_H3_FFN_CHUNK_TOKENS.get()
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor,
+                pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
+        if pre_quantized is not None:
+            return self._forward(hidden_states, pre_quantized)
         tokens = hidden_states.shape[-2] if hidden_states.dim() > 1 else 0
         if (self.chunk_tokens and tokens > self.chunk_tokens and not torch.is_grad_enabled()
                 and not torch.compiler.is_compiling() and hidden_states.numel() == tokens * hidden_states.shape[-1]):
@@ -174,12 +178,29 @@ class MiniMaxH3FeedForward(nn.Module):
             return out
         return self._forward(hidden_states)
 
-    def _forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def _forward(self, hidden_states: torch.Tensor,
+                 pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
         if self.use_mxfp8:
             from fastvideo.layers.mxfp8linear import mxfp8_swiglu_feed_forward
 
             return mxfp8_swiglu_feed_forward(hidden_states, self.fc_in, self.fc_out)
-        hidden_states, _ = self.fc_in(hidden_states)
+        if self.fuse_swiglu and pre_quantized is not None and not torch.compiler.is_compiling():
+            from fastvideo.models.dits.minimax_h3_fusions.swiglu_gemm import try_swiglu_gemm
+
+            fused = try_swiglu_gemm(self, hidden_states, pre_quantized)
+            if fused is not None:
+                return fused
+        if pre_quantized is None:
+            hidden_states, _ = self.fc_in(hidden_states)
+        else:
+            hidden_states = self.fc_in.quant_method.apply(
+                self.fc_in, hidden_states, self.fc_in.bias, pre_quantized=pre_quantized,
+            )
+        if (self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states)
+                and not torch.compiler.is_compiling()):
+            quantized = self._quantized_swiglu_output(hidden_states)
+            if quantized is not None:
+                return quantized
         if self.fuse_swiglu and _can_run_minimax_h3_fusion(hidden_states):
             hidden_states = minimax_h3_swiglu(hidden_states)
         elif self.exact_swiglu and exact_kernels.supports_rowwise(hidden_states):
@@ -190,6 +211,43 @@ class MiniMaxH3FeedForward(nn.Module):
         hidden_states, _ = self.fc_out(hidden_states)
         return hidden_states
 
+    def input_quantization_scale(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        """Static scale for a fused producer, or None when ordinary preparation is needed."""
+        method = getattr(self.fc_in, "quant_method", None)
+        if (self.chunk_tokens or type(method) is not NVFP4QuantizeMethod
+                or not method.wants_prequantized_input() or hidden_states.dtype != torch.bfloat16
+                or hidden_states.shape[-1] % 64 or torch.cuda.get_device_capability(hidden_states.device)[0] < 10):
+            return None
+        scale = method._static_activation_global_sf()
+        if scale is None:
+            scale = getattr(self.fc_in, "_nvfp4_input_global_sf", None)
+        if scale is None and not method._dynamic_activation_scale():
+            scale = method.x_global_sf
+        return scale
+
+    def _quantized_swiglu_output(self, packed: torch.Tensor) -> torch.Tensor | None:
+        """Fuse BF16 SwiGLU with compatible static/unit-scale NVFP4 preparation."""
+        layer = self.fc_out
+        method = getattr(layer, "quant_method", None)
+        if (type(method) is not NVFP4QuantizeMethod or not method.wants_prequantized_input()
+                or packed.dtype != torch.bfloat16 or packed.shape[-1] % 128 or not packed.is_contiguous()
+                or torch.cuda.get_device_capability(packed.device)[0] < 10):
+            return None
+        scale = method._static_activation_global_sf()
+        if scale is None:
+            scale = getattr(layer, "_nvfp4_input_global_sf", None)
+        if scale is None:
+            if method._dynamic_activation_scale():
+                return None
+            scale = method.x_global_sf
+        from fastvideo.models.dits.minimax_h3_fusions.swiglu_quant import swiglu_quantize_nvfp4
+
+        pre = swiglu_quantize_nvfp4(packed, scale)
+        # apply(pre_quantized=...) reads only this view's shape. The BF16
+        # SwiGLU activation itself is never allocated or written.
+        shape_view = packed[..., :packed.shape[-1] // 2]
+        return method.apply(layer, shape_view, layer.bias, pre_quantized=pre)
+
 
 class MiniMaxH3Attention(nn.Module):
     """Full self-attention over one sequence-parallel packed document."""
@@ -351,12 +409,30 @@ class MiniMaxH3Attention(nn.Module):
             return exact_kernels.rope_prefix(hidden_states, rotary_emb)
         return self._apply_rotary_emb(hidden_states, rotary_emb)
 
+    def input_quantization_scale(
+        self, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> torch.Tensor | None:
+        """Allow producer quantization only when this call will use the shared FP4 route."""
+        fine_override = getattr(getattr(self.distributed_attention, "attn_impl", None), "fine_attention_override", None)
+        if (not self._vsa_fp4 or fine_override is not None or rotary_emb is None
+                or torch.is_grad_enabled() or torch.compiler.is_compiling()
+                or not hidden_states.is_cuda or hidden_states.dtype != torch.bfloat16
+                or torch.cuda.get_device_capability(hidden_states.device)[0] < 10):
+            return None
+        meta = get_forward_context().attn_metadata
+        if not isinstance(meta, MiniMaxH3VSAMetadata) or not meta.exempt or meta.pad_tokens:
+            return None
+        from fastvideo.models.dits.minimax_h3_vsa_fp4 import _shared_qkv_gate_scale
+
+        return _shared_qkv_gate_scale(self, hidden_states)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor] | None,
         original_seq_len: int,
         attn_metadata: AttentionMetadata | None = None,
+        pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         fine_override = getattr(getattr(self.distributed_attention, "attn_impl", None), "fine_attention_override", None)
         if (self._vsa_fp4 and fine_override is None and rotary_emb is not None and not torch.is_grad_enabled()
@@ -367,13 +443,22 @@ class MiniMaxH3Attention(nn.Module):
             if isinstance(meta, MiniMaxH3VSAMetadata) and meta.exempt and not meta.pad_tokens:
                 use_fused_rope = self.fuse_qknorm_rope and _can_run_minimax_h3_fusion(hidden_states)
                 if not model_parallel_is_initialized() or get_sp_world_size() == 1:
-                    hidden_states = vsa_fp4_attention(self, hidden_states, rotary_emb, meta, use_fused_rope)
+                    hidden_states = vsa_fp4_attention(self, hidden_states, rotary_emb, meta, use_fused_rope,
+                                                       pre_quantized=pre_quantized)
                 else:
                     hidden_states = vsa_fp4_attention_sp(self, hidden_states, rotary_emb, meta, use_fused_rope,
-                                                         get_sp_group())
+                                                         get_sp_group(), pre_quantized=pre_quantized)
                 with STAGES.span("out_proj"):
-                    hidden_states, _ = self.to_out(hidden_states)
+                    if isinstance(hidden_states, tuple):
+                        shape_view, prepared = hidden_states
+                        hidden_states = self.to_out.quant_method.apply(
+                            self.to_out, shape_view, self.to_out.bias, pre_quantized=prepared,
+                        )
+                    else:
+                        hidden_states, _ = self.to_out(hidden_states)
                 return hidden_states
+        if pre_quantized is not None:
+            raise RuntimeError("Prequantized H3 attention input requires the eager FP4 VSA route")
         if (self._vsa_tile_first and hidden_states.is_cuda and rotary_emb is not None
                 and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
                 and (not model_parallel_is_initialized() or get_sp_world_size() == 1)):
@@ -548,7 +633,7 @@ class MiniMaxH3AdaLayerNormModulation(nn.Module):
             prefix=f"{prefix}.linear",
         )
 
-    def enable_host_cache(self, table: dict | None = None) -> None:
+    def enable_host_cache(self, table: dict | None = None, *, keep_resident: bool = False) -> None:
         """Keep the projection in pinned host memory and cache its output per timestep set.
 
         The modulation is a pure function of the timestep embedding, and few-step checkpoints sample a fixed
@@ -563,7 +648,11 @@ class MiniMaxH3AdaLayerNormModulation(nn.Module):
             self._cache_key = None
             return
         weight, bias = weight.data, bias.data
-        if table is None:
+        if table is None and keep_resident:
+            # Unified-memory GPUs share the host pool. Reuse the original allocation
+            # instead of making a pinned copy and copying it back on every cache miss.
+            self._host_weight, self._host_bias = weight, bias
+        elif table is None:
             self._host_weight = weight.to("cpu").pin_memory()
             self._host_bias = bias.to("cpu").pin_memory()
         else:
@@ -700,7 +789,10 @@ class MiniMaxH3TransformerBlock(nn.Module):
         adaln_indices: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
         original_seq_len: int,
-    ) -> torch.Tensor:
+        *,
+        _attention_prepared: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+        _next_block: MiniMaxH3TransformerBlock | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
         with nvtx_range("minimax_h3.transformer_block.adaln_projection"):
             shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
                 t.to(hidden_states.dtype) for t in self.adaln_proj(temb))
@@ -710,7 +802,21 @@ class MiniMaxH3TransformerBlock(nn.Module):
         use_exact_modulate = (not use_modulate_fusion and self.exact_modulate and adaln_indices.is_contiguous()
                               and exact_kernels.supports_rowwise(hidden_states, shift_msa, scale_msa, gate_msa,
                                                                  shift_mlp, scale_mlp, gate_mlp))
-        if use_modulate_fusion:
+        attn_quantized = _attention_prepared
+        attn_scale = (self.attn.input_quantization_scale(hidden_states, rotary_emb)
+                      if attn_quantized is None and use_modulate_fusion and not torch.compiler.is_compiling() else None)
+        if attn_quantized is not None:
+            norm_hidden_states = hidden_states
+        elif attn_scale is not None:
+            from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import modulate_quantize_nvfp4
+
+            with nvtx_range("minimax_h3.transformer_block.modulate_quant_fusion"):
+                attn_quantized = modulate_quantize_nvfp4(
+                    hidden_states, self.norm1.weight, scale_msa, shift_msa, adaln_indices,
+                    self.norm1.eps, attn_scale,
+                )
+            norm_hidden_states = hidden_states  # Shape only; QKV and the gate consume attn_quantized.
+        elif use_modulate_fusion:
             with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
                 norm_hidden_states = fused_rmsnorm_modulate(
                     hidden_states,
@@ -730,11 +836,26 @@ class MiniMaxH3TransformerBlock(nn.Module):
                 norm_hidden_states = norm_hidden_states * (
                     1.0 + scale_msa.index_select(0, adaln_indices)) + shift_msa.index_select(0, adaln_indices)
         with nvtx_range("minimax_h3.transformer_block.self_attention"):
-            attention_output = self.attn(norm_hidden_states, rotary_emb, original_seq_len)
+            attention_output = (self.attn(norm_hidden_states, rotary_emb, original_seq_len,
+                                          pre_quantized=attn_quantized) if attn_quantized is not None
+                                else self.attn(norm_hidden_states, rotary_emb, original_seq_len))
         # The attention input is dead now. Keeping it until assignment below
         # overlaps three full-width activations during the residual fusion.
         del norm_hidden_states
-        if use_modulate_fusion:
+        ff_quantized = None
+        ff_scale = (self.ff.input_quantization_scale(hidden_states)
+                    if use_modulate_fusion and not torch.compiler.is_compiling() else None)
+        if ff_scale is not None:
+            from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import residual_modulate_quantize_nvfp4
+
+            with nvtx_range("minimax_h3.transformer_block.modulate_quant_fusion"):
+                hidden_states, ff_quantized = residual_modulate_quantize_nvfp4(
+                    hidden_states, attention_output, gate_msa, self.norm2.weight, scale_mlp,
+                    shift_mlp, adaln_indices, self.norm2.eps, ff_scale,
+                )
+            # The following NVFP4 projection consumes only this tensor's shape.
+            norm_hidden_states = hidden_states
+        elif use_modulate_fusion:
             with nvtx_range("minimax_h3.transformer_block.modulate_fusion"):
                 hidden_states, norm_hidden_states = fused_residual_gate_rmsnorm_modulate(
                     hidden_states,
@@ -759,8 +880,23 @@ class MiniMaxH3TransformerBlock(nn.Module):
                     1.0 + scale_mlp.index_select(0, adaln_indices)) + shift_mlp.index_select(0, adaln_indices)
         del attention_output
         with nvtx_range("minimax_h3.transformer_block.feed_forward"), STAGES.span("feed_forward"):
-            feed_forward_output = self.ff(norm_hidden_states)
+            feed_forward_output = (self.ff(norm_hidden_states, pre_quantized=ff_quantized)
+                                   if ff_quantized is not None else self.ff(norm_hidden_states))
         if use_modulate_fusion and not torch.compiler.is_compiling():
+            if _next_block is not None and _next_block.fuse_modulate:
+                next_quant_scale = _next_block.attn.input_quantization_scale(hidden_states, rotary_emb)
+                if next_quant_scale is not None:
+                    from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import (
+                        residual_modulate_quantize_nvfp4,
+                    )
+
+                    next_shift, next_scale, *_ = _next_block.adaln_proj(temb)
+                    return residual_modulate_quantize_nvfp4(
+                        hidden_states, feed_forward_output, gate_mlp, _next_block.norm1.weight,
+                        next_scale.to(hidden_states.dtype), next_shift.to(hidden_states.dtype),
+                        adaln_indices, _next_block.norm1.eps, next_quant_scale,
+                        round_residual_before_norm=True,
+                    )
             return _gated_residual(hidden_states, gate_mlp, adaln_indices, feed_forward_output)
         if use_exact_modulate:
             return exact_kernels.gate_residual(hidden_states, gate_mlp, feed_forward_output, adaln_indices)
@@ -1181,7 +1317,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         object.__setattr__(self, "_splice_late", late)
         self._splice_from_step = int(from_step)
 
-    def enable_adaln_host_cache(self, table_path: str | None = None) -> None:
+    def enable_adaln_host_cache(self, table_path: str | None = None, *, keep_resident: bool = False) -> None:
         """Move every block's AdaLN projection to pinned host memory behind a per-timestep cache.
 
         With ``table_path`` (written by FASTVIDEO_H3_ADALN_DUMP), the cache is prefilled from precomputed
@@ -1193,7 +1329,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
             raw = torch.load(table_path, map_location="cpu")
             tables = {int(i): {ast.literal_eval(k): v for k, v in blk.items()} for i, blk in raw.items()}
         for index, block in enumerate(self.transformer_blocks):
-            block.adaln_proj.enable_host_cache(None if tables is None else tables[index])
+            block.adaln_proj.enable_host_cache(None if tables is None else tables[index], keep_resident=keep_resident)
         self._adaln_host_cache = True
         self._adaln_dumped_entries = -1
 
@@ -1331,6 +1467,18 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 block.adaln_proj._cache_key = key
             self._move_adaln_tables(temb.device)
 
+        # Keep cross-block preparation local to this forward. Compiled/wrapped
+        # blocks retain their normal tensor-only interface.
+        cross_block = (not torch.compiler.is_compiling() and getattr(self, "_adaln_host_cache", False)
+                       and _can_run_minimax_h3_fusion(packed_hidden_states)
+                       and packed_hidden_states.dtype == torch.bfloat16
+                       and all(type(b) is MiniMaxH3TransformerBlock
+                               and getattr(b.forward, "__func__", None) is MiniMaxH3TransformerBlock.forward
+                               and getattr(b, "_compiled_call_impl", None) is None
+                               and not b._forward_pre_hooks and not b._forward_hooks
+                               and b.norm1.weight.device == packed_hidden_states.device
+                               for b in self.transformer_blocks))
+        attention_prepared = None
         # The eager driver owns profiling markers while each block's compiled
         # forward owns the graph that the marker surrounds.
         for block_index, block in enumerate(self.transformer_blocks):
@@ -1338,13 +1486,21 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                 logger.info("H3_MEMORY_BLOCK %d allocated=%.3f GiB reserved=%.3f GiB", block_index,
                             torch.cuda.memory_allocated() / 2**30, torch.cuda.memory_reserved() / 2**30)
             with nvtx_range(f"minimax_h3.transformer_block.{block_index}"), STAGES.span("block_total"):
-                packed_hidden_states = block(
-                    packed_hidden_states,
-                    temb,
-                    adaln_indices,
-                    rotary_emb,
-                    original_seq_len,
-                )
+                if cross_block:
+                    next_block = (self.transformer_blocks[block_index + 1]
+                                  if block_index + 1 < len(self.transformer_blocks) else None)
+                    result = block(
+                        packed_hidden_states, temb, adaln_indices, rotary_emb, original_seq_len,
+                        _attention_prepared=attention_prepared, _next_block=next_block,
+                    )
+                    if isinstance(result, tuple):
+                        packed_hidden_states, attention_prepared = result
+                    else:
+                        packed_hidden_states, attention_prepared = result, None
+                else:
+                    packed_hidden_states = block(
+                        packed_hidden_states, temb, adaln_indices, rotary_emb, original_seq_len,
+                    )
 
         packed_hidden_states = self.norm_out(
             packed_hidden_states,

fastvideo/models/dits/minimax_h3_fusions/block_lists.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Compact VSA tile masks into the native FP4 kernel's descending block lists."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _block_lists(MASK, VALID, IDX, NUM, KV, QUAD,
+                 MB: tl.constexpr, MH: tl.constexpr, MR: tl.constexpr, MC: tl.constexpr,
+                 H: tl.constexpr, T: tl.constexpr, HT: tl.constexpr,
+                 NB: tl.constexpr, BLOCK: tl.constexpr, WRITE_QUAD: tl.constexpr):
+    query = tl.program_id(0)
+    head = tl.program_id(1)
+    batch = tl.program_id(2)
+    lane = tl.arange(0, BLOCK)
+    col = NB - 1 - lane
+    valid_col = lane < NB
+    a, b = col * 2, col * 2 + 1
+    va = tl.load(VALID + a // HT, (a < T * HT) & valid_col, other=0).to(tl.int32) - a % HT * 64
+    vb = tl.load(VALID + b // HT, (b < T * HT) & valid_col, other=0).to(tl.int32) - b % HT * 64
+    va = tl.minimum(tl.maximum(va, 0), 64)
+    vb = tl.minimum(tl.maximum(vb, 0), 64)
+    q0, q1 = query * 2, query * 2 + 1
+    base = batch * MB + head * MH
+    m00 = tl.load(MASK + base + q0 // HT * MR + a // HT * MC,
+                  (q0 < T * HT) & (a < T * HT) & valid_col, other=0)
+    m01 = tl.load(MASK + base + q0 // HT * MR + b // HT * MC,
+                  (q0 < T * HT) & (b < T * HT) & valid_col, other=0)
+    m10 = tl.load(MASK + base + q1 // HT * MR + a // HT * MC,
+                  (q1 < T * HT) & (a < T * HT) & valid_col, other=0)
+    m11 = tl.load(MASK + base + q1 // HT * MR + b // HT * MC,
+                  (q1 < T * HT) & (b < T * HT) & valid_col, other=0)
+    # The native builder pads an odd final query half with a key-zero entry.
+    m10 = tl.where(q1 >= T * HT, a == 0, m10)
+    m11 = tl.where(q1 >= T * HT, b == 0, m11)
+    bits = ((m00 & (va > 0)).to(tl.uint8) + 2 * (m01 & (vb > 0)).to(tl.uint8)
+            + 4 * (m10 & (va > 0)).to(tl.uint8) + 8 * (m11 & (vb > 0)).to(tl.uint8))
+    keep = (bits != 0) & valid_col
+    pos = tl.cumsum(keep.to(tl.int32), 0) - 1
+    offset = ((batch * H + head) * NB + query) * NB
+    tl.store(IDX + offset + lane, 0, valid_col)
+    if WRITE_QUAD:
+        tl.store(QUAD + offset + lane, 0, valid_col)
+    tl.debug_barrier()
+    tl.store(IDX + offset + pos, col, keep)
+    if WRITE_QUAD:
+        tl.store(QUAD + offset + pos, bits, keep)
+    tl.store(NUM + (batch * H + head) * NB + query, tl.sum(keep.to(tl.int32), 0))
+    if batch == 0 and head == 0 and query == 0:
+        tl.store(KV + a, va, valid_col)
+        tl.store(KV + b, vb, valid_col)
+
+
+def tile_mask_to_fp4_blocks(mask: torch.Tensor, tile: int, valid: torch.Tensor):
+    """Match Sage's indices, counts, half validity, and optional quadrant bytes.
+
+    Inputs come from H3's mask builder: square CUDA bool masks and contiguous
+    integer validity counts, with positive tile sizes divisible by 64.
+    """
+    batch, heads, tiles, _ = mask.shape
+    halves = tile // 64
+    blocks = triton.cdiv(tiles * halves, 2)
+    indices = torch.empty((batch, heads, blocks, blocks), device=mask.device, dtype=torch.int32)
+    counts = torch.empty((batch, heads, blocks), device=mask.device, dtype=torch.int32)
+    kv_valid = torch.empty((blocks * 2,), device=mask.device, dtype=torch.int32)
+    quadrants = torch.empty_like(indices, dtype=torch.uint8) if halves == 1 else None
+    _block_lists[(blocks, heads, batch)](
+        mask, valid, indices, counts, kv_valid, quadrants if quadrants is not None else indices,
+        *mask.stride(), heads, tiles, halves, blocks, triton.next_power_of_2(blocks), halves == 1,
+        num_warps=4,
+    )
+    return indices, counts, kv_valid, quadrants

fastvideo/models/dits/minimax_h3_fusions/key_norm_mean.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse compact key normalization with the GB10 ATen mean reduction tree.
+
+For 18,432 rows and 56 head-128 outputs, ATen uses four accumulator lanes,
+four row-reduction warps and 18 CTAs per output. Preserve those additions and
+the BF16 normalization boundary while eliminating the separate full-key read.
+The caller restricts this specialization to the measured one-Spark geometry.
+"""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _gather_normalize_partials(
+    X, W, C, S, Index, Y, Partial, ROWS: tl.constexpr, HEADS: tl.constexpr,
+    INPUT_ROWS: tl.constexpr, ROTARY: tl.constexpr, EPS: tl.constexpr, CTAS: tl.constexpr,
+):
+    cta, head, batch = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    lanes = tl.arange(0, 16)[:, None]
+    cols = tl.arange(0, 128)[None, :]
+    accumulators = tl.full((16, 128), 0, tl.float32)
+    for iteration in range(tl.cdiv(ROWS, 16 * CTAS)):
+        # ATen input_idx = warp + cta * 4, step_input = 4 * CTAS;
+        # its four independent accumulators advance by step_input * 4.
+        row = lanes % 4 + cta * 4 + (lanes // 4) * (4 * CTAS) + iteration * (16 * CTAS)
+        source = tl.load(Index + row, row < ROWS, other=-1)
+        offset = ((batch * INPUT_ROWS + source) * HEADS + head) * 128 + cols
+        x = tl.load(X + offset, (row < ROWS) & (source >= 0), other=0).to(tl.float32)
+        squared = x * x
+        # Original four-warp normalization order, independent of this launch's layout.
+        for shift in tl.static_range(5):
+            other = tl.gather(squared, tl.broadcast_to(cols ^ (16 >> shift), (16, 128)), axis=1)
+            squared = squared + other
+        for shift in tl.static_range(2):
+            other = tl.gather(squared, tl.broadcast_to(cols ^ (64 >> shift), (16, 128)), axis=1)
+            squared = squared + other
+        variance = tl.gather(squared, tl.full((16, 1), 0, tl.int32), axis=1) / 128
+        normalized = x * tl.rsqrt(variance + EPS) * tl.load(W + cols).to(tl.float32)
+        partner = tl.where(cols < ROTARY // 2, cols + ROTARY // 2, cols - ROTARY // 2)
+        partner_normalized = tl.gather(normalized, tl.broadcast_to(partner, (16, 128)), axis=1)
+        rotated = tl.where(cols < ROTARY // 2, -partner_normalized, partner_normalized)
+        rotary_mask = (row < ROWS) & (cols < ROTARY)
+        cos = tl.load(C + row * ROTARY + cols, rotary_mask, other=1).to(tl.float32)
+        sin = tl.load(S + row * ROTARY + cols, rotary_mask, other=0).to(tl.float32)
+        output = tl.where(cols < ROTARY, normalized * cos + rotated * sin, normalized).to(tl.bfloat16)
+        tl.store(Y + ((batch * ROWS + row) * HEADS + head) * 128 + cols, output, row < ROWS)
+        # Mean consumes the stored BF16 values, not the normalization accumulator.
+        accumulators = accumulators + tl.where(row < ROWS, output.to(tl.float32), 0.0)
+
+    warp = tl.arange(0, 4)[:, None]
+    sums = tl.gather(accumulators, tl.broadcast_to(warp, (4, 128)), axis=0)
+    for accumulator in tl.static_range(1, 4):
+        sums = sums + tl.gather(accumulators, tl.broadcast_to(warp + accumulator * 4, (4, 128)), axis=0)
+    s0 = tl.gather(sums, tl.full((1, 128), 0, tl.int32), axis=0)
+    s1 = tl.gather(sums, tl.full((1, 128), 1, tl.int32), axis=0)
+    s2 = tl.gather(sums, tl.full((1, 128), 2, tl.int32), axis=0)
+    s3 = tl.gather(sums, tl.full((1, 128), 3, tl.int32), axis=0)
+    tl.store(Partial + ((batch * CTAS + cta) * HEADS + head) * 128 + cols, (s0 + s2) + (s1 + s3))
+
+
+@triton.jit
+def _finish_mean(Partial, Mean, HEADS: tl.constexpr, ROWS: tl.constexpr, CTAS: tl.constexpr):
+    head, batch = tl.program_id(0), tl.program_id(1)
+    cols = tl.arange(0, 128)
+    s0 = tl.full((128,), 0, tl.float32)
+    s1, s2, s3 = s0, s0, s0
+    for i in tl.static_range(tl.cdiv(CTAS, 4)):
+        s0 = s0 + tl.load(Partial + ((batch * CTAS + i * 4) * HEADS + head) * 128 + cols,
+                          i * 4 < CTAS, other=0)
+        s1 = s1 + tl.load(Partial + ((batch * CTAS + i * 4 + 1) * HEADS + head) * 128 + cols,
+                          i * 4 + 1 < CTAS, other=0)
+        s2 = s2 + tl.load(Partial + ((batch * CTAS + i * 4 + 2) * HEADS + head) * 128 + cols,
+                          i * 4 + 2 < CTAS, other=0)
+        s3 = s3 + tl.load(Partial + ((batch * CTAS + i * 4 + 3) * HEADS + head) * 128 + cols,
+                          i * 4 + 3 < CTAS, other=0)
+    tl.store(Mean + (batch * HEADS + head) * 128 + cols, ((s0 + s2) + (s1 + s3)) * (1.0 / ROWS))
+
+
+def gather_key_norm_mean(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+    eps: float, input_index: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return original BF16 normalized keys and their ATen-equivalent BF16 mean."""
+    if (x.ndim != 4 or x.shape[0] != 1 or x.shape[2:] != (56, 128)
+            or not x.is_cuda or x.dtype != torch.bfloat16 or input_index.numel() != 18432
+            or weight.shape != (128,) or cos.shape != (18432, 96) or sin.shape != cos.shape
+            or input_index.dtype not in (torch.int32, torch.int64)):
+        raise ValueError("Fused key mean requires CUDA BF16 batch-1/head-56, 18432 rows and rotary width 96")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Fused key mean requires eager inference")
+    x, weight, cos, sin, input_index = (t.contiguous() for t in (x, weight, cos, sin, input_index))
+    rows, heads, dim, ctas = 18432, 56, 128, 18
+    out = torch.empty((1, rows, heads, dim), dtype=x.dtype, device=x.device)
+    partial = torch.empty((1, ctas, heads, dim), dtype=torch.float32, device=x.device)
+    mean = torch.empty((1, 1, heads, dim), dtype=x.dtype, device=x.device)
+    _gather_normalize_partials[(ctas, heads, 1)](
+        x, weight, cos, sin, input_index, out, partial, rows, heads, x.shape[1], 96, eps, ctas, num_warps=4,
+    )
+    _finish_mean[(heads, 1)](partial, mean, heads, rows, ctas)
+    return out, mean

fastvideo/models/dits/minimax_h3_fusions/key_quant.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse BF16 key smoothing, Sage FP4 packing, and valid prefix K/V preparation.
+
+The FP4 token permutation and 64-row scale swizzle match fastvideo-kernel's
+scaled_fp4_quant_permute. Both FP4 and FP8 consume the rounded BF16 key.
+"""
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+from .swiglu_quant import _rcp
+from .value_quant import _left_four
+
+
+@triton.jit
+def _smooth_pack(K, V, MEAN, INV, K4, S4, K8, V8, KB: tl.constexpr, KS: tl.constexpr, KH: tl.constexpr,
+                 VB: tl.constexpr, VS: tl.constexpr, VH: tl.constexpr,
+                 B: tl.constexpr, N: tl.constexpr, H: tl.constexpr, VALID: tl.constexpr,
+                 R: tl.constexpr, VALUE_INDEX=None, GATHER_VALUE: tl.constexpr = False,
+                 PREPARE_PREFIX: tl.constexpr = True, COPY_VALUES: tl.constexpr = True,
+                 POOL=None, TILE_VALID=None, POOL_KEYS: tl.constexpr = False,
+                 PACKED_SCALE=None, TILE_INVERSE=None, Remote=None, LOCAL_ROWS: tl.constexpr = 0,
+                 PAIR: tl.constexpr = False, RANK: tl.constexpr = 0, RECEIVED_KEYS: tl.constexpr = False):
+    row = tl.program_id(0).to(tl.int64) * R + tl.arange(0, R)
+    batch = tl.program_id(1)
+    head = tl.program_id(2)
+    col = tl.arange(0, 128)
+    rem = row % 32
+    permuted = row // 32 * 32 + (rem // 8) * 2 + (rem % 8 // 2) * 8 + rem % 2
+    source = row if POOL_KEYS else permuted
+    destination4 = permuted if POOL_KEYS else row
+    mask = row[:, None] < N
+    if RECEIVED_KEYS:
+        packed_row = tl.load(TILE_INVERSE + source, row < N, other=-1)
+        valid_key = mask & (packed_row[:, None] >= 0)
+        peer = packed_row // LOCAL_ROWS
+        packed = (peer * 3 + 1) * LOCAL_ROWS + packed_row % LOCAL_ROWS
+        if PAIR:
+            offset = (LOCAL_ROWS + packed_row % LOCAL_ROWS)[:, None] * H * 128 + head * 128 + col[None, :]
+            pointer = tl.where(peer[:, None] == RANK, K + offset, Remote + offset)
+        else:
+            pointer = K + packed[:, None] * H * 128 + head * 128 + col[None, :]
+        code = tl.load(pointer, valid_key, other=0.0).to(tl.float32)
+        transport_scale = tl.load(PACKED_SCALE + packed * H + head, (row < N) & (packed_row >= 0), other=0)
+        key = (code * transport_scale[:, None]).to(tl.bfloat16).to(tl.float32)
+    else:
+        key = tl.load(K + batch * KB + source[:, None] * KS + head * KH + col[None, :],
+                      mask, other=0).to(tl.float32)
+    if POOL_KEYS:
+        raw = tl.trans(tl.reshape(key, (R // 16, 16, 128)), (1, 2, 0))
+        if R == 64:
+            even, odd = tl.split(tl.reshape(raw, (16, 128, 2, 2)))
+            x0, x2 = tl.split(even)
+            x1, x3 = tl.split(odd)
+            acc = ((x0 + x1) + x2) + x3
+        else:
+            even, odd = tl.split(tl.reshape(raw, (16, 128, 4, 2)))
+            x04, x26 = tl.split(tl.reshape(even, (16, 128, 2, 2)))
+            x15, x37 = tl.split(tl.reshape(odd, (16, 128, 2, 2)))
+            x0, x4 = tl.split(x04)
+            x1, x5 = tl.split(x15)
+            x2, x6 = tl.split(x26)
+            x3, x7 = tl.split(x37)
+            acc = ((((((x0 + x1) + x2) + x3) + x4) + x5) + x6) + x7
+        even, odd = tl.split(tl.reshape(tl.trans(acc), (128, 4, 2, 2)))
+        y0, y2 = tl.split(even)
+        y1, y3 = tl.split(odd)
+        y0, y1 = _left_four(y0, 128), _left_four(y1, 128)
+        y2, y3 = _left_four(y2, 128), _left_four(y3, 128)
+        total = (y0 + y2) + (y1 + y3)
+        tile = tl.program_id(0)
+        count = tl.load(TILE_VALID + tile).to(tl.float32)
+        tl.store(POOL + ((batch * (N // R) + tile) * H + head) * 128 + col, tl.div_rn(total, count))
+    mean = tl.load(MEAN + (batch * H + head) * 128 + col).to(tl.float32)
+    key = (key - mean[None, :]).to(tl.bfloat16).to(tl.float32)
+    groups = tl.reshape(key, (R, 8, 16))
+    maximum = tl.max(tl.abs(groups), 2)
+    sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+    inv4 = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+    q = tl.reshape(groups * inv4[:, :, None], (R, 64, 2))
+    lo, hi = tl.split(q)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q,$1,$2;
+        mov.b32 $0,{q,q,q,q};
+    }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+    c4 = tl.arange(0, 64)
+    tl.store(K4 + ((batch * H + head) * N + destination4[:, None]) * 64 + c4[None, :], packed, mask)
+    g = tl.arange(0, 8)
+    si = ((destination4[:, None] // 64) * 512 + (g[None, :] // 4) * 256 + g[None, :] % 4
+          + (destination4[:, None] % 64 // 16) * 4 + (destination4[:, None] % 16) * 16)
+    tl.store(S4 + (batch * H + head) * N * 8 + si, sf, mask)
+    if PREPARE_PREFIX:
+        destination = tl.load(INV + source, row < N, other=-1).to(tl.int64)
+        valid = mask & (destination[:, None] >= 0)
+        scale8 = tl.max(tl.abs(key), 1) / 448.0
+        inv8 = tl.where(scale8 == 0, 0, 1.0 / scale8)
+        key8 = tl.minimum(tl.maximum(key * inv8[:, None], -448.0), 448.0).to(tl.float8e4nv).to(tl.float32)
+        offset = ((batch * VALID + destination[:, None]) * H + head) * 128 + col[None, :]
+        tl.store(K8 + offset, key8 * scale8[:, None], valid)
+        if COPY_VALUES:
+            value_source = tl.load(VALUE_INDEX + source, row < N, other=-1) if GATHER_VALUE else source
+            value = tl.load(V + batch * VB + value_source[:, None] * VS + head * VH + col[None, :], valid, other=0)
+            tl.store(V8 + offset, value, valid)
+
+
+def _prepare_keys(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor, n_valid: int,
+    value_index: torch.Tensor | None = None, pool_valid: torch.Tensor | None = None, *, tile_elems: int = 64,
+    prefix_value: torch.Tensor | None = None,
+    received_keys: tuple[torch.Tensor | PeerSlices, torch.Tensor, torch.Tensor] | None = None,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    """Return Sage-packed K and FP8-rounded K/BF16 V gathered into prefix order."""
+    if (not k.is_cuda or k.dtype != torch.bfloat16 or v.dtype != k.dtype or v.device != k.device
+            or k.ndim != 4 or v.ndim != 4 or k.shape[0] != v.shape[0] or k.shape[2:] != v.shape[2:]
+            or (value_index is None and k.shape != v.shape) or k.shape[-1] != 128 or k.shape[1] % 128
+            or k.stride(-1) != 1 or v.stride(-1) != 1):
+        raise ValueError("Key quantization requires CUDA BF16 BSHD with 128-wide heads and padded token rows")
+    if (mean.shape != (k.shape[0], 1, k.shape[2], 128) or mean.dtype != k.dtype
+            or mean.device != k.device or not mean.is_contiguous()):
+        raise ValueError("Key quantization requires a contiguous BF16 token mean")
+    if (inverse.shape != (k.shape[1],) or inverse.device != k.device
+            or inverse.dtype not in (torch.int32, torch.int64)):
+        raise ValueError("Key quantization requires an integer inverse row map on the input device")
+    if value_index is not None and (value_index.shape != (k.shape[1],) or value_index.device != v.device
+                                    or value_index.dtype not in (torch.int32, torch.int64)
+                                    or not value_index.is_contiguous()):
+        raise ValueError("Compact prefix values require an integer tile-to-packed row map")
+    if torch.is_grad_enabled():
+        raise RuntimeError("Key quantization is inference-only")
+    batch, rows, heads, _ = k.shape
+    k4 = torch.empty((batch, heads, rows, 64), device=k.device, dtype=torch.uint8)
+    s4 = torch.empty((batch, heads, rows, 8), device=k.device, dtype=torch.float8_e4m3fn)
+    k8 = torch.empty((batch, n_valid, heads, 128), device=k.device, dtype=torch.bfloat16)
+    if prefix_value is not None and (prefix_value.shape != k8.shape or prefix_value.dtype != k8.dtype
+                                     or prefix_value.device != k8.device or not prefix_value.is_contiguous()):
+        raise ValueError("Prepared prefix values must match the contiguous BF16 prefix-key output")
+    v8 = torch.empty_like(k8) if prefix_value is None else prefix_value
+    pooled = None
+    if pool_valid is not None:
+        if tile_elems not in (64, 128):
+            raise ValueError("Key pooling supports 64- and 128-token tiles")
+        if (pool_valid.shape != (rows // tile_elems,) or pool_valid.device != k.device
+                or pool_valid.dtype not in (torch.int32, torch.int64) or not pool_valid.is_contiguous()):
+            raise ValueError("Key pooling requires contiguous integer tile validity counts")
+        pooled = torch.empty((batch, rows // tile_elems, heads, 128), device=k.device, dtype=torch.float32)
+    key_input, remote = k, k
+    transport_scale, tile_inverse, local_rows, rank, pair = None, None, 0, 0, None
+    if received_keys is not None:
+        payload, transport_scale, tile_inverse = received_keys
+        pair = payload if isinstance(payload, PeerSlices) else None
+        template = pair.template if pair is not None else payload
+        if (batch != 1 or pool_valid is not None or prefix_value is None
+                or template.ndim != 5 or template.shape[1] != 3 or template.shape[2] < 1
+                or template.shape[3:] != (heads, 128) or template.dtype != torch.float8_e4m3fn
+                or template.device != k.device or not template.is_contiguous()
+                or transport_scale.shape != template.shape[:-1] or transport_scale.dtype != torch.float32
+                or transport_scale.device != k.device or not transport_scale.is_contiguous()
+                or tile_inverse.shape != (rows,) or tile_inverse.dtype not in (torch.int32, torch.int64)
+                or tile_inverse.device != k.device or not tile_inverse.is_contiguous()):
+            raise ValueError("Received keys require FP8 QKV, FP32 scales, an inverse row map and prepared prefix values")
+        if pair is not None and (template.shape[0] != 2 or pair.rank not in (0, 1)
+                or any(t.shape != template.shape[1:] or t.dtype != template.dtype
+                       or t.device != template.device or not t.is_contiguous() for t in (pair.local, pair.remote))):
+            raise ValueError("Direct-local received keys require matching contiguous two-rank peer slices")
+        key_input = pair.local if pair is not None else template
+        remote = pair.remote if pair is not None else template
+        local_rows, rank = template.shape[2], pair.rank if pair is not None else 0
+    block = tile_elems if pooled is not None else 4
+    _smooth_pack[(triton.cdiv(rows, block), batch, heads)](
+        key_input, v, mean, inverse, k4, s4, k8, v8, *k.stride()[:3], *v.stride()[:3],
+        batch, rows, heads, n_valid, block, VALUE_INDEX=value_index, GATHER_VALUE=value_index is not None,
+        POOL=pooled, TILE_VALID=pool_valid, POOL_KEYS=pooled is not None, COPY_VALUES=prefix_value is None,
+        PACKED_SCALE=transport_scale, TILE_INVERSE=tile_inverse, Remote=remote, LOCAL_ROWS=local_rows,
+        PAIR=pair is not None, RANK=rank, RECEIVED_KEYS=received_keys is not None,
+        num_warps=8 if pooled is not None else 4,
+    )
+    return (k4, s4), k8, v8, pooled
+
+
+def quantize_smoothed_keys(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor, n_valid: int,
+    value_index: torch.Tensor | None = None, *, prefix_value: torch.Tensor | None = None,
+    received_keys: tuple[torch.Tensor | PeerSlices, torch.Tensor, torch.Tensor] | None = None,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]:
+    """Prepare keys and optionally consume values already in dense-prefix order.
+
+    When ``prefix_value`` is supplied, V is used only for metadata validation;
+    its contents are not read. The supplied buffer is returned without a copy.
+    With ``received_keys=(payload, scale, tile_inverse)``, K is metadata only;
+    keys are expanded from the current call's FP8 packets with the same BF16
+    rounding before smoothing. This requires prepared prefix values and batch one.
+    """
+    packed, prefix_key, prefix_value, _ = _prepare_keys(
+        k, v, mean, inverse, n_valid, value_index, prefix_value=prefix_value, received_keys=received_keys,
+    )
+    return packed, prefix_key, prefix_value
+
+
+def pool_quantize_smoothed_keys(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor, n_valid: int,
+    valid: torch.Tensor, value_index: torch.Tensor | None = None, *, tile_elems: int = 64,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Also return exact ATen-order FP32 means of unsmoothed 64- or 128-token tiles."""
+    packed, prefix_key, prefix_value, pooled = _prepare_keys(
+        k, v, mean, inverse, n_valid, value_index, valid, tile_elems=tile_elems,
+    )
+    assert pooled is not None
+    return packed, prefix_key, prefix_value, pooled.permute(0, 2, 1, 3)
+
+
+def quantize_smoothed_keys_only(k: torch.Tensor, mean: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack the rounded BF16 key shift without allocating dense-prefix outputs."""
+    if (not k.is_cuda or k.dtype != torch.bfloat16 or k.ndim != 4 or k.shape[-1] != 128
+            or k.shape[1] % 128 or k.stride(-1) != 1):
+        raise ValueError("Key-only preparation requires CUDA BF16 head-128 with padded token rows")
+    if (mean.shape != (k.shape[0], 1, k.shape[2], 128) or mean.dtype != k.dtype
+            or mean.device != k.device or not mean.is_contiguous()):
+        raise ValueError("Key-only preparation requires a contiguous BF16 token mean")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Key-only preparation requires eager inference")
+    batch, rows, heads, _ = k.shape
+    codes = torch.empty((batch, heads, rows, 64), device=k.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, rows, 8), device=k.device, dtype=torch.float8_e4m3fn)
+    _smooth_pack[(triton.cdiv(rows, 4), batch, heads)](
+        k, k, mean, k, codes, scales, k, k, *k.stride()[:3], *k.stride()[:3],
+        batch, rows, heads, 0, 4, PREPARE_PREFIX=False, num_warps=4,
+    )
+    return codes, scales

fastvideo/models/dits/minimax_h3_fusions/kv_pool.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Share the compact value read across prefix attention, pooling and native FP4 packing."""
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+from .tile_pool import _tile_sum
+
+
+@triton.jit
+def _smooth_pack(K, V, MEAN, INV, K4, S4, K8, V8, KB: tl.constexpr, KS: tl.constexpr, KH: tl.constexpr,
+                 VB: tl.constexpr, VS: tl.constexpr, VH: tl.constexpr,
+                 B: tl.constexpr, N: tl.constexpr, H: tl.constexpr, VALID: tl.constexpr,
+                 R: tl.constexpr, VALUE_INDEX=None, GATHER_VALUE: tl.constexpr = False,
+                 PREPARE_PREFIX: tl.constexpr = True,
+                 POOL=None, TILE_VALID=None, POOL_KEYS: tl.constexpr = False,
+                 VP=None, V4=None, VS4=None, HEAD_FIRST: tl.constexpr = False):
+    row = tl.program_id(0).to(tl.int64) * R + tl.arange(0, R)
+    batch = tl.program_id(1)
+    head = tl.program_id(2)
+    col = tl.arange(0, 128)
+    rem = row % 32
+    permuted = row // 32 * 32 + (rem // 8) * 2 + (rem % 8 // 2) * 8 + rem % 2
+    source = row if POOL_KEYS else permuted
+    destination4 = permuted if POOL_KEYS else row
+    mask = row[:, None] < N
+    key = tl.load(K + batch * KB + source[:, None] * KS + head * KH + col[None, :],
+                  mask, other=0).to(tl.float32)
+    if POOL_KEYS:
+        total = _tile_sum(key, R)
+        tile = tl.program_id(0)
+        count = tl.load(TILE_VALID + tile).to(tl.float32)
+        tl.store(POOL + ((batch * (N // R) + tile) * H + head) * 128 + col, tl.div_rn(total, count))
+    mean = tl.load(MEAN + (batch * H + head) * 128 + col).to(tl.float32)
+    key = (key - mean[None, :]).to(tl.bfloat16).to(tl.float32)
+    groups = tl.reshape(key, (R, 8, 16))
+    maximum = tl.max(tl.abs(groups), 2)
+    sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+    inv4 = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+    q = tl.reshape(groups * inv4[:, :, None], (R, 64, 2))
+    lo, hi = tl.split(q)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q,$1,$2;
+        mov.b32 $0,{q,q,q,q};
+    }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+    c4 = tl.arange(0, 64)
+    tl.store(K4 + ((batch * H + head) * N + destination4[:, None]) * 64 + c4[None, :], packed, mask)
+    g = tl.arange(0, 8)
+    si = ((destination4[:, None] // 64) * 512 + (g[None, :] // 4) * 256 + g[None, :] % 4
+          + (destination4[:, None] % 64 // 16) * 4 + (destination4[:, None] % 16) * 16)
+    tl.store(S4 + (batch * H + head) * N * 8 + si, sf, mask)
+    if PREPARE_PREFIX:
+        destination = tl.load(INV + source, row < N, other=-1).to(tl.int64)
+        valid = mask & (destination[:, None] >= 0)
+        scale8 = tl.max(tl.abs(key), 1) / 448.0
+        inv8 = tl.where(scale8 == 0, 0, 1.0 / scale8)
+        key8 = tl.minimum(tl.maximum(key * inv8[:, None], -448.0), 448.0).to(tl.float8e4nv).to(tl.float32)
+        offset = ((batch * VALID + destination[:, None]) * H + head) * 128 + col[None, :]
+        if HEAD_FIRST:
+            offset = ((batch * H + head) * VALID + destination[:, None]) * 128 + col[None, :]
+        tl.store(K8 + offset, key8 * scale8[:, None], valid)
+        value_source = tl.load(VALUE_INDEX + source, row < N, other=-1) if GATHER_VALUE else source
+        value = tl.load(V + batch * VB + value_source[:, None] * VS + head * VH + col[None, :],
+                        mask & (value_source[:, None] >= 0), other=0)
+        tl.store(V8 + offset, value, valid)
+
+
+    # V was already gathered for dense prefix attention; keep its rounded bits.
+    vf = value.to(tl.float32)
+    vg = tl.reshape(vf, (R // 16, 16, 128))
+    maximum_v = tl.max(tl.abs(vg), 1)
+    sfv = (maximum_v * (1.0 / 6.0)).to(tl.float8e4nv)
+    invv = tl.where(sfv.to(tl.float32) == 0, 0, _rcp(sfv.to(tl.float32)))
+    scaled_v = tl.reshape(vg * invv[:, None, :], (R // 2, 2, 128))
+    vlo, vhi = tl.split(tl.trans(scaled_v, (0, 2, 1)))
+    vcodes = tl.inline_asm_elementwise("""{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q,$1,$2;
+        mov.b32 $0,{q,q,q,q};
+    }""", constraints='=r,f,f', args=[vhi, vlo], dtype=tl.uint8, is_pure=True, pack=1)
+    vr = tl.arange(0, R // 2)
+    voffset = (((batch * H + head) * 128 + col[None, :]) * (N // 2)
+               + tile * (R // 2) + vr[:, None])
+    tl.store(V4 + voffset, vcodes)
+    group_v = tile * (R // 16) + tl.arange(0, R // 16)
+    soffset = ((batch * H + head) * 128 * (N // 16)
+               + (col[None, :] // 64) * 64 * (N // 16)
+               + (group_v[:, None] // 4) * 256 + group_v[:, None] % 4
+               + (col[None, :] % 64 // 16) * 4 + (col[None, :] % 16) * 16)
+    tl.store(VS4 + soffset, sfv)
+    vtotal = _tile_sum(vf, R)
+    tl.store(VP + ((batch * (N // R) + tile) * H + head) * 128 + col, tl.div_rn(vtotal, count))
+
+
+def pool_quantize_prefix_kv(
+    k: torch.Tensor, v: torch.Tensor, mean: torch.Tensor, inverse: torch.Tensor,
+    n_valid: int, valid: torch.Tensor, value_index: torch.Tensor | None, *, tile_elems: int = 64,
+) -> tuple[
+    tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor],
+    tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
+]:
+    """Return exact pooled/native KV plus dense-prefix KV with contiguous head sequences.
+
+    Keys are already normalized in tile order. Values remain in compact projection
+    order with a row map, or padded projection order without one. A row map
+    zero-fills missing source rows; direct inputs retain every padded value. Each value is read
+    once for the prefix BF16 tensor, FP32 tile mean, and transposed FP4 codes/scales.
+    """
+    if (not k.is_cuda or k.ndim != 4 or v.ndim != 4 or k.dtype != torch.bfloat16
+            or v.dtype != k.dtype or v.device != k.device or k.shape[0] != v.shape[0]
+            or k.shape[2:] != v.shape[2:] or k.shape[-1] != 128 or k.shape[1] % 128
+            or tile_elems not in (64, 128) or (value_index is None and v.shape[1] != k.shape[1])
+            or k.stride(-1) != 1 or v.stride(-1) != 1):
+        raise ValueError("Combined KV preparation requires CUDA BF16 BSHD with aligned head-128 keys")
+    batch, rows, heads, _ = k.shape
+    if (mean.shape != (batch, 1, heads, 128) or mean.dtype != k.dtype or mean.device != k.device
+            or not mean.is_contiguous() or not 0 < n_valid <= rows):
+        raise ValueError("Combined KV preparation requires a BF16 key mean and valid prefix extent")
+    for index, size in ((inverse, rows), (value_index, rows), (valid, rows // tile_elems)):
+        if index is None:
+            continue
+        if (index.shape != (size,) or index.device != k.device or not index.is_contiguous()
+                or index.dtype not in (torch.int32, torch.int64)):
+            raise ValueError("Combined KV preparation requires contiguous integer row maps and tile counts")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Combined KV preparation requires eager inference")
+    k4 = torch.empty((batch, heads, rows, 64), device=k.device, dtype=torch.uint8)
+    s4 = torch.empty((batch, heads, rows, 8), device=k.device, dtype=torch.float8_e4m3fn)
+    k8 = torch.empty((batch, heads, n_valid, 128), device=k.device, dtype=torch.bfloat16).transpose(1, 2)
+    v8 = torch.empty_like(k8)
+    kp = torch.empty((batch, rows // tile_elems, heads, 128), device=k.device, dtype=torch.float32)
+    vp = torch.empty_like(kp)
+    v4 = torch.empty((batch, heads, 128, rows // 2), device=k.device, dtype=torch.uint8)
+    vs4 = torch.empty((batch, heads, 128, rows // 16), device=k.device, dtype=torch.float8_e4m3fn)
+    _smooth_pack[(rows // tile_elems, batch, heads)](
+        k, v, mean, inverse, k4, s4, k8, v8, *k.stride()[:3], *v.stride()[:3],
+        batch, rows, heads, n_valid, tile_elems, VALUE_INDEX=value_index, GATHER_VALUE=value_index is not None,
+        POOL=kp, TILE_VALID=valid, POOL_KEYS=True, VP=vp, V4=v4, VS4=vs4, HEAD_FIRST=True, num_warps=4 if tile_elems == 64 else 8,
+    )
+    return ((k4, s4), k8, v8, kp.permute(0, 2, 1, 3)), (vp.permute(0, 2, 1, 3), (v4, vs4))

fastvideo/models/dits/minimax_h3_fusions/modulation_quant.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact residual modulation fused with NVFP4 activation preparation."""
+import torch
+import triton
+import triton.language as tl
+
+from .modulation import (
+    _num_warps,
+    _require_forward_only,
+    _require_triton_cuda,
+    _row_addressable,
+    _validate_contract,
+    _validate_residual_branch,
+)
+from .swiglu_quant import _rcp as rcp
+
+
+@triton.jit
+def _modulate_quant(
+    hidden_out_ptr,
+    Q, S, SF,
+    residual_ptr,
+    branch_ptr,
+    gate_ptr,
+    weight_ptr,
+    scale_ptr,
+    shift_ptr,
+    index_ptr,
+    n_cols,
+    n_index,
+    eps,
+    stride_input_row,
+    stride_gate_row,
+    stride_scale_row,
+    stride_shift_row,
+    BLOCK: tl.constexpr,
+    WIDTH: tl.constexpr,
+    ADD_RESIDUAL: tl.constexpr,
+    ROUND_RESIDUAL: tl.constexpr = False,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.arange(0, BLOCK)
+    mask = cols < n_cols
+    input_offsets = row * stride_input_row + cols
+    table_row = tl.load(index_ptr + row % n_index).to(tl.int64)
+
+    hidden = tl.load(residual_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32)
+    if ADD_RESIDUAL:
+        branch = tl.load(branch_ptr + input_offsets, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(
+            gate_ptr + table_row * stride_gate_row + cols,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        hidden = hidden + gate * branch
+        tl.store(hidden_out_ptr + input_offsets, hidden.to(hidden_out_ptr.dtype.element_ty), mask=mask)
+
+    if ROUND_RESIDUAL:
+        hidden = hidden.to(tl.bfloat16).to(tl.float32)
+    variance = tl.sum(hidden * hidden, axis=0) / n_cols
+    normed = hidden * tl.math.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    scale = tl.load(
+        scale_ptr + table_row * stride_scale_row + cols,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    shift = tl.load(
+        shift_ptr + table_row * stride_shift_row + cols,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    output = normed * weight * (1.0 + scale) + shift
+    y = tl.where(mask, output.to(tl.bfloat16).to(tl.float32), 0.0)
+    grouped = tl.reshape(y, (BLOCK // 16, 16))
+    amax = tl.max(tl.abs(grouped), 1)
+    global_sf = tl.load(SF)
+    sf8 = (global_sf * (amax * rcp(6.0))).to(tl.float8e4nv)
+    inv = tl.where(amax != 0, rcp(sf8.to(tl.float32) * rcp(global_sf)), 0.0)
+    quant = tl.reshape(grouped * inv[:, None], (BLOCK // 2, 2))
+    lo, hi = tl.split(quant)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+        mov.b32 $0, {q,q,q,q};
+    }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+    qc = tl.arange(0, BLOCK // 2)
+    tl.store(Q + row * (WIDTH // 2) + qc, packed, qc < WIDTH // 2)
+    group = tl.arange(0, BLOCK // 16)
+    scale_tile = (row // 128) * (WIDTH // 64) + group // 4
+    si = ((scale_tile * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4
+    tl.store(S + si, sf8.to(tl.uint8, bitcast=True), group < WIDTH // 16)
+
+
+def residual_modulate_quantize_nvfp4(
+    residual: torch.Tensor, branch: torch.Tensor, gate: torch.Tensor, weight: torch.Tensor,
+    scale: torch.Tensor, shift: torch.Tensor, index: torch.Tensor, eps: float,
+    global_scale: torch.Tensor,
+    *, round_residual_before_norm: bool = False,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """Return the updated residual and packed FP4 modulated activation.
+
+    The FP32 residual participates in RMSNorm before its BF16 output cast,
+    matching the existing fused residual/modulation kernel exactly. The
+    modulated activation is rounded to BF16 before FP4 preparation.
+    Across transformer blocks, round the residual to BF16 before RMSNorm too,
+    preserving the separate block output's rounding boundary.
+    """
+    _validate_residual_branch(residual, branch)
+    _validate_contract(residual, weight, (gate, scale, shift), index, eps)
+    _require_forward_only(residual, branch, gate, weight, scale, shift)
+    _require_triton_cuda(residual)
+    width = residual.shape[-1]
+    if residual.dtype != torch.bfloat16 or width % 64:
+        raise ValueError("NVFP4 modulation requires BF16 with a width divisible by 64")
+    if (global_scale.numel() != 1 or global_scale.dtype != torch.float32
+            or global_scale.device != residual.device):
+        raise ValueError("NVFP4 modulation requires a scalar FP32 scale on the input device")
+    flat = residual.reshape(-1, width).contiguous()
+    branch = branch.reshape(-1, width).contiguous()
+    gate, scale, shift = (_row_addressable(t) for t in (gate, scale, shift))
+    rows = flat.shape[0]
+    hidden = torch.empty_like(flat)
+    codes = torch.empty((rows, width // 2), dtype=torch.uint8, device=residual.device)
+    scales = torch.zeros((triton.cdiv(rows, 128) * 128, width // 16), dtype=torch.uint8,
+                         device=residual.device)
+    block = triton.next_power_of_2(width)
+    _modulate_quant[(rows,)](
+        hidden, codes, scales, global_scale, flat, branch, gate, weight.contiguous(),
+        scale, shift, index.contiguous(), width, index.numel(), eps, flat.stride(0),
+        gate.stride(0), scale.stride(0), shift.stride(0), BLOCK=block, WIDTH=width,
+        ADD_RESIDUAL=True, ROUND_RESIDUAL=round_residual_before_norm, num_warps=_num_warps(block),
+    )
+    return hidden.view_as(residual), (codes, scales, global_scale)
+
+
+def modulate_quantize_nvfp4(
+    x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor,
+    index: torch.Tensor, eps: float, global_scale: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """RMSNorm and row-indexed modulation followed by exact BF16-to-NVFP4 preparation."""
+    _validate_contract(x, weight, (scale, shift), index, eps)
+    _require_forward_only(x, weight, scale, shift)
+    _require_triton_cuda(x)
+    width = x.shape[-1]
+    if x.dtype != torch.bfloat16 or width % 64:
+        raise ValueError("NVFP4 modulation requires BF16 with a width divisible by 64")
+    if global_scale.numel() != 1 or global_scale.dtype != torch.float32 or global_scale.device != x.device:
+        raise ValueError("NVFP4 modulation requires a scalar FP32 scale on the input device")
+    flat = x.reshape(-1, width).contiguous()
+    scale, shift = (_row_addressable(t) for t in (scale, shift))
+    rows = flat.shape[0]
+    codes = torch.empty((rows, width // 2), dtype=torch.uint8, device=x.device)
+    scales = torch.zeros((triton.cdiv(rows, 128) * 128, width // 16), dtype=torch.uint8, device=x.device)
+    block = triton.next_power_of_2(width)
+    _modulate_quant[(rows,)](
+        flat, codes, scales, global_scale, flat, flat, scale, weight.contiguous(),
+        scale, shift, index.contiguous(), width, index.numel(), eps, flat.stride(0),
+        scale.stride(0), scale.stride(0), shift.stride(0), BLOCK=block, WIDTH=width,
+        ADD_RESIDUAL=False, num_warps=_num_warps(block),
+    )
+    return codes, scales, global_scale

fastvideo/models/dits/minimax_h3_fusions/prefix_attention.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Optional exact query-tile specialization for the measured GB10 Omni prefix."""
+from functools import lru_cache
+
+import torch
+
+
+@lru_cache(maxsize=1)
+def _native():
+    try:
+        import _fastvideo_h3_prefix
+    except ImportError:
+        return None
+    return _fastvideo_h3_prefix
+
+
+def try_prefix_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor | None:
+    """BHSD operands; keep Torch SDPA outside the validated eager configuration."""
+    if (torch.is_grad_enabled() or torch.compiler.is_compiling()
+            or q.shape != (1, 56, 15360, 128) or k.shape != (1, 56, 29678, 128)
+            or v.shape != k.shape or not q.is_cuda or k.device != q.device or v.device != q.device
+            or q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype
+            or q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1
+            or torch.cuda.get_device_capability(q.device) != (12, 1)
+            or torch.version.git_version != '7661cd9c6b841b62b7f411aa52ec51f05457263b'):
+        return None
+    native = _native()
+    return None if native is None else native.forward(q, k, v)

fastvideo/models/dits/minimax_h3_fusions/prefix_quant.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FP8 prefix-key preparation fused with valid-key/value gathering."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _prefix_kv(K, V, IDX, KO, VO, KB: tl.constexpr, KS: tl.constexpr, KH: tl.constexpr,
+               VB: tl.constexpr, VS: tl.constexpr, VH: tl.constexpr,
+               B: tl.constexpr, N: tl.constexpr, H: tl.constexpr, D: tl.constexpr, R: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64) * R + tl.arange(0, R)
+    col = tl.arange(0, D)
+    batch = row // (N * H)
+    token = row // H % N
+    head = row % H
+    source = tl.load(IDX + token, row < B * N * H, other=0).to(tl.int64)
+    mask = row[:, None] < B * N * H
+    k = tl.load(K + batch[:, None] * KB + source[:, None] * KS + head[:, None] * KH + col[None, :],
+                mask, other=0).to(tl.float32)
+    scale = tl.max(tl.abs(k), 1) / 448.0
+    inv = tl.where(scale == 0, 0, 1.0 / scale)
+    rounded = tl.minimum(tl.maximum(k * inv[:, None], -448.0), 448.0).to(tl.float8e4nv).to(tl.float32)
+    tl.store(KO + row[:, None] * D + col[None, :], rounded * scale[:, None], mask)
+    v = tl.load(V + batch[:, None] * VB + source[:, None] * VS + head[:, None] * VH + col[None, :],
+                mask, other=0)
+    tl.store(VO + row[:, None] * D + col[None, :], v, mask)
+
+
+def quantize_gather_prefix_kv(
+    k: torch.Tensor, v: torch.Tensor, index: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply the existing per-head FP8 round trip to K and gather valid K/V rows.
+
+    K remains BF16 at the SDPA boundary. The output keeps the existing BSHD
+    layout, and each head's amax is independent of the other token rows.
+    """
+    if (not k.is_cuda or k.dtype != torch.bfloat16 or v.dtype != k.dtype or v.device != k.device
+            or k.shape != v.shape or k.ndim != 4 or k.shape[-1] != 128
+            or k.stride(-1) != 1 or v.stride(-1) != 1):
+        raise ValueError("Prefix quantization requires CUDA BF16 BSHD tensors with contiguous 128-wide heads")
+    if index.ndim != 1 or index.device != k.device or index.dtype not in (torch.int32, torch.int64):
+        raise ValueError("Prefix quantization requires a CUDA integer row index")
+    if torch.is_grad_enabled():
+        raise RuntimeError("Prefix quantization is inference-only")
+    batch, _, heads, dim = k.shape
+    rows = index.numel()
+    ko = torch.empty((batch, rows, heads, dim), device=k.device, dtype=k.dtype)
+    vo = torch.empty_like(ko)
+    if rows:
+        _prefix_kv[(triton.cdiv(batch * rows * heads, 4),)](
+            k, v, index, ko, vo, *k.stride()[:3], *v.stride()[:3], batch, rows, heads, dim, 4, num_warps=4,
+        )
+    return ko, vo

fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py

@@ -19,6 +19,7 @@ except ImportError:  # pragma: no cover - exercised only in environments without
 
 
 if HAVE_TRITON:
+    from .swiglu_quant import _rcp
 
     @triton.jit
     def _qknorm_partial_rope_kernel(
@@ -34,42 +35,84 @@ if HAVE_TRITON:
         seq_len,
         eps,
         BLOCK_SIZE: tl.constexpr,
+        ROWS: tl.constexpr,
+        ROWS_PER_PROGRAM: tl.constexpr,
+        codes_ptr=None,
+        scales_ptr=None,
+        PACK_FP4: tl.constexpr = False,
+        START: tl.constexpr = 0,
+        input_index_ptr=None,
+        INPUT_ROWS: tl.constexpr = 0,
+        GATHER_INPUT: tl.constexpr = False,
     ):
         # int64, like the sibling kernels: with int32 program ids,
         # ``row * head_dim`` wraps once the flattened input reaches 2**31
         # elements (H3's 56 heads x 128 head_dim crosses that at
         # batch*seq >= 299_593 tokens per rank) and the loads/stores below
         # become out-of-bounds. ``seq_index`` inherits int64 from ``row``.
-        row = tl.program_id(0).to(tl.int64)
-        seq_index = (row // num_heads) % seq_len
-        cols = tl.arange(0, BLOCK_SIZE)
-        head_mask = cols < head_dim
-        row_offset = row * head_dim
+        # Keep each row's four-warp reduction order while amortizing CTA scheduling.
+        for i in tl.static_range(ROWS_PER_PROGRAM):
+            row = tl.program_id(0).to(tl.int64) * ROWS_PER_PROGRAM + i
+            seq_index = (row // num_heads) % seq_len
+            cols = tl.arange(0, BLOCK_SIZE)
+            head_mask = (cols < head_dim) & (row < ROWS)
+            row_offset = row * head_dim
+            if GATHER_INPUT:
+                source = tl.load(input_index_ptr + seq_index, row < ROWS, other=-1).to(tl.int64)
+                input_offset = (((row // (seq_len * num_heads)) * INPUT_ROWS + source) * num_heads
+                                + row % num_heads) * head_dim
+                input_valid = source >= 0
+            else:
+                input_offset, input_valid = row_offset, True
+            x = tl.load(x_ptr + input_offset + cols, mask=head_mask & input_valid, other=0.0).to(tl.float32)
+            variance = tl.sum(x * x, axis=0) / head_dim
+            inv_rms = tl.math.rsqrt(variance + eps)
+            weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
+            normalized = x * inv_rms * weight
 
-        x = tl.load(x_ptr + row_offset + cols, mask=head_mask, other=0.0).to(tl.float32)
-        variance = tl.sum(x * x, axis=0) / head_dim
-        inv_rms = tl.math.rsqrt(variance + eps)
-        weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
-        normalized = x * inv_rms * weight
+            rotary_mask = (cols < rotary_dim) & (row < ROWS)
+            first_half = cols < half_rotary_dim
+            partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
+            partner_x = tl.load(
+                x_ptr + input_offset + partner_col,
+                mask=rotary_mask & input_valid,
+                other=0.0,
+            ).to(tl.float32)
+            partner_weight = tl.load(weight_ptr + partner_col, mask=rotary_mask, other=0.0).to(tl.float32)
+            partner_normalized = partner_x * inv_rms * partner_weight
+            rotated = tl.where(first_half, -partner_normalized, partner_normalized)
 
-        rotary_mask = cols < rotary_dim
-        first_half = cols < half_rotary_dim
-        partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
-        partner_x = tl.load(
-            x_ptr + row_offset + partner_col,
-            mask=rotary_mask,
-            other=0.0,
-        ).to(tl.float32)
-        partner_weight = tl.load(weight_ptr + partner_col, mask=rotary_mask, other=0.0).to(tl.float32)
-        partner_normalized = partner_x * inv_rms * partner_weight
-        rotated = tl.where(first_half, -partner_normalized, partner_normalized)
-
-        table_offset = seq_index * rotary_dim + cols
-        cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
-        sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
-        rotary_output = normalized * cos + rotated * sin
-        output = tl.where(rotary_mask, rotary_output, normalized)
-        tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+            table_offset = seq_index * rotary_dim + cols
+            cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
+            sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
+            rotary_output = normalized * cos + rotated * sin
+            output = tl.where(rotary_mask, rotary_output, normalized)
+            tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+            if PACK_FP4:
+                rounded = output.to(tl.bfloat16).to(tl.float32)
+                grouped = tl.reshape(rounded, (BLOCK_SIZE // 16, 16))
+                maximum = tl.max(tl.abs(grouped), 1)
+                sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+                inv = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+                values = tl.reshape(grouped * inv[:, None], (BLOCK_SIZE // 2, 2))
+                lo, hi = tl.split(values)
+                packed = tl.inline_asm_elementwise('''{
+                    .reg .b8 q;
+                    cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+                    mov.b32 $0, {q,q,q,q};
+                }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+                batch = row // (seq_len * num_heads)
+                head = row % num_heads
+                c4 = tl.arange(0, BLOCK_SIZE // 2)
+                offset = ((batch * num_heads + head) * (seq_len - START) + seq_index - START) * (head_dim // 2)
+                tl.store(codes_ptr + offset + c4, packed, (row < ROWS) & (seq_index >= START) & (c4 < head_dim // 2))
+                group = tl.arange(0, BLOCK_SIZE // 16)
+                qi = seq_index - START
+                si = ((qi // 64) * (64 * head_dim // 16) + (group // 4) * 256 + group % 4
+                      + (qi % 64 // 16) * 4 + (qi % 16) * 16)
+                scale_offset = (batch * num_heads + head) * (seq_len - START) * (head_dim // 16) + si
+                tl.store(scales_ptr + scale_offset, sf,
+                         (row < ROWS) & (seq_index >= START) & (group < head_dim // 16))
 
 
 _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
@@ -143,7 +186,9 @@ def _fused_qknorm_rope_impl(
     flat_x = x.reshape(-1, head_dim).contiguous()
     flat_out = torch.empty_like(flat_x)
     block_size = 1 << (head_dim - 1).bit_length()
-    _qknorm_partial_rope_kernel[(flat_x.shape[0], )](
+    rows_per_program = (8 if head_dim == 128 and x.dtype == torch.bfloat16
+                        and torch.cuda.get_device_capability(x.device)[0] >= 10 else 1)
+    _qknorm_partial_rope_kernel[(triton.cdiv(flat_x.shape[0], rows_per_program), )](
         flat_out,
         flat_x,
         weight,
@@ -156,6 +201,8 @@ def _fused_qknorm_rope_impl(
         seq_len,
         eps,
         BLOCK_SIZE=block_size,
+        ROWS=flat_x.shape[0],
+        ROWS_PER_PROGRAM=rows_per_program,
         num_warps=4,
     )
     return flat_out.view(batch, seq_len, num_heads, head_dim)
@@ -214,3 +261,74 @@ def fused_qknorm_rope(
 
 
 __all__ = ["HAVE_TRITON", "fused_qknorm_rope"]
+
+
+def qknorm_rope_quantize_fp4(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+    eps: float, start: int = 0,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Keep the BF16 norm/RoPE result and pack its suffix for native Sage FP4.
+
+    The suffix begins on a 128-token boundary. Its codes and 64-row scale
+    swizzle are allocated contiguously, as required by the native attention
+    kernel. BF16 rounding precedes the per-16-element scale calculation.
+    """
+    batch, seq_len, heads, dim, rotary_dim = _validate_inputs(x, weight, cos, sin, eps)
+    if (not x.is_cuda or x.dtype != torch.bfloat16 or dim != 128 or seq_len % 128
+            or start < 0 or start >= seq_len or start % 128):
+        raise ValueError("FP4 Q packing requires CUDA BF16 head-128 and aligned sequence/suffix lengths")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("FP4 Q packing requires eager inference")
+    x, weight, cos, sin = (t.contiguous() for t in (x, weight, cos, sin))
+    out = torch.empty_like(x)
+    codes = torch.empty((batch, heads, seq_len - start, 64), device=x.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, seq_len - start, 8), device=x.device, dtype=torch.float8_e4m3fn)
+    _qknorm_partial_rope_kernel[(triton.cdiv(batch * seq_len * heads, 4),)](
+        out, x, weight, cos, sin, dim, rotary_dim, rotary_dim // 2, heads, seq_len, eps,
+        BLOCK_SIZE=128, ROWS=batch * seq_len * heads, ROWS_PER_PROGRAM=4,
+        codes_ptr=codes, scales_ptr=scales, PACK_FP4=True, START=start, num_warps=4,
+    )
+    return out, (codes, scales)
+
+
+def gather_qknorm_rope(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float,
+    input_index: torch.Tensor, *, fp4_start: int | None = None,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+    """Normalize compact projected rows while gathering them into padded tile order.
+
+    Negative indices denote zero projection rows. Pooling and sparse attention
+    still receive every original tile row; only the linear projections omit
+    padding. An optional aligned suffix is packed directly for native FP4 Q.
+    """
+    if x.ndim != 4 or input_index.ndim != 1:
+        raise ValueError("Gathered normalization requires BSHD input and a one-dimensional row map")
+    rows = input_index.numel()
+    shape_view = x[:, :1].expand(-1, rows, -1, -1)
+    batch, _, heads, dim, rotary_dim = _validate_inputs(shape_view, weight, cos, sin, eps)
+    if (not x.is_cuda or x.dtype != torch.bfloat16 or dim != 128 or rows % 128
+            or input_index.device != x.device or input_index.dtype not in (torch.int32, torch.int64)
+            or not input_index.is_contiguous()):
+        raise ValueError("Gathered normalization requires CUDA BF16 head-128 and an aligned integer row map")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Gathered normalization requires eager inference")
+    quantized = fp4_start is not None
+    start = 0 if fp4_start is None else fp4_start
+    if start < 0 or start >= rows or start % 128:
+        raise ValueError("The FP4 query suffix must begin at a valid 128-token boundary")
+    x, weight, cos, sin = (t.contiguous() for t in (x, weight, cos, sin))
+    out = torch.empty((batch, rows, heads, dim), device=x.device, dtype=x.dtype)
+    codes, scales = out, out
+    if quantized:
+        codes = torch.empty((batch, heads, rows - start, 64), device=x.device, dtype=torch.uint8)
+        scales = torch.empty((batch, heads, rows - start, 8), device=x.device, dtype=torch.float8_e4m3fn)
+    from .qknorm_rope_vector import qknorm_rope_vector_kernel
+
+    unroll = 16 if quantized else 8
+    qknorm_rope_vector_kernel[(triton.cdiv(batch * rows * heads, unroll),)](
+        out, x, weight, cos, sin, dim, rotary_dim, rotary_dim // 2, heads, rows, eps,
+        BLOCK_SIZE=128, ROWS=batch * rows * heads, ROWS_PER_PROGRAM=unroll,
+        codes_ptr=codes, scales_ptr=scales, PACK_FP4=quantized, START=start,
+        input_index_ptr=input_index, INPUT_ROWS=x.shape[1], GATHER_INPUT=True, num_warps=4,
+    )
+    return out, (codes, scales) if quantized else None

fastvideo/models/dits/minimax_h3_fusions/qknorm_rope_vector.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Vectorized BF16 head-128 normalization with the original four-warp sum order.
+
+The scalar-row kernel reduces within each 32-lane warp first, then combines
+its four warp sums. A plain vectorized ``tl.sum`` chooses another tree and
+changes BF16 rounding at a few inputs. Explicit butterfly permutations keep
+that original order while processing multiple rows in registers together.
+"""
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+
+
+@triton.jit
+def qknorm_rope_vector_kernel(
+    out_ptr,
+    x_ptr,
+    weight_ptr,
+    cos_ptr,
+    sin_ptr,
+    head_dim: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    half_rotary_dim: tl.constexpr,
+    num_heads: tl.constexpr,
+    seq_len: tl.constexpr,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+    ROWS: tl.constexpr,
+    ROWS_PER_PROGRAM: tl.constexpr,
+    codes_ptr=None,
+    scales_ptr=None,
+    PACK_FP4: tl.constexpr = False,
+    START: tl.constexpr = 0,
+    input_index_ptr=None,
+    INPUT_ROWS: tl.constexpr = 0,
+    GATHER_INPUT: tl.constexpr = False,
+):
+    # int64, like the sibling kernels: with int32 program ids,
+    # ``row * head_dim`` wraps once the flattened input reaches 2**31
+    # elements (H3's 56 heads x 128 head_dim crosses that at
+    # batch*seq >= 299_593 tokens per rank) and the loads/stores below
+    # become out-of-bounds. ``seq_index`` inherits int64 from ``row``.
+    tl.static_assert(BLOCK_SIZE == 128 and head_dim == 128)
+    row = tl.program_id(0).to(tl.int64) * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)[:, None]
+    seq_index = (row // num_heads) % seq_len
+    cols = tl.arange(0, BLOCK_SIZE)[None, :]
+    head_mask = (cols < head_dim) & (row < ROWS)
+    row_offset = row * head_dim
+    if GATHER_INPUT:
+        source = tl.load(input_index_ptr + seq_index, row < ROWS, other=-1).to(tl.int64)
+        input_offset = (((row // (seq_len * num_heads)) * INPUT_ROWS + source) * num_heads
+                        + row % num_heads) * head_dim
+        input_valid = source >= 0
+    else:
+        input_offset, input_valid = row_offset, True
+    x = tl.load(x_ptr + input_offset + cols, mask=head_mask & input_valid, other=0.0).to(tl.float32)
+    # Warp-local offsets 16..1, then the four-warp tree at offsets 64, 32.
+    squared = x * x
+    for shift in tl.static_range(5):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (16 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    for shift in tl.static_range(2):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (64 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    variance = tl.gather(squared, tl.full((ROWS_PER_PROGRAM, 1), 0, tl.int32), axis=1) / head_dim
+    inv_rms = tl.math.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
+    normalized = x * inv_rms * weight
+
+    rotary_mask = (cols < rotary_dim) & (row < ROWS)
+    first_half = cols < half_rotary_dim
+    partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
+    # Reuse the identical normalized partner instead of loading and normalizing it again.
+    partner_normalized = tl.gather(
+        normalized, tl.broadcast_to(partner_col, (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1,
+    )
+    rotated = tl.where(first_half, -partner_normalized, partner_normalized)
+
+    table_offset = seq_index * rotary_dim + cols
+    cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
+    sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
+    rotary_output = normalized * cos + rotated * sin
+    output = tl.where(rotary_mask, rotary_output, normalized)
+    tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+    if PACK_FP4:
+        rounded = output.to(tl.bfloat16).to(tl.float32)
+        grouped = tl.reshape(rounded, (ROWS_PER_PROGRAM, BLOCK_SIZE // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 2)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inv = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+        values = tl.reshape(grouped * inv[:, :, None], (ROWS_PER_PROGRAM, BLOCK_SIZE // 2, 2))
+        lo, hi = tl.split(values)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        batch = row // (seq_len * num_heads)
+        head = row % num_heads
+        c4 = tl.arange(0, BLOCK_SIZE // 2)[None, :]
+        offset = ((batch * num_heads + head) * (seq_len - START) + seq_index - START) * (head_dim // 2)
+        tl.store(codes_ptr + offset + c4, packed, (row < ROWS) & (seq_index >= START) & (c4 < head_dim // 2))
+        group = tl.arange(0, BLOCK_SIZE // 16)[None, :]
+        qi = seq_index - START
+        si = ((qi // 64) * (64 * head_dim // 16) + (group // 4) * 256 + group % 4
+              + (qi % 64 // 16) * 4 + (qi % 16) * 16)
+        scale_offset = (batch * num_heads + head) * (seq_len - START) * (head_dim // 16) + si
+        tl.store(scales_ptr + scale_offset, sf,
+                 (row < ROWS) & (seq_index >= START) & (group < head_dim // 16))
+
+

fastvideo/models/dits/minimax_h3_fusions/quant_tile.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Reuse one packed-input quantization for tile-order QKV and packed-order gate."""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _tile_quantized(Q, S, Index, OQ, OS, WIDTH: tl.constexpr, ROWS: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    src = tl.load(Index + row).to(tl.int64)
+    val = tl.load(Q + src * (WIDTH // 2) + col, (src >= 0) & (col < WIDTH // 2), other=0)
+    tl.store(OQ + row * (WIDTH // 2) + col, val, col < WIDTH // 2)
+    group = tl.program_id(1) * (BLOCK // 8) + tl.arange(0, BLOCK // 8)
+    src_scale = ((((src // 128) * (WIDTH // 64) + group // 4) * 32 + src % 32) * 4 + (src // 32) % 4) * 4 + group % 4
+    dst_scale = ((((row // 128) * (WIDTH // 64) + group // 4) * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4
+    scale = tl.load(S + src_scale, (src >= 0) & (group < WIDTH // 16), other=0)
+    tl.store(OS + dst_scale, scale, group < WIDTH // 16)
+
+
+def tile_quantized(q: torch.Tensor, s: torch.Tensor, inverse: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather NVFP4 rows, including 128x4-swizzled scales; -1 indices produce zero padding."""
+    rows, width = inverse.numel(), q.shape[-1] * 2
+    oq = torch.empty((rows, width // 2), device=q.device, dtype=q.dtype)
+    os = torch.empty((rows, width // 16), device=s.device, dtype=s.dtype)
+    _tile_quantized[(rows, triton.cdiv(width // 2, 512))](q, s, inverse, oq, os, width, rows, 512)
+    return oq, os
+

fastvideo/models/dits/minimax_h3_fusions/query_pool.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Prepare exact pooled/FP4 queries while retaining BF16 only for dense prefix attention."""
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+from .tile_pool import _tile_sum
+
+
+@triton.jit
+def _query_pool_kernel(
+    out_ptr,
+    x_ptr,
+    weight_ptr,
+    cos_ptr,
+    sin_ptr,
+    head_dim: tl.constexpr,
+    rotary_dim: tl.constexpr,
+    half_rotary_dim: tl.constexpr,
+    num_heads: tl.constexpr,
+    seq_len: tl.constexpr,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+    ROWS: tl.constexpr,
+    ROWS_PER_PROGRAM: tl.constexpr,
+    TILE_START: tl.constexpr = 0,
+    codes_ptr=None,
+    scales_ptr=None,
+    PACK_FP4: tl.constexpr = False,
+    START: tl.constexpr = 0,
+    input_index_ptr=None,
+    INPUT_ROWS: tl.constexpr = 0,
+    GATHER_INPUT: tl.constexpr = False,
+    pool_ptr=None,
+    valid_ptr=None,
+    PREFIX: tl.constexpr = 0,
+    PREFIX_ONLY: tl.constexpr = False,
+    PREFIX_FP8: tl.constexpr = False,
+):
+    # int64, like the sibling kernels: with int32 program ids,
+    # ``row * head_dim`` wraps once the flattened input reaches 2**31
+    # elements (H3's 56 heads x 128 head_dim crosses that at
+    # batch*seq >= 299_593 tokens per rank) and the loads/stores below
+    # become out-of-bounds. ``seq_index`` inherits int64 from ``row``.
+    tl.static_assert(BLOCK_SIZE == 128 and head_dim == 128)
+    tile = tl.program_id(0).to(tl.int64) + TILE_START
+    current_head = tl.program_id(1)
+    current_batch = tl.program_id(2)
+    seq_index = tile * ROWS_PER_PROGRAM + tl.arange(0, ROWS_PER_PROGRAM)[:, None]
+    row = (current_batch * seq_len + seq_index) * num_heads + current_head
+    cols = tl.arange(0, BLOCK_SIZE)[None, :]
+    head_mask = (cols < head_dim) & (row < ROWS)
+    row_offset = row * head_dim
+    if GATHER_INPUT:
+        source = tl.load(input_index_ptr + seq_index, row < ROWS, other=-1).to(tl.int64)
+        input_offset = (((row // (seq_len * num_heads)) * INPUT_ROWS + source) * num_heads
+                        + row % num_heads) * head_dim
+        input_valid = source >= 0
+    else:
+        input_offset, input_valid = row_offset, True
+    x = tl.load(x_ptr + input_offset + cols, mask=head_mask & input_valid, other=0.0).to(tl.float32)
+    # Warp-local offsets 16..1, then the four-warp tree at offsets 64, 32.
+    squared = x * x
+    for shift in tl.static_range(5):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (16 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    for shift in tl.static_range(2):
+        other = tl.gather(squared, tl.broadcast_to(cols ^ (64 >> shift), (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1)
+        squared = squared + other
+    variance = tl.gather(squared, tl.full((ROWS_PER_PROGRAM, 1), 0, tl.int32), axis=1) / head_dim
+    inv_rms = tl.math.rsqrt(variance + eps)
+    weight = tl.load(weight_ptr + cols, mask=head_mask, other=0.0).to(tl.float32)
+    normalized = x * inv_rms * weight
+
+    rotary_mask = (cols < rotary_dim) & (row < ROWS)
+    first_half = cols < half_rotary_dim
+    partner_col = tl.where(first_half, cols + half_rotary_dim, cols - half_rotary_dim)
+    # Reuse the identical normalized partner instead of loading and normalizing it again.
+    partner_normalized = tl.gather(
+        normalized, tl.broadcast_to(partner_col, (ROWS_PER_PROGRAM, BLOCK_SIZE)), axis=1,
+    )
+    rotated = tl.where(first_half, -partner_normalized, partner_normalized)
+
+    table_offset = seq_index * rotary_dim + cols
+    cos = tl.load(cos_ptr + table_offset, mask=rotary_mask, other=1.0).to(tl.float32)
+    sin = tl.load(sin_ptr + table_offset, mask=rotary_mask, other=0.0).to(tl.float32)
+    rotary_output = normalized * cos + rotated * sin
+    output = tl.where(rotary_mask, rotary_output, normalized)
+    if PREFIX_ONLY:
+        output_offset = ((current_batch * PREFIX + seq_index) * num_heads + current_head) * head_dim
+        if PREFIX_FP8:
+            if tile * ROWS_PER_PROGRAM < PREFIX:
+                prefix_value = output.to(tl.bfloat16).to(tl.float32)
+                scale8 = tl.max(tl.abs(prefix_value), 1) / 448.0
+                inv8 = tl.where(scale8 == 0, 0, 1.0 / scale8)
+                prefix_value = tl.minimum(tl.maximum(prefix_value * inv8[:, None], -448.0), 448.0)
+                prefix_value = prefix_value.to(tl.float8e4nv).to(tl.float32) * scale8[:, None]
+                tl.store(out_ptr + output_offset + cols, prefix_value, mask=head_mask & (seq_index < PREFIX))
+        else:
+            tl.store(out_ptr + output_offset + cols, output.to(out_ptr.dtype.element_ty),
+                     mask=head_mask & (seq_index < PREFIX))
+    else:
+        tl.store(out_ptr + row_offset + cols, output.to(out_ptr.dtype.element_ty), mask=head_mask)
+    rounded_pool = output.to(tl.bfloat16).to(tl.float32)
+    total = _tile_sum(rounded_pool, ROWS_PER_PROGRAM)
+    count = tl.load(valid_ptr + tile).to(tl.float32)
+    pool_offset = ((current_batch * (seq_len // ROWS_PER_PROGRAM) + tile) * num_heads + current_head) * 128
+    tl.store(pool_ptr + pool_offset + tl.arange(0, 128), tl.div_rn(total, count))
+    if PACK_FP4:
+        rounded = output.to(tl.bfloat16).to(tl.float32)
+        grouped = tl.reshape(rounded, (ROWS_PER_PROGRAM, BLOCK_SIZE // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 2)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inv = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+        values = tl.reshape(grouped * inv[:, :, None], (ROWS_PER_PROGRAM, BLOCK_SIZE // 2, 2))
+        lo, hi = tl.split(values)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        batch = row // (seq_len * num_heads)
+        head = row % num_heads
+        c4 = tl.arange(0, BLOCK_SIZE // 2)[None, :]
+        offset = ((batch * num_heads + head) * (seq_len - START) + seq_index - START) * (head_dim // 2)
+        tl.store(codes_ptr + offset + c4, packed, (row < ROWS) & (seq_index >= START) & (c4 < head_dim // 2))
+        group = tl.arange(0, BLOCK_SIZE // 16)[None, :]
+        qi = seq_index - START
+        si = ((qi // 64) * (64 * head_dim // 16) + (group // 4) * 256 + group % 4
+              + (qi % 64 // 16) * 4 + (qi % 16) * 16)
+        scale_offset = (batch * num_heads + head) * (seq_len - START) * (head_dim // 16) + si
+        tl.store(scales_ptr + scale_offset, sf,
+                 (row < ROWS) & (seq_index >= START) & (group < head_dim // 16))
+
+
+
+
+def prepare_pooled_queries(
+    x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float,
+    input_index: torch.Tensor | None, valid: torch.Tensor, prefix_rows: int, *, tile_elems: int = 64,
+    prefix_fp8: bool = False, split_prefix: bool = False,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+    """Return BF16 prefix queries, FP4 video queries, and exact FP32 tile means.
+
+    The native suffix starts at the preceding 128-token boundary, so its FP4
+    data also includes a half-prefix block when the prefix ends on tile 64.
+    An absent row map reads padded projection rows directly.
+    Omni tile-128 callers may also fuse the existing prefix FP8 round trip.
+    Splitting the prefix and suffix launches specializes away unused packing
+    and prefix stores while retaining every consumed value and tile sum.
+    Pooling and native codes always use the original BF16-normalized queries.
+    No BF16 video-query intermediate is needed by either attention branch.
+    """
+    from .qknorm_rope import _validate_inputs
+
+    if x.ndim != 4 or (input_index is not None and input_index.ndim != 1):
+        raise ValueError("Pooled queries require BSHD input and a one-dimensional row map")
+    rows = x.shape[1] if input_index is None else input_index.numel()
+    shape_view = x[:, :1].expand(-1, rows, -1, -1)
+    batch, _, heads, dim, rotary_dim = _validate_inputs(shape_view, weight, cos, sin, eps)
+    if (not x.is_cuda or x.dtype != torch.bfloat16 or dim != 128 or rows % 128
+            or tile_elems not in (64, 128) or (prefix_fp8 and tile_elems != 128)
+            or (input_index is not None and (input_index.device != x.device
+                or input_index.dtype not in (torch.int32, torch.int64) or not input_index.is_contiguous()))):
+        raise ValueError("Pooled queries require CUDA BF16 head-128 and an aligned integer row map")
+    if (valid.shape != (rows // tile_elems,) or valid.device != x.device
+            or valid.dtype not in (torch.int32, torch.int64) or not valid.is_contiguous()
+            or prefix_rows <= 0 or prefix_rows >= rows or prefix_rows % tile_elems):
+        raise ValueError("Pooled queries require tile validity and a nonempty aligned prefix")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Pooled query preparation requires eager inference")
+    if split_prefix and (tile_elems != 128 or not prefix_fp8):
+        raise ValueError("Split query preparation requires tile-128 FP8 prefix queries")
+    x, weight, cos, sin = (t.contiguous() for t in (x, weight, cos, sin))
+    start = prefix_rows // 128 * 128
+    prefix = torch.empty((batch, prefix_rows, heads, dim), device=x.device, dtype=x.dtype)
+    codes = torch.empty((batch, heads, rows - start, 64), device=x.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, rows - start, 8), device=x.device, dtype=torch.float8_e4m3fn)
+    pooled = torch.empty((batch, rows // tile_elems, heads, dim), device=x.device, dtype=torch.float32)
+    launches = ((0, rows // tile_elems, True, prefix_fp8, prefix_rows),)
+    if split_prefix:
+        prefix_tiles = prefix_rows // tile_elems
+        launches = ((0, prefix_tiles, False, True, prefix_rows),
+                    (prefix_tiles, rows // tile_elems - prefix_tiles, True, False, 0))
+    for first_tile, tiles, pack, fp8, keep_prefix in launches:
+        _query_pool_kernel[(tiles, heads, batch)](
+            prefix, x, weight, cos, sin, dim, rotary_dim, rotary_dim // 2, heads, rows, eps,
+            BLOCK_SIZE=128, ROWS=batch * rows * heads, ROWS_PER_PROGRAM=tile_elems, TILE_START=first_tile,
+            codes_ptr=codes, scales_ptr=scales, PACK_FP4=pack, START=start,
+            input_index_ptr=input_index, INPUT_ROWS=x.shape[1], GATHER_INPUT=input_index is not None,
+            pool_ptr=pooled, valid_ptr=valid, PREFIX=keep_prefix, PREFIX_ONLY=True, PREFIX_FP8=fp8, num_warps=8,
+        )
+    return prefix, (codes, scales), pooled.permute(0, 2, 1, 3)

fastvideo/models/dits/minimax_h3_fusions/sp_exchange.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Dequantize an H3 Ulysses exchange directly into the VSA tile layout."""
+
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+from .swiglu_quant import _rcp as rcp
+
+
+@triton.jit
+def _unpack_fp8_tiles(Payload, Scale, Inverse, Out, LOCAL_ROWS: tl.constexpr,
+                      TILE_ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr, BLOCK: tl.constexpr,
+                      Codes=None, QueryScales=None, PACK_QUERY: tl.constexpr = False, START: tl.constexpr = 0):
+    tile_row = tl.program_id(0).to(tl.int64)
+    projection = tl.program_id(1)
+    column = tl.arange(0, BLOCK)
+    packed_row = tl.load(Inverse + tile_row).to(tl.int64)
+    source = packed_row // LOCAL_ROWS
+    local_row = packed_row % LOCAL_ROWS
+    source_row = (source * 3 + projection) * LOCAL_ROWS + local_row
+    valid = (packed_row >= 0) & (column < HEADS * DIM)
+    value = tl.load(Payload + source_row * HEADS * DIM + column, valid, other=0.0).to(tl.float32)
+    scale = tl.load(Scale + source_row * HEADS + column // DIM, valid, other=0)
+    decoded = (value * scale).to(tl.bfloat16)
+    destination = (projection * TILE_ROWS + tile_row) * HEADS * DIM + column
+    tl.store(Out + destination, decoded, column < HEADS * DIM)
+    if PACK_QUERY and projection == 0 and tile_row >= START:
+        grouped = tl.reshape(decoded.to(tl.float32), (BLOCK // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 1)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inverse = tl.where(sf.to(tl.float32) == 0, 0, rcp(sf.to(tl.float32)))
+        values = tl.reshape(grouped * inverse[:, None], (BLOCK // 2, 2))
+        lo, hi = tl.split(values)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        qrow = tile_row - START
+        qrows: tl.constexpr = TILE_ROWS - START
+        code_col = tl.arange(0, BLOCK // 2)
+        head, local_code = code_col // (DIM // 2), code_col % (DIM // 2)
+        tl.store(Codes + (head * qrows + qrow) * (DIM // 2) + local_code,
+                  packed, code_col < HEADS * DIM // 2)
+        group = tl.arange(0, BLOCK // 16)
+        head, local_group = group // (DIM // 16), group % (DIM // 16)
+        swizzled = ((qrow // 64) * (64 * DIM // 16) + (local_group // 4) * 256 + local_group % 4
+                    + (qrow % 64 // 16) * 4 + (qrow % 16) * 16)
+        tl.store(QueryScales + head * qrows * (DIM // 16) + swizzled, sf, group < HEADS * DIM // 16)
+
+
+def unpack_fp8_tiles(payload: torch.Tensor, scale: torch.Tensor, inverse: torch.Tensor,
+                     output: torch.Tensor, *, query_start: int | None = None
+                     ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Map `[W, 3, local_rows, Hs, D]` FP8 receives to `[3, tile_rows, Hs, D]` BF16.
+
+    `inverse[tile_row]` is the global packed row, or -1 for VSA padding. Every
+    output element is written, including padding, so a reused buffer cannot
+    retain stale values. The float32 multiply and BF16 rounding match the
+    separate unpack and index-copy operations. An optional aligned query suffix
+    also receives native Sage FP4 codes and scales after that BF16 rounding.
+    """
+    _, projections, local_rows, heads, dim = payload.shape
+    assert projections == 3 and payload.dtype == torch.float8_e4m3fn
+    assert scale.dtype == torch.float32 and output.dtype == torch.bfloat16
+    assert payload.is_contiguous() and scale.is_contiguous() and output.is_contiguous()
+    assert output.shape == (3, inverse.numel(), heads, dim)
+    if query_start is not None:
+        if dim != 128 or inverse.numel() % 128 or query_start % 128 or not 0 <= query_start < inverse.numel():
+            raise ValueError("Fused query packing requires head-128 and an aligned token suffix")
+        query_rows = inverse.numel() - query_start
+        codes = torch.empty((1, heads, query_rows, dim // 2), device=payload.device, dtype=torch.uint8)
+        query_scales = torch.empty((1, heads, query_rows, dim // 16), device=payload.device, dtype=torch.float8_e4m3fn)
+    else:
+        codes, query_scales = output, output
+    block = triton.next_power_of_2(heads * dim)
+    _unpack_fp8_tiles[(inverse.numel(), 3)](
+        payload, scale, inverse, output, local_rows, inverse.numel(), heads, dim, block,
+        Codes=codes, QueryScales=query_scales, PACK_QUERY=query_start is not None,
+        START=query_start if query_start is not None else 0, num_warps=8 if block >= 2048 else 4,
+    )
+    return output if query_start is None else (output, (codes, query_scales))
+
+
+@triton.jit
+def _unpack_fp8_gate(Payload, Scale, Compressed, Gate, RowTile, Out, Codes, Scales, GlobalScale,
+                     ROWS: tl.constexpr, HEADS: tl.constexpr, WORLD: tl.constexpr,
+                     TILES: tl.constexpr, DIM: tl.constexpr, BLOCK: tl.constexpr, QUANTIZED: tl.constexpr, CHUNK: tl.constexpr, CHUNKS: tl.constexpr,
+                     Remote=None, PAIR: tl.constexpr = False, RANK: tl.constexpr = 0):
+    row = tl.program_id(0).to(tl.int64)
+    column = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    head, dim = column // DIM, column % DIM
+    source, local_head = head // HEADS, head % HEADS
+    source_head = (source * ROWS + row) * HEADS + local_head
+    valid = column < WORLD * HEADS * DIM
+    if PAIR:
+        local_offset = (row * HEADS + local_head) * DIM + dim
+        pointer = tl.where(source == RANK, Payload + local_offset, Remote + local_offset)
+    else:
+        pointer = Payload + source_head * DIM + dim
+    value = tl.load(pointer, valid, other=0.0).to(tl.float32)
+    scale = tl.load(Scale + source_head, valid, other=0.0)
+    # Retain the materialized unpack's BF16 rounding before applying the gate.
+    decoded = (value * scale).to(tl.bfloat16).to(tl.float32)
+    tile = tl.load(RowTile + row).to(tl.int64)
+    compressed = tl.load(Compressed + (head * TILES + tile) * DIM + dim, valid, other=0.0).to(tl.float32)
+    full_head = source * HEADS * CHUNKS + CHUNK * HEADS + local_head
+    destination = row * WORLD * HEADS * CHUNKS * DIM + full_head * DIM + dim
+    gate = tl.load(Gate + destination, valid, other=0.0).to(tl.float32)
+    if QUANTIZED:
+        width: tl.constexpr = WORLD * HEADS * CHUNKS * DIM
+        rounded = tl.where(valid, tl.fma(compressed, gate, decoded).to(tl.bfloat16).to(tl.float32), 0.0)
+        grouped = tl.reshape(rounded, (BLOCK // 16, 16))
+        maximum = tl.max(tl.abs(grouped), 1)
+        global_scale = tl.load(GlobalScale)
+        fp4_scale = (global_scale * (maximum * rcp(6.0))).to(tl.float8e4nv)
+        inverse = tl.where(maximum != 0, rcp(fp4_scale.to(tl.float32) * rcp(global_scale)), 0.0)
+        quantized = tl.reshape(grouped * inverse[:, None], (BLOCK // 2, 2))
+        lo, hi = tl.split(quantized)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        code_column = tl.program_id(1) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+        code_head = code_column // (DIM // 2)
+        full_code_head = (code_head // HEADS) * HEADS * CHUNKS + CHUNK * HEADS + code_head % HEADS
+        full_code_column = full_code_head * (DIM // 2) + code_column % (DIM // 2)
+        tl.store(Codes + row * (width // 2) + full_code_column, packed, code_column < WORLD * HEADS * DIM // 2)
+        group = tl.program_id(1) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+        group_head = group // (DIM // 16)
+        full_group_head = (group_head // HEADS) * HEADS * CHUNKS + CHUNK * HEADS + group_head % HEADS
+        full_group = full_group_head * (DIM // 16) + group % (DIM // 16)
+        scale_tile = (row // 128) * (width // 64) + full_group // 4
+        scale_index = ((scale_tile * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + full_group % 4
+        tl.store(Scales + scale_index, fp4_scale.to(tl.uint8, bitcast=True), group < WORLD * HEADS * DIM // 16)
+    else:
+        tl.store(Out + destination, tl.fma(compressed, gate, decoded), valid)
+
+
+def unpack_fp8_gate(payload: torch.Tensor | PeerSlices, scale: torch.Tensor, compressed: torch.Tensor,
+                     gate: torch.Tensor, row_tile: torch.Tensor, global_scale: torch.Tensor | None = None,
+                     *, head_chunk: int = 0, head_chunks: int = 1, output: torch.Tensor | None = None,
+                     fp4_output: tuple[torch.Tensor, torch.Tensor] | None = None,
+                     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Unpack head groups and apply the local gate without a BF16 intermediate.
+
+    Receives are `[W, rows, Hs, D]`; the gathered compression branch is
+    `[W, Hs, tiles, D]`, and the local gate/result are `[rows, W*Hs, D]`.
+    A chunked exchange writes only its corresponding heads in a shared full
+    output, preserving rank-major then chunk-major head order without copies.
+    """
+    pair = payload if isinstance(payload, PeerSlices) else None
+    payload = pair.template if pair is not None else payload
+    world, rows, heads, dim = payload.shape
+    assert payload.dtype == torch.float8_e4m3fn and scale.dtype == torch.float32
+    assert compressed.dtype == gate.dtype == torch.bfloat16
+    assert all(t.is_contiguous() for t in (payload, scale, compressed, gate, row_tile))
+    assert head_chunks > 0 and 0 <= head_chunk < head_chunks
+    assert gate.shape == (rows, world * heads * head_chunks, dim)
+    assert compressed.shape[:2] == (world, heads) and compressed.shape[-1] == dim
+    quantized = global_scale is not None
+    if quantized:
+        if output is not None or (head_chunks > 1 and fp4_output is None):
+            raise ValueError("Chunked FP4 preparation requires shared codes/scales and no BF16 output")
+        width = world * heads * head_chunks * dim
+        if width % 64 or dim % 16:
+            raise ValueError("Fused SP gate quantization requires a width divisible by 64")
+        if global_scale.numel() != 1 or global_scale.dtype != torch.float32 or global_scale.device != gate.device:
+            raise ValueError("Fused SP gate quantization requires a scalar FP32 scale on the input device")
+        if fp4_output is None:
+            codes = torch.empty((rows, width // 2), device=gate.device, dtype=torch.uint8)
+            fp4_scales = torch.zeros((triton.cdiv(rows, 128) * 128, width // 16), device=gate.device, dtype=torch.uint8)
+        else:
+            codes, fp4_scales = fp4_output
+            assert codes.shape == (rows, width // 2)
+            assert fp4_scales.shape == (triton.cdiv(rows, 128) * 128, width // 16)
+            assert all(t.dtype == torch.uint8 and t.device == gate.device and t.is_contiguous()
+                       for t in fp4_output)
+        output = gate  # Unused pointer when the kernel writes FP4 codes directly.
+    else:
+        if output is None:
+            output = torch.empty_like(gate)
+        assert output.shape == gate.shape and output.dtype == gate.dtype and output.is_contiguous()
+        codes, fp4_scales, global_scale = output, output, output
+    _unpack_fp8_gate[(rows, triton.cdiv(world * heads * dim, 512))](
+        pair.local if pair is not None else payload, scale, compressed, gate, row_tile, output, codes, fp4_scales, global_scale,
+        rows, heads, world, compressed.shape[2], dim, 512, quantized, head_chunk, head_chunks,
+        Remote=pair.remote if pair is not None else payload, PAIR=pair is not None, RANK=pair.rank if pair is not None else 0,
+    )
+    return (codes, fp4_scales, global_scale) if quantized else output

fastvideo/models/dits/minimax_h3_fusions/sp_head_pipeline.py

@@ -0,0 +1,250 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Overlap independent H3 attention head groups with FP8 Ulysses exchanges."""
+
+import torch
+
+from fastvideo.distributed.async_collectives import (
+    PeerSlices, all_gather_tensor_async, all_to_all_tensor_async, exchange_peer_tensor_async,
+)
+
+from fastvideo.attention.backends.video_sparse_attn_h3 import _build_block_mask, _pool_tiles
+from fastvideo.models.dits import minimax_h3_vsa_fp4 as base
+from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+from fastvideo.models.dits.minimax_h3_fusions.sp_qkv_pack import pack_normalized_qkv_fp8
+
+
+def _start_exchange(payload, scale, group):
+    if payload.shape[0] == 2 and payload.shape[-1] == 128:
+        packet = exchange_peer_tensor_async(payload, group)
+    else:
+        packet = all_to_all_tensor_async(payload.view(torch.uint8), group)
+    return packet, all_to_all_tensor_async(scale, group, align_bytes=16)
+
+
+def _finish_exchange(packet):
+    payload, scale = packet
+    received = payload.wait()
+    if not isinstance(received, PeerSlices):
+        received = received.view(torch.float8_e4m3fn)
+    return received, scale.wait()
+
+
+def _output_quant_method(attn):
+    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+    method = getattr(attn.to_out, "quant_method", None)
+    if (type(method) is NVFP4QuantizeMethod and method.wants_prequantized_input()
+            and method.uses_unit_activation_scale(attn.to_out) and attn.attention_head_dim % 64 == 0):
+        return method
+    return None
+
+
+def _prepare_tiles(api, layout, payload, scale, query_start):
+    if (payload.shape[-1] == 128 and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(payload.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4_permute", "scale_and_quant_fp4_transpose"))):
+        return layout.tiles_from_fp8(payload, scale, query_start=query_start)
+    return layout.tiles_from_fp8(payload, scale), None
+
+
+def _prepare_values(api, value, valid, tile, logical):
+    if (tile == 64 and logical == value.shape[1] and value.shape[-1] == 128 and value.stride(-1) == 1
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(value.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_permute"))):
+        from fastvideo.models.dits.minimax_h3_fusions.value_quant import pool_quantize_values
+        return pool_quantize_values(value, valid)
+    return None, None
+
+
+def _prepare_received(api, layout, payload, scale, query_start, valid, *, prefix_values=False, direct_keys=False):
+    template = payload.template if isinstance(payload, PeerSlices) else payload
+    if (layout.tile == 64 and layout.n_tiles * layout.tile == layout.rows and template.shape[-1] == 128
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(template.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_permute",
+                                                    "scale_and_quant_fp4_transpose"))):
+        from .sp_receive import unpack_pool_quantize_fp8
+        inverse, qkv = layout.fp8_buffers(template)
+        prefix_value, prefix = None, None
+        # The dense branch consumes gathered values and the sparse branch
+        # consumes native V codes, so raw BF16 V need not be materialized.
+        if prefix_values:
+            valid_index, prefix_inverse = layout.valid_keys_from_tiles(valid, api.HALF_N)
+            prefix_value = torch.empty((1, valid_index.numel(), *qkv.shape[2:]),
+                                       device=qkv.device, dtype=qkv.dtype)
+            prefix = (prefix_inverse, prefix_value)
+        # Native FP4 query codes supply the suffix; the raw query suffix is
+        # used only for shape/dtype metadata. Keep a full block beyond the
+        # aligned start to cover a dense prefix ending halfway through it.
+        pooled, prepared_value, prepared_query = unpack_pool_quantize_fp8(
+            payload, scale, inverse, valid, qkv, query_start,
+            query_rows=min(query_start + base._BLOCK, layout.rows), prefix_values=prefix, store_keys=not direct_keys,
+        )
+        return qkv, prepared_query, pooled, prepared_value, prefix_value
+    if isinstance(payload, PeerSlices):
+        payload = payload.materialize()
+    qkv, prepared_query = _prepare_tiles(api, layout, payload, scale, query_start)
+    return qkv, prepared_query, None, None, None
+
+
+def _pack_output(out, layout, world, rows):
+    if isinstance(out, tuple):
+        from .sp_output_pack import pack_split_fp8
+        return pack_split_fp8(out, layout.untile, world, rows)
+    return base._pack_seq_fp8(out, layout.untile, world, rows)
+
+
+def attention_pipelined(attn, hidden_states, rotary_emb, meta, sp_group, chunks, pre_quantized=None):
+    """Pipeline head groups so communication overlaps independent attention work."""
+    world, rank = sp_group.world_size, sp_group.rank_in_group
+    rows, heads, dim = hidden_states.shape[1], attn.num_attention_heads, attn.attention_head_dim
+    assert heads % (world * chunks) == 0
+    assert attn.to_gate_compress is not None and attn._gate_active()
+    layouts = getattr(meta, '_h3_fp4_pipeline_layouts', None)
+    if layouts is None or len(layouts) != chunks:
+        layouts = [base._SPTileLayout(meta, rank, rows) for _ in range(chunks)]
+        meta._h3_fp4_pipeline_layouts = layouts
+    options = getattr(attn, '_vsa_fp4_options', {})
+    early_sends = (options.get('sp_projection_overlap', False) and world == 2 and dim == 128
+                   and hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
+                   and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+                   and base._shared_qkv_gate_scale(attn, hidden_states) is not None)
+    input_gather = False
+    if early_sends and options.get('sp_input_gather', False):
+        from .sp_input_gather import supports_input_gather
+        input_gather = supports_input_gather(attn, hidden_states)
+    if input_gather:
+        from .sp_input_gather import project_gathered_inputs
+        with base.STAGES.span('pipe_project_pack'):
+            packets, gate = project_gathered_inputs(attn, hidden_states, rotary_emb, sp_group, chunks, pre_quantized)
+    elif early_sends:
+        from .sp_projection import project_and_exchange
+        with base.STAGES.span('pipe_project_pack'):
+            packets, gate_pre = project_and_exchange(attn, hidden_states, rotary_emb, sp_group.device_group,
+                                                     chunks, pre_quantized)
+    else:
+        with base.STAGES.span('pipe_project_pack'):
+            shared = base._shared_qkv_gate_projections(attn, hidden_states, pre_quantized=pre_quantized)
+            if shared is None:
+                projected = base._shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states)
+                gate_pre = None
+            else:
+                projected, gate_pre = shared
+            query, key, value = (t.unflatten(-1, (heads, dim)) for t in projected)
+            cos, sin = (t.to(query.dtype) for t in rotary_emb)
+            payload, scale = pack_normalized_qkv_fp8(
+                query[0],
+                key[0],
+                value[0],
+                attn.norm_q.weight,
+                attn.norm_k.weight,
+                cos,
+                sin,
+                attn.norm_q.eps,
+                attn.norm_k.eps,
+                world,
+                chunks,
+            )
+        packets = [_start_exchange(payload[chunk], scale[chunk], sp_group.device_group) for chunk in range(chunks)]
+    if not input_gather:
+        with base.STAGES.span('gate_proj'):
+            if gate_pre is None:
+                gate, _ = attn.to_gate_compress(hidden_states)
+            else:
+                linear = attn.to_gate_compress
+                gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
+    vbs = meta.variable_block_sizes
+    api = base._api()
+    pending = []
+    for chunk, layout in enumerate(layouts):
+        prefix_mode, prefix_rows = base._prefix_settings(attn, meta, layout.n_tiles)
+        direct_keys = bool(options.get('sp_direct_keys', False) and options.get('sp_pooled_key_mean', False)
+                           and options.get('smooth_k') is True and prefix_rows and prefix_mode == 'fp8')
+        with base.STAGES.span('pipe_receive'):
+            received, received_scale = _finish_exchange(packets[chunk])
+            qkv, prepared_query, pooled_qkv, prepared_value, prefix_value = _prepare_received(
+                api, layout, received, received_scale, prefix_rows // base._BLOCK * base._BLOCK, vbs,
+                prefix_values=bool(prefix_rows and prefix_mode == 'fp8' and options.get('smooth_k') is True),
+                direct_keys=direct_keys,
+            )
+        q_t, k_t, v_t = qkv[0:1], qkv[1:2], qkv[2:3]
+        logical = layout.n_tiles * layout.tile
+        with base.STAGES.span('pipe_attention'):
+            q_pool = pooled_qkv[0:1] if pooled_qkv is not None else _pool_tiles(q_t[:, :logical], vbs, layout.tile)
+            k_pool = pooled_qkv[1:2] if pooled_qkv is not None else _pool_tiles(k_t[:, :logical], vbs, layout.tile)
+            scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) / (dim**0.5)
+            sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity
+            mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans,
+                                     meta.span_sparsities)
+            idx, num, valid, quad = base._tile_mask_to_fp4_blocks(api, mask, layout.tile, vbs)
+            if pooled_qkv is None:
+                pooled, prepared_value = _prepare_values(api, v_t, vbs, layout.tile, logical)
+            else:
+                pooled = pooled_qkv[2:3]
+            prepared_keys = None
+            if ((options.get('sp_pooled_key_mean', False) or prefix_value is not None)
+                    and options.get('smooth_k') is True
+                    and pooled_qkv is not None and prefix_rows and prefix_mode == 'fp8'):
+                from .key_quant import quantize_smoothed_keys
+                from .sp_key_mean import key_mean_from_tiles
+                valid_index = layout.valid_key_index(valid, api.HALF_N)
+                inverse = layout.valid_key_inverse(valid, api.HALF_N)
+                mean = (key_mean_from_tiles(pooled_qkv[1:2], vbs, logical)
+                        if options.get('sp_pooled_key_mean', False) else k_t.mean(dim=1, keepdim=True))
+                key_packet = None
+                if direct_keys and prefix_value is not None:
+                    template = received.template if isinstance(received, PeerSlices) else received
+                    tile_inverse, _ = layout.fp8_buffers(template)
+                    key_packet = (received, received_scale, tile_inverse)
+                prepared_keys = quantize_smoothed_keys(k_t, v_t, mean, inverse, valid_index.numel(),
+                                                       prefix_value=prefix_value, received_keys=key_packet)
+            out = base._fp4_fine_attention(api,
+                                           q_t,
+                                           k_t,
+                                           v_t,
+                                           idx,
+                                           num,
+                                           valid,
+                                           quad,
+                                           logical=logical,
+                                           prefix_rows=prefix_rows,
+                                           precision=prefix_mode,
+                                           options=options,
+                                           layout=layout, prepared_query=prepared_query, prepared_value=prepared_value,
+                                           prepared_keys=prepared_keys, return_split=True)
+        with base.STAGES.span('pipe_output_pack'):
+            out_payload, out_scale = _pack_output(out, layout, world, rows)
+        # Queue this output before computing the next head group. The pending
+        # handles retain their inputs until the final wait and gate application.
+        out_packet = _start_exchange(out_payload, out_scale, sp_group.device_group)
+        with base.STAGES.span('gate_compress'):
+            if pooled is None:
+                pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile)
+            compressed = torch.matmul(torch.softmax(scores, dim=-1), pooled)[0].to(qkv.dtype).contiguous()
+        gathered = all_gather_tensor_async(compressed, sp_group.device_group, world)
+        pending.append((out_packet, gathered))
+    full_gate = gate[0].view(rows, heads, dim)
+    method = _output_quant_method(attn)
+    if method is None:
+        full_output = torch.empty_like(full_gate)
+    else:
+        width = heads * dim
+        codes = torch.empty((rows, width // 2), device=gate.device, dtype=torch.uint8)
+        scales = torch.zeros((((rows + 127) // 128) * 128, width // 16), device=gate.device, dtype=torch.uint8)
+    for chunk, (packet, gathered) in enumerate(pending):
+        with base.STAGES.span('pipe_output_wait_gate'):
+            received, received_scale = _finish_exchange(packet)
+            compression = gathered.wait()
+            if method is None:
+                unpack_fp8_gate(received, received_scale, compression, full_gate, layouts[chunk].local_row_tile,
+                                 head_chunk=chunk, head_chunks=chunks, output=full_output)
+            else:
+                unpack_fp8_gate(received, received_scale, compression, full_gate, layouts[chunk].local_row_tile,
+                                 method.x_global_sf, head_chunk=chunk, head_chunks=chunks, fp4_output=(codes, scales))
+    if method is not None:
+        return gate, (codes, scales, method.x_global_sf)
+    return full_output.flatten(1, 2).unsqueeze(0)

fastvideo/models/dits/minimax_h3_fusions/sp_input_gather.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Gather exact packed inputs before head-sharded two-rank H3 projections.
+
+The packed inputs already carry the shared unit activation scale. Only weight
+views are sharded; each row/head computes the same Q/K/V projection, norm/RoPE,
+and transport FP8 rounding as the sequence-sharded producer. Scale reassembly
+handles the unaligned join between independently padded 128-row scale tiles.
+"""
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import all_gather_tensor_async
+from fastvideo.distributed.projection_exchange import completed_projection_packets
+from fastvideo.layers.quantization.nvfp4_config import _mm_fp4, _mm_fp4_backend
+
+from .sp_projection import allocate
+from .sp_qkv_ordered import _normalize
+
+
+@triton.jit
+def _join_scale_rows(Source, Out, ROWS: tl.constexpr, PADDED: tl.constexpr, GROUPS: tl.constexpr,
+                     TOTAL_PADDED: tl.constexpr, BLOCK: tl.constexpr):
+    index = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    row, group = index // GROUPS, index % GROUPS
+    rank, local_row = row // ROWS, row % ROWS
+    local_tile = (local_row // 128) * (GROUPS // 4) + group // 4
+    local_index = ((local_tile * 32 + local_row % 32) * 4 + (local_row // 32) % 4) * 4 + group % 4
+    value = tl.load(Source + rank * PADDED * GROUPS + local_index, row < 2 * ROWS, other=0)
+    tile = (row // 128) * (GROUPS // 4) + group // 4
+    destination = ((tile * 32 + row % 32) * 4 + (row // 32) % 4) * 4 + group % 4
+    tl.store(Out + destination, value, index < TOTAL_PADDED * GROUPS)
+
+
+def join_scale_rows(source: torch.Tensor, rows: int) -> torch.Tensor:
+    assert source.ndim == 3 and source.shape[0] == 2 and source.is_contiguous()
+    assert source.dtype == torch.uint8 and source.shape[2] % 4 == 0
+    padded, groups = source.shape[1:]
+    assert padded == triton.cdiv(rows, 128) * 128
+    total_padded = triton.cdiv(rows * 2, 128) * 128
+    output = torch.empty((total_padded, groups), device=source.device, dtype=source.dtype)
+    _join_scale_rows[(triton.cdiv(total_padded * groups, 512), )](
+        source, output, rows, padded, groups, total_padded, 512, num_warps=4)
+    return output
+
+
+@triton.jit
+def _pack_gathered(X, Weight, Cos, Sin, Payload, Scale, LOCAL_ROWS: tl.constexpr, HEADS: tl.constexpr,
+                   CHUNKS: tl.constexpr, ROTARY: tl.constexpr, EPS: tl.constexpr, PROJECTION: tl.constexpr,
+                   BLOCK_ROWS: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
+    row, head = row_head // HEADS, row_head % HEADS
+    column = tl.arange(0, 128)
+    valid = row[:, None] < LOCAL_ROWS * 2
+    if PROJECTION == 2:
+        values = tl.load(X + row_head[:, None] * 128 + column[None, :], valid, other=0.0).to(tl.float32)
+    else:
+        values = _normalize(X, Weight, Cos, Sin, row, head, HEADS, LOCAL_ROWS * 2, 128, ROTARY, EPS, 128)
+    maximum = tl.max(tl.abs(values), 1)
+    scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    chunk_heads: tl.constexpr = HEADS // CHUNKS
+    chunk = head // chunk_heads
+    source = row // LOCAL_ROWS
+    destination = ((
+        (chunk * 2 + source) * 3 + PROJECTION) * LOCAL_ROWS + row % LOCAL_ROWS) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * 128 + column[None, :], quantized, valid)
+    tl.store(Scale + destination, scale, row < LOCAL_ROWS * 2)
+
+
+def pack_gathered(x, payload, scale, projection, weight=None, cosine=None, sine=None, epsilon=0.0) -> None:
+    chunks, world, projections, rows, heads, dim = payload.shape
+    assert world == 2 and projections == 3 and dim == 128
+    assert x.shape == (2 * rows, chunks * heads * dim) and x.is_contiguous()
+    assert payload.is_contiguous() and scale.is_contiguous()
+    if projection == 2:
+        weight, cosine, sine = x, x, x
+        rotary = 0
+    else:
+        assert weight.shape == (128, ) and cosine.shape == sine.shape
+        assert cosine.shape[0] == rows * 2 and cosine.is_contiguous() and sine.is_contiguous()
+        rotary = cosine.shape[1]
+    block_rows = 16 if projection == 2 else 8
+    _pack_gathered[(triton.cdiv(2 * rows * chunks * heads, block_rows), )](
+        x, weight, cosine, sine, payload, scale, rows, chunks * heads, chunks,
+        rotary, epsilon, projection, block_rows, num_warps=4 if projection == 2 else 8)
+
+
+def supports_input_gather(attn, hidden: torch.Tensor) -> bool:
+    """Keep other shapes/backends on the established projection exchange."""
+    if (not hidden.is_cuda or hidden.dtype != torch.bfloat16
+            or hidden.shape != (1, 7430, 5376) or attn.num_attention_heads != 56
+            or torch.cuda.get_device_capability(hidden.device) != (12, 1)
+            or _mm_fp4_backend() != "cudnn"):
+        return False
+    for linear in (attn.to_q, attn.to_k, attn.to_v):
+        weight = getattr(linear, "_nvfp4_weight", None)
+        scale = getattr(linear, "_nvfp4_weight_scale", None)
+        if (weight is None or scale is None or weight.shape != (7168, 2688)
+                or scale.shape != (7168, 336) or linear.bias is not None
+                or weight.dtype != torch.uint8 or scale.dtype != torch.uint8
+                or weight.device != hidden.device or scale.device != hidden.device
+                or not weight.is_contiguous() or not scale.is_contiguous()):
+            return False
+    return True
+
+
+def project_gathered_inputs(attn, hidden, rotary, sp_group, chunks, pre_quantized=None):
+    """Return completed head packets and the gate computed while inputs travel.
+
+    The caller checks eager, two-rank, head-128 inference and the common unit
+    activation-scale contract through _shared_qkv_gate_scale. No activations
+    or packets survive this call except through the returned consumer handles.
+    """
+    linears = (attn.to_q, attn.to_k, attn.to_v)
+    pre = pre_quantized if pre_quantized is not None else linears[0].quant_method.quantize_input(hidden)
+    a, sa, input_scale = pre
+    a, sa = a.reshape(-1, a.shape[-1]), sa.reshape(-1, sa.shape[-1])
+    cosine, sine = (item.to(hidden.dtype) for item in rotary)
+    pending = [all_gather_tensor_async(value, sp_group.device_group, 2) for value in (a, sa, cosine, sine)]
+    linear = attn.to_gate_compress
+    gate = linear.quant_method.apply(linear, hidden, linear.bias, pre_quantized=pre)
+    gathered_a, gathered_sa, gathered_cosine, gathered_sine = (item.wait() for item in pending)
+    full_a = gathered_a.flatten(0, 1)
+    full_sa = join_scale_rows(gathered_sa, a.shape[0])
+    full_cosine, full_sine = gathered_cosine.flatten(0, 1), gathered_sine.flatten(0, 1)
+    rows, heads, rank = hidden.shape[1], attn.num_attention_heads, sp_group.rank_in_group
+    payload, scale = allocate(rows, heads, 2, chunks, hidden.device)
+    for projection, linear in enumerate(linears):
+        half = linear._nvfp4_weight.shape[0] // 2
+        weight = linear._nvfp4_weight.narrow(0, rank * half, half)
+        weight_scale = linear._nvfp4_weight_scale.narrow(0, rank * half, half)
+        alpha = (linear._nvfp4_alpha / input_scale if hasattr(linear, "_nvfp4_alpha")
+                 else 1.0 / (input_scale * linear._weight_global_sf))
+        value = _mm_fp4(full_a, weight.T, full_sa, weight_scale.T, alpha, torch.bfloat16, None, backend="cudnn")
+        if projection == 2:
+            pack_gathered(value, payload, scale, projection)
+        else:
+            norm = attn.norm_q if projection == 0 else attn.norm_k
+            pack_gathered(value, payload, scale, projection, norm.weight, full_cosine, full_sine, norm.eps)
+    return completed_projection_packets(payload, scale, rank), gate

fastvideo/models/dits/minimax_h3_fusions/sp_key_mean.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Recover global key means from the SP receive's complete tile statistics.
+
+Every tile contributes its FP32 mean times its valid count, and the total is
+normalized by the padded sequence length, matching the zero-padded key tensor.
+This changes floating-point reduction order and may alter the generated sample.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _key_mean(Pool, Counts, Mean, HEAD_STRIDE: tl.constexpr, TILE_STRIDE: tl.constexpr,
+              TILES: tl.constexpr, SEQ: tl.constexpr, ROWS: tl.constexpr, COLS: tl.constexpr):
+    head = tl.program_id(0)
+    rows = tl.arange(0, ROWS)
+    cols = tl.program_id(1) * COLS + tl.arange(0, COLS)
+    pooled = tl.load(Pool + head * HEAD_STRIDE + rows[:, None] * TILE_STRIDE + cols[None, :],
+                     rows[:, None] < TILES, other=0.0)
+    count = tl.load(Counts + rows, rows < TILES, other=0).to(tl.float32)
+    pooled = tl.where(count[:, None] > 0, pooled, 0.0)
+    total = tl.sum(pooled * count[:, None], 0)
+    mean = total * (1.0 / SEQ)
+    tl.store(Mean + head * 128 + cols, mean)
+
+
+def key_mean_from_tiles(pool: torch.Tensor, counts: torch.Tensor, sequence: int) -> torch.Tensor:
+    """Return BF16 ``[1, 1, H, 128]`` means from FP32 ``[1, H, tiles, 128]``.
+
+    Counts must be the same divisors used to produce ``pool``. Padding is
+    included in ``sequence`` and contributes zero, just as in ``key.mean``.
+    """
+    batch, heads, tiles, dim = pool.shape
+    if (batch != 1 or dim != 128 or sequence != tiles * 64 or pool.stride(-1) != 1
+            or pool.dtype != torch.float32 or counts.shape != (tiles,)
+            or counts.dtype not in (torch.int32, torch.int64) or not pool.is_cuda
+            or counts.device != pool.device):
+        raise ValueError("Pooled SP key means require CUDA tile-64 FP32 statistics with 128 channels")
+    result = torch.empty((1, 1, heads, dim), device=pool.device, dtype=torch.bfloat16)
+    _key_mean[(heads, 4)](
+        pool, counts, result, pool.stride(1), pool.stride(2), tiles, sequence,
+        triton.next_power_of_2(tiles), 32, num_warps=4,
+    )
+    return result

fastvideo/models/dits/minimax_h3_fusions/sp_output_pack.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pack SP attention's split prefix/video results without a concatenation buffer."""
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _pack_split(Prefix, Video, Untile, Payload, Scale, LENGTH: tl.constexpr, ROWS: tl.constexpr, H: tl.constexpr,
+                D: tl.constexpr, PREFIX: tl.constexpr, VIDEO: tl.constexpr, START: tl.constexpr, BLOCK_H: tl.constexpr,
+                BLOCK_D: tl.constexpr, PSH: tl.constexpr, PSS: tl.constexpr, VSH: tl.constexpr, VSS: tl.constexpr):
+    row_head = tl.program_id(0) * BLOCK_H + tl.arange(0, BLOCK_H)
+    column = tl.arange(0, BLOCK_D)
+    row, head = row_head // H, row_head % H
+    tile_row = tl.load(Untile + row, row < LENGTH, other=-1)
+    valid = (row_head[:, None] < ROWS * H) & (tile_row[:, None] >= 0) & (column[None, :] < D)
+    is_prefix = tile_row < PREFIX
+    prefix = tl.load(Prefix + head[:, None] * PSH + tile_row[:, None] * PSS + column[None, :],
+                     valid & is_prefix[:, None],
+                     other=0).to(tl.float32)
+    video = tl.load(Video + head[:, None] * VSH + (tile_row[:, None] - START) * VSS + column[None, :],
+                    valid & ~is_prefix[:, None],
+                    other=0).to(tl.float32)
+    value = tl.where(is_prefix[:, None], prefix, video)
+    scale = tl.maximum(tl.max(tl.abs(value), 1) / 448.0, 1e-12)
+    quant = (value / scale[:, None]).to(tl.float8e4nv)
+    tl.store(Payload + row_head[:, None] * D + column[None, :], quant,
+             (row_head[:, None] < ROWS * H) & (column[None, :] < D))
+    tl.store(Scale + row_head, scale, row_head < ROWS * H)
+
+
+def pack_split_fp8(parts, untile, world, rows, block=4) -> tuple[torch.Tensor, torch.Tensor]:
+    prefix, video, start = parts
+    _, heads, _, dim = video.shape
+    assert prefix.dtype == video.dtype == torch.bfloat16 and prefix.shape[0] == video.shape[0] == 1
+    assert prefix.shape[1] == heads and prefix.shape[3] == dim and prefix.stride(-1) == video.stride(-1) == 1
+    assert untile.is_contiguous() and untile.dtype in (torch.int32, torch.int64)
+    assert prefix.device == video.device == untile.device and untile.numel() <= world * rows
+    assert 0 <= start <= prefix.shape[2] and video.is_cuda
+    payload = torch.empty((world, rows, heads, dim), device=video.device, dtype=torch.float8_e4m3fn)
+    scale = torch.empty(payload.shape[:-1], device=video.device, dtype=torch.float32)
+    _pack_split[(triton.cdiv(world * rows * heads, block), )](prefix,
+                                                              video,
+                                                              untile,
+                                                              payload,
+                                                              scale,
+                                                              untile.numel(),
+                                                              world * rows,
+                                                              heads,
+                                                              dim,
+                                                              prefix.shape[2],
+                                                              video.shape[2],
+                                                              start,
+                                                              block,
+                                                              triton.next_power_of_2(dim),
+                                                              prefix.stride(1),
+                                                              prefix.stride(2),
+                                                              video.stride(1),
+                                                              video.stride(2),
+                                                              num_warps=4)
+    return payload, scale
+
+

fastvideo/models/dits/minimax_h3_fusions/sp_projection.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact per-projection QKV packing for overlap with two-rank SP traffic."""
+
+import torch
+import triton
+import triton.language as tl
+
+from .sp_qkv_ordered import _normalize
+from .sp_qkv_pack import _pack_value_fp8
+
+
+@triton.jit
+def _pack_one(X, Weight, Cos, Sin, Payload, Scale, ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+              ROTARY: tl.constexpr, WORLD: tl.constexpr, CHUNKS: tl.constexpr, EPS: tl.constexpr,
+              PROJECTION: tl.constexpr, BLOCK_ROWS: tl.constexpr, BLOCK_DIM: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
+    row, head = row_head // HEADS, row_head % HEADS
+    col = tl.arange(0, BLOCK_DIM)
+    values = _normalize(X, Weight, Cos, Sin, row, head, HEADS, ROWS, DIM, ROTARY, EPS, BLOCK_DIM)
+    maximum = tl.max(tl.abs(values), 1)
+    scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    chunk = head % local_heads // chunk_heads
+    destination = ((
+        (chunk * WORLD + head // local_heads) * 3 + PROJECTION) * ROWS + row) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * DIM + col[None, :], quantized,
+             (row[:, None] < ROWS) & (col[None, :] < DIM))
+    tl.store(Scale + destination, scale, row < ROWS)
+
+
+def allocate(rows, heads, world, chunks, device) -> tuple[torch.Tensor, torch.Tensor]:
+    payload = torch.empty((chunks, world, 3, rows, heads // (world * chunks), 128),
+                          device=device,
+                          dtype=torch.float8_e4m3fn)
+    return payload, torch.empty(payload.shape[:-1], device=device, dtype=torch.float32)
+
+
+def pack_one(x, weight, cosine, sine, payload, scale, epsilon, projection) -> None:
+    chunks, world, _, rows, local_heads, dim = payload.shape
+    heads = local_heads * world * chunks
+    _pack_one[(triton.cdiv(rows * heads, 8), )](
+        x,
+        weight,
+        cosine,
+        sine,
+        payload,
+        scale,
+        rows,
+        heads,
+        dim,
+        cosine.shape[1],
+        world,
+        chunks,
+        epsilon,
+        projection,
+        8,
+        128,
+        num_warps=8,
+    )
+
+
+def pack_value(x, payload, scale) -> None:
+    chunks, world, _, rows, local_heads, dim = payload.shape
+    heads = local_heads * world * chunks
+    _pack_value_fp8[(triton.cdiv(rows * heads, 16), )](
+        x,
+        payload,
+        scale,
+        rows,
+        heads,
+        dim,
+        world,
+        chunks,
+        16,
+        128,
+        num_warps=4,
+    )
+
+
+def project_and_exchange(attn, hidden, rotary, group, chunks, pre_quantized):
+    """Start Q traffic before K's GEMM, and K traffic before V's GEMM.
+
+    The caller checks that the three NVFP4 projections and gate can share
+    packed input. Normalization, RoPE and FP8 division preserve the ordinary
+    packer's arithmetic and reduction order. Nothing is reused across calls.
+    """
+    from fastvideo.distributed.projection_exchange import ProjectionExchange
+
+    linears = (attn.to_q, attn.to_k, attn.to_v)
+    pre = pre_quantized if pre_quantized is not None else linears[0].quant_method.quantize_input(hidden)
+    rows, heads = hidden.shape[1], attn.num_attention_heads
+    cosine, sine = (item.to(hidden.dtype) for item in rotary)
+    payload, scale = allocate(rows, heads, 2, chunks, hidden.device)
+    exchange = ProjectionExchange(payload, scale, group)
+    for projection, linear in enumerate(linears):
+        value = linear.quant_method.apply(linear, hidden, linear.bias, pre_quantized=pre)
+        if projection == 2:
+            pack_value(value, payload, scale)
+        else:
+            norm = attn.norm_q if projection == 0 else attn.norm_k
+            pack_one(value, norm.weight, cosine, sine, payload, scale, norm.eps, projection)
+        exchange.start(projection)
+    return exchange.packets(), pre

fastvideo/models/dits/minimax_h3_fusions/sp_qkv_ordered.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Parallel SP QK rows with the original four-warp RMS addition order.
+
+The standalone head-128 normalization reduces 32 lanes in each warp first,
+then combines four warp sums. Explicit XOR gathers retain that addition tree
+when multiple heads share a CTA; a plain two-dimensional reduction changes it.
+"""
+import triton
+import triton.language as tl
+
+@triton.jit
+def _normalize(X, Weight, Cos, Sin, row, head, HEADS: tl.constexpr, ROWS: tl.constexpr, DIM: tl.constexpr,
+               ROTARY: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr):
+    col = tl.arange(0, BLOCK)
+    offset = (row[:, None] * HEADS + head[:, None]) * DIM
+    mask = (row[:, None] < ROWS) & (col[None, :] < DIM)
+    value = tl.load(X + offset + col[None, :], mask, other=0.0).to(tl.float32)
+    squared = value * value
+    indices = tl.broadcast_to(col[None, :], squared.shape)
+    for shift in tl.static_range(5):
+        squared = squared + tl.gather(squared, indices ^ (16 >> shift), 1)
+    for shift in tl.static_range(2):
+        squared = squared + tl.gather(squared, indices ^ (64 >> shift), 1)
+    variance = tl.sum(tl.where(col[None, :] == 0, squared, 0.0), 1) / DIM
+    inv = tl.math.rsqrt(variance + EPS)
+    weight = tl.load(Weight + col, col < DIM, other=0.0).to(tl.float32)
+    normalized = value * inv[:, None] * weight[None, :]
+    partner = tl.where(col < ROTARY // 2, col + ROTARY // 2, col - ROTARY // 2)
+    rot_mask = (row[:, None] < ROWS) & (col[None, :] < ROTARY)
+    pv = tl.load(X + offset + partner[None, :], rot_mask, other=0.0).to(tl.float32)
+    pw = tl.load(Weight + partner, col < ROTARY, other=0.0).to(tl.float32)
+    pn = pv * inv[:, None] * pw[None, :]
+    rotated = tl.where(col[None, :] < ROTARY // 2, -pn, pn)
+    cosine = tl.load(Cos + row[:, None] * ROTARY + col[None, :], rot_mask, other=1.0).to(tl.float32)
+    sine = tl.load(Sin + row[:, None] * ROTARY + col[None, :], rot_mask, other=0.0).to(tl.float32)
+    return tl.where(col[None, :] < ROTARY, normalized * cosine + rotated * sine,
+                    normalized).to(tl.bfloat16).to(tl.float32)
+
+
+@triton.jit
+def _pack_qk(Q, K, QW, KW, Cos, Sin, Payload, Scale, ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+             ROTARY: tl.constexpr, WORLD: tl.constexpr, CHUNKS: tl.constexpr, QEPS: tl.constexpr, KEPS: tl.constexpr,
+             BLOCK_ROWS: tl.constexpr, BLOCK_DIM: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)).to(tl.int64)
+    row, head = row_head // HEADS, row_head % HEADS
+    col = tl.arange(0, BLOCK_DIM)
+    projection = tl.program_id(1)
+    if projection == 0:
+        values = _normalize(Q, QW, Cos, Sin, row, head, HEADS, ROWS, DIM, ROTARY, QEPS, BLOCK_DIM)
+    else:
+        values = _normalize(K, KW, Cos, Sin, row, head, HEADS, ROWS, DIM, ROTARY, KEPS, BLOCK_DIM)
+    maximum = tl.max(tl.abs(values), 1)
+    scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    chunk = head % local_heads // chunk_heads
+    destination = ((
+        (chunk * WORLD + head // local_heads) * 3 + projection) * ROWS + row) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * DIM + col[None, :], quantized,
+             (row[:, None] < ROWS) & (col[None, :] < DIM))
+    tl.store(Scale + destination, scale, row < ROWS)
+

fastvideo/models/dits/minimax_h3_fusions/sp_qkv_pack.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse H3 local Q/K normalization and RoPE with FP8 Ulysses packing."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _normalized_rotary(X, Weight, Cos, Sin, row, head, HEADS: tl.constexpr, DIM: tl.constexpr,
+                        ROTARY: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr):
+    column = tl.arange(0, BLOCK)
+    offset = (row * HEADS + head) * DIM
+    value = tl.load(X + offset + column, column < DIM, other=0.0).to(tl.float32)
+    variance = tl.sum(value * value, axis=0) / DIM
+    inverse_rms = tl.math.rsqrt(variance + EPS)
+    weight = tl.load(Weight + column, column < DIM, other=0.0).to(tl.float32)
+    normalized = value * inverse_rms * weight
+    partner = tl.where(column < ROTARY // 2, column + ROTARY // 2, column - ROTARY // 2)
+    partner_value = tl.load(X + offset + partner, column < ROTARY, other=0.0).to(tl.float32)
+    partner_weight = tl.load(Weight + partner, column < ROTARY, other=0.0).to(tl.float32)
+    partner_normalized = partner_value * inverse_rms * partner_weight
+    rotated = tl.where(column < ROTARY // 2, -partner_normalized, partner_normalized)
+    cosine = tl.load(Cos + row * ROTARY + column, column < ROTARY, other=1.0).to(tl.float32)
+    sine = tl.load(Sin + row * ROTARY + column, column < ROTARY, other=0.0).to(tl.float32)
+    result = tl.where(column < ROTARY, normalized * cosine + rotated * sine, normalized)
+    # The former standalone norm/RoPE kernel materialized BF16 here.
+    return result.to(tl.bfloat16).to(tl.float32)
+
+
+@triton.jit
+def _pack_normalized_qk(Q, K, QWeight, KWeight, Cos, Sin, Payload, Scale,
+                         ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+                         ROTARY: tl.constexpr, WORLD: tl.constexpr, CHUNKS: tl.constexpr,
+                         QEPS: tl.constexpr, KEPS: tl.constexpr, BLOCK: tl.constexpr, ROWS_PER_PROGRAM: tl.constexpr):
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    # Preserve each head's four-warp reduction while amortizing CTA scheduling.
+    for index in tl.static_range(ROWS_PER_PROGRAM):
+        row_head = tl.program_id(0).to(tl.int64) * ROWS_PER_PROGRAM + index
+        if ROWS * HEADS % ROWS_PER_PROGRAM == 0 or row_head < ROWS * HEADS:
+            row, head = row_head // HEADS, row_head % HEADS
+            projection = tl.program_id(1)
+            column = tl.arange(0, BLOCK)
+            if projection == 0:
+                values = _normalized_rotary(Q, QWeight, Cos, Sin, row, head, HEADS, DIM, ROTARY, QEPS, BLOCK)
+            else:
+                values = _normalized_rotary(K, KWeight, Cos, Sin, row, head, HEADS, DIM, ROTARY, KEPS, BLOCK)
+            maximum = tl.max(tl.where(column < DIM, tl.abs(values), 0.0), axis=0)
+            scale = tl.maximum(maximum * (1.0 / 448.0), 1e-12)
+            quantized = (values / scale).to(tl.float8e4nv)
+            chunk = (head % local_heads) // chunk_heads
+            destination = (((chunk * WORLD + head // local_heads) * 3 + projection) * ROWS + row) * chunk_heads + head % chunk_heads
+            tl.store(Payload + destination * DIM + column, quantized, column < DIM)
+            tl.store(Scale + destination, scale)
+
+
+
+@triton.jit
+def _pack_value_fp8(V, Payload, Scale, ROWS: tl.constexpr, HEADS: tl.constexpr, DIM: tl.constexpr,
+                    WORLD: tl.constexpr, CHUNKS: tl.constexpr, BLOCK_HEADS: tl.constexpr, BLOCK_DIM: tl.constexpr):
+    row_head = (tl.program_id(0) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)).to(tl.int64)
+    column = tl.arange(0, BLOCK_DIM)
+    valid = (row_head[:, None] < ROWS * HEADS) & (column[None, :] < DIM)
+    values = tl.load(V + row_head[:, None] * DIM + column[None, :], valid, other=0.0).to(tl.float32)
+    scale = tl.maximum(tl.max(tl.abs(values), axis=1) * (1.0 / 448.0), 1e-12)
+    quantized = (values / scale[:, None]).to(tl.float8e4nv)
+    row, head = row_head // HEADS, row_head % HEADS
+    local_heads: tl.constexpr = HEADS // WORLD
+    chunk_heads: tl.constexpr = local_heads // CHUNKS
+    chunk = (head % local_heads) // chunk_heads
+    destination = (((chunk * WORLD + head // local_heads) * 3 + 2) * ROWS + row) * chunk_heads + head % chunk_heads
+    tl.store(Payload + destination[:, None] * DIM + column[None, :], quantized, valid)
+    tl.store(Scale + destination, scale, row_head < ROWS * HEADS)
+
+
+def pack_normalized_qkv_fp8(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
+                            q_weight: torch.Tensor, k_weight: torch.Tensor,
+                            cos: torch.Tensor, sin: torch.Tensor, q_eps: float, k_eps: float,
+                            world: int, head_chunks: int = 1) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack local `[rows,H,D]` after BF16 norm/RoPE rounding.
+
+    With multiple chunks, write `[chunks,W,3,rows,H/W/chunks,D]` directly so
+    each chunk's all-to-all buffer is contiguous without another memory pass.
+    """
+    rows, heads, dim = query.shape
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert q_weight.dtype == k_weight.dtype == cos.dtype == sin.dtype == query.dtype
+    assert query.shape == key.shape == value.shape and head_chunks > 0 and heads % (world * head_chunks) == 0
+    assert all(t.is_contiguous() for t in (query, key, value, q_weight, k_weight, cos, sin))
+    assert q_weight.shape == k_weight.shape == (dim,)
+    assert cos.shape == sin.shape and cos.shape[0] == rows and 0 < cos.shape[1] <= dim
+    assert cos.shape[1] % 2 == 0
+    payload = torch.empty((head_chunks, world, 3, rows, heads // (world * head_chunks), dim), device=query.device, dtype=torch.float8_e4m3fn)
+    scales = torch.empty((head_chunks, world, 3, rows, heads // (world * head_chunks)), device=query.device, dtype=torch.float32)
+    rows_per_program = 8 if dim == 128 and torch.cuda.get_device_capability(query.device)[0] >= 10 else 1
+    if dim == 128 and torch.cuda.get_device_capability(query.device)[0] >= 10:
+        from .sp_qkv_ordered import _pack_qk
+        _pack_qk[(triton.cdiv(rows * heads, 8), 2)](
+            query, key, q_weight, k_weight, cos, sin, payload, scales,
+            rows, heads, dim, cos.shape[1], world, head_chunks, q_eps, k_eps, 8, 128, num_warps=8,
+        )
+    else:
+        _pack_normalized_qk[(triton.cdiv(rows * heads, rows_per_program), 2)](
+            query, key, q_weight, k_weight, cos, sin, payload, scales,
+            rows, heads, dim, cos.shape[1], world, head_chunks, q_eps, k_eps,
+            triton.next_power_of_2(dim), rows_per_program, num_warps=4,
+        )
+    _pack_value_fp8[(triton.cdiv(rows * heads, 16),)](
+        value, payload, scales, rows, heads, dim, world, head_chunks, 16, triton.next_power_of_2(dim), num_warps=4,
+    )
+    return (payload[0], scales[0]) if head_chunks == 1 else (payload, scales)

fastvideo/models/dits/minimax_h3_fusions/sp_receive.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Share SP FP8 receive reads across BF16 expansion, tile pooling and Q/V FP4 packing."""
+import torch
+import triton
+import triton.language as tl
+
+from fastvideo.distributed.async_collectives import PeerSlices
+
+from .value_quant import _left_four
+from .swiglu_quant import _rcp
+
+@triton.jit
+def _receive_all(Payload, Scale, Inverse, Valid, Out, Pool, Codes, Scales, QCodes, QScales, ROWS: tl.constexpr,
+                 S: tl.constexpr, H: tl.constexpr, D: tl.constexpr, BLOCK: tl.constexpr, START: tl.constexpr,
+                 Remote=None, PAIR: tl.constexpr = False, RANK: tl.constexpr = 0,
+                 QUERY_ROWS: tl.constexpr = -1, PrefixOut=None, PrefixInverse=None,
+                 PREFIX_VALUES: tl.constexpr = False, STORE_KEYS: tl.constexpr = True):
+    tile = tl.program_id(1)
+    projection = tl.program_id(2)
+    col = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    row = tl.arange(0, 16)
+    head = col // D
+    acc = tl.full((16, BLOCK), 0, tl.float32)
+    for i in tl.static_range(4):
+        token = tile * 64 + i * 16 + row
+        source = tl.load(Inverse + token)
+        packed = (source // ROWS * 3 + projection) * ROWS + source % ROWS
+        valid = (col[None, :] < H * D) & (source[:, None] >= 0)
+        if PAIR:
+            local_offset = (projection * ROWS + source % ROWS)[:, None] * H * D + col[None, :]
+            pointer = tl.where((source // ROWS)[:, None] == RANK, Payload + local_offset, Remote + local_offset)
+        else:
+            pointer = Payload + packed[:, None] * H * D + col[None, :]
+        value = tl.load(pointer, valid, other=0.0).to(tl.float32)
+        scale = tl.load(Scale + packed[:, None] * H + head[None, :], valid, other=0)
+        x = (value * scale).to(tl.bfloat16).to(tl.float32)
+        store_mask = (col[None, :] < H * D) & (STORE_KEYS | (projection != 1))
+        if QUERY_ROWS >= 0:
+            store_mask = store_mask & ((projection != 0) | (token[:, None] < QUERY_ROWS))
+        if PREFIX_VALUES and projection == 2:
+            destination = tl.load(PrefixInverse + token)
+            tl.store(PrefixOut + destination[:, None] * H * D + col[None, :], x,
+                     (col[None, :] < H * D) & (destination[:, None] >= 0))
+        else:
+            tl.store(Out + (projection * S + token[:, None]) * H * D + col[None, :], x, store_mask)
+        acc = acc + x
+        if projection == 2:
+            maximum = tl.max(tl.abs(x), 0)
+            sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+            inverse = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+            scaled = x * inverse[None, :]
+            pair = tl.trans(tl.reshape(scaled, (8, 2, BLOCK)), (0, 2, 1))
+            lo, hi = tl.split(pair)
+            codes = tl.inline_asm_elementwise('''{
+                .reg .b8 q;
+                cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+                mov.b32 $0, {q,q,q,q};
+            }''',
+                                              constraints='=r,f,f',
+                                              args=[hi, lo],
+                                              dtype=tl.uint8,
+                                              is_pure=True,
+                                              pack=1)
+            qrow = tl.arange(0, 8)
+            offset = (col[None, :] * (S // 2) + tile * 32 + i * 8 + qrow[:, None])
+            tl.store(Codes + offset, codes, col[None, :] < H * D)
+            group = tile * 4 + i
+            dim = col % D
+            offset = (head * D * (S // 16) + (dim // 64) * 64 * (S // 16) + (group // 4) * 256 + group % 4 +
+                      (dim % 64 // 16) * 4 + (dim % 16) * 16)
+            tl.store(Scales + offset, sf, col < H * D)
+        elif projection == 0:
+            groups = tl.reshape(x, (16, BLOCK // 16, 16))
+            maximum = tl.max(tl.abs(groups), 2)
+            sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+            inverse = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+            quant = tl.reshape(groups * inverse[:, :, None], (16, BLOCK // 2, 2))
+            lo, hi = tl.split(quant)
+            packed_q = tl.inline_asm_elementwise('''{
+                .reg .b8 q;
+                cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+                mov.b32 $0, {q,q,q,q};
+            }''',
+                                                 constraints='=r,f,f',
+                                                 args=[hi, lo],
+                                                 dtype=tl.uint8,
+                                                 is_pure=True,
+                                                 pack=1)
+            qrow = token - START
+            qseq: tl.constexpr = S - START
+            qcol = tl.program_id(0) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+            qhead, qdim = qcol // (D // 2), qcol % (D // 2)
+            tl.store(QCodes + (qhead[None, :] * qseq + qrow[:, None]) * (D // 2) + qdim[None, :], packed_q,
+                     (qrow[:, None] >= 0) & (qcol[None, :] < H * D // 2))
+            group = tl.program_id(0) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+            qhead, qdim = group // (D // 16), group % (D // 16)
+            offset = ((qrow[:, None] // 64) * 512 + (qdim[None, :] // 4) * 256 + qdim[None, :] % 4 +
+                      (qrow[:, None] % 64 // 16) * 4 + (qrow[:, None] % 16) * 16)
+            tl.store(QScales + qhead[None, :] * qseq * (D // 16) + offset, sf,
+                     (qrow[:, None] >= 0) & (group[None, :] < H * D // 16))
+    even, odd = tl.split(tl.reshape(tl.trans(acc), (BLOCK, 4, 2, 2)))
+    y0, y2 = tl.split(even)
+    y1, y3 = tl.split(odd)
+    y0, y1 = _left_four(y0, BLOCK), _left_four(y1, BLOCK)
+    y2, y3 = _left_four(y2, BLOCK), _left_four(y3, BLOCK)
+    total = (y0 + y2) + (y1 + y3)
+    count = tl.load(Valid + tile).to(tl.float32)
+    mean = tl.div_rn(total, count)
+    tl.store(Pool + (projection * (S // 64) + tile) * H * D + col, mean, col < H * D)
+
+
+def unpack_pool_quantize_fp8(
+    payload: torch.Tensor | PeerSlices, scale: torch.Tensor, inverse: torch.Tensor,
+    valid: torch.Tensor, output: torch.Tensor, query_start: int, *, query_rows: int | None = None,
+    prefix_values: tuple[torch.Tensor, torch.Tensor] | None = None, store_keys: bool = True,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+    """Expand FP8 QKV and prepare exact tile means plus native Q/V FP4 inputs.
+
+    By default every tile row is written, including zero padding. When the
+    caller consumes native FP4 queries, ``query_rows`` may limit BF16 query
+    stores to the dense prefix. The omitted suffix is untouched and must not
+    be read; pooling and native query codes still cover every row. K/V stores
+    remain complete by default. ``prefix_values=(inverse, output)`` instead
+    writes BF16 values directly into dense-prefix order; the entire raw V
+    tensor is then untouched and must not be read. Native V codes and pooling
+    still cover every row. Pooling keeps ATen's
+    tile-64 accumulation order and reads the rounded BF16 expansion. Query
+    scales use Sage's native suffix layout; value scales use its transposed
+    layout. With ``store_keys=False``, raw K is untouched; the caller must
+    prepare keys directly from these packets. Prefix values use the
+    caller's inverse map, whose entries are -1 or in-bounds output row indices.
+    """
+    pair = payload if isinstance(payload, PeerSlices) else None
+    payload = pair.template if pair is not None else payload
+    start, block = query_start, 256
+    _, _, rows, heads, dim = payload.shape
+    assert payload.dtype == torch.float8_e4m3fn and scale.dtype == torch.float32
+    assert payload.shape[1] == 3 and dim == 128 and output.dtype == torch.bfloat16
+    assert inverse.numel() % 128 == 0 and 0 <= start < inverse.numel() and start % 128 == 0
+    assert output.shape == (3, inverse.numel(), heads, dim) and scale.shape == payload.shape[:-1]
+    assert valid.shape == (inverse.numel() // 64,)
+    assert all(t.is_contiguous() for t in (payload, scale, inverse, valid, output))
+    assert inverse.dtype in (torch.int32, torch.int64) and valid.dtype in (torch.int32, torch.int64)
+    assert all(t.device == payload.device for t in (scale, inverse, valid, output))
+    assert payload.is_cuda and not torch.is_grad_enabled()
+    prefix_inverse, prefix_output = (None, None) if prefix_values is None else prefix_values
+    if prefix_values is not None:
+        if (prefix_inverse.shape != inverse.shape or prefix_inverse.dtype not in (torch.int32, torch.int64)
+                or prefix_inverse.device != payload.device or not prefix_inverse.is_contiguous()
+                or prefix_output.ndim != 4 or prefix_output.shape[0] != 1
+                or prefix_output.shape[2:] != (heads, dim) or prefix_output.dtype != torch.bfloat16
+                or prefix_output.device != payload.device or not prefix_output.is_contiguous()):
+            raise ValueError("Prefix values require a tile inverse and contiguous BF16 [1, valid, heads, 128] output")
+    seq = inverse.numel()
+    if query_rows is not None and not 0 <= query_rows <= seq:
+        raise ValueError("BF16 query rows must lie within the tile sequence")
+    pool = torch.empty((3, seq // 64, heads, dim), device=payload.device, dtype=torch.float32)
+    codes = torch.empty((1, heads, dim, seq // 2), device=payload.device, dtype=torch.uint8)
+    scales = torch.empty((1, heads, dim, seq // 16), device=payload.device, dtype=torch.float8_e4m3fn)
+    qcodes = torch.empty((1, heads, seq - start, dim // 2), device=payload.device, dtype=torch.uint8)
+    qscales = torch.empty((1, heads, seq - start, dim // 16), device=payload.device, dtype=torch.float8_e4m3fn)
+    _receive_all[(triton.cdiv(heads * dim, block), seq // 64, 3)](
+        pair.local if pair is not None else payload, scale, inverse, valid, output, pool, codes, scales, qcodes, qscales,
+        rows, seq, heads, dim, block, start, Remote=pair.remote if pair is not None else payload,
+        PAIR=pair is not None, RANK=pair.rank if pair is not None else 0,
+        QUERY_ROWS=-1 if query_rows is None else query_rows, PrefixOut=prefix_output, PrefixInverse=prefix_inverse,
+        PREFIX_VALUES=prefix_values is not None, STORE_KEYS=store_keys, num_warps=4,
+    )
+    return pool.permute(0, 2, 1, 3), (codes, scales), (qcodes, qscales)

fastvideo/models/dits/minimax_h3_fusions/swiglu_gemm.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact eager GB10 FFN-up GEMM, BF16 SwiGLU, and NVFP4 output preparation.
+
+Only immutable weights are cached. Their value/gate columns are paired in
+memory so each GEMM tile can produce complete SwiGLU quantization groups.
+The checkpoint layout and files are unchanged. Both original BF16 rounding
+boundaries precede the original block-16 scale/code conversion.
+"""
+from functools import lru_cache
+
+import torch
+
+
+@lru_cache(maxsize=1)
+def _backend():
+    try:
+        import _fastvideo_h3_swiglu_gemm
+    except ImportError:
+        return None
+    return _fastvideo_h3_swiglu_gemm
+
+
+def pair_weight_columns(weight: torch.Tensor, scales: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pair value/gate rows of packed weights and their 128x4-swizzled scales."""
+    n, k = weight.shape[0], weight.shape[1] * 2
+    values = torch.arange(n // 2, device=weight.device)
+    permutation = torch.stack((values, values + n // 2), dim=1).reshape(-1)
+    paired_weight = weight[permutation].contiguous().T
+    plain_scales = scales.reshape(n // 128, k // 64, 32, 4, 4).permute(0, 3, 2, 1, 4).reshape(n, k // 16)
+    paired_scales = (plain_scales[permutation].reshape(n // 128, 4, 32, k // 64, 4)
+                     .permute(0, 3, 2, 1, 4).contiguous().reshape_as(scales).T)
+    return paired_weight, paired_scales
+
+
+def _weight_key(weight: torch.Tensor, scales: torch.Tensor) -> tuple:
+    # Packed inference buffers may have no version counter. Their storage is
+    # immutable; replacement or migration changes the identity/storage key.
+    return tuple((id(t), t.data_ptr(), t.device, None if t.is_inference() else t._version)
+                 for t in (weight, scales))
+
+
+def try_swiglu_gemm(ff, hidden_states: torch.Tensor, pre_quantized) -> torch.Tensor | None:
+    """Return None outside the measured T2AV/Omni single/pair-Spark NVFP4 shapes."""
+    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod, _mm_fp4_row_chunks
+
+    if (pre_quantized is None or torch.is_grad_enabled() or torch.compiler.is_compiling()
+            or not hidden_states.is_cuda or hidden_states.dtype != torch.bfloat16
+            or torch.cuda.get_device_capability(hidden_states.device) != (12, 1)
+            or type(getattr(ff.fc_in, 'quant_method', None)) is not NVFP4QuantizeMethod
+            or type(getattr(ff.fc_out, 'quant_method', None)) is not NVFP4QuantizeMethod
+            or ff.fc_in.bias is not None):
+        return None
+    a, a_scale, input_scale = pre_quantized
+    a = a.reshape(-1, a.shape[-1])
+    a_scale = a_scale.reshape(-1, a_scale.shape[-1])
+    weight = ff.fc_in._nvfp4_weight
+    weight_scales = ff.fc_in._nvfp4_weight_scale
+    if (a.shape not in ((14859, 2688), (7430, 2688), (29678, 2688)) or weight.shape != (28672, 2688)
+            or not ff.fc_out.quant_method.wants_prequantized_input()):
+        return None
+    backend = _backend()
+    if backend is None:
+        return None
+    method = ff.fc_out.quant_method
+    output_scale = method._static_activation_global_sf()
+    if output_scale is None:
+        output_scale = getattr(ff.fc_out, '_nvfp4_input_global_sf', None)
+    if output_scale is None:
+        if method._dynamic_activation_scale():
+            return None
+        output_scale = method.x_global_sf
+    key = _weight_key(weight, weight_scales)
+    if getattr(ff.fc_in, '_h3_swiglu_weight_key', None) != key:
+        paired_weight, paired_scales = pair_weight_columns(weight, weight_scales)
+        ff.fc_in.register_buffer('_h3_swiglu_paired_weight', paired_weight, persistent=False)
+        ff.fc_in.register_buffer('_h3_swiglu_paired_scales', paired_scales, persistent=False)
+        ff.fc_in._h3_swiglu_weight_key = key
+    paired_weight = ff.fc_in._h3_swiglu_paired_weight
+    paired_scales = ff.fc_in._h3_swiglu_paired_scales
+    if hasattr(ff.fc_in, '_nvfp4_alpha'):
+        alpha = ff.fc_in._nvfp4_alpha / input_scale
+    else:
+        alpha = 1.0 / (input_scale * ff.fc_in._weight_global_sf)
+    rows, width = a.shape[0], weight.shape[0] // 2
+    q = torch.empty((rows, width // 2), device=a.device, dtype=torch.uint8)
+    scales = torch.empty(((rows + 127) // 128 * 128, width // 16), device=a.device, dtype=torch.uint8)
+    start = 0
+    parts = 8 if rows == 29678 else (2 if rows == 7430 else 4)
+    for chunk, chunk_scale in _mm_fp4_row_chunks(a, a_scale, parts):
+        end = start + chunk.shape[0]
+        backend.run(chunk, paired_weight, chunk_scale, paired_scales, alpha, output_scale,
+                    q[start:end], scales[start:(end + 127) // 128 * 128])
+        start = end
+    # apply(pre_quantized=...) reads only hidden_states' leading shape.
+    return method.apply(ff.fc_out, hidden_states, ff.fc_out.bias, pre_quantized=(q, scales, output_scale))

fastvideo/models/dits/minimax_h3_fusions/swiglu_quant.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+"""H3 SwiGLU and exact FlashInfer-compatible NVFP4 activation quantization.
+
+The SwiGLU result is rounded to BF16 before quantization, exactly as in the
+separate kernels. Scales use FlashInfer's default fast-math recipe and 128x4
+swizzle. Fusing the two passes avoids a full-width BF16 intermediate.
+"""
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _rcp(x):
+    return tl.inline_asm_elementwise(
+        "rcp.approx.ftz.f32 $0, $1;", constraints="=f,f", args=[x],
+        dtype=tl.float32, is_pure=True, pack=1,
+    )
+
+
+@triton.jit
+def _swiglu_quant(X, SF, Q, S, ROWS: tl.constexpr, WIDTH: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    valid = (row < ROWS) & (col < WIDTH)
+    value = tl.load(X + row * WIDTH * 2 + col, valid, other=0).to(tl.float32)
+    gate = tl.load(X + row * WIDTH * 2 + WIDTH + col, valid, other=0).to(tl.float32)
+    y = (value * (gate * tl.sigmoid(gate))).to(tl.bfloat16).to(tl.float32)
+    grouped = tl.reshape(y, (BLOCK // 16, 16))
+    amax = tl.max(tl.abs(grouped), 1)
+    global_sf = tl.load(SF)
+    scale = (global_sf * (amax * _rcp(6.0))).to(tl.float8e4nv)
+    inv = tl.where(amax != 0, _rcp(scale.to(tl.float32) * _rcp(global_sf)), 0.0)
+    quantized = tl.reshape(grouped * inv[:, None], (BLOCK // 2, 2))
+    lo, hi = tl.split(quantized)
+    packed = tl.inline_asm_elementwise('''{
+        .reg .b8 q;
+        cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+        mov.b32 $0, {q,q,q,q};
+    }''', constraints='=r,f,f', args=[hi,lo], dtype=tl.uint8, is_pure=True, pack=1)
+    qcol = tl.program_id(1) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+    tl.store(Q + row * (WIDTH // 2) + qcol, packed, (row < ROWS) & (qcol < WIDTH // 2))
+    group = tl.program_id(1) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+    sidx = (((((row // 128) * (WIDTH // 64) + group // 4) * 32 + row % 32) * 4
+             + (row // 32) % 4) * 4 + group % 4)
+    tl.store(S + sidx, scale.to(tl.uint8, bitcast=True), group < WIDTH // 16)
+
+
+
+def swiglu_quantize_nvfp4(
+    x: torch.Tensor, global_scale: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Quantize value-first SwiGLU to packed FP4 codes and padded 128x4 scales.
+
+    Requires contiguous CUDA BF16 inputs with a hidden width divisible by 64,
+    a scalar FP32 global scale, and native Blackwell FP4 conversion support.
+    The caller selects this only for eager inference with a compatible NVFP4
+    output projection; other paths retain the ordinary SwiGLU implementation.
+    """
+    if not x.is_cuda or x.dtype != torch.bfloat16 or not x.is_contiguous():
+        raise ValueError("NVFP4 SwiGLU requires contiguous CUDA BF16 input")
+    if x.ndim < 2 or x.shape[-1] == 0 or x.shape[-1] % 128:
+        raise ValueError("NVFP4 SwiGLU requires a packed width divisible by 128")
+    if global_scale.numel() != 1 or global_scale.dtype != torch.float32 or global_scale.device != x.device:
+        raise ValueError("NVFP4 SwiGLU requires a scalar FP32 scale on the input device")
+    if torch.is_grad_enabled():
+        raise RuntimeError("NVFP4 SwiGLU is inference-only")
+    rows, width = x.numel() // x.shape[-1], x.shape[-1] // 2
+    quantized = torch.empty((rows, width // 2), device=x.device, dtype=torch.uint8)
+    scales = torch.empty((triton.cdiv(rows, 128) * 128, width // 16), device=x.device, dtype=torch.uint8)
+    if rows:
+        block = 8192
+        _swiglu_quant[(scales.shape[0], triton.cdiv(width, block))](
+            x, global_scale, quantized, scales, rows, width, block, num_warps=8,
+        )
+    return quantized, scales, global_scale

fastvideo/models/dits/minimax_h3_fusions/tile_pool.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ATen-compatible tile sums for the fused attention producers."""
+import triton
+import triton.language as tl
+
+from .value_quant import _left_four
+
+
+@triton.jit
+def _tile_sum(values, R: tl.constexpr):
+    """Sum 64 or 128 rows in the same order as the eager FP32 tile reduction."""
+    raw = tl.trans(tl.reshape(values, (R // 16, 16, 128)), (1, 2, 0))
+    if R == 64:
+        even, odd = tl.split(tl.reshape(raw, (16, 128, 2, 2)))
+        x0, x2 = tl.split(even)
+        x1, x3 = tl.split(odd)
+        acc = ((x0 + x1) + x2) + x3
+    else:
+        even, odd = tl.split(tl.reshape(raw, (16, 128, 4, 2)))
+        x04, x26 = tl.split(tl.reshape(even, (16, 128, 2, 2)))
+        x15, x37 = tl.split(tl.reshape(odd, (16, 128, 2, 2)))
+        x0, x4 = tl.split(x04)
+        x1, x5 = tl.split(x15)
+        x2, x6 = tl.split(x26)
+        x3, x7 = tl.split(x37)
+        acc = ((((((x0 + x1) + x2) + x3) + x4) + x5) + x6) + x7
+    even, odd = tl.split(tl.reshape(tl.trans(acc), (128, 4, 2, 2)))
+    y0, y2 = tl.split(even)
+    y1, y3 = tl.split(odd)
+    y0, y1 = _left_four(y0, 128), _left_four(y1, 128)
+    y2, y3 = _left_four(y2, 128), _left_four(y3, 128)
+    return (y0 + y2) + (y1 + y3)

fastvideo/models/dits/minimax_h3_fusions/untile_gate.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fuse VSA's packed-order gather and compression-gate addcmul."""
+
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp as rcp
+
+
+@triton.jit
+def _untile_gate_kernel(O, P, C, G, Untile, RowTile, Y, Q, S, SF,
+                        OB: tl.constexpr, OH: tl.constexpr, OR: tl.constexpr,
+                        PB: tl.constexpr, PH: tl.constexpr, PR: tl.constexpr,
+                        CB: tl.constexpr, CT: tl.constexpr, CH: tl.constexpr,
+                        PREFIX: tl.constexpr, START: tl.constexpr,
+                        ROWS: tl.constexpr, HD: tl.constexpr, D: tl.constexpr,
+                        BLOCK: tl.constexpr, QUANTIZED: tl.constexpr):
+    row = tl.program_id(0)
+    batch = tl.program_id(2)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    head, dim = col // D, col % D
+    src = tl.load(Untile + row).to(tl.int64)
+    tile = tl.load(RowTile + row).to(tl.int64)
+    if PREFIX > 0:
+        if src < PREFIX:
+            a = tl.load(P + batch * PB + head * PH + src * PR + dim, col < HD, other=0).to(tl.float32)
+        else:
+            a = tl.load(O + batch * OB + head * OH + (src - START) * OR + dim, col < HD, other=0).to(tl.float32)
+    else:
+        a = tl.load(O + batch * OB + head * OH + src * OR + dim, col < HD, other=0).to(tl.float32)
+    b = tl.load(C + batch * CB + tile * CT + head * CH + dim, col < HD, other=0).to(tl.float32)
+    offset = (batch * ROWS + row).to(tl.int64) * HD + col
+    gate = tl.load(G + offset, col < HD, other=0).to(tl.float32)
+    # ATen's BF16 addcmul uses float32 fused multiply-add and one final cast.
+    if QUANTIZED:
+        y = tl.where(col < HD, tl.fma(b, gate, a).to(tl.bfloat16).to(tl.float32), 0.0)
+        grouped = tl.reshape(y, (BLOCK // 16, 16))
+        amax = tl.max(tl.abs(grouped), 1)
+        global_sf = tl.load(SF)
+        scale = (global_sf * (amax * rcp(6.0))).to(tl.float8e4nv)
+        inv = tl.where(amax != 0, rcp(scale.to(tl.float32) * rcp(global_sf)), 0.0)
+        quant = tl.reshape(grouped * inv[:, None], (BLOCK // 2, 2))
+        lo, hi = tl.split(quant)
+        packed = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints="=r,f,f", args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        packed_row = (batch * ROWS + row).to(tl.int64)
+        qc = tl.program_id(1) * (BLOCK // 2) + tl.arange(0, BLOCK // 2)
+        tl.store(Q + packed_row * (HD // 2) + qc, packed, qc < HD // 2)
+        group = tl.program_id(1) * (BLOCK // 16) + tl.arange(0, BLOCK // 16)
+        scale_tile = (packed_row // 128) * (HD // 64) + group // 4
+        si = ((scale_tile * 32 + packed_row % 32) * 4 + (packed_row // 32) % 4) * 4 + group % 4
+        tl.store(S + si, scale.to(tl.uint8, bitcast=True), group < HD // 16)
+    else:
+        tl.store(Y + offset, tl.fma(b, gate, a), col < HD)
+
+
+def untile_gate(out: torch.Tensor | tuple[torch.Tensor, torch.Tensor, int], compressed: torch.Tensor,
+                gate: torch.Tensor, untile: torch.Tensor, row_tile: torch.Tensor,
+                global_scale: torch.Tensor | None = None,
+                ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Combine attention and compression in packed order, optionally preparing NVFP4 directly."""
+    batch, rows, hd = gate.shape
+    quantized = global_scale is not None
+    if quantized:
+        if gate.dtype != torch.bfloat16 or hd % 64:
+            raise ValueError("Fused gate quantization requires BF16 and a width divisible by 64")
+        if (global_scale.numel() != 1 or global_scale.dtype != torch.float32
+                or global_scale.device != gate.device):
+            raise ValueError("Fused gate quantization requires a scalar FP32 scale on the input device")
+        codes = torch.empty((batch * rows, hd // 2), device=gate.device, dtype=torch.uint8)
+        scales = torch.zeros((triton.cdiv(batch * rows, 128) * 128, hd // 16),
+                             device=gate.device, dtype=torch.uint8)
+        result = gate  # Unused output pointer in the quantized specialization.
+    else:
+        result = torch.empty_like(gate)
+        codes, scales, global_scale = result, result, result
+    if isinstance(out, tuple):
+        prefix, out, start = out
+        prefix_rows = prefix.shape[2]
+    else:
+        prefix, prefix_rows, start = out, 0, 0
+    _untile_gate_kernel[(rows, triton.cdiv(hd, 512), batch)](
+        out, prefix, compressed, gate, untile, row_tile, result, codes, scales, global_scale,
+        *out.stride()[:3], *prefix.stride()[:3], *compressed.stride()[:3],
+        prefix_rows, start, rows, hd, out.shape[-1], 512, quantized,
+    )
+    return (codes, scales, global_scale) if quantized else result

fastvideo/models/dits/minimax_h3_fusions/value_quant.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Share the value read between exact tile pooling and transposed Sage FP4 packing.
+
+For BF16 tile-64/128 pooling, ATen uses four accumulators on each of four warp
+rows. Keep that addition order, including the final (0+2)+(1+3) reduction,
+instead of a tree sum. Quantization independently reduces the maximum of
+16 consecutive tokens and uses Sage's transposed 64-row scale layout.
+"""
+import torch
+import triton
+import triton.language as tl
+
+from .swiglu_quant import _rcp
+
+
+@triton.jit
+def _left_four(x, BLOCK: tl.constexpr):
+    even, odd = tl.split(tl.reshape(x, (BLOCK, 2, 2)))
+    a0, a2 = tl.split(even)
+    a1, a3 = tl.split(odd)
+    return ((a0 + a1) + a2) + a3
+
+
+@triton.jit
+def _pool_and_quantize(X, VALID, POOL, CODES, SCALES,
+                       SB: tl.constexpr, SS: tl.constexpr, SH: tl.constexpr,
+                       S: tl.constexpr, H: tl.constexpr, D: tl.constexpr, BLOCK: tl.constexpr, WIDE: tl.constexpr,
+                       INPUT_INDEX=None, GATHER_INPUT: tl.constexpr = False, TILE: tl.constexpr = 64):
+    tile, batch = tl.program_id(1), tl.program_id(2)
+    col = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    if WIDE:
+        tile, batch, col = tile.to(tl.int64), batch.to(tl.int64), col.to(tl.int64)
+    row = tl.arange(0, 16)
+    head, dim = col // D, col % D
+    acc = tl.full((16, BLOCK), 0, tl.float32)
+    for i in tl.static_range(TILE // 16):
+        token = tile * TILE + i * 16 + row
+        source = tl.load(INPUT_INDEX + token) if GATHER_INPUT else token
+        x = tl.load(X + batch * SB + source[:, None] * SS + head[None, :] * SH + dim[None, :],
+                    (col[None, :] < H * D) & (source[:, None] >= 0), other=0).to(tl.float32)
+        acc = acc + x
+        maximum = tl.max(tl.abs(x), 0)
+        sf = (maximum * (1.0 / 6.0)).to(tl.float8e4nv)
+        inverse = tl.where(sf.to(tl.float32) == 0, 0, _rcp(sf.to(tl.float32)))
+        scaled = x * inverse[None, :]
+        pair = tl.trans(tl.reshape(scaled, (8, 2, BLOCK)), (0, 2, 1))
+        lo, hi = tl.split(pair)
+        codes = tl.inline_asm_elementwise('''{
+            .reg .b8 q;
+            cvt.rn.satfinite.e2m1x2.f32 q, $1, $2;
+            mov.b32 $0, {q,q,q,q};
+        }''', constraints='=r,f,f', args=[hi, lo], dtype=tl.uint8, is_pure=True, pack=1)
+        qrow = tl.arange(0, 8)
+        offset = (((batch * H + head[None, :]) * D + dim[None, :]) * (S // 2)
+                  + tile * (TILE // 2) + i * 8 + qrow[:, None])
+        tl.store(CODES + offset, codes, col[None, :] < H * D)
+        group = tile * (TILE // 16) + i
+        offset = ((batch * H + head) * D * (S // 16) + (dim // 64) * 64 * (S // 16)
+                  + (group // 4) * 256 + group % 4 + (dim % 64 // 16) * 4 + (dim % 16) * 16)
+        tl.store(SCALES + offset, sf, col < H * D)
+    even, odd = tl.split(tl.reshape(tl.trans(acc), (BLOCK, 4, 2, 2)))
+    y0, y2 = tl.split(even)
+    y1, y3 = tl.split(odd)
+    y0, y1 = _left_four(y0, BLOCK), _left_four(y1, BLOCK)
+    y2, y3 = _left_four(y2, BLOCK), _left_four(y3, BLOCK)
+    total = (y0 + y2) + (y1 + y3)
+    count = tl.load(VALID + tile).to(tl.float32)
+    mean = tl.div_rn(total, count)
+    tl.store(POOL + ((batch * (S // TILE) + tile) * H * D) + col, mean, col < H * D)
+
+
+def pool_quantize_values(
+    value: torch.Tensor, valid: torch.Tensor, input_index: torch.Tensor | None = None, *, tile_elems: int = 64,
+) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Return FP32 tile means and contiguous native FP4 V codes/scales.
+
+    The caller uses 64- or 128-token tiles, with all logical rows present. Values
+    must have unit feature stride and 128-wide heads; token rows and heads
+    may be strided. This is an eager, inference-only Blackwell path.
+    """
+    if (value.ndim != 4 or not value.is_cuda or value.dtype != torch.bfloat16
+            or value.shape[-1] != 128 or value.stride(-1) != 1
+            or (input_index is None and value.shape[1] % 128)):
+        raise ValueError("Value preparation requires CUDA BF16 BSHD, head-128 and padded token rows")
+    if torch.is_grad_enabled() or torch.compiler.is_compiling():
+        raise RuntimeError("Value preparation requires eager inference")
+    if tile_elems not in (64, 128):
+        raise ValueError("Value preparation supports 64- and 128-token tiles")
+    batch, seq_len, heads, dim = value.shape
+    if input_index is not None:
+        if (input_index.ndim != 1 or input_index.numel() % 128 or input_index.device != value.device
+                or input_index.dtype not in (torch.int32, torch.int64) or not input_index.is_contiguous()):
+            raise ValueError("Compact values require a contiguous aligned integer row map")
+        seq_len = input_index.numel()
+    if (valid.shape != (seq_len // tile_elems,) or valid.device != value.device or not valid.is_contiguous()
+            or valid.dtype not in (torch.int32, torch.int64)):
+        raise ValueError("Value preparation requires contiguous integer tile validity counts")
+    pooled = torch.empty((batch, seq_len // tile_elems, heads, dim), device=value.device, dtype=torch.float32)
+    codes = torch.empty((batch, heads, dim, seq_len // 2), device=value.device, dtype=torch.uint8)
+    scales = torch.empty((batch, heads, dim, seq_len // 16), device=value.device, dtype=torch.float8_e4m3fn)
+    # Keep wide indexing for large/strided batches without charging small clips
+    # for 64-bit vector arithmetic on every packed value.
+    span = 1 + sum((size - 1) * stride for size, stride in zip(value.shape, value.stride(), strict=True))
+    block = 128 if input_index is not None else 32
+    span = max(span, batch * seq_len * heads * dim)
+    _pool_and_quantize[(triton.cdiv(heads * dim, block), seq_len // tile_elems, batch)](
+        value, valid, pooled, codes, scales, *value.stride()[:3], seq_len, heads, dim, block, span >= 2**31,
+        INPUT_INDEX=input_index, GATHER_INPUT=input_index is not None, TILE=tile_elems, num_warps=4,
+    )
+    return pooled.permute(0, 2, 1, 3), (codes, scales)

fastvideo/models/dits/minimax_h3_vsa_fp4.py

@@ -19,6 +19,7 @@ with one gather before ``to_out``.
 from __future__ import annotations
 
 import math
+from collections.abc import Callable
 from typing import Any
 
 import torch
@@ -28,6 +29,16 @@ import fastvideo.envs as envs
 from fastvideo.attention.backends.video_sparse_attn_h3 import (MiniMaxH3VSAMetadata, _build_block_mask, _pool_tiles)
 
 
+def _tile_mask_to_fp4_blocks(api: Any, mask: torch.Tensor, tile: int, valid: torch.Tensor):
+    if (mask.is_cuda and mask.dtype == torch.bool and mask.shape[-1] == mask.shape[-2]
+            and tile > 0 and tile % 64 == 0 and valid.is_contiguous()
+            and not torch.compiler.is_compiling()):
+        from fastvideo.models.dits.minimax_h3_fusions.block_lists import tile_mask_to_fp4_blocks
+
+        return tile_mask_to_fp4_blocks(mask, tile, valid)
+    return api.vsa_tile_mask_to_fp4_blocks(mask, tile, valid)
+
+
 def _fp4_kernel_kwargs(two_level_p: bool | None = None) -> dict[str, object]:
     """Quantization options for the FP4 video kernel.
 
@@ -42,7 +53,10 @@ def _fp4_kernel_kwargs(two_level_p: bool | None = None) -> dict[str, object]:
 
 
 def _sparse_fp4_attention(api: Any, *args: Any, two_level_p: bool | None = None,
-                          smooth_k: bool | None = None) -> torch.Tensor:
+                          smooth_k: bool | None = None,
+                          prepared_key: tuple[torch.Tensor, torch.Tensor] | None = None,
+                          prepared_query: tuple[torch.Tensor, torch.Tensor] | None = None,
+                          prepared_value: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
     """Run the block-sparse FP4 kernel on lists built by vsa_tile_mask_to_fp4_blocks.
 
     Those lists are in range by construction, so skip the host-sync validation where the installed
@@ -52,8 +66,30 @@ def _sparse_fp4_attention(api: Any, *args: Any, two_level_p: bool | None = None,
     quantizes K. The shift is exact under softmax, so it only changes quantization error.
     """
     smooth = envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.get() if smooth_k is None else smooth_k
+    key = args[1]
+    if (smooth and prepared_key is None and key.is_cuda and key.dtype == torch.bfloat16
+            and key.shape[-1] == 128 and key.shape[1] % 128 == 0 and key.stride(-1) == 1
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(key.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_transpose"))):
+        from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys_only
+
+        prepared_key = quantize_smoothed_keys_only(key, key.mean(dim=1, keepdim=True))
+        smooth = False
     if smooth:
-        args = (args[0], args[1] - args[1].mean(dim=1, keepdim=True), *args[2:])
+        args = (args[0], key - key.mean(dim=1, keepdim=True), *args[2:])
+    if prepared_key is not None or prepared_query is not None or prepared_value is not None:
+        query, key, value, indices, counts, valid, quadrants = args
+        delta = api._zero_delta_s(query.shape[0], query.shape[2], key.shape[1], query.device)
+        single_level = _fp4_kernel_kwargs(two_level_p).get("single_level_p_quant", True)
+        return api.blockscaled_fp4_attn_sparse(
+            prepared_query if prepared_query is not None else api.scale_and_quant_fp4(query.transpose(1, 2)),
+            prepared_key if prepared_key is not None else api.scale_and_quant_fp4_permute(key.transpose(1, 2)),
+            prepared_value if prepared_value is not None else api.scale_and_quant_fp4_transpose(value.transpose(1, 2)),
+            delta, key.shape[1],
+            indices, counts, valid, quadrants, False, query.dtype == torch.bfloat16, single_level,
+        )[0]
     try:
         return api.sageattn_blackwell_sparse_bshd(*args, validate=False, **_fp4_kernel_kwargs(two_level_p))
     except TypeError as error:
@@ -143,12 +179,30 @@ class _KeyIndexCache:
     tile: int
     n_tiles: int
     _valid_key_index: torch.Tensor | None = None
+    _valid_key_inverse: torch.Tensor | None = None
 
     def valid_key_index(self, kv_valid: torch.Tensor, half: int) -> torch.Tensor:
         if self._valid_key_index is None:
             self._valid_key_index = _valid_key_index(kv_valid, half, self.n_tiles * self.tile)
         return self._valid_key_index
 
+    def valid_keys_from_tiles(self, tile_valid: torch.Tensor, half: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare prefix maps before mask selection, splitting 128-token tile counts into native halves."""
+        native_valid = tile_valid
+        if self._valid_key_index is None and self.tile != half:
+            offsets = torch.arange(self.tile // half, device=tile_valid.device) * half
+            native_valid = (tile_valid[:, None] - offsets[None, :]).clamp(0, half).flatten()
+        index = self.valid_key_index(native_valid, half)
+        return index, self.valid_key_inverse(native_valid, half)
+
+    def valid_key_inverse(self, kv_valid: torch.Tensor, half: int) -> torch.Tensor:
+        if self._valid_key_inverse is None:
+            index = self.valid_key_index(kv_valid, half)
+            self._valid_key_inverse = torch.full((self.n_tiles * self.tile,), -1,
+                                                 device=index.device, dtype=torch.int64)
+            self._valid_key_inverse[index] = torch.arange(index.numel(), device=index.device)
+        return self._valid_key_inverse
+
 
 class _TileLayout(_KeyIndexCache):
     """Per-step tile-order state shared by every block of one forward."""
@@ -164,6 +218,20 @@ class _TileLayout(_KeyIndexCache):
         self.cos = cos.new_zeros((self.rows, cos.shape[-1])).index_copy_(0, self.untile, cos)
         self.sin = sin.new_zeros((self.rows, sin.shape[-1])).index_copy_(0, self.untile, sin)
         self._buf: torch.Tensor | None = None
+        self._quant_inverse: torch.Tensor | None = None
+
+    def gather_quantized(self, quantized: torch.Tensor, scales: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Reorder per-row FP4 codes and scales without expanding them back to BF16."""
+        from fastvideo.models.dits.minimax_h3_fusions.quant_tile import tile_quantized
+
+        return tile_quantized(quantized, scales, self.packed_row_inverse())
+
+    def packed_row_inverse(self) -> torch.Tensor:
+        """Geometry-only map from padded tile rows to compact projection rows."""
+        if self._quant_inverse is None:
+            self._quant_inverse = torch.full((self.rows,), -1, device=self.untile.device, dtype=torch.int64)
+            self._quant_inverse[self.untile] = torch.arange(self.untile.numel(), device=self.untile.device)
+        return self._quant_inverse
 
     def gather_in(self, x: torch.Tensor) -> torch.Tensor:
         """Packed ``[B, L, C]`` -> tile-ordered ``[B, rows, C]``; pad rows stay zero.
@@ -207,6 +275,48 @@ def _shared_input_projections(linears: tuple[Any, ...], x: torch.Tensor) -> list
     return [m.apply(linear, x, linear.bias, pre_quantized=pre) for m, linear in zip(methods, linears, strict=True)]
 
 
+def _shared_qkv_gate_scale(attn: Any, x: torch.Tensor) -> torch.Tensor | None:
+    """Return the common unit scale when QKV and the gate can share packed input."""
+    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+    linears = (attn.to_q, attn.to_k, attn.to_v, attn.to_gate_compress)
+    if attn.to_gate_compress is None or x.shape[0] != 1 or x.shape[-1] % 64 or not attn._gate_active():
+        return None
+    if not all(type(getattr(linear, "quant_method", None)) is NVFP4QuantizeMethod
+               and linear.quant_method.wants_prequantized_input()
+               and linear.quant_method.uses_unit_activation_scale(linear) for linear in linears):
+        return None
+    return linears[0].quant_method.x_global_sf
+
+
+def _shared_qkv_gate_projections(
+    attn: Any, x: torch.Tensor, layout: _TileLayout | None = None,
+    pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+):
+    """Share packed activation preparation across QKV and the compression gate.
+
+    FP4 groups never cross rows, so tile reordering preserves every code and
+    scale. SP projects local packed rows directly and omits the tile layout.
+    Compatible fused modulation producers can supply the packed input.
+    """
+    if _shared_qkv_gate_scale(attn, x) is None:
+        if pre_quantized is not None:
+            raise RuntimeError("Prequantized H3 input requires compatible QKV and gate projections")
+        return None
+    linears = (attn.to_q, attn.to_k, attn.to_v, attn.to_gate_compress)
+    pre = linears[0].quant_method.quantize_input(x) if pre_quantized is None else pre_quantized
+    qkv_pre, shape_view = pre, x
+    if layout is not None:
+        q, scales = layout.gather_quantized(pre[0], pre[1])
+        qkv_pre = (q, scales, pre[2])
+        # apply(pre_quantized=...) only reads x.shape. This view supplies the
+        # tile-order shape without allocating or scattering a full BF16 tensor.
+        shape_view = x[:, :1].expand(-1, layout.rows, -1)
+    projected = [linear.quant_method.apply(linear, shape_view, linear.bias, pre_quantized=qkv_pre)
+                 for linear in linears[:3]]
+    return projected, pre
+
+
 @torch.compile(dynamic=True, fullgraph=True)
 def _fp8_per_token(x: torch.Tensor) -> torch.Tensor:
     """``[..., D]`` BF16 -> per-(row, head) FP8 E4M3 -> BF16, the prefix Q/K precision of the QK study.
@@ -232,7 +342,10 @@ def _valid_key_index(kv_valid: torch.Tensor, half: int, logical: int) -> torch.T
 
 
 def _prefix_dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, prefix_rows: int,
-                            logical: int, valid_index: torch.Tensor, precision: str) -> torch.Tensor:
+                            logical: int, valid_index: torch.Tensor, precision: str,
+                            key_is_smoothed: bool = False,
+                            prepared_kv: tuple[torch.Tensor, torch.Tensor] | None = None,
+                            query_is_fp8: bool = False) -> torch.Tensor:
     """Dense attention for the prefix query rows, ``[B, H, prefix_rows, D]``.
 
     VSA gives every prefix (non-video) query tile every key column, so those rows need no
@@ -246,12 +359,32 @@ def _prefix_dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch
     import torch.nn.functional as F
 
     q, k, v = query[:, :prefix_rows], key[:, :logical], value[:, :logical]
-    k = k - k.mean(dim=1, keepdim=True)
-    if precision == "fp8":
-        q, k = _fp8_per_token(q), _fp8_per_token(k)
+    if not key_is_smoothed:
+        k = k - k.mean(dim=1, keepdim=True)
     # Attend only the valid keys. Gathering them keeps SDPA on its unmasked (flash) kernel; a boolean
     # mask over all rows falls back to the memory-efficient kernel, 3-4x slower at OmniRef prefix sizes.
-    k, v = k.index_select(1, valid_index), v.index_select(1, valid_index)
+    if prepared_kv is not None:
+        q = _fp8_per_token(q) if precision == "fp8" and not query_is_fp8 else q
+        k, v = prepared_kv
+    elif (precision == "fp8" and k.is_cuda and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and k.dtype == torch.bfloat16 and v.dtype == k.dtype
+            and k.shape[-1] == 128 and k.stride(-1) == 1 and v.stride(-1) == 1
+            and torch.cuda.get_device_capability(k.device)[0] >= 10):
+        from fastvideo.models.dits.minimax_h3_fusions.prefix_quant import quantize_gather_prefix_kv
+
+        q = _fp8_per_token(q) if not query_is_fp8 else q
+        k, v = quantize_gather_prefix_kv(k, v, valid_index)
+    else:
+        if precision == "fp8":
+            q = _fp8_per_token(q) if not query_is_fp8 else q
+            k = _fp8_per_token(k)
+        k, v = k.index_select(1, valid_index), v.index_select(1, valid_index)
+    if prefix_rows == 15360 and precision == "fp8":
+        from fastvideo.models.dits.minimax_h3_fusions.prefix_attention import try_prefix_attention
+
+        specialized = try_prefix_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2))
+        if specialized is not None:
+            return specialized
     return F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
                                           scale=1.0 / math.sqrt(q.shape[-1]))
 
@@ -259,10 +392,20 @@ def _prefix_dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch
 def _fp4_fine_attention(api: Any, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
                         q2k_idx: torch.Tensor, q2k_num: torch.Tensor, kv_valid: torch.Tensor,
                         q2k_quad: torch.Tensor | None, *, logical: int, prefix_rows: int, precision: str,
-                        options: dict[str, Any], layout: _KeyIndexCache | None = None) -> torch.Tensor:
+                        options: dict[str, Any], layout: _KeyIndexCache | None = None,
+                        return_split: bool = False,
+                        prepared_query: tuple[torch.Tensor, torch.Tensor] | None = None,
+                        prepared_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+                        prefix_query_is_fp8: bool = False,
+                        value_index: torch.Tensor | None = None,
+                        prepared_keys: tuple[
+                            tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor
+                        ] | None = None,
+                        ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, int]:
     numerics = {"two_level_p": options.get("two_level_p"), "smooth_k": options.get("smooth_k")}
     if not prefix_rows:
-        return _sparse_fp4_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad, **numerics)
+        return _sparse_fp4_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad,
+                                      prepared_query=prepared_query, prepared_value=prepared_value, **numerics)
     if not getattr(api, "KEY_LENGTH_SIZING", False):
         raise RuntimeError("the MiniMax-H3 FP4 prefix split needs a fastvideo-kernel build whose "
                            "sageattn_blackwell_sparse_bshd sizes by key length; rebuild fastvideo-kernel "
@@ -271,15 +414,52 @@ def _fp4_fine_attention(api: Any, query: torch.Tensor, key: torch.Tensor, value:
         valid_index = layout.valid_key_index(kv_valid, api.HALF_N)
     else:
         valid_index = _valid_key_index(kv_valid, api.HALF_N, logical)
+    smooth = numerics["smooth_k"]
+    if smooth is None:
+        smooth = envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.get()
+    shared_smoothing = bool(smooth and logical == key.shape[1])
+    prepared_key, prepared_kv = None, None
+    if prepared_keys is not None:
+        prepared_key, prefix_key, prefix_value = prepared_keys
+        prepared_kv = (prefix_key, prefix_value)
+        numerics["smooth_k"] = False
+    elif (shared_smoothing and precision == "fp8" and layout is not None and key.is_cuda
+            and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and key.dtype == torch.bfloat16 and value.dtype == key.dtype and key.shape[-1] == 128
+            and key.stride(-1) == 1 and value.stride(-1) == 1
+            and torch.cuda.get_device_capability(key.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_transpose"))):
+        from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys
+
+        inverse = layout.valid_key_inverse(kv_valid, api.HALF_N)
+        prepared_key, prefix_key, prefix_value = quantize_smoothed_keys(
+            key, value, key.mean(dim=1, keepdim=True), inverse, valid_index.numel(), value_index=value_index,
+        )
+        prepared_kv = (prefix_key, prefix_value)
+        numerics["smooth_k"] = False
+    elif shared_smoothing:
+        # Both branches reduce the same rows with the same strides and dtype.
+        # Keep one rounded BF16 subtraction result instead of repeating it.
+        key = key - key.mean(dim=1, keepdim=True)
+        numerics["smooth_k"] = False
     with STAGES.span("prefix_attention"):
-        prefix_out = _prefix_dense_attention(query, key, value, prefix_rows, logical, valid_index, precision)
+        prefix_out = _prefix_dense_attention(query, key, value, prefix_rows, logical, valid_index, precision,
+                                             key_is_smoothed=shared_smoothing, prepared_kv=prepared_kv,
+                                             query_is_fp8=prefix_query_is_fp8)
     # A tile-64 prefix can end halfway through a native 128-row block. Keep that whole query block and its
     # original quadrant mask, then discard the overlapping prefix half and trailing padding from its output.
     head = prefix_rows // _BLOCK
     start = head * _BLOCK
+    # With prepared FP4 queries, the raw suffix supplies only dtype/device/head
+    # metadata; query may hold just the BF16 prefix without a video intermediate.
     video_out = _sparse_fp4_attention(api, query[:, start:], key, value,
                                       q2k_idx[:, :, head:].contiguous(), q2k_num[:, :, head:].contiguous(), kv_valid,
-                                      None if q2k_quad is None else q2k_quad[:, :, head:].contiguous(), **numerics)
+                                      None if q2k_quad is None else q2k_quad[:, :, head:].contiguous(),
+                                      prepared_key=prepared_key, prepared_query=prepared_query,
+                                      prepared_value=prepared_value, **numerics)
+    if return_split:
+        return prefix_out, video_out, start
     return torch.cat((prefix_out, video_out[:, :, prefix_rows - start:logical - start]), dim=2)
 
 
@@ -353,54 +533,188 @@ def vsa_tile_first_attention(attn: Any, hidden_states: torch.Tensor,
 
 
 def vsa_fp4_attention(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor],
-                      meta: MiniMaxH3VSAMetadata, use_fused_rope: bool) -> torch.Tensor:
+                      meta: MiniMaxH3VSAMetadata, use_fused_rope: bool,
+                      pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+                      ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
     """Attention core for ``MiniMaxH3Attention``; returns the pre-``to_out`` ``[B, L, H*D]``."""
     api = _api()
     layout = _layout_for(meta, rotary_emb)
     heads, dim = attn.num_attention_heads, attn.attention_head_dim
+    prefix_mode, prefix_rows = _prefix_settings(attn, meta, layout.n_tiles)
+    sim_fp8 = envs.FASTVIDEO_H3_SIM_SP_FP8.get()
+    prepared_query, q_pooled = None, None
+    prepared_key_mean = None
+    prefix_query_is_fp8 = False
+    smooth = getattr(attn, "_vsa_fp4_options", {}).get("smooth_k")
+    smooth = envs.FASTVIDEO_H3_VSA_FP4_K_SMOOTH.get() if smooth is None else smooth
+    # The measured cuDNN/GB10 route preserves every projected row while omitting
+    # zero padding from QKV GEMMs. Other backends keep the tiled projection path.
+    compact = (use_fused_rope and not sim_fp8 and layout.tile == 64
+               and layout.rows == layout.n_tiles * layout.tile and dim == 128
+               and hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
+               and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+               and (not prefix_rows or (prefix_mode == "fp8" and smooth))
+               and envs.FASTVIDEO_NVFP4_MM_BACKEND.get() == "cudnn"
+               and torch.cuda.get_device_capability(hidden_states.device) == (12, 1)
+               and _shared_qkv_gate_scale(attn, hidden_states) is not None
+               and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                      "scale_and_quant_fp4", "scale_and_quant_fp4_permute",
+                                                      "scale_and_quant_fp4_transpose")))
+    input_index = layout.packed_row_inverse() if compact else None
     with STAGES.span("qkv_proj_rope"):
-        x_tiles = layout.gather_in(hidden_states)
-        query, key, value = (t.unflatten(-1, (heads, dim))
-                             for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles))
-        if use_fused_rope:
+        shared = _shared_qkv_gate_projections(attn, hidden_states, None if compact else layout, pre_quantized)
+        if shared is None:
+            x_tiles = layout.gather_in(hidden_states)
+            projected = _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), x_tiles)
+            gate_pre = None
+        else:
+            projected, gate_pre = shared
+        query, key, value = (t.unflatten(-1, (heads, dim)) for t in projected)
+        del projected, shared
+        if compact:
+            from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import gather_qknorm_rope
+
+            cos, sin = layout.cos.to(query.dtype), layout.sin.to(query.dtype)
+            if prefix_rows:
+                from fastvideo.models.dits.minimax_h3_fusions.query_pool import prepare_pooled_queries
+
+                query, prepared_query, q_pooled = prepare_pooled_queries(
+                    query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, input_index,
+                    meta.variable_block_sizes, prefix_rows,
+                )
+            else:
+                query, prepared_query = gather_qknorm_rope(
+                    query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, input_index, fp4_start=0,
+                )
+            if prefix_rows and key.shape[0] == 1 and heads == 56 and input_index.numel() == 18432 and cos.shape[1] == 96:
+                from fastvideo.models.dits.minimax_h3_fusions.key_norm_mean import gather_key_norm_mean
+
+                key, prepared_key_mean = gather_key_norm_mean(
+                    key, attn.norm_k.weight, cos, sin, attn.norm_k.eps, input_index,
+                )
+            else:
+                key, _ = gather_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps, input_index)
+        elif use_fused_rope:
             from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
             cos, sin = layout.cos.to(query.dtype), layout.sin.to(query.dtype)
-            query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
+            if (not sim_fp8 and query.is_cuda and query.dtype == torch.bfloat16
+                    and dim == 128 and query.shape[1] % 128 == 0
+                    and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+                    and torch.cuda.get_device_capability(query.device)[0] >= 10
+                    and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                            "scale_and_quant_fp4", "scale_and_quant_fp4_permute",
+                                                            "scale_and_quant_fp4_transpose"))):
+                from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import qknorm_rope_quantize_fp4
+
+                if (layout.tile == 128 and query.shape[1] == layout.n_tiles * layout.tile
+                        and prefix_mode == "fp8" and prefix_rows
+                        and torch.cuda.get_device_capability(query.device) == (12, 1)):
+                    from fastvideo.models.dits.minimax_h3_fusions.query_pool import prepare_pooled_queries
+
+                    query, prepared_query, q_pooled = prepare_pooled_queries(
+                        query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, None,
+                        meta.variable_block_sizes, prefix_rows, tile_elems=128, prefix_fp8=True, split_prefix=True,
+                    )
+                    prefix_query_is_fp8 = True
+                else:
+                    query, prepared_query = qknorm_rope_quantize_fp4(
+                        query, attn.norm_q.weight, cos, sin, attn.norm_q.eps, prefix_rows // _BLOCK * _BLOCK,
+                    )
+            else:
+                query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
             key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps)
         else:
             rope = (layout.cos, layout.sin)
             query = attn._apply_rotary_emb(attn.norm_q(query), rope)
             key = attn._apply_rotary_emb(attn.norm_k(key), rope)
 
-    sim_fp8 = envs.FASTVIDEO_H3_SIM_SP_FP8.get()
     if sim_fp8:
         query, key, value = (_fp8_roundtrip(t) for t in (query, key, value))
 
     vbs = meta.variable_block_sizes
     logical = layout.n_tiles * layout.tile
+    prepared_keys, k_pooled = None, None
+    v_pooled, prepared_value = None, None
+    fused_gate = (attn.to_gate_compress is not None and attn._gate_active()
+                  and not sim_fp8 and query.is_cuda and query.dtype == torch.bfloat16)
+    omni_key_pool = (layout.tile == 128 and logical == key.shape[1] and prefix_mode == "fp8" and smooth
+                     and not sim_fp8 and key.is_cuda and key.dtype == value.dtype == torch.bfloat16
+                     and dim == 128 and key.stride(-1) == value.stride(-1) == 1
+                     and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+                     and torch.cuda.get_device_capability(key.device)[0] >= 10
+                     and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                             "scale_and_quant_fp4", "scale_and_quant_fp4_transpose")))
+    if prefix_rows and (compact or omni_key_pool):
+        from fastvideo.models.dits.minimax_h3_fusions.key_quant import pool_quantize_smoothed_keys
+
+        valid_index, inverse = layout.valid_keys_from_tiles(vbs, api.HALF_N)
+        key_mean = key.mean(dim=1, keepdim=True) if prepared_key_mean is None else prepared_key_mean
+        if fused_gate and (compact or torch.cuda.get_device_capability(key.device) == (12, 1)):
+            from fastvideo.models.dits.minimax_h3_fusions.kv_pool import pool_quantize_prefix_kv
+
+            key_result, value_result = pool_quantize_prefix_kv(
+                key, value, key_mean, inverse, valid_index.numel(), vbs, input_index,
+                tile_elems=layout.tile,
+            )
+            packed_key, prefix_key, prefix_value, k_pooled = key_result
+            v_pooled, prepared_value = value_result
+        else:
+            packed_key, prefix_key, prefix_value, k_pooled = pool_quantize_smoothed_keys(
+                key, value, key_mean, inverse, valid_index.numel(), vbs,
+                value_index=input_index, tile_elems=layout.tile,
+            )
+        prepared_keys = (packed_key, prefix_key, prefix_value)
     with STAGES.span("select_mask"):
-        q_pooled = _pool_tiles(query[:, :logical], vbs, layout.tile)
-        k_pooled = _pool_tiles(key[:, :logical], vbs, layout.tile)
+        if q_pooled is None:
+            q_pooled = _pool_tiles(query[:, :logical], vbs, layout.tile)
+        if k_pooled is None:
+            k_pooled = _pool_tiles(key[:, :logical], vbs, layout.tile)
         scores = torch.matmul(q_pooled, k_pooled.transpose(-2, -1)) / (dim**0.5)
         sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity
         mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans,
                                  meta.span_sparsities)
-        q2k_idx, q2k_num, kv_valid, q2k_quad = api.vsa_tile_mask_to_fp4_blocks(mask, layout.tile, vbs)
-    prefix_mode, prefix_rows = _prefix_settings(attn, meta, layout.n_tiles)
+        q2k_idx, q2k_num, kv_valid, q2k_quad = _tile_mask_to_fp4_blocks(api, mask, layout.tile, vbs)
+    if (prepared_value is None and fused_gate and layout.tile in (64, 128) and (compact or logical == value.shape[1]) and dim == 128
+            and value.stride(-1) == 1 and not torch.is_grad_enabled() and not torch.compiler.is_compiling()
+            and torch.cuda.get_device_capability(value.device)[0] >= 10
+            and all(hasattr(api, name) for name in ("blockscaled_fp4_attn_sparse", "_zero_delta_s",
+                                                    "scale_and_quant_fp4", "scale_and_quant_fp4_permute"))):
+        from fastvideo.models.dits.minimax_h3_fusions.value_quant import pool_quantize_values
+
+        v_pooled, prepared_value = pool_quantize_values(value, vbs, input_index=input_index, tile_elems=layout.tile)
     with STAGES.span("fp4_attention"):
         out = _fp4_fine_attention(api, query, key, value, q2k_idx, q2k_num, kv_valid, q2k_quad,
                                   logical=logical, prefix_rows=prefix_rows, precision=prefix_mode,
-                                  options=getattr(attn, "_vsa_fp4_options", {}), layout=layout)
-    with STAGES.span("out_untile"):
-        out = out.transpose(1, 2).index_select(1, layout.untile)  # [B, L, H, D], packed order
+                                  options=getattr(attn, "_vsa_fp4_options", {}), layout=layout,
+                                  return_split=fused_gate, prepared_query=prepared_query, prepared_value=prepared_value,
+                                  prefix_query_is_fp8=prefix_query_is_fp8,
+                                  value_index=input_index, prepared_keys=prepared_keys)
+    if not fused_gate:
+        with STAGES.span("out_untile"):
+            out = out.transpose(1, 2).index_select(1, layout.untile)  # [B, L, H, D], packed order
     if sim_fp8:
         out = _fp8_roundtrip(out)
 
     if attn.to_gate_compress is not None and attn._gate_active():
         with STAGES.span("gate_compress"):
-            gate, _ = attn.to_gate_compress(hidden_states)
-            v_pooled = _pool_tiles(value[:, :logical], vbs, layout.tile)
-            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled).permute(0, 2, 1, 3).to(out.dtype)
+            if gate_pre is None:
+                gate, _ = attn.to_gate_compress(hidden_states)
+            else:
+                linear = attn.to_gate_compress
+                gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
+            if v_pooled is None:
+                v_pooled = _pool_tiles(value[:, :logical], vbs, layout.tile)
+            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled).permute(0, 2, 1, 3).to(query.dtype)
+            if fused_gate:
+                from fastvideo.models.dits.minimax_h3_fusions.untile_gate import untile_gate
+                from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+                method = getattr(attn.to_out, "quant_method", None)
+                if (type(method) is NVFP4QuantizeMethod and method.wants_prequantized_input()
+                        and method.uses_unit_activation_scale(attn.to_out) and gate.shape[-1] % 64 == 0):
+                    prepared = untile_gate(out, out_c, gate, layout.untile, layout.row_tile, method.x_global_sf)
+                    return gate, prepared
+                return untile_gate(out, out_c, gate, layout.untile, layout.row_tile)
             gate = gate.unflatten(-1, (heads, dim))
             chunk = envs.FASTVIDEO_H3_GATE_CHUNK_ROWS.get()
             if chunk > 0 and out.shape[1] > chunk:
@@ -490,12 +804,23 @@ def _apply_gate(out: torch.Tensor, out_c: torch.Tensor, row_tile: torch.Tensor,
     return out + out_c.index_select(0, row_tile) * gate
 
 
-def _all_to_all(payload: torch.Tensor, scale: torch.Tensor, group: Any) -> tuple[torch.Tensor, torch.Tensor]:
+def _all_to_all(payload: torch.Tensor, scale: torch.Tensor, group: Any,
+                overlap: Callable[[], None] | None = None) -> tuple[torch.Tensor, torch.Tensor]:
     import torch.distributed as dist
     recv = torch.empty_like(payload)
     recv_scale = torch.empty_like(scale)
-    dist.all_to_all_single(recv.view(torch.uint8), payload.view(torch.uint8), group=group)
-    dist.all_to_all_single(recv_scale, scale, group=group)
+    if overlap is None:
+        dist.all_to_all_single(recv.view(torch.uint8), payload.view(torch.uint8), group=group)
+        dist.all_to_all_single(recv_scale, scale, group=group)
+    else:
+        # Keep the current compute stream independent of NCCL until the local
+        # projection has been enqueued. wait() then orders consumers of recv.
+        payload_work = dist.all_to_all_single(recv.view(torch.uint8), payload.view(torch.uint8),
+                                              group=group, async_op=True)
+        scale_work = dist.all_to_all_single(recv_scale, scale, group=group, async_op=True)
+        overlap()
+        payload_work.wait()
+        scale_work.wait()
     return recv, recv_scale
 
 
@@ -513,6 +838,7 @@ class _SPTileLayout(_KeyIndexCache):
         # Rows past the sequence are SP padding; their outputs are discarded.
         self.local_row_tile = row_tile[local.clamp_max(self.seq_len - 1)]
         self._buf: torch.Tensor | None = None
+        self._inverse: torch.Tensor | None = None
 
     def tiles_from(self, qkv: torch.Tensor) -> torch.Tensor:
         """``[3, L, Hs, D]`` packed -> ``[3, rows, Hs, D]`` tile order with zero padding (reused buffer)."""
@@ -521,9 +847,26 @@ class _SPTileLayout(_KeyIndexCache):
             self._buf = qkv.new_zeros(shape)
         return self._buf.index_copy_(1, self.untile, qkv)
 
+    def fp8_buffers(self, payload: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return the invariant row map and scratch tensor overwritten by each receive."""
+        if self._inverse is None:
+            self._inverse = torch.full((self.rows,), -1, device=payload.device, dtype=torch.int64)
+            self._inverse[self.untile] = torch.arange(self.seq_len, device=payload.device)
+        shape = (3, self.rows, *payload.shape[3:])
+        if self._buf is None or self._buf.shape != shape:
+            self._buf = torch.empty(shape, dtype=torch.bfloat16, device=payload.device)
+        return self._inverse, self._buf
+
+    def tiles_from_fp8(self, payload: torch.Tensor, scale: torch.Tensor, *, query_start: int | None = None):
+        from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_tiles
+
+        inverse, output = self.fp8_buffers(payload)
+        return unpack_fp8_tiles(payload, scale, inverse, output, query_start=query_start)
+
 
 def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor],
-                         meta: MiniMaxH3VSAMetadata, use_fused_rope: bool, sp_group: Any) -> torch.Tensor:
+                         meta: MiniMaxH3VSAMetadata, use_fused_rope: bool, sp_group: Any,
+                         pre_quantized: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
     """Ulysses-SP attention core on local sequence rows ``[1, rows, C]``; returns pre-``to_out`` ``[1, rows, H*D]``."""
     import torch.distributed as dist
 
@@ -531,29 +874,65 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup
     world, rank = sp_group.world_size, sp_group.rank_in_group
     heads, dim = attn.num_attention_heads, attn.attention_head_dim
     local_rows = hidden_states.shape[1]
+    chunks = getattr(attn, "_vsa_fp4_options", {}).get("sp_head_chunks", 1)
+    if heads % (world * chunks):
+        raise ValueError("H3 SP attention heads must divide evenly across ranks and head chunks")
+    if (chunks > 1 and use_fused_rope and hidden_states.is_cuda and hidden_states.dtype == torch.bfloat16
+            and attn.to_gate_compress is not None and attn._gate_active()):
+        from fastvideo.models.dits.minimax_h3_fusions.sp_head_pipeline import attention_pipelined
+        return attention_pipelined(attn, hidden_states, rotary_emb, meta, sp_group, chunks, pre_quantized=pre_quantized)
     layout = getattr(meta, "_h3_fp4_sp_layout", None)
     if layout is None:
         layout = _SPTileLayout(meta, rank, local_rows)
         meta._h3_fp4_sp_layout = layout  # type: ignore[attr-defined]
 
     with STAGES.span("qkv_proj_rope"):
-        query, key, value = (t.unflatten(-1, (heads, dim))
-                             for t in _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states))
+        shared = _shared_qkv_gate_projections(attn, hidden_states, pre_quantized=pre_quantized)
+        if shared is None:
+            projected = _shared_input_projections((attn.to_q, attn.to_k, attn.to_v), hidden_states)
+            gate_pre = None
+        else:
+            projected, gate_pre = shared
+        query, key, value = (t.unflatten(-1, (heads, dim)) for t in projected)
+        del projected, shared
+        fused_transport = (use_fused_rope and query.is_cuda and query.dtype == torch.bfloat16
+                           and all(t.is_contiguous() for t in (query, key, value)))
         if use_fused_rope:
-            from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
             cos, sin = rotary_emb[0].to(query.dtype), rotary_emb[1].to(query.dtype)
-            query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
-            key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps)
+            if not fused_transport:
+                from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
+                query = fused_qknorm_rope(query, attn.norm_q.weight, cos, sin, attn.norm_q.eps)
+                key = fused_qknorm_rope(key, attn.norm_k.weight, cos, sin, attn.norm_k.eps)
         else:
             query = attn._apply_rotary_emb(attn.norm_q(query), rotary_emb)
             key = attn._apply_rotary_emb(attn.norm_k(key), rotary_emb)
 
+    gate = None
+
+    def project_gate() -> None:
+        nonlocal gate
+        with STAGES.span("gate_proj"):
+            if gate_pre is None:
+                gate, _ = attn.to_gate_compress(hidden_states)
+            else:
+                linear = attn.to_gate_compress
+                gate = linear.quant_method.apply(linear, hidden_states, linear.bias, pre_quantized=gate_pre)
+
+    gate_active = attn.to_gate_compress is not None and attn._gate_active()
     with STAGES.span("qkv_pack"):
-        payload, scale = _pack_heads_fp8(query[0], key[0], value[0], world)
+        if fused_transport:
+            from fastvideo.models.dits.minimax_h3_fusions.sp_qkv_pack import pack_normalized_qkv_fp8
+            payload, scale = pack_normalized_qkv_fp8(
+                query[0], key[0], value[0], attn.norm_q.weight, attn.norm_k.weight,
+                cos, sin, attn.norm_q.eps, attn.norm_k.eps, world,
+            )
+        else:
+            payload, scale = _pack_heads_fp8(query[0], key[0], value[0], world)
     with STAGES.span("qkv_all_to_all"):
-        payload, scale = _all_to_all(payload, scale, sp_group.device_group)
+        payload, scale = _all_to_all(payload, scale, sp_group.device_group,
+                                     overlap=project_gate if gate_active else None)
     with STAGES.span("qkv_unpack_tile"):
-        qkv = layout.tiles_from(_unpack_seq_fp8(payload, scale, layout.seq_len))  # [3, R, Hs, D]
+        qkv = layout.tiles_from_fp8(payload, scale)  # [3, R, Hs, D]
     q_t, k_t, v_t = qkv[0:1], qkv[1:2], qkv[2:3]
 
     vbs = meta.variable_block_sizes
@@ -565,28 +944,43 @@ def vsa_fp4_attention_sp(attn: Any, hidden_states: torch.Tensor, rotary_emb: tup
         sparsity = 0.0 if attn._layer_idx in meta.dense_layers else meta.VSA_sparsity
         mask = _build_block_mask(scores, meta.num_prefix_tiles, sparsity, meta.exempt, meta.video_tile_spans,
                                  meta.span_sparsities)
-        q2k_idx, q2k_num, kv_valid, q2k_quad = api.vsa_tile_mask_to_fp4_blocks(mask, layout.tile, vbs)
+        q2k_idx, q2k_num, kv_valid, q2k_quad = _tile_mask_to_fp4_blocks(api, mask, layout.tile, vbs)
     with STAGES.span("fp4_attention"):
         out_bhsd = _fp4_fine_attention(api, q_t, k_t, v_t, q2k_idx, q2k_num, kv_valid, q2k_quad,
                                        logical=logical, prefix_rows=prefix_rows, precision=prefix_mode,
                                        options=getattr(attn, "_vsa_fp4_options", {}), layout=layout)
 
+    out_c, gathered = None, None
+
+    def compress_gate() -> None:
+        nonlocal out_c, gathered
+        with STAGES.span("gate_compress"):
+            v_pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile)
+            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)[0].to(qkv.dtype)  # [Hs, n_tiles, D]
+            gathered = torch.empty((world, *out_c.shape), dtype=out_c.dtype, device=out_c.device)
+
     with STAGES.span("out_pack"):
         payload, scale = _pack_seq_fp8(out_bhsd, layout.untile, world, local_rows)
     with STAGES.span("out_all_to_all"):
-        payload, scale = _all_to_all(payload, scale, sp_group.device_group)
-    with STAGES.span("out_unpack"):
-        out = _unpack_heads_fp8(payload, scale)  # [rows, H, D]
-
-    if attn.to_gate_compress is not None and attn._gate_active():
-        with STAGES.span("gate_compress"):
-            v_pooled = _pool_tiles(v_t[:, :logical], vbs, layout.tile)
-            out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled)[0].to(out.dtype)  # [Hs, n_tiles, D]
-            gathered = torch.empty((world, *out_c.shape), dtype=out_c.dtype, device=out_c.device)
+        payload, scale = _all_to_all(payload, scale, sp_group.device_group,
+                                     overlap=compress_gate if gate_active else None)
+    if gate_active:
+        assert out_c is not None and gathered is not None
         with STAGES.span("gate_all_gather"):
             dist.all_gather_into_tensor(gathered, out_c.contiguous(), group=sp_group.device_group)
         with STAGES.span("gate_apply"):
-            out_c_all = gathered.flatten(0, 1).transpose(0, 1)  # [n_tiles, H, D]
-            gate, _ = attn.to_gate_compress(hidden_states)
-            out = _apply_gate(out, out_c_all, layout.local_row_tile, gate[0].unflatten(-1, (heads, dim)))
+            from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+            from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+            assert gate is not None
+            local_gate = gate[0].unflatten(-1, (heads, dim))
+            method = getattr(attn.to_out, "quant_method", None)
+            if (type(method) is NVFP4QuantizeMethod and method.wants_prequantized_input()
+                    and method.uses_unit_activation_scale(attn.to_out) and gate.shape[-1] % 64 == 0):
+                prepared = unpack_fp8_gate(payload, scale, gathered, local_gate, layout.local_row_tile, method.x_global_sf)
+                return gate, prepared
+            out = unpack_fp8_gate(payload, scale, gathered, local_gate, layout.local_row_tile)
+    else:
+        with STAGES.span("out_unpack"):
+            out = _unpack_heads_fp8(payload, scale)  # [rows, H, D]
     return out.flatten(1, 2).unsqueeze(0)

fastvideo/models/loader/component_loader.py

@@ -849,14 +849,57 @@ class VAELoader(ComponentLoader):
         # The swap runs where the VAE was built (the CPU, see ``_h3_nvfp4_decoder_path``); a pre-packed
         # file never allocates the quantized linears' full-precision weights, a bf16 file packs on the GPU.
         device = next(vae.parameters()).device
-        metadata = apply_nvfp4_decoder_checkpoint(vae, load_nvfp4_decoder_checkpoint(path, device), freeze=False)
+        checkpoint = load_nvfp4_decoder_checkpoint(path, device)
+        student_layers = int(checkpoint["metadata"].get("student_layers", loaded_blocks))
+        if student_layers > loaded_blocks:
+            # A complete decoder export can also replace a shallower light VAE.
+            # Build its topology from the host geometry; strict loading below
+            # requires every exported weight and rejects incompatible widths.
+            from fastvideo.models.vaes.minimax_h3_video import MiniMaxH3VideoViTDecoder3d
+
+            arch = vae.config
+            with torch.device(device):
+                vae.decoder = MiniMaxH3VideoViTDecoder3d(
+                    in_channels=arch.latent_channels,
+                    out_channels=arch.out_channels,
+                    patch_size=vae.spatial_compression_ratio,
+                    patch_size_t=vae.temporal_compression_ratio,
+                    num_layers=student_layers,
+                    num_attention_heads=arch.decoder_num_attention_heads,
+                    attention_head_dim=arch.decoder_attention_head_dim,
+                    num_register_tokens=arch.decoder_num_register_tokens,
+                    ffn_mult=arch.decoder_ffn_mult,
+                    rope_theta=arch.decoder_rope_theta,
+                    rope_dim_ratio=arch.decoder_rope_dim_ratio,
+                    norm_eps=arch.decoder_norm_eps,
+                ).eval()
+        metadata = apply_nvfp4_decoder_checkpoint(vae, checkpoint, freeze=False)
         vae.to(target_device)
         if target_device.type == "cuda":
             # Frozen layers hold only packed FP4 weights; packing needs the CUDA FP4 kernels.
             freeze_nvfp4_linears(vae.decoder)
-        logger.info("MiniMax-H3 VAE decoder: NVFP4 %s (%d of %d blocks, act_scale=%s), bf16 decode autocast", path,
+        logger.info("MiniMax-H3 VAE decoder: NVFP4 %s (%d blocks, original %d, act_scale=%s), bf16 decode autocast", path,
                     len(vae.decoder.transformer_blocks), loaded_blocks, metadata.get("act_scale"))
 
+    @staticmethod
+    def _quantize_h3_dense_decoder(vae: nn.Module) -> None:
+        """Quantize the original full-depth decoder in memory; leave the checkpoint untouched."""
+        from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import (
+            convert_decoder_to_nvfp4,
+            freeze_nvfp4_linears_on_device,
+        )
+
+        # Fused-state caching reads input_amax._version. Even an outer inference
+        # context must create ordinary version-tracked buffers during conversion.
+        with torch.inference_mode(False), torch.no_grad():
+            names = convert_decoder_to_nvfp4(vae.decoder, compute_dtype=torch.bfloat16, act_scale="unit")
+            freeze_nvfp4_linears_on_device(vae.decoder)
+        vae.requires_grad_(False)
+        vae.decode_autocast_dtype = torch.bfloat16
+        logger.info("MiniMax-H3 VAE decoder: runtime NVFP4 quantization of %d original block linears "
+                    "(%d blocks, unit activation scale), bf16 decode autocast", len(names),
+                    len(vae.decoder.transformer_blocks))
+
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the VAE based on the model path, and inference args."""
         config = get_diffusers_config(model=model_path)
@@ -867,6 +910,10 @@ class VAELoader(ComponentLoader):
         assert class_name is not None, (
             "Model config does not contain a _class_name attribute. Only diffusers format is supported.")
         fastvideo_args.model_paths["vae"] = model_path
+        quantize_dense = (class_name == "AutoencoderKLMiniMaxH3"
+                          and getattr(fastvideo_args, "h3_vae_nvfp4_quantize_dense", False))
+        if quantize_dense and self._h3_nvfp4_decoder_path(class_name, fastvideo_args):
+            raise ValueError("h3_vae_nvfp4_quantize_dense and an exported NVFP4 decoder are mutually exclusive")
 
         from fastvideo.platforms import current_platform
 
@@ -993,9 +1040,9 @@ class VAELoader(ComponentLoader):
                 int8_convrot_overlay_to_apply,
             )
             int8_convrot_path = int8_convrot_overlay_to_apply(model_path)
-            if int8_convrot_path is not None and self._h3_nvfp4_decoder_path(class_name, fastvideo_args):
+            if int8_convrot_path is not None and (quantize_dense or self._h3_nvfp4_decoder_path(class_name, fastvideo_args)):
                 logger.info("Skipping MiniMax-H3 INT8 ConvRot VAE overlay %s: the NVFP4 decoder replaces the "
-                            "decoder", int8_convrot_path)
+                            "decoder execution", int8_convrot_path)
                 int8_convrot_path = None
             safetensors_list = dense_vae_safetensors(safetensors_list)
         if not safetensors_list:
@@ -1041,6 +1088,8 @@ class VAELoader(ComponentLoader):
         nvfp4_decoder = self._h3_nvfp4_decoder_path(class_name, fastvideo_args)
         if nvfp4_decoder:
             self._apply_h3_nvfp4_decoder(vae, nvfp4_decoder, target_device)
+        elif quantize_dense:
+            self._quantize_h3_dense_decoder(vae)
         if (class_name == "AutoencoderKLWan" and getattr(vae.config, "use_light_vae", False)
                 and target_device.type == "cuda" and hasattr(vae, "optimize_memory_format")):
             vae.optimize_memory_format()
@@ -1332,7 +1381,12 @@ class TransformerLoader(ComponentLoader):
             adaln_table = envs.FASTVIDEO_H3_ADALN_TABLE.get() or None
             adaln_host_cache = (fastvideo_args.inference_mode and not fastvideo_args.use_fsdp_inference
                                 and (envs.FASTVIDEO_H3_ADALN_CACHE.get() or adaln_table is not None))
-            layerwise_load = layerwise_load or adaln_host_cache
+            from fastvideo.platforms import current_platform
+            adaln_resident_cache = (adaln_host_cache and adaln_table is None
+                                    and not fastvideo_args.dit_layerwise_offload
+                                    and current_platform.has_unified_memory())
+            # Tables skip the projections; unified-memory caches reuse the device allocation.
+            layerwise_load = layerwise_load or (adaln_host_cache and adaln_table is None and not adaln_resident_cache)
             model = maybe_load_fsdp_model(
                 model_cls=model_cls,
                 init_params={
@@ -1381,8 +1435,9 @@ class TransformerLoader(ComponentLoader):
         model = model.eval()
 
         if adaln_host_cache and hasattr(model, "enable_adaln_host_cache"):
-            model.enable_adaln_host_cache(adaln_table)
+            model.enable_adaln_host_cache(adaln_table, keep_resident=adaln_resident_cache)
             logger.info("AdaLN modulation: %s", "precomputed tables from " + adaln_table if adaln_table
+                        else "resident per-timestep cache" if adaln_resident_cache
                         else "projections in pinned host memory behind a per-timestep cache")
         if layerwise_load and not fastvideo_args.dit_layerwise_offload:
             model = model.to(get_local_torch_device())

fastvideo/models/vaes/minimax_h3_parallel.py

@@ -43,7 +43,7 @@ output is bitwise equal to serial ``decode_to_pixels``.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 import torch
 
@@ -194,6 +194,8 @@ def decode_to_pixels_parallel(
     output: torch.Tensor | None,
     group: "GroupCoordinator",
     strategy: str = DEFAULT_DECODE_GATHER_STRATEGY,
+    *,
+    idle_callback: Callable[[], None] | None = None,
 ) -> torch.Tensor | None:
     """Chunk-parallel ``decode_to_pixels`` across a sequence-parallel group.
 
@@ -202,10 +204,14 @@ def decode_to_pixels_parallel(
     group's first rank, which assembles bitwise the serial
     ``decode_to_pixels`` result into ``output``. Only the first rank passes
     ``output`` (validated exactly like the serial API); other ranks pass
-    ``None`` and receive ``None``.
+    ``None`` and receive ``None``. An optional callback runs once on a rank
+    with no chunk in the final round (single-item batches only). It must not
+    issue collectives; the normal round rendezvous follows it on every rank.
     """
     if strategy not in DECODE_GATHER_STRATEGIES:
         raise ValueError(f"Unknown parallel-decode strategy {strategy!r}; expected one of {DECODE_GATHER_STRATEGIES}.")
+    if idle_callback is not None and z.shape[0] != 1:
+        raise ValueError("An idle decode callback requires a single-item batch.")
     is_leader = group.rank_in_group == 0
     if is_leader:
         if output is None:
@@ -226,7 +232,7 @@ def decode_to_pixels_parallel(
                 slice_output = output[batch_index:batch_index + 1] if output is not None else None
                 _decode_single_parallel(vae, z_slice, slice_output, group, strategy)
         else:
-            _decode_single_parallel(vae, z, output, group, strategy)
+            _decode_single_parallel(vae, z, output, group, strategy, idle_callback=idle_callback)
     finally:
         # Drain the leader's async chunk copies before the caller (or an
         # exception handler) can read or release the pinned buffer.
@@ -386,6 +392,8 @@ def _decode_single_parallel(
     output: torch.Tensor | None,
     group: "GroupCoordinator",
     strategy: str,
+    *,
+    idle_callback: Callable[[], None] | None = None,
 ) -> None:
     if _decode_single_tile_parallel(vae, z, output, group, strategy):
         return
@@ -413,6 +421,9 @@ def _decode_single_parallel(
         for round_index in range(_num_rounds(num_chunks, world_size)):
             chunk_index = round_index * world_size + rank
             if chunk_index >= num_chunks:
+                if idle_callback is not None:
+                    idle_callback()
+                    idle_callback = None
                 segment = torch.zeros(segment_shape, dtype=segment_dtype, device=z.device)
             elif round_index == 0 and first_segment is not None:
                 segment = first_segment

fastvideo/pipelines/basic/minimax_h3/minimax_h3_pipeline.py

@@ -740,8 +740,10 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
                 audio_scheduler=audio_scheduler,
             ),
         )
-        self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=None if use_taeh3 else vae))
-        self.add_stage("audio_decoding_stage", MiniMaxH3AudioDecodingStage(audio_vae=audio_vae))
+        audio_stage = MiniMaxH3AudioDecodingStage(audio_vae=audio_vae)
+        self.add_stage("video_decoding_stage",
+                       MiniMaxH3VideoDecodingStage(vae=None if use_taeh3 else vae, audio_stage=audio_stage))
+        self.add_stage("audio_decoding_stage", audio_stage)
         self._denoise_stages_ready = True
 
     def _add_stages(self, fastvideo_args: FastVideoArgs, *, ref2va: bool) -> None:

fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_decoding.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from typing import Any
+from collections.abc import Callable
 
 import torch
 
@@ -28,10 +29,12 @@ from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import current_platform
-from fastvideo.utils import allocate_cpu_tensor_with_pin_fallback
+from fastvideo.utils import PRECISION_TO_TYPE, allocate_cpu_tensor_with_pin_fallback
 
 logger = init_logger(__name__)
 
+_IDLE_AUDIO_KEY = "minimax_h3_idle_audio"
+
 
 def _layout(batch: ForwardBatch) -> MiniMaxH3PackedLayout:
     layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY)
@@ -65,9 +68,12 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
 
     performance_component_metric = "vae_decode_time_s"
 
-    def __init__(self, vae: AutoencoderKLMiniMaxH3 | None) -> None:
+    def __init__(self,
+                 vae: AutoencoderKLMiniMaxH3 | None,
+                 audio_stage: MiniMaxH3AudioDecodingStage | None = None) -> None:
         super().__init__()
         self.vae = vae
+        self.audio_stage = audio_stage
         self._output_buffer: tuple[tuple[tuple[int, ...], bool], torch.Tensor] | None = None
         self._unified_memory: bool | None = None
 
@@ -94,6 +100,31 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
             self._output_buffer = (key, buffer)
         return self._output_buffer[1]
 
+    def _idle_audio_callback(self, batch: ForwardBatch, args: FastVideoArgs, group: Any,
+                             device: torch.device) -> Callable[[], None] | None:
+        """Use the last idle slot of an odd two-rank chunk plan for audio.
+
+        The result lives only on this request. The audio stage transfers it
+        after video assembly, keeping its normal output ownership and cleanup.
+        """
+        audio_stage, vae = self.audio_stage, self.vae
+        if audio_stage is None or vae is None or group.world_size != 2 or batch.raw_latent_shape[0] != 1:
+            return None
+        _, chunks, _ = vae._temporal_decode_plan(batch.raw_latent_shape[2])
+        if chunks <= 1 or chunks % 2 == 0:
+            return None
+        batch.extra[_IDLE_AUDIO_KEY] = None
+        if group.rank_in_group != 1:
+            return None
+        # Capture the outer stage context before video enables its own autocast.
+        enabled, dtype = torch.is_autocast_enabled(device.type), torch.get_autocast_dtype(device.type)
+
+        def decode_audio() -> None:
+            with torch.autocast(device_type=device.type, enabled=enabled, dtype=dtype):
+                batch.extra[_IDLE_AUDIO_KEY] = audio_stage.decode_audio(batch, args)
+
+        return decode_audio
+
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         result = VerificationResult()
         result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none)
@@ -109,6 +140,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         """Decode H3 video latents into normalized CPU pixels."""
+        batch.extra.pop(_IDLE_AUDIO_KEY, None)
         placeholder = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
         sp_group, is_output_rank, parallel = _decode_participation(fastvideo_args, fastvideo_args.vae_parallel_decode)
         if not is_output_rank and not parallel:
@@ -140,12 +172,15 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
                 return batch
             expected = taeh3_decoded_pixel_shape(tuple(latents.shape))
             logger.info("MiniMax-H3 video decode: TAEH3 preview (%s -> %s)", tuple(latents.shape), expected)
+            decode_precision = (fastvideo_args.pipeline_config.vae_decode_precision
+                                or fastvideo_args.pipeline_config.vae_precision)
             with nvtx_range("minimax_h3.taeh3"):
                 pixels = decode_ncthw_latents_taeh3(
                     latents,
                     device=device,
                     checkpoint_path=getattr(fastvideo_args, "taeh3_checkpoint", None),
                     chunk_size=int(getattr(fastvideo_args, "taeh3_chunk_size", 5) or 5),
+                    dtype=PRECISION_TO_TYPE[decode_precision],
                 )
             batch.output = pixels.float().cpu() if is_output_rank else placeholder
             if is_output_rank and tuple(batch.output.shape) != expected:
@@ -175,6 +210,7 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
             # Attribute the streamed decoder computation while retaining
             # per-chunk device-to-host transfer and pinned-buffer reuse.
             # fp16 for the release decoders; NVFP4 decoders were trained under bf16.
+            idle_callback = self._idle_audio_callback(batch, fastvideo_args, sp_group, device) if parallel else None
             autocast_dtype = getattr(self.vae, "decode_autocast_dtype", torch.float16)
             with (
                     nvtx_range("minimax_h3.vae"),
@@ -184,7 +220,15 @@ class MiniMaxH3VideoDecodingStage(PipelineStage):
                     strategy = fastvideo_args.vae_parallel_decode_strategy or DEFAULT_DECODE_GATHER_STRATEGY
                     logger.info("MiniMax-H3 VAE decode: sequence-parallel chunks across %d ranks (%s)",
                                 sp_group.world_size, strategy)
-                    decode_to_pixels_parallel(self.vae, latents, output, sp_group, strategy=strategy)
+                    if idle_callback is None:
+                        decode_to_pixels_parallel(self.vae, latents, output, sp_group, strategy=strategy)
+                    else:
+                        decode_to_pixels_parallel(self.vae,
+                                                  latents,
+                                                  output,
+                                                  sp_group,
+                                                  strategy=strategy,
+                                                  idle_callback=idle_callback)
                 else:
                     self.vae.decode_to_pixels(latents, output)
             if is_output_rank and output is not None and output.dtype == torch.float32:
@@ -225,15 +269,26 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
     @torch.no_grad()
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         """Decode H3 audio latents into a stereo CPU waveform."""
-        # Audio decode is sub-second, so preserve the serial path's output
-        # ownership (global rank zero, or the first denoise rank under the
-        # MiniMax-H3 encoder split).
-        if model_parallel_is_initialized() and not h3_is_output_worker(fastvideo_args):
+        produced = None
+        if _IDLE_AUDIO_KEY in batch.extra:
+            group = get_sp_group()
+            produced = group.broadcast_object(batch.extra.pop(_IDLE_AUDIO_KEY), src=1)
+        is_output = not model_parallel_is_initialized() or h3_is_output_worker(fastvideo_args)
+        if not is_output:
             batch.extra["audio"] = torch.empty((0, 2), device="cpu", dtype=torch.float32)
             batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
-            self._clear_runtime(batch)
-            return batch
+        else:
+            # A tile-parallel decoder may not use a temporal idle slot. All
+            # ranks observe None in that case and retain the serial audio path.
+            waveform, sample_rate = produced if produced is not None else self.decode_audio(batch, fastvideo_args)
+            batch.extra["audio"] = waveform
+            batch.extra["audio_sample_rate"] = sample_rate
+        self._clear_runtime(batch)
+        return batch
 
+    @torch.no_grad()
+    def decode_audio(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> tuple[torch.Tensor, int]:
+        """Synthesize this request without changing its latents or stage state."""
         layout = _layout(batch)
         if batch.audio_latents is None:
             raise ValueError("MiniMax-H3 audio latents are missing at decode.")
@@ -246,10 +301,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
         try:
             latents = self.audio_vae.denormalize_latents(latents.to(device=device, dtype=torch.float32))
             if fastvideo_args.output_type == "latent":
-                batch.extra["audio"] = latents.detach().float().cpu()
-                batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
-                self._clear_runtime(batch)
-                return batch
+                return latents.detach().float().cpu(), self.audio_vae.sampling_rate
 
             # The range isolates waveform synthesis from packing and runtime
             # cleanup so the audio decoder has one stable timeline boundary.
@@ -258,10 +310,7 @@ class MiniMaxH3AudioDecodingStage(PipelineStage):
             if decoded.ndim != 3 or decoded.shape[0] != 2 or decoded.shape[1] != 1:
                 raise ValueError("MiniMax-H3 audio VAE must decode stereo channels as two mono batch items; "
                                  f"got {tuple(decoded.shape)}.")
-            batch.extra["audio"] = decoded[:, 0].transpose(0, 1).contiguous().cpu()
-            batch.extra["audio_sample_rate"] = self.audio_vae.sampling_rate
-            self._clear_runtime(batch)
-            return batch
+            return decoded[:, 0].transpose(0, 1).contiguous().cpu(), self.audio_vae.sampling_rate
         finally:
             if fastvideo_args.vae_cpu_offload:
                 pinned_offload.unload(self.audio_vae)

fastvideo/tests/ops/quantization/test_nvfp4_row_split.py

@@ -9,8 +9,8 @@ from fastvideo.layers.quantization import nvfp4_config
 from fastvideo.layers.quantization.nvfp4_config import _mm_fp4_row_chunks
 
 
-@pytest.mark.parametrize("rows", [35328, 35000, 16384, 20481])
-@pytest.mark.parametrize("parts", [1, 2, 4])
+@pytest.mark.parametrize("rows", [35328, 35000, 16384, 20481, 14859])
+@pytest.mark.parametrize("parts", [1, 2, 4, 8])
 def test_row_chunks_cover_rows_in_order(rows: int, parts: int) -> None:
     a = torch.arange(rows * 2, dtype=torch.float32).view(rows, 2)
     scale = torch.arange(rows + (-rows) % 128, dtype=torch.float32).view(-1, 1)
@@ -55,3 +55,39 @@ def test_large_cudnn_call_splits_only_when_opted_in(monkeypatch, policy, expecte
         nvfp4_config._mm_fp4(torch.empty(16384, 32), torch.empty(16, 64), torch.empty(16384, 1),
                               torch.empty(16, 1), None, torch.bfloat16, None, backend="cudnn")
     assert calls == [expected]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4 direct-row parity requires CUDA")
+def test_compact_h3_direct_rows_match_whole_cudnn(monkeypatch):
+    """Real GB10 GEMMs, including non-aligned final rows and the public dispatch."""
+    if torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip("This measured dispatch is specific to GB10")
+    import flashinfer
+    from fastvideo import envs
+
+    torch.manual_seed(7331)
+    with torch.inference_mode():
+        a = torch.randn((14859, 5376), device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn((28672, 5376), device="cuda", dtype=torch.bfloat16)
+        sf = torch.ones((), device="cuda", dtype=torch.float32)
+        qa, sa = flashinfer.nvfp4_quantize(a, sf)
+        qw, sw = flashinfer.nvfp4_quantize(weight, sf)
+        qa, qw = qa.view(torch.uint8), qw.view(torch.uint8)
+        expected = flashinfer.mm_fp4(qa, qw.T, sa, sw.T, sf, torch.bfloat16, backend="cudnn")
+        calls = []
+        direct = nvfp4_config._mm_fp4_direct_rows
+
+        def recorded(*args):
+            calls.append(args[-1])
+            return direct(*args)
+
+        monkeypatch.setattr(nvfp4_config, "_mm_fp4_direct_rows", recorded)
+        with envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.override("table"):
+            actual = nvfp4_config._mm_fp4(qa, qw.T, sa, sw.T, sf, torch.bfloat16, None, backend="cudnn")
+        assert calls == [8]
+        assert torch.equal(actual, expected)
+        calls.clear()
+        with envs.FASTVIDEO_NVFP4_MM_ROW_SPLIT.override("off"):
+            whole = nvfp4_config._mm_fp4(qa, qw.T, sa, sw.T, sf, torch.bfloat16, None, backend="cudnn")
+        assert not calls
+        assert torch.equal(whole, expected)

fastvideo/tests/ops/test_minimax_h3_prefix.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact native prefix parity across query/key tail masks and the Omni geometry."""
+import pytest
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.nn.functional import scaled_dot_product_attention
+
+
+# Enough independent query blocks keep Torch on its original unsplit K traversal.
+# Low-occupancy shapes use Torch's split-K heuristic and are deliberately not dispatched.
+@pytest.mark.parametrize('shape', [(2, 7, 519, 777), (2, 56, 256, 384), (2, 56, 257, 513), (1, 56, 15360, 29678)])
+@pytest.mark.parametrize('magnitude', [1.0, 8.0])
+@torch.inference_mode()
+def test_prefix_preserves_sdpa_output_bytes(shape, magnitude):
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip('requires GB10')
+    native = pytest.importorskip('_fastvideo_h3_prefix')
+    from fastvideo.models.dits.minimax_h3_fusions.prefix_attention import try_prefix_attention
+
+    batch, heads, nq, nk = shape
+    torch.manual_seed(nq + nk)
+    q = (torch.randn((batch, nq, heads, 128), device='cuda', dtype=torch.bfloat16) * magnitude).transpose(1, 2)
+    k = torch.randn((batch, heads, nk, 128), device='cuda', dtype=torch.bfloat16) * magnitude
+    v = torch.randn_like(k)
+    with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
+        expected = scaled_dot_product_attention(q, k, v)
+    actual = native.forward(q, k, v)
+    assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
+    selected = try_prefix_attention(q, k, v)
+    if shape == (1, 56, 15360, 29678):
+        assert selected is not None
+        assert torch.equal(selected.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
+    else:
+        assert selected is None
+
+
+@pytest.mark.parametrize('shape', [(1, 4, 256, 384), (1, 7, 257, 513)])
+@torch.inference_mode()
+def test_low_occupancy_shapes_keep_torch_splitkv_dispatch(shape):
+    from fastvideo.models.dits.minimax_h3_fusions.prefix_attention import try_prefix_attention
+
+    batch, heads, nq, nk = shape
+    q = torch.empty((batch, heads, nq, 128), dtype=torch.bfloat16)
+    k = torch.empty((batch, heads, nk, 128), dtype=torch.bfloat16)
+    assert try_prefix_attention(q, k, k) is None

fastvideo/tests/ops/test_minimax_h3_swiglu_gemm.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GB10 native epilogue parity, including padded-scheduler store bounds."""
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires one SM121 GPU")
+
+
+@pytest.mark.parametrize("rows", [1, 127, 128, 129, 1025, 7430])
+@pytest.mark.parametrize("magnitude", [1.0, 8.0])
+@torch.inference_mode()
+def test_epilogue_matches_bf16_and_quantization_boundaries(rows, magnitude):
+    if torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip("native SM121 epilogue")
+    native = pytest.importorskip('_fastvideo_h3_swiglu_gemm')
+    import flashinfer
+    from flashinfer import SfLayout, nvfp4_quantize
+    from fastvideo.models.dits.minimax_h3_fusions.swiglu import minimax_h3_swiglu
+    from fastvideo.models.dits.minimax_h3_fusions.swiglu_gemm import pair_weight_columns
+    from fastvideo.models.dits.minimax_h3_fusions.swiglu_quant import swiglu_quantize_nvfp4
+
+    torch.manual_seed(rows)
+    n, k = 512, 256
+    unit = torch.ones((), device='cuda', dtype=torch.float32)
+    a = torch.randn(rows, k, device='cuda', dtype=torch.bfloat16) * magnitude
+    w = torch.randn(n, k, device='cuda', dtype=torch.bfloat16)
+    qa, sa = nvfp4_quantize(a, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
+    qw, sw = nvfp4_quantize(w, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
+    bf16 = flashinfer.mm_fp4(qa, qw.T, sa, sw.T, unit, torch.bfloat16, backend='cudnn')
+    expected_q, expected_s, _ = swiglu_quantize_nvfp4(bf16, unit)
+    expected_y = minimax_h3_swiglu(bf16)
+    expected_amax = expected_y.float().abs().reshape(rows, n // 32, 16).amax(-1)
+    paired_w, paired_s = pair_weight_columns(qw, sw)
+    qn, sn = expected_q.numel(), expected_s.numel()
+    qstorage = torch.full((qn + 1048576,), 165, device='cuda', dtype=torch.uint8)
+    sstorage = torch.full((sn + 1048576,), 165, device='cuda', dtype=torch.uint8)
+    q = qstorage[:qn].view_as(expected_q)
+    s = sstorage[:sn].view_as(expected_s)
+    amax = torch.empty_like(expected_s, dtype=torch.float32)
+    y = torch.empty_like(expected_y)
+    for raster, swizzle in [(0, 1), (0, 4), (2, 8)]:
+        native.run(qa, paired_w, sa, paired_s, unit, unit, q, s, amax, y, raster=raster, swizzle=swizzle)
+        assert torch.equal(q, expected_q)
+        assert torch.equal(s, expected_s)
+        assert torch.equal(y.view(torch.int16), expected_y.view(torch.int16))
+        assert torch.equal(amax[:rows], expected_amax)
+        assert bool((qstorage[qn:] == 165).all())
+        assert bool((sstorage[sn:] == 165).all())
+
+
+@torch.inference_mode()
+def test_omni_eight_chunks_preserve_tail_scales_and_output_guards():
+    if torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip("native SM121 epilogue")
+    native = pytest.importorskip('_fastvideo_h3_swiglu_gemm')
+    import flashinfer
+    from flashinfer import SfLayout, nvfp4_quantize
+    from fastvideo.layers.quantization.nvfp4_config import _mm_fp4_row_chunks
+    from fastvideo.models.dits.minimax_h3_fusions.swiglu_gemm import pair_weight_columns
+    from fastvideo.models.dits.minimax_h3_fusions.swiglu_quant import swiglu_quantize_nvfp4
+
+    torch.manual_seed(831)
+    rows, n, k = 29678, 512, 256
+    unit = torch.ones((), device='cuda', dtype=torch.float32)
+    a = torch.randn(rows, k, device='cuda', dtype=torch.bfloat16)
+    w = torch.randn(n, k, device='cuda', dtype=torch.bfloat16)
+    qa, sa = nvfp4_quantize(a, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
+    qw, sw = nvfp4_quantize(w, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
+    bf16 = flashinfer.mm_fp4(qa, qw.T, sa, sw.T, unit, torch.bfloat16, backend='cudnn')
+    expected_q, expected_s, _ = swiglu_quantize_nvfp4(bf16, unit)
+    paired_w, paired_s = pair_weight_columns(qw, sw)
+    qn, sn = expected_q.numel(), expected_s.numel()
+    qstorage = torch.full((qn + 1048576,), 165, device='cuda', dtype=torch.uint8)
+    sstorage = torch.full((sn + 1048576,), 165, device='cuda', dtype=torch.uint8)
+    q, scales = qstorage[:qn].view_as(expected_q), sstorage[:sn].view_as(expected_s)
+    start = 0
+    for chunk, chunk_scale in _mm_fp4_row_chunks(qa, sa, 8):
+        end = start + chunk.shape[0]
+        native.run(chunk, paired_w, chunk_scale, paired_s, unit, unit,
+                   q[start:end], scales[start:(end + 127) // 128 * 128])
+        start = end
+    assert start == rows
+    assert torch.equal(q, expected_q) and torch.equal(scales, expected_s)
+    assert bool((qstorage[qn:] == 165).all()) and bool((sstorage[sn:] == 165).all())

fastvideo/tests/stages/test_minimax_h3_vae_streaming.py

@@ -422,3 +422,55 @@ def test_encoder_split_decodes_in_the_denoise_group_on_its_first_rank(monkeypatc
                         lambda: SimpleNamespace(is_first_rank=True, world_size=1, rank_in_group=0, ranks=[0]))
     result = MiniMaxH3VideoDecodingStage(EncoderRankVAE()).forward(ForwardBatch(data_type="video"), args)
     assert result.output.shape == (0, 3, 0, 0, 0)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+@pytest.mark.parametrize("use_idle_slot", [True, False])
+def test_idle_audio_preserves_request_and_output_ownership(monkeypatch, rank, use_idle_slot) -> None:
+    from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import MiniMaxH3AudioDecodingStage
+
+    calls = []
+    waveform = torch.arange(16, dtype=torch.float32).reshape(8, 2)
+    audio = MiniMaxH3AudioDecodingStage(SimpleNamespace(sampling_rate=32000))
+
+    def decode_audio(batch, args):
+        assert not torch.is_autocast_enabled("cpu")
+        assert batch.latents is not None and batch.audio_latents is not None
+        calls.append(batch)
+        return waveform.clone(), 32000
+
+    audio.decode_audio = decode_audio
+    group = SimpleNamespace(world_size=2, rank_in_group=rank)
+
+    def broadcast(value, src):
+        assert src == 1
+        if rank == 1 and use_idle_slot:
+            torch.testing.assert_close(value[0], waveform, atol=0, rtol=0)
+        return (waveform.clone(), 32000) if use_idle_slot else None
+
+    group.broadcast_object = broadcast
+    monkeypatch.setattr(minimax_h3_decoding, "get_sp_group", lambda: group)
+    monkeypatch.setattr(minimax_h3_decoding, "model_parallel_is_initialized", lambda: True)
+    monkeypatch.setattr(minimax_h3_decoding, "h3_is_output_worker", lambda args: rank == 0)
+    video = MiniMaxH3VideoDecodingStage(SimpleNamespace(_temporal_decode_plan=lambda frames: (0, 3, 20)), audio)
+    saved = []
+    for _ in range(2):
+        batch = ForwardBatch(data_type="video", raw_latent_shape=(1, 4, 13, 4, 4),
+                             latents=torch.ones(2, 4), audio_latents=torch.ones(2, 4))
+        callback = video._idle_audio_callback(batch, SimpleNamespace(), group, torch.device("cpu"))
+        assert (callback is not None) == (rank == 1)
+        if callback is not None and use_idle_slot:
+            # Video's autocast must not change audio's arithmetic.
+            with torch.autocast("cpu", dtype=torch.bfloat16):
+                callback()
+        assert batch.latents is not None and batch.audio_latents is not None
+        result = audio.forward(batch, SimpleNamespace())
+        saved.append(result.extra["audio"])
+        assert result.latents is None and result.audio_latents is None
+        assert not any(key.startswith("minimax_h3_") for key in result.extra)
+        assert result.extra["audio_sample_rate"] == 32000
+        assert result.extra["audio"].shape == ((8, 2) if rank == 0 else (0, 2))
+    assert len(calls) == (2 if (rank == 1 and use_idle_slot) or (rank == 0 and not use_idle_slot) else 0)
+    if rank == 0:
+        assert saved[0].data_ptr() != saved[1].data_ptr()
+        torch.testing.assert_close(saved[0], waveform, atol=0, rtol=0)

fastvideo/tests/transformers/test_minimax_h3_attention_preparation.py

@@ -0,0 +1,397 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Byte parity for native Blackwell attention input preparation."""
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+    reason="requires Blackwell native FP4 conversion",
+)
+
+
+@pytest.mark.parametrize("batch,heads,seq,start", [(1, 1, 128, 0), (2, 3, 512, 128), (1, 56, 1024, 384)])
+def test_norm_rope_query_codes_and_suffix_scales_are_exact(batch, heads, seq, start):
+    from attn_qat_infer import api
+    from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import (
+        fused_qknorm_rope, qknorm_rope_quantize_fp4,
+    )
+    torch.manual_seed(3581)
+    x = torch.randn((batch, seq, heads, 128), device="cuda", dtype=torch.bfloat16)
+    x[:, :3] = 0
+    weight = torch.randn((128,), device="cuda", dtype=x.dtype)
+    angles = torch.randn((seq, 48), device="cuda")
+    angles = torch.cat((angles, angles), -1)
+    cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
+    with torch.inference_mode():
+        expected = fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+        codes, scales = api.scale_and_quant_fp4(expected[:, start:].transpose(1, 2))
+        actual, (actual_codes, actual_scales) = qknorm_rope_quantize_fp4(x, weight, cos, sin, 1e-6, start)
+    assert torch.equal(expected, actual)
+    assert torch.equal(codes, actual_codes)
+    assert torch.equal(scales.view(torch.uint8), actual_scales.view(torch.uint8))
+    assert actual_codes.is_contiguous() and actual_scales.is_contiguous()
+
+
+@pytest.mark.parametrize("batch,heads", [(1, 3), (1, 56), (2, 3)])
+@pytest.mark.parametrize("rotary_dim", [96, 128])
+def test_omni_identity_gather_preserves_padded_qk_and_query_packing(batch, heads, rotary_dim):
+    from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import (
+        fused_qknorm_rope, gather_qknorm_rope, qknorm_rope_quantize_fp4,
+    )
+
+    torch.manual_seed(2459)
+    rows = 512
+    x = torch.randn((batch, rows, heads, 128), device="cuda", dtype=torch.bfloat16)
+    exponents = torch.randint(-12, 8, x.shape, device=x.device)
+    x = torch.ldexp(x.float(), exponents).to(x.dtype)
+    x[:, 17:29] = 0  # Preserve the padded rows too, including interleaved padding.
+    weight = torch.randn(128, device=x.device, dtype=x.dtype)
+    cos, sin = (torch.randn((rows, rotary_dim), device=x.device, dtype=x.dtype) for _ in range(2))
+    identity = torch.arange(rows, device=x.device, dtype=torch.int64)
+    with torch.inference_mode():
+        expected = fused_qknorm_rope(x, weight, cos, sin, 1e-6)
+        actual, no_pack = gather_qknorm_rope(x, weight, cos, sin, 1e-6, identity)
+        assert no_pack is None
+        assert torch.equal(expected.view(torch.uint8), actual.view(torch.uint8))
+        expected, expected_pack = qknorm_rope_quantize_fp4(x, weight, cos, sin, 1e-6, 128)
+        actual, actual_pack = gather_qknorm_rope(x, weight, cos, sin, 1e-6, identity, fp4_start=128)
+    assert torch.equal(expected.view(torch.uint8), actual.view(torch.uint8))
+    assert actual_pack is not None
+    for old, new in zip(expected_pack, actual_pack, strict=True):
+        assert torch.equal(old.view(torch.uint8), new.view(torch.uint8))
+
+
+@pytest.mark.parametrize("tile,tiles", [(64, 1), (64, 5), (64, 288), (128, 5), (256, 3)])
+@pytest.mark.parametrize("density", [0.0, 0.5, 1.0])
+def test_native_block_lists_preserve_order_quadrants_and_partial_tiles(tile, tiles, density):
+    from attn_qat_infer import api
+    from fastvideo.models.dits.minimax_h3_fusions.block_lists import tile_mask_to_fp4_blocks
+
+    torch.manual_seed(842)
+    mask = torch.rand((2, 3, tiles, tiles), device="cuda") < density
+    valid = torch.randint(0, tile + 1, (tiles,), device="cuda", dtype=torch.int32)
+    valid[0] = tile
+    expected = api.vsa_tile_mask_to_fp4_blocks(mask, tile, valid, validate=False)
+    actual = tile_mask_to_fp4_blocks(mask, tile, valid)
+    for a, b in zip(actual, expected, strict=True):
+        if a is None:
+            assert b is None
+        else:
+            assert torch.equal(a, b)
+
+
+@pytest.mark.parametrize("smooth", [False, True])
+@pytest.mark.parametrize("prepared", ["query", "value", "both"])
+def test_prepared_inputs_keep_native_key_permutation_without_prepared_keys(smooth, prepared):
+    from attn_qat_infer import api
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _sparse_fp4_attention
+
+    torch.manual_seed(918)
+    q, k, v = (torch.randn((1, 256, 2, 128), device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    mask = torch.ones((1, 2, 4, 4), device="cuda", dtype=torch.bool)
+    valid = torch.full((4,), 64, device="cuda", dtype=torch.int32)
+    lists = api.vsa_tile_mask_to_fp4_blocks(mask, 64, valid, validate=False)
+    with torch.inference_mode():
+        expected = _sparse_fp4_attention(api, q, k, v, *lists, smooth_k=smooth)
+        kwargs = {}
+        if prepared in ("query", "both"):
+            kwargs["prepared_query"] = api.scale_and_quant_fp4(q.transpose(1, 2))
+        if prepared in ("value", "both"):
+            kwargs["prepared_value"] = api.scale_and_quant_fp4_transpose(v.transpose(1, 2))
+        actual = _sparse_fp4_attention(api, q, k, v, *lists, smooth_k=smooth, **kwargs)
+    assert torch.equal(actual, expected)
+
+
+@pytest.mark.parametrize("batch,heads,seq,strided", [(1, 1, 128, False), (2, 3, 512, True), (1, 56, 1024, False)])
+@pytest.mark.parametrize("wide", [False, True])
+@pytest.mark.parametrize("tile", [64, 128])
+@pytest.mark.parametrize("large_range", [False, True])
+def test_value_pool_and_native_transposed_quantization_are_bit_exact(
+    batch, heads, seq, strided, wide, tile, large_range,
+):
+    from attn_qat_infer import api
+    from fastvideo.attention.backends.video_sparse_attn_h3 import _pool_tiles
+    from fastvideo.models.dits.minimax_h3_fusions.value_quant import _pool_and_quantize, pool_quantize_values
+
+    torch.manual_seed(5517)
+    x = torch.randn((batch, seq, heads * (2 if strided else 1), 128), device="cuda", dtype=torch.bfloat16)
+    if strided:
+        x = x[:, :, ::2]
+    if large_range:
+        exponents = torch.randint(-20, 9, x.shape, device=x.device)
+        x.copy_(torch.ldexp(x.float(), exponents).to(x.dtype))
+    x[:, :16] = 0
+    valid = torch.randint(1, tile + 1, (seq // tile,), device="cuda", dtype=torch.int32)
+    with torch.inference_mode():
+        expected = _pool_tiles(x, valid, tile)
+        codes, scales = api.scale_and_quant_fp4_transpose(x.transpose(1, 2))
+        actual, (actual_codes, actual_scales) = pool_quantize_values(x, valid, tile_elems=tile)
+        if wide:
+            _pool_and_quantize[(heads * 4, seq // tile, batch)](
+                x, valid, actual, actual_codes, actual_scales, *x.stride()[:3],
+                seq, heads, 128, 32, True, TILE=tile, num_warps=4,
+            )
+    assert torch.equal(expected.contiguous().view(torch.int32), actual.contiguous().view(torch.int32))
+    assert torch.equal(codes, actual_codes)
+    assert torch.equal(scales.view(torch.uint8), actual_scales.view(torch.uint8))
+
+
+@pytest.mark.parametrize("batch,heads", [(1, 3), (2, 4), (1, 56)])
+@pytest.mark.parametrize("rotary_dim", [32, 96, 128])
+def test_compact_projection_consumers_match_zero_padded_inputs(batch, heads, rotary_dim):
+    from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys, pool_quantize_smoothed_keys
+    from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import (
+        fused_qknorm_rope, gather_qknorm_rope, qknorm_rope_quantize_fp4,
+    )
+    from fastvideo.models.dits.minimax_h3_fusions.value_quant import pool_quantize_values
+
+    torch.manual_seed(7619)
+    rows, packed_rows = 512, 337
+    compact = torch.randn((batch, packed_rows, heads, 128), device="cuda", dtype=torch.bfloat16)
+    positions = torch.randperm(rows, device="cuda")[:packed_rows]
+    inverse = torch.full((rows,), -1, device="cuda", dtype=torch.int64)
+    inverse[positions] = torch.arange(packed_rows, device="cuda")
+    tiled = torch.zeros((batch, rows, heads, 128), device="cuda", dtype=compact.dtype)
+    tiled.index_copy_(1, positions, compact)
+    valid = (inverse >= 0).reshape(-1, 64).sum(-1).to(torch.int32)
+    weight = torch.randn(128, device="cuda", dtype=compact.dtype)
+    cos, sin = (torch.randn((rows, rotary_dim), device="cuda", dtype=compact.dtype) for _ in range(2))
+    with torch.inference_mode():
+        expected = fused_qknorm_rope(tiled, weight, cos, sin, 1e-6)
+        actual, no_pack = gather_qknorm_rope(compact, weight, cos, sin, 1e-6, inverse)
+        assert no_pack is None and torch.equal(expected, actual)
+        expected, expected_pack = qknorm_rope_quantize_fp4(tiled, weight, cos, sin, 1e-6, 128)
+        actual, actual_pack = gather_qknorm_rope(compact, weight, cos, sin, 1e-6, inverse, fp4_start=128)
+        assert torch.equal(expected, actual)
+        assert torch.equal(expected_pack[0], actual_pack[0])
+        assert torch.equal(expected_pack[1].view(torch.uint8), actual_pack[1].view(torch.uint8))
+        expected_pool, expected_pack = pool_quantize_values(tiled, valid)
+        actual_pool, actual_pack = pool_quantize_values(compact, valid, input_index=inverse)
+        assert torch.equal(expected_pool.contiguous().view(torch.int32), actual_pool.contiguous().view(torch.int32))
+        assert torch.equal(expected_pack[0], actual_pack[0])
+        assert torch.equal(expected_pack[1].view(torch.uint8), actual_pack[1].view(torch.uint8))
+        valid_rows = (inverse >= 0).nonzero().flatten()
+        prefix_inverse = torch.full_like(inverse, -1)
+        prefix_inverse[valid_rows] = torch.arange(packed_rows, device="cuda")
+        mean = tiled.mean(dim=1, keepdim=True)
+        expected_k, expected_k8, expected_v8 = quantize_smoothed_keys(tiled, tiled, mean, prefix_inverse, packed_rows)
+        actual_k, actual_k8, actual_v8 = quantize_smoothed_keys(
+            tiled, compact, mean, prefix_inverse, packed_rows, value_index=inverse,
+        )
+        pooled_k, pooled_k8, pooled_v8, key_pool = pool_quantize_smoothed_keys(
+            tiled, compact, mean, prefix_inverse, packed_rows, valid, value_index=inverse,
+        )
+        assert torch.equal(key_pool.contiguous().view(torch.int32), expected_pool.contiguous().view(torch.int32))
+        assert torch.equal(pooled_k[0], expected_k[0])
+        assert torch.equal(pooled_k[1].view(torch.uint8), expected_k[1].view(torch.uint8))
+        assert torch.equal(pooled_k8, expected_k8)
+        assert torch.equal(pooled_v8, expected_v8)
+    assert torch.equal(expected_k[0], actual_k[0])
+    assert torch.equal(expected_k[1].view(torch.uint8), actual_k[1].view(torch.uint8))
+    assert torch.equal(expected_k8, actual_k8)
+    assert torch.equal(expected_v8, actual_v8)
+
+
+@pytest.mark.parametrize("batch,rows,heads", [(1, 128, 3), (2, 512, 4), (1, 1024, 56)])
+def test_key_only_preparation_matches_separate_bf16_shift_and_native_packing(batch, rows, heads):
+    from attn_qat_infer import api
+    from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys_only
+
+    torch.manual_seed(153)
+    key = torch.randn((batch, rows, heads * 2, 128), device="cuda", dtype=torch.bfloat16)[:, :, ::2]
+    with torch.inference_mode():
+        mean = key.mean(dim=1, keepdim=True)
+        expected = api.scale_and_quant_fp4_permute((key - mean).transpose(1, 2))
+        actual = quantize_smoothed_keys_only(key, mean)
+    assert torch.equal(expected[0], actual[0])
+    assert torch.equal(expected[1].view(torch.uint8), actual[1].view(torch.uint8))
+
+
+def test_fused_key_only_attention_matches_original_native_entrypoint():
+    from attn_qat_infer import api
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _sparse_fp4_attention
+
+    torch.manual_seed(255)
+    q, k, v = (torch.randn((1, 256, 2, 128), device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    mask = torch.ones((1, 2, 4, 4), device="cuda", dtype=torch.bool)
+    valid = torch.full((4,), 64, device="cuda", dtype=torch.int32)
+    lists = api.vsa_tile_mask_to_fp4_blocks(mask, 64, valid, validate=False)
+    with torch.inference_mode():
+        expected = api.sageattn_blackwell_sparse_bshd(q, k - k.mean(dim=1, keepdim=True), v, *lists,
+                                                     validate=False, single_level_p_quant=False)
+        actual = _sparse_fp4_attention(api, q, k, v, *lists, smooth_k=True, two_level_p=True)
+    assert torch.equal(expected, actual)
+
+
+@pytest.mark.parametrize("batch,heads,seq,strided", [(1, 1, 128, False), (2, 3, 512, True), (1, 56, 1024, False)])
+@pytest.mark.parametrize("tile", [64, 128])
+@pytest.mark.parametrize("large_range", [False, True])
+def test_key_pool_and_prefix_preparation_preserve_native_bits(batch, heads, seq, strided, tile, large_range):
+    from attn_qat_infer import api
+    from fastvideo.attention.backends.video_sparse_attn_h3 import _pool_tiles
+    from fastvideo.models.dits.minimax_h3_fusions.key_quant import pool_quantize_smoothed_keys
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _fp8_per_token
+
+    torch.manual_seed(5518)
+    k, v = (torch.randn((batch, seq, heads * (2 if strided else 1), 128), device="cuda", dtype=torch.bfloat16)
+            for _ in range(2))
+    if strided:
+        k, v = k[:, :, ::2], v[:, :, ::2]
+    if large_range:
+        exponents = torch.randint(-20, 9, k.shape, device=k.device)
+        k.copy_(torch.ldexp(k.float(), exponents).to(k.dtype))
+    valid = torch.randint(1, tile + 1, (seq // tile,), device=k.device, dtype=torch.int32)
+    positions = (torch.arange(tile, device=k.device)[None, :] < valid[:, None]).flatten().nonzero().flatten()
+    inverse = torch.full((seq,), -1, device=k.device, dtype=torch.int64)
+    inverse[positions] = torch.arange(positions.numel(), device=k.device)
+    k[:, inverse < 0] = 0
+    v[:, inverse < 0] = 0
+    with torch.inference_mode():
+        mean = k.mean(dim=1, keepdim=True)
+        shifted = k - mean
+        expected_pool = _pool_tiles(k, valid, tile)
+        expected_pack = api.scale_and_quant_fp4_permute(shifted.transpose(1, 2))
+        expected_key = _fp8_per_token(shifted.index_select(1, positions))
+        expected_value = v.index_select(1, positions)
+        actual_pack, actual_key, actual_value, actual_pool = pool_quantize_smoothed_keys(
+            k, v, mean, inverse, positions.numel(), valid, tile_elems=tile,
+        )
+    assert torch.equal(expected_pool.contiguous().view(torch.int32), actual_pool.contiguous().view(torch.int32))
+    assert torch.equal(expected_pack[0], actual_pack[0])
+    assert torch.equal(expected_pack[1].view(torch.uint8), actual_pack[1].view(torch.uint8))
+    assert torch.equal(expected_key, actual_key)
+    assert torch.equal(expected_value, actual_value)
+
+
+@pytest.mark.parametrize("tile", [64, 128])
+def test_prefix_maps_from_tile_validity_include_both_native_halves(tile):
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _KeyIndexCache
+
+    counts = torch.tensor([tile, 0, 1, tile - 1, 63, 64], device="cuda", dtype=torch.int32)
+    layout = _KeyIndexCache()
+    layout.tile, layout.n_tiles = tile, counts.numel()
+    expected = (torch.arange(tile, device=counts.device)[None, :] < counts[:, None]).flatten().nonzero().flatten()
+    actual, inverse = layout.valid_keys_from_tiles(counts, 64)
+    assert torch.equal(expected, actual)
+    assert torch.equal(inverse[expected], torch.arange(expected.numel(), device=counts.device))
+    assert torch.equal((inverse >= 0).nonzero().flatten(), expected)
+    again, inverse_again = layout.valid_keys_from_tiles(counts, 64)
+    assert actual.data_ptr() == again.data_ptr()
+    assert inverse.data_ptr() == inverse_again.data_ptr()
+
+@pytest.mark.parametrize("batch,heads,rotary_dim", [(1, 3, 32), (2, 7, 96), (1, 56, 128)])
+@pytest.mark.parametrize("prefix_tiles", [1, 2, 3])
+@pytest.mark.parametrize("tile", [64, 128])
+@pytest.mark.parametrize("compact", [False, True])
+def test_pooled_queries_match_full_materialization(batch, heads, rotary_dim, prefix_tiles, tile, compact):
+    from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import gather_qknorm_rope, qknorm_rope_quantize_fp4
+    from fastvideo.models.dits.minimax_h3_fusions.query_pool import prepare_pooled_queries
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _fp8_per_token, _pool_tiles
+
+    torch.manual_seed(6521)
+    rows, packed_rows = 512, 337 if compact else 512
+    prefix_rows = prefix_tiles * tile
+    x = torch.randn((batch, packed_rows, heads * 2, 128), device="cuda", dtype=torch.bfloat16)[:, :, ::2]
+    positions = torch.randperm(rows, device="cuda")[:packed_rows]
+    inverse = torch.full((rows,), -1, device="cuda", dtype=torch.int64)
+    inverse[positions] = torch.arange(packed_rows, device="cuda")
+    valid = (inverse >= 0).reshape(-1, tile).sum(-1).to(torch.int32)
+    if not compact:
+        inverse = None
+        valid = torch.randint(1, tile + 1, (rows // tile,), device="cuda", dtype=torch.int32)
+        # Direct padded inputs must retain all row values, including invalid rows.
+        x.copy_(torch.ldexp(x.float(), torch.randint(-20, 9, x.shape, device=x.device)).to(x.dtype))
+    weight = torch.randn(128, device="cuda", dtype=x.dtype)
+    cos, sin = (torch.randn((rows, rotary_dim), device="cuda", dtype=x.dtype) for _ in range(2))
+    with torch.inference_mode():
+        if compact:
+            full, packed = gather_qknorm_rope(
+                x, weight, cos, sin, 1e-6, inverse, fp4_start=prefix_rows // 128 * 128,
+            )
+        else:
+            full, packed = qknorm_rope_quantize_fp4(x, weight, cos, sin, 1e-6, prefix_rows // 128 * 128)
+        expected_pool = _pool_tiles(full, valid, tile)
+        prefix, actual_packed, pool = prepare_pooled_queries(
+            x, weight, cos, sin, 1e-6, inverse, valid, prefix_rows, tile_elems=tile,
+        )
+    assert prefix.shape == (batch, prefix_rows, heads, 128)
+    assert torch.equal(prefix, full[:, :prefix_rows])
+    assert torch.equal(actual_packed[0], packed[0])
+    assert torch.equal(actual_packed[1].view(torch.uint8), packed[1].view(torch.uint8))
+    assert torch.equal(pool.contiguous().view(torch.int32), expected_pool.contiguous().view(torch.int32))
+    if tile == 128:
+        for split_prefix in (False, True):
+            with torch.inference_mode():
+                prefix8, packed8, pool8 = prepare_pooled_queries(
+                    x, weight, cos, sin, 1e-6, inverse, valid, prefix_rows, tile_elems=tile,
+                    prefix_fp8=True, split_prefix=split_prefix,
+                )
+                expected_prefix8 = _fp8_per_token(full[:, :prefix_rows])
+            assert torch.equal(prefix8, expected_prefix8)
+            assert torch.equal(packed8[0], packed[0])
+            assert torch.equal(packed8[1].view(torch.uint8), packed[1].view(torch.uint8))
+            assert torch.equal(pool8.contiguous().view(torch.int32), expected_pool.contiguous().view(torch.int32))
+
+
+
+@pytest.mark.parametrize("batch,heads", [(1, 3), (2, 7), (1, 56)])
+@pytest.mark.parametrize("large_range", [False, True])
+@pytest.mark.parametrize("tile", [64, 128])
+@pytest.mark.parametrize("compact", [False, True])
+def test_combined_kv_preserves_separate_prefix_pool_and_native_bytes(batch, heads, large_range, tile, compact):
+    from fastvideo.models.dits.minimax_h3_fusions.key_quant import pool_quantize_smoothed_keys
+    from fastvideo.models.dits.minimax_h3_fusions.kv_pool import pool_quantize_prefix_kv
+    from fastvideo.models.dits.minimax_h3_fusions.value_quant import pool_quantize_values
+
+    torch.manual_seed(7713)
+    rows, compact_rows = 512, 337 if compact else 512
+    key = torch.randn((batch, rows, heads * 2, 128), device="cuda", dtype=torch.bfloat16)[:, :, ::2]
+    value = torch.randn((batch, compact_rows, heads * 2, 128), device="cuda", dtype=torch.bfloat16)[:, :, ::2]
+    if large_range:
+        key.mul_(torch.exp2(torch.randint(-12, 12, key.shape, device="cuda").float()).bfloat16())
+        value.mul_(torch.exp2(torch.randint(-12, 12, value.shape, device="cuda").float()).bfloat16())
+    source = torch.randint(compact_rows, (rows,), device="cuda", dtype=torch.int64)
+    source[torch.rand(rows, device="cuda") < 0.2] = -1
+    if compact:
+        key[:, source < 0] = 0
+    valid_index = (source >= 0).nonzero().flatten()
+    inverse = torch.full_like(source, -1)
+    inverse[valid_index] = torch.arange(valid_index.numel(), device="cuda")
+    valid = (source >= 0).reshape(-1, tile).sum(1).to(torch.int32)
+    if not compact:
+        source = None
+    mean = key.mean(dim=1, keepdim=True)
+    with torch.inference_mode():
+        expected_key = pool_quantize_smoothed_keys(
+            key, value, mean, inverse, valid_index.numel(), valid, value_index=source, tile_elems=tile,
+        )
+        expected_value = pool_quantize_values(value, valid, input_index=source, tile_elems=tile)
+        actual_key, actual_value = pool_quantize_prefix_kv(
+            key, value, mean, inverse, valid_index.numel(), valid, source, tile_elems=tile,
+        )
+    for actual, expected in ((actual_key[0][0], expected_key[0][0]),
+                             (actual_key[0][1], expected_key[0][1]),
+                             (actual_key[1], expected_key[1]), (actual_key[2], expected_key[2]),
+                             (actual_key[3], expected_key[3]), (actual_value[0], expected_value[0]),
+                             (actual_value[1][0], expected_value[1][0]),
+                             (actual_value[1][1], expected_value[1][1])):
+        assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
+    assert actual_key[1].transpose(1, 2).is_contiguous()
+    assert actual_key[2].transpose(1, 2).is_contiguous()
+
+
+@pytest.mark.parametrize("batch", [1, 2])
+def test_preconverted_prefix_queries_preserve_dense_attention(batch):
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _fp8_per_token, _prefix_dense_attention
+
+    torch.manual_seed(4917)
+    q, k, v = (torch.randn((batch, 256, 3, 128), device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    positions = torch.arange(211, device="cuda")
+    with torch.inference_mode():
+        k = k - k.mean(dim=1, keepdim=True)
+        prepared = (_fp8_per_token(k[:, :211]).contiguous(), v[:, :211].contiguous())
+        expected = _prefix_dense_attention(q, k, v, 128, 256, positions, "fp8",
+                                            key_is_smoothed=True, prepared_kv=prepared)
+        actual = _prefix_dense_attention(_fp8_per_token(q[:, :128]), k, v, 128, 256, positions, "fp8",
+                                          key_is_smoothed=True, prepared_kv=prepared, query_is_fp8=True)
+    assert torch.equal(actual, expected)

fastvideo/tests/transformers/test_minimax_h3_fused_quantization.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Blackwell byte parity for H3 producer/FP4 quantization fusions."""
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+    reason="requires a Blackwell GPU with native FP4 conversion",
+)
+
+
+@pytest.mark.parametrize("batch,rows,heads,prefix", [(1, 129, 3, 0), (2, 257, 4, 64), (1, 383, 56, 192)])
+def test_quantized_gate_matches_separate_gather_addcmul_and_quantization(batch, rows, heads, prefix):
+    from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
+    from fastvideo.models.dits.minimax_h3_fusions.untile_gate import untile_gate
+
+    torch.manual_seed(935)
+    tiled = 512
+    start = prefix // 128 * 128
+    full = torch.randn((batch, heads, tiled, 128), device="cuda", dtype=torch.bfloat16)
+    output = (full[:, :, :prefix], full[:, :, start:], start) if prefix else full
+    compressed = torch.randn((batch, heads, tiled // 64, 128), device="cuda", dtype=full.dtype).permute(0, 2, 1, 3)
+    gate = torch.randn((batch, rows, heads * 128), device="cuda", dtype=full.dtype)
+    index = torch.randperm(tiled, device="cuda")[:rows]
+    scale = torch.ones((), device="cuda")
+    layout, _, _ = _require_flashinfer()
+    with torch.inference_mode():
+        expected = full.transpose(1, 2).index_select(1, index)
+        expected.addcmul_(compressed.index_select(1, index // 64), gate.unflatten(-1, (heads, 128)))
+        q, s = _nvfp4_quantize(expected.flatten(0, 1).flatten(1), scale, sfLayout=layout.layout_128x4)
+        actual_q, actual_s, actual_scale = untile_gate(output, compressed, gate, index, index // 64, scale)
+    assert torch.equal(q, actual_q)
+    assert torch.equal(s, actual_s)
+    assert actual_scale is scale
+
+
+@pytest.mark.parametrize("batch,rows,width,global_scale", [(1, 129, 64, 1.0), (2, 257, 5376, 0.7079)])
+def test_quantized_residual_modulation_matches_existing_fusion(batch, rows, width, global_scale):
+    from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
+    from fastvideo.models.dits.minimax_h3_fusions.modulation import fused_residual_gate_rmsnorm_modulate
+    from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import residual_modulate_quantize_nvfp4
+
+    torch.manual_seed(831)
+    residual, branch = (torch.randn((batch, rows, width), device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    gate, scale, shift = (torch.randn((3, width), device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    weight = torch.randn((width,), device="cuda", dtype=torch.bfloat16)
+    index = torch.randint(0, 3, (rows,), device="cuda")
+    sf = torch.tensor(global_scale, device="cuda")
+    args = residual, branch, gate, weight, scale, shift, index, 1e-5
+    layout, _, _ = _require_flashinfer()
+    with torch.inference_mode():
+        expected_hidden, normalized = fused_residual_gate_rmsnorm_modulate(*args)
+        q, s = _nvfp4_quantize(normalized.reshape(-1, width), sf, sfLayout=layout.layout_128x4)
+        hidden, (actual_q, actual_s, actual_sf) = residual_modulate_quantize_nvfp4(*args, sf)
+    assert torch.equal(hidden, expected_hidden)
+    assert torch.equal(q, actual_q)
+    assert torch.equal(s, actual_s)
+    assert actual_sf is sf
+
+
+@pytest.mark.parametrize("batch,rows,width", [(1, 129, 64), (2, 257, 5376)])
+def test_quantized_modulation_matches_existing_fusion(batch, rows, width):
+    from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
+    from fastvideo.models.dits.minimax_h3_fusions.modulation import fused_rmsnorm_modulate
+    from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import modulate_quantize_nvfp4
+
+    torch.manual_seed(834)
+    x = torch.randn((batch, rows, width), device="cuda", dtype=torch.bfloat16)
+    scale, shift = (torch.randn((3, width), device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    weight = torch.randn((width,), device="cuda", dtype=torch.bfloat16)
+    index = torch.randint(0, 3, (rows,), device="cuda")
+    sf = torch.ones((), device="cuda")
+    args = x, weight, scale, shift, index, 1e-5
+    layout, _, _ = _require_flashinfer()
+    with torch.inference_mode():
+        normalized = fused_rmsnorm_modulate(*args)
+        q, s = _nvfp4_quantize(normalized.reshape(-1, width), sf, sfLayout=layout.layout_128x4)
+        actual_q, actual_s, actual_sf = modulate_quantize_nvfp4(*args, sf)
+    assert torch.equal(q, actual_q)
+    assert torch.equal(s, actual_s)
+    assert actual_sf is sf
+
+
+@pytest.mark.parametrize("batch,rows,width", [(1, 129, 64), (2, 257, 5376)])
+def test_next_block_preparation_preserves_bf16_residual_rounding(batch, rows, width):
+    from fastvideo.models.dits.minimax_h3 import _gated_residual
+    from fastvideo.models.dits.minimax_h3_fusions.modulation_quant import (
+        modulate_quantize_nvfp4, residual_modulate_quantize_nvfp4,
+    )
+
+    torch.manual_seed(9183)
+    residual, branch = (torch.randn((batch, rows, width), device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    gate, scale, shift = (torch.randn((3, width), device="cuda", dtype=residual.dtype) for _ in range(3))
+    weight = torch.randn(width, device="cuda", dtype=residual.dtype)
+    index = torch.randint(0, 3, (rows,), device="cuda")
+    global_scale = torch.ones((), device="cuda")
+    with torch.inference_mode():
+        expected_hidden = _gated_residual(residual, gate, index, branch)
+        expected = modulate_quantize_nvfp4(expected_hidden, weight, scale, shift, index, 1e-5, global_scale)
+        hidden, actual = residual_modulate_quantize_nvfp4(
+            residual, branch, gate, weight, scale, shift, index, 1e-5, global_scale,
+            round_residual_before_norm=True,
+        )
+    assert torch.equal(hidden, expected_hidden)
+    assert torch.equal(actual[0], expected[0])
+    assert torch.equal(actual[1], expected[1])
+    assert actual[2] is global_scale

fastvideo/tests/transformers/test_minimax_h3_key_norm_mean.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact GB10 normalization and ATen mean parity at the specialized H3 geometry."""
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 1),
+    reason="requires the measured GB10 reduction geometry",
+)
+
+
+@pytest.mark.parametrize("exponent,strided", [(0, False), (6, False), (12, False), (6, True)])
+def test_key_normalization_and_mean_match_original(exponent, strided):
+    from fastvideo.models.dits.minimax_h3_fusions.key_norm_mean import gather_key_norm_mean
+    from fastvideo.models.dits.minimax_h3_fusions.qknorm_rope import fused_qknorm_rope
+
+    torch.manual_seed(29017 + exponent)
+    rows, source_rows, heads, dim = 18432, 14859, 56, 128
+    width = dim * 2 if strided else dim
+    x = torch.randn((1, source_rows, heads, width), device="cuda", dtype=torch.bfloat16)
+    if exponent:
+        x *= torch.exp2(torch.randint(-exponent, exponent + 1, x.shape, device=x.device).to(x.dtype))
+    if strided:
+        x = x[..., ::2]
+    index = torch.randperm(rows, device=x.device).to(torch.int64)
+    index[index >= source_rows] = -1
+    weight = torch.randn(dim, device=x.device, dtype=x.dtype)
+    angles = torch.randn((rows, 48), device=x.device)
+    angles = torch.cat((angles, angles), -1)
+    cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
+    with torch.inference_mode():
+        expanded = x.index_select(1, index.clamp(min=0))
+        expanded[:, index < 0] = 0
+        expected = fused_qknorm_rope(expanded, weight, cos, sin, 1e-6)
+        expected_mean = expected.mean(dim=1, keepdim=True)
+        actual, actual_mean = gather_key_norm_mean(x, weight, cos, sin, 1e-6, index)
+    assert torch.equal(actual, expected)
+    assert torch.equal(actual_mean, expected_mean)
+    assert actual.is_contiguous() and actual_mean.is_contiguous()

fastvideo/tests/transformers/test_minimax_h3_key_quant.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SageAttention FP4 bytes and prefix FP8 bytes after fused key smoothing."""
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+    reason="requires a Blackwell GPU and the native FP4 attention extension",
+)
+
+
+@pytest.mark.parametrize("batch,tokens,heads", [(1, 384, 56), (2, 256, 3)])
+def test_smoothed_key_packing_matches_separate_sage_and_prefix_preparation(batch, tokens, heads):
+    from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _api, _fp8_per_token
+
+    torch.manual_seed(1362)
+    key, value = (torch.randn((batch, tokens, heads, 128), device="cuda", dtype=torch.bfloat16) for _ in range(2))
+    key[:, 0] = 0
+    index = torch.randperm(tokens, device="cuda")[:tokens - 31]
+    inverse = torch.full((tokens,), -1, device="cuda", dtype=torch.int64)
+    inverse[index] = torch.arange(index.numel(), device="cuda")
+    with torch.inference_mode():
+        mean = key.mean(dim=1, keepdim=True)
+        smooth = key - mean
+        expected_codes, expected_scales = _api().scale_and_quant_fp4_permute(smooth.transpose(1, 2))
+        expected_key = _fp8_per_token(smooth).index_select(1, index)
+        expected_value = value.index_select(1, index)
+        (codes, scales), prefix_key, prefix_value = quantize_smoothed_keys(key, value, mean, inverse, index.numel())
+    assert torch.equal(codes, expected_codes)
+    assert torch.equal(scales.view(torch.uint8), expected_scales.view(torch.uint8))
+    assert torch.equal(prefix_key, expected_key)
+    assert torch.equal(prefix_value, expected_value)

fastvideo/tests/transformers/test_minimax_h3_prefix_quant.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CUDA byte parity for combined FP8 prefix preparation and K/V gather."""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+    reason="requires a Blackwell GPU",
+)
+
+
+@pytest.mark.parametrize("batch,heads,rows,strided", [(1, 56, 257, False), (2, 3, 19, True)])
+def test_prefix_quant_gather_matches_existing_fp8_roundtrip(batch, heads, rows, strided):
+    from fastvideo.models.dits.minimax_h3_fusions.prefix_quant import quantize_gather_prefix_kv
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _fp8_per_token
+
+    torch.manual_seed(9941)
+    key, value = (torch.randn((batch, rows * 2, heads, 128), device="cuda", dtype=torch.bfloat16)
+                  for _ in range(2))
+    if strided:
+        key, value = key[:, ::2], value[:, ::2]
+    key[:, 0] = 0
+    key[:, 1] *= 1e-12
+    index = torch.randperm(key.shape[1], device="cuda")[:rows - 1]
+    with torch.inference_mode():
+        expected_key = _fp8_per_token(key).index_select(1, index)
+        expected_value = value.index_select(1, index)
+        actual_key, actual_value = quantize_gather_prefix_kv(key, value, index)
+    assert torch.equal(actual_key, expected_key)
+    assert torch.equal(actual_value, expected_value)
+    assert actual_key.is_contiguous() and actual_value.is_contiguous()

fastvideo/tests/transformers/test_minimax_h3_sp_input_gather.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact FP4 input gathering and source-ordered head packets on Blackwell."""
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and Triton")
+
+
+def _swizzle(logical):
+    rows, groups = logical.shape
+    return logical.reshape(rows // 128, 4, 32, groups // 4, 4).permute(0, 3, 2, 1, 4).contiguous().reshape(rows, groups)
+
+
+@pytest.mark.parametrize("rows", [1, 127, 128, 129, 7430])
+@pytest.mark.parametrize("groups", [4, 336])
+def test_join_scale_rows_preserves_unaligned_sources_and_zeros_padding(rows, groups):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_input_gather import join_scale_rows
+    torch.manual_seed(217)
+    padded, total = (rows + 127) // 128 * 128, (2 * rows + 127) // 128 * 128
+    logical = torch.randint(0, 256, (2, padded, groups), dtype=torch.uint8, device="cuda")
+    source = torch.stack([_swizzle(part) for part in logical])
+    expected = torch.zeros((total, groups), dtype=torch.uint8, device="cuda")
+    expected[:2 * rows] = logical[:, :rows].reshape(2 * rows, groups)
+    assert torch.equal(join_scale_rows(source, rows), _swizzle(expected))
+
+
+@pytest.mark.parametrize("rows", [1, 127, 7430])
+@pytest.mark.parametrize("projection", [0, 1, 2])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_gathered_head_packets_match_original_exchange(rows, projection, rank):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_input_gather import pack_gathered
+    from fastvideo.models.dits.minimax_h3_fusions.sp_projection import allocate, pack_one, pack_value
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("ordered QK packing requires Blackwell")
+    torch.manual_seed(29176)
+    values = torch.randn((2, rows, 56, 128), device="cuda", dtype=torch.bfloat16)
+    values[:, ::127] = 0
+    weight = torch.randn((128,), device="cuda", dtype=torch.bfloat16)
+    weight[::17] = 0
+    cosine, sine = [torch.randn((2, rows, 96), device="cuda", dtype=torch.bfloat16) for _ in range(2)]
+    gathered = values[:, :, rank * 28:(rank + 1) * 28].contiguous().reshape(2 * rows, 3584)
+    payload, scale = allocate(rows, 56, 2, 7, values.device)
+    pack_gathered(gathered, payload, scale, projection, weight, cosine.flatten(0, 1), sine.flatten(0, 1), 1e-5)
+    for source in range(2):
+        expected, expected_scale = allocate(rows, 56, 2, 7, values.device)
+        if projection == 2:
+            pack_value(values[source], expected, expected_scale)
+        else:
+            pack_one(values[source], weight, cosine[source], sine[source], expected, expected_scale, 1e-5, projection)
+        assert torch.equal(payload[:, source, projection].contiguous().view(torch.uint8),
+                           expected[:, rank, projection].contiguous().view(torch.uint8))
+        assert torch.equal(scale[:, source, projection], expected_scale[:, rank, projection])
+
+
+def test_gathered_input_head_gemm_matches_sequence_projection():
+    from flashinfer import SfLayout, nvfp4_quantize
+    from fastvideo.layers.quantization.nvfp4_config import _mm_fp4
+    from fastvideo.models.dits.minimax_h3_fusions.sp_input_gather import join_scale_rows
+    if torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip("production route is restricted to GB10")
+    torch.manual_seed(62071)
+    rows, inner, columns = 7430, 5376, 7168
+    unit = torch.ones((), device="cuda", dtype=torch.float32)
+    x = torch.randn((2, rows, inner), device="cuda", dtype=torch.bfloat16)
+    x[1] *= 8
+    x[:, ::127] = 0
+    w = torch.randn((columns, inner), device="cuda", dtype=torch.bfloat16) / inner**.5
+    global_scale = torch.full((), 13721.5, device="cuda", dtype=torch.float32)
+    b, sb = nvfp4_quantize(w, global_scale, sfLayout=SfLayout.layout_128x4, do_shuffle=False)
+    inputs = [nvfp4_quantize(part, unit, sfLayout=SfLayout.layout_128x4, do_shuffle=False) for part in x]
+    full_a = torch.cat([p[0] for p in inputs])
+    full_sa = join_scale_rows(torch.stack([p[1] for p in inputs]), rows)
+    expected = [_mm_fp4(a, b.T, sa, sb.T, 1.0 / global_scale, torch.bfloat16, None, backend="cudnn")
+                for a, sa in inputs]
+    for rank in (0, 1):
+        start, stop = rank * 3584, (rank + 1) * 3584
+        output = _mm_fp4(full_a, b[start:stop].T, full_sa, sb[start:stop].T,
+                         1.0 / global_scale, torch.bfloat16, None, backend="cudnn")
+        target = torch.cat([part[:, start:stop] for part in expected])
+        assert torch.equal(output.view(torch.int16), target.view(torch.int16))
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_completed_packets_order_other_stream_and_retain_source(rank):
+    from fastvideo.distributed.projection_exchange import completed_projection_packets
+    producer, consumer = torch.cuda.Stream(), torch.cuda.Stream()
+    with torch.cuda.stream(producer):
+        payload = torch.empty((7, 2, 3, 127, 4, 128), device="cuda", dtype=torch.float8_e4m3fn)
+        scale = torch.empty(payload.shape[:-1], device="cuda", dtype=torch.float32)
+        torch.cuda._sleep(5_000_000)
+        payload.view(torch.uint8).fill_(37)
+        scale.fill_(2.5)
+        packets = completed_projection_packets(payload, scale, rank)
+    with torch.cuda.stream(consumer):
+        received = [(p.wait().materialize(), s.wait().clone()) for p, s in packets]
+    consumer.synchronize()
+    del payload, scale
+    for payload, scale in received:
+        assert torch.all(payload.view(torch.uint8) == 37)
+        assert torch.all(scale == 2.5)
+
+
+def test_input_gather_guard_rejects_other_shapes_and_backend(monkeypatch):
+    from fastvideo.models.dits.minimax_h3_fusions import sp_input_gather
+    if torch.cuda.get_device_capability() != (12, 1):
+        pytest.skip("production route is restricted to GB10")
+    hidden = torch.empty((1, 7430, 5376), device="cuda", dtype=torch.bfloat16)
+    linear = SimpleNamespace(_nvfp4_weight=torch.empty((7168, 2688), dtype=torch.uint8, device="cuda"),
+                             _nvfp4_weight_scale=torch.empty((7168, 336), dtype=torch.uint8, device="cuda"),
+                             bias=None)
+    attn = SimpleNamespace(num_attention_heads=56, to_q=linear, to_k=linear, to_v=linear)
+    monkeypatch.setattr(sp_input_gather, "_mm_fp4_backend", lambda: "cudnn")
+    assert sp_input_gather.supports_input_gather(attn, hidden)
+    assert not sp_input_gather.supports_input_gather(attn, hidden[:, :-1])
+    monkeypatch.setattr(sp_input_gather, "_mm_fp4_backend", lambda: "cutlass")
+    assert not sp_input_gather.supports_input_gather(attn, hidden)
+    monkeypatch.setattr(sp_input_gather, "_mm_fp4_backend", lambda: "cudnn")
+    linear.bias = torch.zeros((7168,), device="cuda")
+    assert not sp_input_gather.supports_input_gather(attn, hidden)

fastvideo/tests/transformers/test_minimax_h3_sp_key_mean.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Pooled key reductions retain the complete padded-sequence mean."""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA Triton reductions")
+
+
+@pytest.mark.parametrize("tiles,heads,scale", [(6, 2, 0.01), (6, 7, 10.0), (254, 4, 1.0), (256, 4, 0.1)])
+def test_pooled_key_mean_matches_float64_reference_within_bf16_rounding(tiles, heads, scale):
+    from fastvideo.attention.backends.video_sparse_attn_h3 import _pool_tiles
+    from fastvideo.models.dits.minimax_h3_fusions.sp_key_mean import key_mean_from_tiles
+
+    torch.manual_seed(511)
+    sequence = tiles * 64
+    key = (torch.randn((1, sequence, heads, 128), device="cuda") * scale).to(torch.bfloat16)
+    counts = torch.randint(0, 65, (tiles,), device="cuda", dtype=torch.int32)
+    counts[0] = 0
+    # Irregular holes exercise partial tiles, not only power-of-two counts.
+    positions = torch.arange(64, device="cuda")[None, :]
+    keep = positions < counts[:, None]
+    key.view(1, tiles, 64, heads, 128).masked_fill_(~keep[None, :, :, None, None], 0)
+    with torch.inference_mode():
+        pool = _pool_tiles(key, counts, 64).float()
+        actual = key_mean_from_tiles(pool, counts, sequence)
+        reference = key.double().mean(dim=1, keepdim=True)
+    # The final BF16 rounding dominates FP32 tile-reduction error. Permit one
+    # BF16 spacing plus a small cancellation allowance, without requiring the
+    # different FP32 reduction trees to select identical final BF16 values.
+    error = (actual.double() - reference).abs()
+    assert (error <= reference.abs() * 0.008 + scale * 1e-7).all()
+    assert actual.shape == (1, 1, heads, 128) and actual.dtype == torch.bfloat16

fastvideo/tests/transformers/test_minimax_h3_sp_qkv_pack.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SP transport preparation must preserve normalized BF16 values and FP8 codes/scales."""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and Triton")
+
+
+@pytest.mark.parametrize("rows,heads,dim,rotary,world", [
+    (17, 2, 64, 48, 1),
+    (33, 6, 128, 96, 2),
+    (47, 9, 128, 128, 3),
+    (7430, 56, 128, 96, 2),
+])
+def test_normalized_qkv_transport_matches_separate_bf16_rope_and_fp8_pack(rows, heads, dim, rotary, world):
+    from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
+    from fastvideo.models.dits.minimax_h3_fusions.sp_qkv_pack import pack_normalized_qkv_fp8
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _pack_heads_fp8
+
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("native FP8 conversion requires Hopper or newer")
+    torch.manual_seed(823)
+    query, key, value = [torch.randn((1, rows, heads, dim), device="cuda", dtype=torch.bfloat16) for _ in range(3)]
+    q_weight, k_weight = [torch.randn(dim, device="cuda", dtype=torch.bfloat16) for _ in range(2)]
+    theta = torch.randn((rows, rotary), device="cuda")
+    cos, sin = theta.cos().bfloat16(), theta.sin().bfloat16()
+    query[:, 0].zero_()
+    key[:, -1].zero_()
+    value[:, 0].zero_()
+    with torch.no_grad():
+        expected = _pack_heads_fp8(fused_qknorm_rope(query, q_weight, cos, sin, 1e-6)[0],
+                                   fused_qknorm_rope(key, k_weight, cos, sin, 1e-5)[0], value[0], world)
+        actual = pack_normalized_qkv_fp8(query[0], key[0], value[0], q_weight, k_weight, cos, sin, 1e-6, 1e-5, world)
+    for candidate, reference in zip(actual, expected, strict=True):
+        assert torch.equal(candidate.view(torch.uint8), reference.view(torch.uint8))
+
+
+@pytest.mark.parametrize("rows,heads,dim,rotary,world,chunks", [
+    (17, 12, 64, 48, 2, 3),
+    (33, 8, 128, 96, 2, 2),
+    (7430, 56, 128, 96, 2, 2),
+    (7430, 56, 128, 96, 2, 4),
+    (7430, 56, 128, 96, 2, 7),
+    (7430, 56, 128, 96, 2, 14),
+    (7430, 56, 128, 96, 2, 28),
+])
+def test_chunked_qkv_transport_matches_separate_pack_and_permute(rows, heads, dim, rotary, world, chunks):
+    from fastvideo.models.dits.minimax_h3_fusions import fused_qknorm_rope
+    from fastvideo.models.dits.minimax_h3_fusions.sp_qkv_pack import pack_normalized_qkv_fp8
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _pack_heads_fp8
+
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("native FP8 conversion requires Hopper or newer")
+    torch.manual_seed(324)
+    q, k, v = [torch.randn((1, rows, heads, dim), device="cuda", dtype=torch.bfloat16) for _ in range(3)]
+    qw, kw = [torch.randn(dim, device="cuda", dtype=torch.bfloat16) for _ in range(2)]
+    theta = torch.randn((rows, rotary), device="cuda")
+    cos, sin = theta.cos().bfloat16(), theta.sin().bfloat16()
+    with torch.no_grad():
+        p, s = _pack_heads_fp8(fused_qknorm_rope(q, qw, cos, sin, 1e-6)[0],
+                               fused_qknorm_rope(k, kw, cos, sin, 1e-5)[0], v[0], world)
+        expected = (p.view(world, 3, rows, chunks, heads // (world * chunks), dim).permute(3, 0, 1, 2, 4, 5).contiguous(),
+                    s.view(world, 3, rows, chunks, heads // (world * chunks)).permute(3, 0, 1, 2, 4).contiguous())
+        actual = pack_normalized_qkv_fp8(q[0], k[0], v[0], qw, kw, cos, sin, 1e-6, 1e-5, world, chunks)
+    for candidate, reference in zip(actual, expected, strict=True):
+        assert torch.equal(candidate.view(torch.uint8), reference.view(torch.uint8))
+
+
+@pytest.mark.parametrize("rows,heads,rotary,world,chunks", [
+    (17, 4, 32, 2, 1), (33, 8, 96, 2, 2), (47, 12, 128, 3, 2),
+    (7430, 56, 96, 2, 4), (7430, 56, 96, 2, 7), (7430, 56, 96, 2, 14),
+])
+def test_separate_projection_packing_matches_joint_qkv_bytes(rows, heads, rotary, world, chunks):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_projection import allocate, pack_one, pack_value
+    from fastvideo.models.dits.minimax_h3_fusions.sp_qkv_pack import pack_normalized_qkv_fp8
+
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("ordered parallel QK producer requires Blackwell")
+    torch.manual_seed(6039)
+    q, k, v = [torch.randn((rows, heads, 128), device="cuda", dtype=torch.bfloat16) for _ in range(3)]
+    qw, kw = [torch.randn(128, device="cuda", dtype=torch.bfloat16) for _ in range(2)]
+    theta = torch.randn((rows, rotary), device="cuda")
+    cos, sin = theta.cos().bfloat16(), theta.sin().bfloat16()
+    q[0].zero_()
+    k[-1].zero_()
+    v[0].zero_()
+    with torch.no_grad():
+        expected = pack_normalized_qkv_fp8(q, k, v, qw, kw, cos, sin, 1e-6, 1e-5, world, chunks)
+        payload, scale = allocate(rows, heads, world, chunks, q.device)
+        pack_one(q, qw, cos, sin, payload, scale, 1e-6, 0)
+        pack_one(k, kw, cos, sin, payload, scale, 1e-5, 1)
+        pack_value(v, payload, scale)
+    actual = (payload[0], scale[0]) if chunks == 1 else (payload, scale)
+    for candidate, reference in zip(actual, expected, strict=True):
+        assert torch.equal(candidate.view(torch.uint8), reference.view(torch.uint8))

fastvideo/tests/transformers/test_minimax_h3_sp_receive_preparation.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+"""SP fused receives retain native quantization, pooling, padding and head order."""
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+    reason="requires Blackwell native FP4 conversion",
+)
+
+
+@pytest.mark.parametrize("world,rows,heads,tiled,start", [
+    (1, 73, 2, 128, 0), (2, 121, 7, 384, 128), (3, 83, 3, 384, 0), (2, 7430, 7, 16384, 768),
+    (2, 7430, 4, 16384, 768), (2, 7430, 2, 16384, 768), (2, 7430, 1, 16384, 768),
+])
+@pytest.mark.parametrize("peer_rank", [None, 0, 1])
+@pytest.mark.parametrize("compact_queries", [False, True])
+@pytest.mark.parametrize("compact_values", [False, True])
+def test_receive_means_and_native_query_value_quantization_are_byte_exact(
+        world, rows, heads, tiled, start, peer_rank, compact_queries, compact_values):
+    if peer_rank is not None and world != 2:
+        pytest.skip("direct-local payloads require two ranks")
+    from attn_qat_infer import api
+    from fastvideo.attention.backends.video_sparse_attn_h3 import _pool_tiles
+    from fastvideo.models.dits.minimax_h3_fusions.sp_receive import unpack_pool_quantize_fp8
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _unpack_seq_fp8
+
+    torch.manual_seed(3927)
+    dim, length = 128, world * rows - 1
+    payload = torch.randn((world, 3, rows, heads, dim), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand(payload.shape[:-1], device="cuda") * 0.5
+    payload[:, :, 0] = 0
+    untile = torch.randperm(tiled, device="cuda")[:length]
+    inverse = torch.full((tiled,), -1, device="cuda", dtype=torch.int64)
+    inverse[untile] = torch.arange(length, device="cuda")
+    valid = (inverse >= 0).view(-1, 64).sum(1).clamp_min(1).to(torch.int32)
+    expected = torch.zeros((3, tiled, heads, dim), device="cuda", dtype=torch.bfloat16)
+    actual = torch.full_like(expected, float("nan"))
+    query_rows = min(start + 128, tiled) if compact_queries else None
+    positions = (inverse >= 0).nonzero().flatten()
+    prefix_inverse = torch.full_like(inverse, -1)
+    prefix_inverse[positions] = torch.arange(positions.numel(), device="cuda")
+    prefix_value = torch.empty((1, positions.numel(), heads, dim), device="cuda", dtype=torch.bfloat16)
+    prefix = (prefix_inverse, prefix_value) if compact_values else None
+    with torch.inference_mode():
+        expected.index_copy_(1, untile, _unpack_seq_fp8(payload, scales, length))
+        pool = torch.cat([_pool_tiles(expected[i:i + 1], valid, 64) for i in range(3)])
+        q_quant = api.scale_and_quant_fp4(expected[0:1, start:].transpose(1, 2))
+        v_quant = api.scale_and_quant_fp4_transpose(expected[2:3].transpose(1, 2))
+        received = payload
+        if peer_rank is not None:
+            from fastvideo.distributed.async_collectives import PeerSlices
+            received = PeerSlices(payload[peer_rank], payload[1 - peer_rank].clone(), torch.empty_like(payload), peer_rank)
+            assert torch.equal(received.materialize().view(torch.uint8), payload.view(torch.uint8))
+        got_pool, got_v, got_q = unpack_pool_quantize_fp8(
+            received, scales, inverse, valid, actual, start, query_rows=query_rows, prefix_values=prefix)
+    live_queries = tiled if query_rows is None else query_rows
+    assert torch.equal(actual[0, :live_queries].view(torch.uint8), expected[0, :live_queries].view(torch.uint8))
+    assert torch.isnan(actual[0, live_queries:]).all()
+    assert torch.equal(actual[1].view(torch.uint8), expected[1].view(torch.uint8))
+    if compact_values:
+        assert torch.isnan(actual[2]).all()
+        assert torch.equal(prefix_value, expected[2:3, positions])
+    else:
+        assert torch.equal(actual[2].view(torch.uint8), expected[2].view(torch.uint8))
+    assert torch.equal(got_pool.contiguous().view(torch.uint8), pool.contiguous().view(torch.uint8))
+    for candidate, reference in zip((*got_q, *got_v), (*q_quant, *v_quant), strict=True):
+        assert torch.equal(candidate.view(torch.uint8), reference.view(torch.uint8))
+
+
+@pytest.mark.parametrize("world,rows,heads,prefix,start,logical", [
+    (1, 79, 2, 64, 0, 128), (2, 123, 3, 128, 128, 384),
+    (3, 77, 7, 192, 128, 384), (2, 7430, 7, 832, 768, 16384),
+    (2, 7430, 4, 832, 768, 16384), (2, 7430, 2, 832, 768, 16384), (2, 7430, 1, 832, 768, 16384),
+])
+@pytest.mark.parametrize("strided", [False, True])
+def test_split_output_pack_matches_concat_gather_and_fp8_rounding(world, rows, heads, prefix, start, logical, strided):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_output_pack import pack_split_fp8
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _pack_seq_fp8
+
+    torch.manual_seed(907)
+    p = torch.randn((1, heads, prefix, 128), device="cuda", dtype=torch.bfloat16)
+    v = torch.randn((1, heads, logical - start, 128), device="cuda", dtype=torch.bfloat16)
+    if strided:
+        p = p.transpose(1, 2).contiguous().transpose(1, 2)
+        v = v.transpose(1, 2).contiguous().transpose(1, 2)
+    untile = torch.randperm(logical, device="cuda")[:world * rows - 1]
+    with torch.inference_mode():
+        expected = _pack_seq_fp8(torch.cat((p, v[:, :, prefix - start:]), dim=2), untile, world, rows)
+        actual = pack_split_fp8((p, v, start), untile, world, rows)
+    for candidate, reference in zip(actual, expected, strict=True):
+        assert torch.equal(candidate.view(torch.uint8), reference.view(torch.uint8))
+
+
+@pytest.mark.parametrize("prefix", [0, 64, 128, 192])
+def test_native_attention_does_not_read_omitted_query_suffix(prefix):
+    """Exercise the real native kernel with a poisoned raw query suffix."""
+    if torch.cuda.get_device_capability()[0] != 12:
+        pytest.skip("requires consumer Blackwell sparse FP4 attention")
+    from attn_qat_infer import api
+    from fastvideo.models.dits import minimax_h3_vsa_fp4 as impl
+    from fastvideo.models.dits.minimax_h3_fusions.sp_receive import unpack_pool_quantize_fp8
+
+    torch.manual_seed(713)
+    rows, heads, tiled = 191, 4, 384
+    start = prefix // 128 * 128
+    payload = torch.randn((2, 3, rows, heads, 128), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand(payload.shape[:-1], device="cuda") * 0.1
+    inverse = torch.arange(tiled, device="cuda", dtype=torch.int64)
+    inverse[2 * rows - 1:] = -1
+    valid = (inverse >= 0).view(-1, 64).sum(1).to(torch.int32)
+    full = torch.empty((3, tiled, heads, 128), device="cuda", dtype=torch.bfloat16)
+    compact = torch.full_like(full, float("nan"))
+    with torch.inference_mode():
+        unpack_pool_quantize_fp8(payload, scales, inverse, valid, full, start)
+        _, prepared_value, prepared_query = unpack_pool_quantize_fp8(
+            payload, scales, inverse, valid, compact, start, query_rows=min(start + 128, tiled))
+        mask = torch.eye(tiled // 64, device="cuda", dtype=torch.bool).expand(1, heads, -1, -1).clone()
+        mask[:, :, :prefix // 64] = True
+        mask[:, :, :, :prefix // 64] = True
+        lists = api.vsa_tile_mask_to_fp4_blocks(mask, 64, valid)
+        kwargs = dict(logical=tiled, prefix_rows=prefix, precision="fp8",
+                      options={"two_level_p": True, "smooth_k": True},
+                      prepared_query=prepared_query, prepared_value=prepared_value)
+        expected = impl._fp4_fine_attention(api, full[0:1], full[1:2], full[2:3], *lists, **kwargs)
+        actual = impl._fp4_fine_attention(api, compact[0:1], compact[1:2], compact[2:3], *lists, **kwargs)
+    assert actual.isfinite().all()
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+
+
+@pytest.mark.parametrize("prefix", [64, 128, 192])
+@pytest.mark.parametrize("pooled_mean", [False, True])
+@pytest.mark.parametrize("direct_keys", [False, True])
+@pytest.mark.parametrize("peer_rank", [None, 0, 1])
+def test_native_attention_accepts_direct_prefix_values_with_poisoned_raw_values(prefix, pooled_mean, direct_keys, peer_rank):
+    """The real native/prefix attention must not read omitted BF16 values."""
+    if torch.cuda.get_device_capability()[0] != 12:
+        pytest.skip("requires consumer Blackwell sparse FP4 attention")
+    from attn_qat_infer import api
+    from fastvideo.models.dits import minimax_h3_vsa_fp4 as impl
+    from fastvideo.models.dits.minimax_h3_fusions.sp_receive import unpack_pool_quantize_fp8
+    from fastvideo.models.dits.minimax_h3_fusions.key_quant import quantize_smoothed_keys
+    from fastvideo.models.dits.minimax_h3_fusions.sp_key_mean import key_mean_from_tiles
+
+    torch.manual_seed(1729)
+    rows, heads, tiled = 191, 4, 384
+    start = prefix // 128 * 128
+    payload = torch.randn((2, 3, rows, heads, 128), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand(payload.shape[:-1], device="cuda") * 0.1
+    inverse = torch.arange(tiled, device="cuda", dtype=torch.int64)
+    inverse[2 * rows - 1:] = -1
+    valid = (inverse >= 0).view(-1, 64).sum(1).to(torch.int32)
+    full = torch.empty((3, tiled, heads, 128), device="cuda", dtype=torch.bfloat16)
+    compact = torch.full_like(full, float("nan"))
+    layout = impl._KeyIndexCache()
+    layout.tile, layout.n_tiles = 64, tiled // 64
+    received = payload
+    if peer_rank is not None:
+        from fastvideo.distributed.async_collectives import PeerSlices
+        received = PeerSlices(payload[peer_rank], payload[1 - peer_rank], payload, peer_rank)
+    with torch.inference_mode():
+        index, prefix_inverse = layout.valid_keys_from_tiles(valid, api.HALF_N)
+        prefix_value = torch.empty((1, index.numel(), heads, 128), device="cuda", dtype=torch.bfloat16)
+        unpack_pool_quantize_fp8(payload, scales, inverse, valid, full, start)
+        pool, prepared_value, prepared_query = unpack_pool_quantize_fp8(
+            received, scales, inverse, valid, compact, start, query_rows=min(start + 128, tiled),
+            prefix_values=(prefix_inverse, prefix_value), store_keys=not direct_keys)
+        mean = key_mean_from_tiles(pool[1:2], valid, tiled) if pooled_mean else full[1:2].mean(dim=1, keepdim=True)
+        expected_keys = quantize_smoothed_keys(full[1:2], full[2:3], mean, prefix_inverse, index.numel())
+        keys = quantize_smoothed_keys(compact[1:2], compact[2:3], mean, prefix_inverse, index.numel(),
+                                      prefix_value=prefix_value,
+                                      received_keys=(received, scales, inverse) if direct_keys else None)
+        if direct_keys:
+            assert torch.isnan(compact[1]).all()
+        assert keys[2].data_ptr() == prefix_value.data_ptr()
+        for candidate, reference in zip((*keys[0], *keys[1:]), (*expected_keys[0], *expected_keys[1:]), strict=True):
+            assert torch.equal(candidate.view(torch.uint8), reference.view(torch.uint8))
+        assert torch.isnan(compact[2]).all()
+        mask = torch.eye(tiled // 64, device="cuda", dtype=torch.bool).expand(1, heads, -1, -1).clone()
+        mask[:, :, :prefix // 64] = True
+        mask[:, :, :, :prefix // 64] = True
+        lists = api.vsa_tile_mask_to_fp4_blocks(mask, 64, valid)
+        kwargs = dict(logical=tiled, prefix_rows=prefix, precision="fp8", layout=layout,
+                      options={"two_level_p": True, "smooth_k": True},
+                      prepared_query=prepared_query, prepared_value=prepared_value)
+        expected = impl._fp4_fine_attention(api, full[0:1], full[1:2], full[2:3], *lists,
+                                             prepared_keys=expected_keys, **kwargs)
+        actual = impl._fp4_fine_attention(api, compact[0:1], compact[1:2], compact[2:3], *lists,
+                                           prepared_keys=keys, **kwargs)
+    assert actual.isfinite().all()
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))

fastvideo/tests/transformers/test_minimax_h3_sp_unpack.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Exact parity for FP8 Ulysses receive unpacking on a CUDA GPU with Triton."""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA and Triton")
+
+
+@pytest.mark.parametrize("world,local_rows,heads,dim,trim", [
+    (1, 17, 2, 64, 0),
+    (2, 33, 3, 128, 1),
+    (3, 47, 19, 128, 2),
+    (2, 7430, 28, 128, 1),
+])
+def test_fp8_receive_tiles_match_separate_unpack_and_scatter(world, local_rows, heads, dim, trim):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_tiles
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _unpack_seq_fp8
+
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("native FP8 conversion requires Hopper or newer")
+    torch.manual_seed(178)
+    sequence_length = world * local_rows - trim
+    tile_rows = ((sequence_length + 127) // 128 + 1) * 128
+    payload = torch.randn((world, 3, local_rows, heads, dim), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand((world, 3, local_rows, heads), device="cuda", dtype=torch.float32).clamp_min_(1e-12)
+    untile = torch.randperm(tile_rows, device="cuda")[:sequence_length]
+    inverse = torch.full((tile_rows,), -1, device="cuda", dtype=torch.int64)
+    inverse[untile] = torch.arange(sequence_length, device="cuda")
+    # Ensure special finite FP8 values and distinct source-rank boundaries are covered.
+    payload[:, :, 0, 0, :4] = torch.tensor([0.0, -0.0, 448.0, -448.0], device="cuda").to(payload.dtype)
+    expected = torch.zeros((3, tile_rows, heads, dim), device="cuda", dtype=torch.bfloat16)
+    expected.index_copy_(1, untile, _unpack_seq_fp8(payload, scales, sequence_length))
+    output = torch.full_like(expected, float("nan"))
+    actual = unpack_fp8_tiles(payload, scales, inverse, output)
+    assert actual.data_ptr() == output.data_ptr()
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+    # Reuse must overwrite old values as well as all padding rows.
+    scales.mul_(0.25)
+    expected.zero_().index_copy_(1, untile, _unpack_seq_fp8(payload, scales, sequence_length))
+    unpack_fp8_tiles(payload, scales, inverse, output)
+    assert torch.equal(output.view(torch.uint8), expected.view(torch.uint8))
+
+
+@pytest.mark.parametrize("world,rows,heads,dim,tiles", [
+    (1, 17, 2, 64, 4),
+    (2, 33, 3, 128, 7),
+    (3, 47, 19, 128, 11),
+    (2, 7430, 28, 128, 256),
+])
+@pytest.mark.parametrize("peer_rank", [None, 0, 1])
+def test_fp8_output_gate_matches_separate_unpack_and_compiled_gate(world, rows, heads, dim, tiles, peer_rank):
+    if peer_rank is not None and world != 2:
+        pytest.skip("direct-local payloads require two ranks")
+    from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _apply_gate, _unpack_heads_fp8
+
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("native FP8 conversion requires Hopper or newer")
+    torch.manual_seed(735)
+    payload = torch.randn((world, rows, heads, dim), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand((world, rows, heads), device="cuda").clamp_min_(1e-12)
+    compressed = torch.randn((world, heads, tiles, dim), device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn((rows, world * heads, dim), device="cuda", dtype=torch.bfloat16)
+    row_tile = torch.randint(tiles, (rows,), device="cuda")
+    payload[:, 0, 0, :4] = torch.tensor([0.0, -0.0, 448.0, -448.0], device="cuda").to(payload.dtype)
+    expected = _apply_gate(_unpack_heads_fp8(payload, scales), compressed.flatten(0, 1).transpose(0, 1),
+                           row_tile, gate)
+    received = payload
+    if peer_rank is not None:
+        from fastvideo.distributed.async_collectives import PeerSlices
+        received = PeerSlices(payload[peer_rank], payload[1 - peer_rank].clone(), torch.empty_like(payload), peer_rank)
+    actual = unpack_fp8_gate(received, scales, compressed, gate, row_tile)
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+
+
+@pytest.mark.parametrize("world,rows,heads,dim,global_scale", [
+    (1, 17, 2, 64, 1.0),
+    (2, 129, 3, 128, 0.7079),
+    (3, 257, 19, 128, 1.0),
+    (2, 7430, 28, 128, 0.7079),
+])
+def test_fp8_output_gate_quantization_matches_separate_fp4(world, rows, heads, dim, global_scale):
+    from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
+    from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _apply_gate, _unpack_heads_fp8
+
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("native FP4 conversion requires Blackwell")
+    torch.manual_seed(156)
+    tiles = 256
+    payload = torch.randn((world, rows, heads, dim), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand((world, rows, heads), device="cuda").clamp_min_(1e-12)
+    compressed = torch.randn((world, heads, tiles, dim), device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn((rows, world * heads, dim), device="cuda", dtype=torch.bfloat16)
+    row_tile = torch.randint(tiles, (rows,), device="cuda")
+    sf = torch.tensor(global_scale, device="cuda")
+    layout, _, _ = _require_flashinfer()
+    with torch.inference_mode():
+        expected = _apply_gate(_unpack_heads_fp8(payload, scales), compressed.flatten(0, 1).transpose(0, 1),
+                               row_tile, gate)
+        q, s = _nvfp4_quantize(expected.flatten(1), sf, sfLayout=layout.layout_128x4)
+        actual_q, actual_s, actual_sf = unpack_fp8_gate(payload, scales, compressed, gate, row_tile, sf)
+    assert torch.equal(q, actual_q)
+    assert torch.equal(s, actual_s)
+    assert actual_sf is sf
+
+
+@pytest.mark.parametrize("world,rows,heads,dim,chunks", [(2, 33, 3, 128, 2), (3, 17, 2, 64, 3), (2, 7430, 14, 128, 2)])
+def test_chunked_output_gate_writes_full_rank_major_head_order(world, rows, heads, dim, chunks):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _apply_gate, _unpack_heads_fp8
+
+    if torch.cuda.get_device_capability()[0] < 9:
+        pytest.skip("native FP8 conversion requires Hopper or newer")
+    torch.manual_seed(943)
+    tiles = 256
+    full_heads = heads * chunks
+    payload = torch.randn((world, rows, full_heads, dim), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand((world, rows, full_heads), device="cuda").clamp_min_(1e-12)
+    compressed = torch.randn((world, full_heads, tiles, dim), device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn((rows, world * full_heads, dim), device="cuda", dtype=torch.bfloat16)
+    row_tile = torch.randint(tiles, (rows,), device="cuda")
+    expected = _apply_gate(_unpack_heads_fp8(payload, scales), compressed.flatten(0, 1).transpose(0, 1), row_tile, gate)
+    output = torch.full_like(expected, float("nan"))
+    for chunk in range(chunks):
+        start, stop = chunk * heads, (chunk + 1) * heads
+        result = unpack_fp8_gate(payload[:, :, start:stop].contiguous(), scales[:, :, start:stop].contiguous(),
+                                 compressed[:, start:stop].contiguous(), gate, row_tile, head_chunk=chunk,
+                                 head_chunks=chunks, output=output)
+        assert result.data_ptr() == output.data_ptr()
+    assert torch.equal(output.view(torch.uint8), expected.view(torch.uint8))
+
+
+@pytest.mark.parametrize("world,rows,heads,dim,chunks", [
+    (2, 33, 3, 128, 2), (3, 17, 2, 64, 3), (2, 7430, 14, 128, 2), (2, 257, 7, 128, 4),
+    (2, 7430, 4, 128, 7), (2, 7430, 2, 128, 14), (2, 7430, 1, 128, 28),
+])
+@pytest.mark.parametrize("peer_rank", [None, 0, 1])
+def test_chunked_fp4_output_matches_full_unpack_gate_and_quantization(world, rows, heads, dim, chunks, peer_rank):
+    if peer_rank is not None and world != 2:
+        pytest.skip("direct-local payloads require two ranks")
+    from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
+    from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_gate
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _apply_gate, _unpack_heads_fp8
+
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("native FP4 conversion requires Blackwell")
+    torch.manual_seed(245)
+    tiles, full_heads = 256, heads * chunks
+    payload = torch.randn((world, rows, full_heads, dim), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand((world, rows, full_heads), device="cuda").clamp_min_(1e-12)
+    compressed = torch.randn((world, full_heads, tiles, dim), device="cuda", dtype=torch.bfloat16)
+    gate = torch.randn((rows, world * full_heads, dim), device="cuda", dtype=torch.bfloat16)
+    row_tile = torch.randint(tiles, (rows,), device="cuda")
+    sf = torch.tensor(0.7079, device="cuda")
+    layout, _, _ = _require_flashinfer()
+    width = world * full_heads * dim
+    codes = torch.empty((rows, width // 2), dtype=torch.uint8, device="cuda")
+    fp4_scales = torch.zeros((((rows + 127) // 128) * 128, width // 16), dtype=torch.uint8, device="cuda")
+    with torch.inference_mode():
+        decoded = _apply_gate(_unpack_heads_fp8(payload, scales), compressed.flatten(0, 1).transpose(0, 1), row_tile, gate)
+        expected_q, expected_s = _nvfp4_quantize(decoded.flatten(1), sf, sfLayout=layout.layout_128x4)
+        for chunk in range(chunks):
+            start, stop = chunk * heads, (chunk + 1) * heads
+            received = payload[:, :, start:stop].contiguous()
+            if peer_rank is not None:
+                from fastvideo.distributed.async_collectives import PeerSlices
+                received = PeerSlices(received[peer_rank], received[1 - peer_rank].clone(), torch.empty_like(received), peer_rank)
+            result = unpack_fp8_gate(received, scales[:, :, start:stop].contiguous(),
+                                     compressed[:, start:stop].contiguous(), gate, row_tile, sf, head_chunk=chunk,
+                                     head_chunks=chunks, fp4_output=(codes, fp4_scales))
+            assert result[0].data_ptr() == codes.data_ptr()
+            assert result[1].data_ptr() == fp4_scales.data_ptr()
+    assert torch.equal(codes, expected_q)
+    assert torch.equal(fp4_scales, expected_s)
+
+
+@pytest.mark.parametrize("world,rows,heads,start", [(1, 17, 2, 0), (2, 33, 7, 128), (3, 47, 3, 0),
+                                                  (2, 7430, 7, 768)])
+def test_fp8_query_unpack_matches_bf16_scatter_and_native_fp4(world, rows, heads, start):
+    from fastvideo.models.dits.minimax_h3_fusions.sp_exchange import unpack_fp8_tiles
+    from fastvideo.models.dits.minimax_h3_vsa_fp4 import _api, _unpack_seq_fp8
+
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("native FP4 conversion requires Blackwell")
+    torch.manual_seed(357)
+    dim, sequence = 128, world * rows - 1
+    tile_rows = ((sequence + 127) // 128 + 1) * 128
+    payload = torch.randn((world, 3, rows, heads, dim), device="cuda").to(torch.float8_e4m3fn)
+    scales = torch.rand((world, 3, rows, heads), device="cuda").clamp_min_(1e-12)
+    untile = torch.randperm(tile_rows, device="cuda")[:sequence]
+    inverse = torch.full((tile_rows,), -1, device="cuda", dtype=torch.int64)
+    inverse[untile] = torch.arange(sequence, device="cuda")
+    expected = torch.zeros((3, tile_rows, heads, dim), device="cuda", dtype=torch.bfloat16)
+    expected.index_copy_(1, untile, _unpack_seq_fp8(payload, scales, sequence))
+    native = _api().scale_and_quant_fp4(expected[0:1, start:].transpose(1, 2))
+    output = torch.full_like(expected, float("nan"))
+    actual, packed = unpack_fp8_tiles(payload, scales, inverse, output, query_start=start)
+    assert actual.data_ptr() == output.data_ptr()
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+    for candidate, reference in zip(packed, native, strict=True):
+        assert torch.equal(candidate.view(torch.uint8), reference.view(torch.uint8))

fastvideo/tests/transformers/test_minimax_h3_swiglu_quant.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Blackwell-only byte parity for the fused H3 SwiGLU/NVFP4 preparation."""
+import pytest
+import torch
+
+
+@pytest.mark.parametrize("rows,width,scale", [(129, 64, 1.0), (257, 14336, 0.1395), (33, 256, 91.75)])
+def test_swiglu_quantization_matches_flashinfer(rows: int, width: int, scale: float) -> None:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("native FP4 conversion requires a Blackwell GPU")
+    pytest.importorskip("flashinfer")
+    from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
+    from fastvideo.models.dits.minimax_h3_fusions.swiglu import minimax_h3_swiglu
+    from fastvideo.models.dits.minimax_h3_fusions.swiglu_quant import swiglu_quantize_nvfp4
+
+    torch.manual_seed(321)
+    x = torch.randn((1, rows, width * 2), device="cuda", dtype=torch.bfloat16) * 3
+    x[:, 0] = 0
+    global_scale = torch.tensor(scale, device="cuda", dtype=torch.float32)
+    layout, _, _ = _require_flashinfer()
+    with torch.inference_mode():
+        expected_codes, expected_scales = _nvfp4_quantize(
+            minimax_h3_swiglu(x).view(rows, width), global_scale, sfLayout=layout.layout_128x4,
+        )
+        codes, scales, returned_scale = swiglu_quantize_nvfp4(x, global_scale)
+    assert torch.equal(codes, expected_codes)
+    assert torch.equal(scales, expected_scales)
+    assert returned_scale is global_scale

fastvideo/tests/transformers/test_minimax_h3_untile_gate.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CUDA parity for H3's packed-order compression gate and resident AdaLN cache."""
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU and Triton")
+
+
+@pytest.mark.parametrize("batch,rows,heads", [(1, 127, 4), (2, 257, 6), (1, 14859, 56)])
+def test_untile_gate_matches_aten_addcmul(batch, rows, heads):
+    from fastvideo.models.dits.minimax_h3_fusions.untile_gate import untile_gate
+
+    torch.manual_seed(751)
+    tiled_rows = ((rows + 127) // 128) * 128
+    out = torch.randn((batch, heads, tiled_rows, 128), device="cuda", dtype=torch.bfloat16)
+    compressed = torch.randn((batch, heads, tiled_rows // 64, 128), device="cuda", dtype=out.dtype)
+    compressed = compressed.permute(0, 2, 1, 3)
+    gate = torch.randn((batch, rows, heads * 128), device="cuda", dtype=out.dtype)
+    untile = torch.randperm(tiled_rows, device="cuda")[:rows]
+    row_tile = untile // 64
+    expected = out.transpose(1, 2).index_select(1, untile)
+    expected.addcmul_(compressed.index_select(1, row_tile), gate.unflatten(-1, (heads, 128)))
+    actual = untile_gate(out, compressed, gate, untile, row_tile)
+    assert torch.equal(actual, expected.flatten(2, 3))
+
+
+def test_resident_adaln_cache_preserves_projection_and_allocation():
+    from fastvideo.models.dits.minimax_h3 import MiniMaxH3AdaLayerNormModulation
+
+    torch.manual_seed(319)
+    mod = MiniMaxH3AdaLayerNormModulation(16, 12).to(device="cuda", dtype=torch.bfloat16)
+    # ReplicatedLinear leaves initialization to the loader.
+    with torch.no_grad():
+        mod.linear.weight.normal_()
+        mod.linear.bias.normal_()
+    inputs = [torch.randn((3, 16), device="cuda", dtype=torch.bfloat16) for _ in range(2)]
+    with torch.inference_mode():
+        expected = [mod(x) for x in inputs]
+        weight_ptr = mod.linear.weight.data_ptr()
+        mod.enable_host_cache(keep_resident=True)
+        assert mod._host_weight.data_ptr() == weight_ptr
+        for _ in range(2):
+            for key, x in enumerate(inputs):
+                mod._cache_key = key
+                assert all(torch.equal(a, b) for a, b in zip(mod(x), expected[key], strict=True))
+        assert len(mod._modulation_cache) == 2
+
+
+@pytest.mark.parametrize("prefix_rows", [64, 128, 192])
+def test_split_untile_gate_matches_concatenation(prefix_rows):
+    from fastvideo.models.dits.minimax_h3_fusions.untile_gate import untile_gate
+
+    torch.manual_seed(457)
+    batch, heads, rows, tiled_rows = 2, 3, 319, 512
+    start = prefix_rows // 128 * 128
+    prefix = torch.randn((batch, heads, prefix_rows, 128), device="cuda", dtype=torch.bfloat16)
+    video = torch.randn((batch, heads, tiled_rows - start, 128), device="cuda", dtype=torch.bfloat16)
+    compressed = torch.randn((batch, heads, tiled_rows // 64, 128), device="cuda", dtype=torch.bfloat16)
+    compressed = compressed.permute(0, 2, 1, 3)
+    gate = torch.randn((batch, rows, heads * 128), device="cuda", dtype=torch.bfloat16)
+    index = torch.randperm(tiled_rows, device="cuda")[:rows]
+    combined = torch.cat((prefix, video[:, :, prefix_rows - start:]), dim=2)
+    expected = combined.transpose(1, 2).index_select(1, index)
+    expected.addcmul_(compressed.index_select(1, index // 64), gate.unflatten(-1, (heads, 128)))
+    actual = untile_gate((prefix, video, start), compressed, gate, index, index // 64)
+    assert torch.equal(actual, expected.flatten(2, 3))
+
+
+@pytest.mark.parametrize("width,scale", [(64, 1.0), (5376, 0.7079)])
+def test_quantized_tile_matches_flashinfer_quantization(width, scale):
+    if torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("NVFP4 quantization requires Blackwell")
+    from fastvideo.layers.quantization.nvfp4_config import _nvfp4_quantize, _require_flashinfer
+    from fastvideo.models.dits.minimax_h3_fusions.quant_tile import tile_quantized
+
+    torch.manual_seed(319)
+    layout, _, _ = _require_flashinfer()
+    x = torch.randn((129, width), device="cuda", dtype=torch.bfloat16)
+    global_sf = torch.tensor(scale, device="cuda")
+    index = torch.randperm(384, device="cuda")[:129]
+    inverse = torch.full((384,), -1, device="cuda", dtype=torch.int64)
+    inverse[index] = torch.arange(129, device="cuda")
+    tiled = x.new_zeros((384, width)).index_copy_(0, index, x)
+    expected = _nvfp4_quantize(tiled, global_sf, sfLayout=layout.layout_128x4)
+    codes, scales = _nvfp4_quantize(x, global_sf, sfLayout=layout.layout_128x4)
+    actual = tile_quantized(codes, scales, inverse)
+    assert torch.equal(actual[0], expected[0])
+    assert torch.equal(actual[1], expected[1])

fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py

@@ -487,7 +487,7 @@ def _write_deploy_checkpoint(path, student_layers: int, base_layers: int, heads:
     return decoder
 
 
-def _load_h3_vae(vae_dir, cpu: bool = True, decoder: str | None = None):
+def _load_h3_vae(vae_dir, cpu: bool = True, decoder: str | None = None, quantize_dense: bool = False):
     from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig
     from fastvideo.configs.pipelines import PipelineConfig
     from fastvideo.fastvideo_args import FastVideoArgs
@@ -495,7 +495,8 @@ def _load_h3_vae(vae_dir, cpu: bool = True, decoder: str | None = None):
 
     args = FastVideoArgs(model_path=str(vae_dir),
                          pipeline_config=PipelineConfig(vae_config=MiniMaxH3VideoVAEConfig(), vae_precision="fp32"),
-                         h3_vae_nvfp4_decoder=decoder)
+                         h3_vae_nvfp4_decoder=decoder,
+                         h3_vae_nvfp4_quantize_dense=quantize_dense)
     args.vae_cpu_offload = cpu  # load on CPU
     return VAELoader().load(str(vae_dir), args)
 
@@ -666,12 +667,13 @@ def _fake_packed_checkpoint(vae, act_scale="unit"):
     }
 
 
-def test_loader_applies_packed_safetensors_without_master_weights(tmp_path):
+@pytest.mark.parametrize("loaded_layers", [5, 2])
+def test_loader_applies_packed_safetensors_without_master_weights(tmp_path, loaded_layers):
     import fastvideo.envs as envs
     from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import write_nvfp4_decoder_safetensors
     from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import keep_evenly_spaced_blocks, nvfp4_linears
 
-    _write_vae_dir(tmp_path / "vae", decoder_layers=5)
+    _write_vae_dir(tmp_path / "vae", decoder_layers=loaded_layers)
     student = _tiny_h3_vae(_tiny_vae_arch(5))
     keep_evenly_spaced_blocks(student.decoder, 3)
     checkpoint = _fake_packed_checkpoint(student)
@@ -796,3 +798,34 @@ def test_packed_export_round_trip_is_bit_identical(tmp_path, fused, act_scale):
     latents = torch.randn(2, 24, 3, 8, 8, device="cuda")
     with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
         assert torch.equal(loaded.decoder(latents), reference(latents))
+
+
+@pytest.mark.parametrize("outer_inference", [False, True])
+def test_loader_quantizes_original_dense_decoder_without_depth_change(tmp_path, outer_inference):
+    import fastvideo.envs as envs
+    from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import nvfp4_linears
+
+    _write_vae_dir(tmp_path / "vae", decoder_layers=5)
+    with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(None), envs.FASTVIDEO_H3_VAE_INT8_OVERLAY.override(False):
+        dense = _load_h3_vae(tmp_path / "vae")
+        with torch.inference_mode(outer_inference):
+            quantized = _load_h3_vae(tmp_path / "vae", quantize_dense=True)
+    assert len(quantized.decoder.transformer_blocks) == 5
+    assert quantized.decode_autocast_dtype == torch.bfloat16
+    layers = nvfp4_linears(quantized.decoder)
+    assert len(layers) == 30 and all(layer.act_scale == "unit" for layer in layers)
+    assert all(layer.input_amax._version >= 0 for layer in layers)
+    assert not any(parameter.requires_grad for parameter in quantized.parameters())
+    original = dense.state_dict()
+    converted = quantized.state_dict()
+    for name, value in original.items():
+        assert torch.equal(value, converted[name]), name
+
+
+def test_loader_rejects_dense_quantization_with_exported_decoder(tmp_path):
+    import fastvideo.envs as envs
+
+    _write_vae_dir(tmp_path / "vae", decoder_layers=5)
+    with envs.FASTVIDEO_H3_VAE_NVFP4_DECODER.override(None):
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            _load_h3_vae(tmp_path / "vae", decoder=str(tmp_path / "unused.pt"), quantize_dense=True)

fastvideo/tests/vaes/test_minimax_h3_parallel_vae.py

@@ -443,3 +443,27 @@ def test_encode_keyframe_routes_through_the_tile_split(monkeypatch, tile_paralle
     results = group.run(lambda rank: vae.encode_keyframe(keyframe).latent_dist.parameters)
     for moments in results:
         assert_close(moments, expected, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.parametrize("latent_frames", [12, 13])
+@torch.inference_mode()
+def test_idle_callback_preserves_decode_and_runs_only_on_idle_rank(latent_frames: int) -> None:
+    vae = _tiny_vae()
+    latents = torch.randn(1, 4, latent_frames, 4, 4)
+    expected = torch.empty(vae.decoded_pixel_shape(latents.shape), dtype=torch.float32)
+    vae.decode_to_pixels(latents, expected)
+    _, chunks, _ = vae._temporal_decode_plan(latent_frames)
+    group = _ThreadedFakeGroup(2)
+    calls = [0, 0]
+
+    def rank_main(rank):
+        output = torch.empty_like(expected) if rank == 0 else None
+
+        def callback():
+            calls[rank] += 1
+
+        return decode_to_pixels_parallel(vae, latents, output, group, idle_callback=callback)
+
+    result = group.run(rank_main)
+    assert_close(result[0], expected, atol=0, rtol=0)
+    assert calls == [0, chunks % 2]