Sol-Attn and Spark for MiniMax-H3

This package ports the Sol-Attn integration and the spark_reblock and spark_reweight primitives used by MiniMax-H3-Sparse. It includes local Sol-Attn kernels, so inference does not depend on a sibling checkout. No SVOO, SVG, SVG2, SVG-EAR, Radial, or experimental compensation methods are included.

Install

From the MiniMax-H3 repository, in the same environment as your H3 pipeline:

python -m pip install -e '.[cuda]'

The pipeline still needs the H3-capable diffusers build and model dependencies listed in the repository's main README and requirements.txt. CUDA execution requires a matching PyTorch/CUDA installation. Plain Sol-Attn selects the native SM80 Triton backend or the SM90/SM100/SM120 CuTe backend before launch; other NVIDIA GPUs with compute capability at least 8.0 use the portable Triton backend. Spark-H3's Diffusers integration contains architecture dispatch for SM80, SM90, SM100, and SM120 and raises before reblocking on every other architecture. Production kernel support is complete for SM80 and SM120. SM90 and SM100 remain development targets until compatible GPU resources become available for implementation and validation. Spark requires CUDA and Triton for its supported execution paths. Sol-Attn requires contiguous BF16 [batch, tokens, heads, 128] tensors. These are forward/inference kernels.

H3 pipeline integration

from h3_sparse_attention import install_h3_sol_attn

# pipe is an already loaded MiniMaxH3Pipeline or MiniMaxH3ModularPipeline.
# MiniMaxH3Scheduler includes the terminal sigma=0 in num_inference_steps,
# but that terminal point does not run the transformer.
num_denoise_steps = 19
num_inference_steps = num_denoise_steps + 1

with install_h3_sol_attn(
    pipe.transformer,
    num_denoise_steps=num_denoise_steps,
    warmup_mode="warmup_ratio",
    warmup_ratio=0.2,
    sol_tau=1.0,
    sol_thresh_type="diag",
    sol_kv_splits=1,
    dense_layers=1,
    sol_force_local_blocks=True,
) as attention:
    result = pipe(..., num_inference_steps=num_inference_steps)
    stats = attention.summary()

The context manager installs attention processors and a layout pre-hook, then restores them even if inference raises. Call attention.reset() before a second pipeline invocation inside the same context. With the default MiniMaxH3Scheduler sigma grid, prefer the installer's num_denoise_steps parameter, which directly counts actual denoising iterations and Transformer evaluations. The legacy installer parameter num_inference_steps remains supported and is interpreted as the sigma-grid count. For the pipeline, set num_inference_steps = num_denoise_steps + 1 because the grid includes the terminal 0. With warmup_mode="warmup_ratio", warmup is computed from actual denoising calls as ceil(num_denoise_steps * warmup_ratio), matching ComfyUI's model-evaluation convention. With warmup_mode="warmup_steps", warmup_steps gives the exact dense prefix length. Both modes cap the result at num_denoise_steps; the defaults are ratio mode, warmup_ratio=0.2, and warmup_steps=4. Thus 19 denoising steps with the default ratio produce four dense evaluations. The legacy warmup_percent parameter remains supported when none of the new warmup parameters are supplied. The first transformer layer remains dense by default.

For fixed Top-K routing, sol_exact_block_radius=0 forces only the self block exact, while 1 forces the self block and its two immediate neighbors. These blocks are added outside the Top-K budget, matching sol_force_local_blocks. With reblocking this option requires landmark_tree_v2_layout_reuse="q_from_k" or "k_from_q"; independently reordered Q/K blocks have no shared locality.

Only the generated target-video grid is sparse. Conditioning video, text, and audio are packed after the target as exact K/V sinks; their query rows are computed with dense attention. If the video boundary splits a 64-token block, that final mixed block is also computed densely. The adapter supports H3's packed batch size one and rejects external attention masks. It preserves Q/K normalization, RoPE, fused or separate QKV projections, and the output projection.

sol_video_tail_mode="dense" is the default: the incomplete final video block is included in the dense suffix. Set sol_video_tail_mode="pad" to run all video query rows through sparse attention instead. Pad mode fills only the non-video Q slots in that physical block with the mean Q of its real video rows during sparse execution, preserving the block routing centroid. Real context Q rows are still evaluated once by dense attention, and the complete, unmodified K/V sequence remains visible to both paths. The packed context suffix must contain enough slots to complete the 64-row block. The ComfyUI Spark node currently supports only the dense-tail choice.

Spark inference integration

install_h3_spark_attn enables reblocking and reweighting together in the H3 attention path. It follows the updated MiniMax-H3-Sparse installer defaults:

from h3_sparse_attention import install_h3_spark_attn

# MiniMaxH3Scheduler includes the terminal sigma=0 in num_inference_steps,
# but that terminal point does not run the transformer.
num_denoise_steps = 19
num_inference_steps = num_denoise_steps + 1

with install_h3_spark_attn(
    pipe.transformer,
    num_denoise_steps=num_denoise_steps,
    warmup_mode="warmup_steps",
    warmup_steps=4,
    dense_layers=1,
    topk_ratio=0.1,
) as attention:
    result = pipe(..., num_inference_steps=num_inference_steps)
    stats = attention.summary()

The installer is also exported from h3_sparse_attention.spark. It uses the same reversible processor installation, dense warmup, dense first layer, and exact context handling as the Sol installer. Call attention.reset() before another pipeline invocation in the same context. Reset/removal releases the cached reblocking plans and query topology.

Default Spark settings match the source:

Sequences shorter than 8192 packed tokens use the original dense attention; min_tokens=0 disables this size gate.
Native mean Top-K routing with ratio 0.1, gemm_radix cutoffs, and packed_external_no_route_qk execution. This exports the packed route once and avoids recomputing route QK inside the attention kernel. Set sol_route_topk_execution="threshold" explicitly for the historical compiled Table 4 path. The packed route is currently implemented on SM120; SM80, SM90, and SM100 resolve this default to threshold and report that fallback in the plugin summary.
Landmark-v2 Q/K reblocking using opposite second moments, cosine distance, fanout 16 with power_of_two_fanout scheduling, 32 midpoint landmarks, and no temporal grouping.
The midpoint direction builder defaults to the historical two-stage calculation for bitwise compatibility with the Table 4 Spark run. Set landmark_tree_v2_midpoint_direction_mode="fused" to opt into the newer, faster direction kernel; it can change near-tie reblock routes and outputs. The ComfyUI Spark node exposes the same choice as its optional midpoint_direction_mode input, also defaulting to legacy.
One global representative per head, formed by averaging all target-video queries after reblocking; the default stored anchor dtype is BF16.
Global-anchor K/V and log-mass summaries merged with exact attention. The former target-189 representatives remain available only as an explicit ablation and are not the default Spark-H3 configuration.
The positional local band is disabled after reblocking.

The approximate branch defaults to per-query processing. Set sol_tail_granularity="block" in install_h3_spark_attn(...) to share one approximate softmax/output across each 64-row query block, as in ComfyUI's block-granularity mode. Exact attention still uses the real query rows, and Top-K/reblock/global-reweight settings do not change. This opt-in mode uses the existing exact route/export plus a separate Triton block-tail merge; it is a functional compatibility path, not the fused query-mode fast path. For Top-K it supports threshold, packed_external, and packed_external_no_route_qk routing, but not fused routing. The default "query" mode is unchanged.

For reweight numeric ablations, the Diffusers installer also accepts sol_global_anchor_dtype="bfloat16", sol_reweight_summary_math="comfy_fp32", and sol_reweight_logmass_key="pre_round". The defaults are BF16 global anchor, Tensor-Core summary, and log-mass subtraction using the stored BF16 summary key. comfy_fp32 uses FP32 lane-wise FMA for anchor–key scores and FP32 softmax/weighted K/V reduction before BF16 summary storage; pre_round subtracts the unrounded FP32 weighted key in log-mass. These two switches can be varied independently of anchor dtype and of query/block tail granularity. The CUDA and Triton reduction orders are not bitwise identical. ComfyUI Sol's downstream INT8 attention consumption is a separate backend choice, not part of reweight, and is not emulated by these switches.

For a component ablation, sol_reweight_components accepts "full" (default), "weights_only", "bias_only", or "none". The weight component replaces each physical key block's ordinary mean K/V with anchor-softmax-weighted K/V. The bias component replaces the ordinary log(block_length) mass with logsumexp(anchor · K) - anchor · active_summary_K, where active_summary_K is weighted or ordinary according to the first switch. Thus bias-only is recentered against its own mean K rather than incorrectly reusing the full path's weighted-K shift. These four arms retain the same Top-K route, reblock, query granularity and exact branch; "none" is a within-Spark summary baseline, not the original Sol-Attn kernel.

Select landmark_tree_v2_fanout_mode="arbitrary_fanout" for balanced arbitrary child counts. landmark_tree_v2_fanout is an alias for landmark_tree_v2_children; landmark_tree_v2_root_fanout and landmark_tree_v2_final_fanout independently control the first nonfinal and final rounds. Both default to inheriting the ordinary fanout (16 by default). The default power_of_two_fanout uses the original strict power-of-two policy in every round, including final nodes; its final fanout must inherit or equal the ordinary fanout (16/16/16 by default). Select power_of_two_arbitrary_final explicitly for power-of-two nonfinal rounds followed by arbitrary final splits. arbitrary_fanout uses balanced arbitrary child counts throughout. Both arbitrary modes allow a separate final fanout. Arbitrary final splits can change the resulting token grouping compared with the former strictly power-of-two tree; it is not a bitwise-equivalent kernel optimization. Small-node kernel fusion remains available to all three modes. Automatic fusion includes eight-child splits up to 512 tokens, alongside two-/four-child splits up to the configured fused-node size limit (1024 by default). The fusion enable switch and size limit still apply to every case. Reweighting follows the actual reblocking hierarchy. Cosine scoring defaults to normalized FP16; H3_LMV2_COS_PRECISION selects other precision modes.

Sol and Spark always avoid redundant Q/K/V layout copies; the former H3_SOL_LAYOUT_FAST comparison switch is no longer used. Optional sol_route_global_weighted_mean=True enables global weighted Top-K routing; sol_route_global_weighted_side selects both, query, or key.

Set sol_route_topk_ratio=None to use native Sol tau routing with Spark reweighting (sol_virtual_query_route_score="native_mean"). SM120 reweighting now skips unused ordinary value sums. Spark evaluates only complete video query blocks by default; pad mode additionally evaluates the incomplete video tail sparsely with query-only padding. All real context query rows are evaluated once by dense attention with the complete K/V set.

Configuration overrides are passed as keyword arguments, for example sol_route_topk_ratio=0.2 or sol_virtual_query_levels_up=2. Setting levels-up automatically clears the default target-block setting; specifying both is an error.

Plain Sol on SM80 retains its pointer-based Triton kernels. Spark BF16 on SM80 and SM120 uses a CuTe fused virtual-query mainloop automatically above 8192 packed tokens; smaller inputs use the streamed exact/skipped implementation. The SM80 mainloop follows the SM120 structure: CTA-local native-threshold or fixed-budget fused Top-K routing, warp exact-index compaction, two-stage cp.async K/V streaming, and a single FP32 online-softmax/output accumulator for exact and approximate tiles. Set sol_route_topk_execution="fused" to select Top-K inside each CTA and remove the external cutoff tensor/pass. The SM80 implementation reuses the Q shared-memory tile for route scores after Q is register resident, so fused Top-K does not increase shared-memory usage. Its exact FP32 radix boundary selection distributes score scans over all four CTA warps, and the summary producer writes AK/AV directly in the N-major layout consumed by cp.async, avoiding per-layer transpose copies. Kernel failures propagate to the caller; there is no error-triggered backend fallback. See the kernel performance matrix for the current cross-architecture benchmark format and locally measured results. H3_SPARK_REWEIGHT_FUSED=0 or 1 selects the source's fallback/fused path on SM80/SM90/SM100/SM120 for testing. On SM80, fused mode supports native-threshold and CTA-local fused Top-K routes; unsupported external/hybrid route modes raise directly. This code does not change environment variables.

Plain install_h3_sol_attn retains its base Sol defaults. The standalone Spark primitives below remain available. Unrelated attention methods and the source's other compensation pipelines are not exposed by this port.

Direct Sol-Attn

from sol_attn import sol_attn

out = sol_attn(q, k, v, tau=1.0, thresh_type="diag",
               sink_start=video_tokens, sink_tokens=q.shape[1] - video_tokens)

The kernel's sinks make K/V exact. Use the H3 adapter above when context query rows must also be computed densely.

Spark primitives

from h3_sparse_attention.spark import spark_reblock, spark_reweight

# samples: [..., T, D], in raster order for grid_shape=(frames, height, width).
blocks = spark_reblock(samples, grid_shape=(frames, height, width))
# blocks.permutation / blocks.inverse_permutation: [..., T]
# blocks.block_indices: [..., T // 64, 64]
# blocks.excluded_indices holds the T % 64 remainder.

# a: [B, P, H, 128] query representatives; k, v: [B, T, H, 128].
ak, av, lm = spark_reweight(a, k, v)
# ak, av: [B, P, H, ceil(T / 64), 128]
# lm:     [B, P, H, ceil(T / 64)] (FP32 log-mass correction)

spark_reblock is a direct alias for recursive_landmark_tree_v2_blocks. Defaults match the source: 64-token leaves, cosine distance, eight children, 32 midpoint landmarks, flat initial order, single-token groups, and stable parent order. CPU supports a PyTorch reference path; CUDA takes FP16/BF16. The output includes remainder rows in the permutation and its inverse. The PreparedLandmarkTreeV2Permutation class remains available from h3_sparse_attention.landmark_tree_v2 for repeated CUDA graph execution.

spark_reweight is a direct alias for virtual_summaries. It requires matching CUDA FP16/BF16 inputs, contiguous K/V, and contiguous anchor head/feature dimensions. For each anchor and physical 64-token K/V block, it computes softmax-weighted keys and values and lm = logsumexp(a @ k / sqrt(128)) - a @ ak / sqrt(128). Partial final blocks exclude padding. It does not choose query anchors, select exact routes, or merge attention outputs. Its output allocation scales with B * P * H * ceil(T / 64) * 128; choose the anchor count accordingly. spark_block remains a compatibility alias for spark_reblock.

Validation and provenance

OMP_NUM_THREADS=1 python -m pytest tests -q

Tests cover CPU reblocking/reference parity, CUDA reblocking, FP16/BF16 reweighting against explicit softmax (including partial blocks and strided anchors), exact Sol sinks, packed target selection, and H3 processor cleanup. CUDA graph replay and the Triton fallback are also covered. GPU tests require a CUDA device; H3 integration tests require H3 diffusers. Validation results for the current sync are recorded in PORT_MANIFEST.json. Spark tests additionally cover installer defaults, temporal topology, fused and streamed all-exact parity, and skipped-block reweighting against an explicit softmax oracle. A synthetic default-Spark case matched the updated source bitwise; its settings are recorded in PORT_MANIFEST.json. Full model video generation was not run.

PORT_MANIFEST.json records source commits and source-file hashes. The Sol code is copied from the Sana-sol-engine working tree used by MiniMax-H3-Sparse, including its local kernel patches. Required landmark helpers were extracted without the other attention methods. The file named mahalanobis_kmeans.py contains shared Hilbert-order and metric-factor utilities, not the clustering method. Exact-branch kernels are extracted from the source compensation modules without importing their compensation methods. See sol_attn/THIRD_PARTY_NOTICES.md and its included license for kernel notices.