h3_sparse_attention/README.md

# 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:

```bash
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

```python
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](https://huggingface.co/docs/diffusers/main/api/schedulers/minimax_h3),
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:

```python
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](../docs/kernel_performance_matrix.md) 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

```python
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

```python
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

```bash
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.