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;

  • measured latency: the attention launch runs about 9-15 % faster on both GB202 SKUs (RTX 5090 and RTX PRO 6000 Blackwell; 220 registers instead of 246), the complete call proportionally less on short plans where pre-processing dominates.

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.

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