flashinfer.diffusion_ops.minimax_h3_sm120_varlen_attention_nvfp4_nodelta

flashinfer.diffusion_ops.minimax_h3_sm120_varlen_attention_nvfp4_nodelta(q: Tensor, k: Tensor, v: Tensor, cu_seqlens: Tensor, out: Tensor | None = None, *, cu_seqlens_host: Sequence[int] | None = None, softmax_scale: float | None = None) → Tensor

Opt-in uncompensated variant of minimax_h3_sm120_varlen_attention_nvfp4() (same contract, same signature, same workspace and pre-processing launches; only the attention kernel differs).

The default route quantizes each 128-row Q block after removing its block mean and adds the removed term qm K^T back to the scores inside the kernel, so its softmax equals the uncentred one up to FP4 rounding. This route keeps the centring but does not add the term back: the scores of every Q block carry a per-key offset -(qm . k_j), which the softmax does not cancel (unlike the per-row constant of the K mean removal). It is a deliberate precision / latency trade-off that is NOT part of the SageAttention3 recipe:

  • measured error vs the FP32 oracle on Gaussian inputs: relative L2 about 0.209 vs 0.191 for the default route (+9 %), max-abs about 0.064 vs 0.060; the FP4 block-scaled tolerance atol = 1.0, rtol = 0.1 still holds on every element, and the tests bound the relative L2 at 1.15x the default route’s on the same inputs;

  • the deviation grows on short segments: the dropped term is the mean over a 128-row Q block, and a partial block of n rows has a block mean about sqrt(128 / n) larger. Measured relative-L2 ratio vs the default route for a single segment of N tokens (56 heads, Gaussian inputs): 1.10 at N >= 128 (full blocks), 1.13 at 100, 1.18 at 64, 1.23 at 51, 1.32 at 32, 1.64 at 16 (max-abs 0.79, still inside the FP4 tolerance). In a packed stream the aggregate error is dominated by the short segments, whose outputs are larger per element. Segments of fewer than 8 tokens exceed the FP4 tolerance on this route (max-abs 1.2-2.7 at 2-6 tokens; a 1-token segment is exact by construction), so do not use it for streams that contain such segments – the default route stays within tolerance there;

  • latency: the attention program differs (ATTENTION_NODELTA / ATTENTION_WS_NODELTA, which skip the compensation term); its warp-specialised form is taken from a mean segment of 1536 tokens with no upper bound, and on parts of at least 170 SMs also for a plan of 2048 to 16384 padded tokens; on parts of 188 SMs the overlapped form’s two attention launches are warp-specialised too (ATTENTION_WS_NODELTA_LEAD / ATTENTION_WS_NODELTA_TAIL, or for plans of long segments ATTENTION_WS_NODELTA_WSIL_LEAD / ATTENTION_WS_NODELTA_WSIL_TAIL), behind the producers PREPROCESS_HEAD_MAJOR_KA (the outputs of PREPROCESS_HEAD_MAJOR; a K item waits for its head’s K sums in a second warp).

Use the default route unless the extra error has been validated for the model in question. Inputs, output, cu_seqlens handling, workspace and caching are exactly those of the default route.

Concurrency: safe on any number of streams, host threads, CUDA graphs and MPS clients. Exec graphs instantiated more than once from one captured graph (or cloned as child graphs) share this op’s workspace and must not run concurrently.

CUDA graphs: the first call for a given cu_seqlens_host builds the plan’s device tables from host lists, a pageable host-to-device copy that a stream capture does not allow. Call the op once for that plan outside the capture before capturing it; later calls reuse the tables, and the capture then records kernel launches only. This holds with CUDA_MODULE_LOADING=LAZY as well: once the tables exist, the first launch of each program may happen inside the capture.

Parameters:

  • q (torch.Tensor) – Contiguous bfloat16 CUDA tensors of shape [tokens, heads, 128] (packed THD layout; MiniMax-H3 uses 56 heads).

  • k (torch.Tensor) – Contiguous bfloat16 CUDA tensors of shape [tokens, heads, 128] (packed THD layout; MiniMax-H3 uses 56 heads).

  • v (torch.Tensor) – Contiguous bfloat16 CUDA tensors of shape [tokens, heads, 128] (packed THD layout; MiniMax-H3 uses 56 heads).

  • cu_seqlens (torch.Tensor) – int32 tensor of shape [segments + 1] on the same device with cu_seqlens[0] = 0, non-decreasing entries and cu_seqlens[-1] = tokens. Segments shorter than 8 tokens are outside this route’s FP4 tolerance (see above).

  • out (Optional[torch.Tensor]) – Optional pre-allocated output of the same shape/dtype as q; allocated when omitted.

  • cu_seqlens_host (Optional[Sequence[int]]) – Host copy of cu_seqlens (avoids a synchronizing device-to-host copy). The segment plan is cached per (cu_seqlens, heads, device) and shared with the default route.

  • softmax_scale (Optional[float]) – Softmax scale; defaults to 1 / sqrt(128).

Returns:

bfloat16 [tokens, heads, 128] attention output of the uncompensated route.

Return type:

torch.Tensor