flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4

flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4(query: Tensor, swa_kv_cache: Tensor, workspace_buffer: Tensor, sparse_indices: Tensor | None = None, compressed_kv_cache: Tensor | None = None, sparse_topk_lens: Tensor | None = None, seq_lens: Tensor | None = None, out: Tensor | None = None, bmm1_scale: float | Tensor = 1.0, bmm2_scale: float | Tensor = 1.0, sinks: Tensor | None = None, kv_layout: Literal['HND', 'NHD'] = 'HND', cum_seq_lens_q: Tensor | None = None, max_q_len: int | None = None, enable_pdl: bool | None = None, swa_topk_lens: Tensor | None = None, extra_sparse_indices: Tensor | None = None, extra_sparse_topk_lens: Tensor | None = None, backend: Literal['auto', 'trtllm-gen', 'cute-dsl', 'sparse'] = 'auto', hca_swa_indices: Tensor | None = None, hca_compressed_block_tables: Tensor | None = None, hca_seq_lens: Tensor | None = None, hca_is_causal: bool = True, hca_use_persistent: bool = False, hca_sparse_indices_format: Literal['compressed-page-aligned'] | None = None, remapped_sparse_indices_buffer: Tensor | None = None, sparse_indices_are_storage_offsets: bool | None = None, dsv4_inv_rope_cos_sin_cache: Tensor | None = None, dsv4_output_scale: Tensor | None = None, *, kv_cache_format: Literal['fp8', 'nvfp4'] = 'fp8') → Tensor | Tuple[Tensor, Tensor]

Decode DeepSeek V4 sparse MLA.

The implementation is selected from the query device architecture.

On SM100/SM103, this calls the TRTLLM-GEN DeepSeek V4 sparse MLA kernels. The query and both KV pools use head dim 512. The query may be BF16 or per-tensor FP8 E4M3 and the default output is BF16. When dsv4_inv_rope_cos_sin_cache is provided, the fixed TRTLLM-GEN RopeQuant epilogue instead writes group-major FP8 E4M3 values and packed UE8M0 scales. The first 128 columns of sparse_indices are SWA entries into swa_kv_cache; remaining columns are compressed entries into compressed_kv_cache. sparse_topk_lens gives the total active sparse length for each query token and must include the fixed 128 SWA entries. seq_lens provides the original KV sequence length for the SWA validity window.

On SM120/SM121, this calls the packed sparse backend. swa_kv_cache is the required packed uint8 SWA pool: 584 bytes per token for FP8 or 384 bytes per token for group-16 NVFP4, selected by kv_cache_format. sparse_indices and swa_topk_lens describe the active SWA segment. To add a compressed segment, pass compressed_kv_cache as another pool in the same format and pass extra_sparse_indices with extra_sparse_topk_lens. The SM120/SM121 path accepts BF16 query tensors and produces BF16 output.

With backend="cute-dsl" on SM100/SM103, this calls the DeepSeek V4 HCA kernel. Its SWA stream consumes arbitrary physical token-row indices, matching the TRTLLM-GEN dynamic-token-sparse ABI, while its compressed stream consumes physical page IDs. HCA currently accepts dense FP8 E4M3 query/KV tensors and produces BF16 output.

Parameters:

  • query (torch.Tensor) – Dense query input [batch_size, q_len_per_request, num_heads, 512] or varlen query input [sum_q, num_heads, 512] when cum_seq_lens_q is provided. SM100/SM103 accepts BF16 or FP8 E4M3; SM120/SM121 accepts BF16.

  • swa_kv_cache (torch.Tensor) – SWA KV cache. TRTLLM-GEN uses head dim 512; SM120 sparse uses an opaque packed uint8 record with last dimension 584 (FP8) or 384 (NVFP4). Layout follows kv_layout.

  • workspace_buffer (torch.Tensor) – Byte workspace used by TRTLLM-GEN or HCA split-K reduction. The TRTLLM-GEN multi-CTA KV counters are managed in a separate internal buffer.

  • sparse_indices (Optional[torch.Tensor]) – TRTLLM-GEN combined sparse table, or the SM120 sparse SWA segment. Pass None for the explicit backend="cute-dsl" metadata path. Combined HCA tables whose compressed segment is a canonical page expansion may instead be converted by setting hca_sparse_indices_format="compressed-page-aligned". SWA entries remain arbitrary absolute token rows in this mode.

  • compressed_kv_cache (Optional[torch.Tensor]) – Primary/compressed KV cache in the same backend layout as swa_kv_cache. Required by trtllm-gen and HCA, and by SM120 sparse when extra_sparse_indices is provided.

  • sparse_topk_lens (Optional[torch.Tensor]) – Flattened total sparse MLA top-k lengths in query-token order, shape [sum_q]. Values must already include the fixed 128 SWA entries, matching TRTLLM-GEN sparseMlaTopkLengths. For TRTLLM-GEN they must not exceed sparse_indices.shape[-1]. HCA also requires this tensor; there it describes the visible window-plus-compressed slot count.

  • seq_lens (Optional[torch.Tensor]) – Original KV sequence lengths, shape [batch_size] INT32. Required by trtllm-gen and by compressed-page-aligned HCA metadata conversion.

  • out (Optional[torch.Tensor]) – Optional preallocated output. The default path expects the same shape as query and BF16 dtype. RopeQuant expects FP8 E4M3 shape [sum_q, 16, 4096] with strides (4096, sum_q * 4096, 1). If omitted, FlashInfer allocates the appropriate layout.

  • bmm1_scale (Union[float, torch.Tensor]) – Fused per-tensor scale for QK and softmax. Tensor form must be FP32. HCA currently accepts only a Python float.

  • bmm2_scale (Union[float, torch.Tensor]) – Fused per-tensor scale for VO. Tensor form must be FP32. HCA currently accepts only a Python float.

  • sinks (Optional[torch.Tensor]) – Optional attention sink logits, shape [num_heads] FP32.

  • kv_layout (Literal["HND", "NHD"]) – Layout of both KV pools.

  • cum_seq_lens_q (Optional[torch.Tensor]) – Cumulative query lengths for varlen query input, shape [batch_size + 1] INT32. When provided, dynamic top-k lengths are consumed in flattened query-token order. On SM120/SM121 (sparse backend) the query is processed per token and per-token sparse indices fully determine visibility, so cum_seq_lens_q and max_q_len are accepted for call-site compatibility but do not affect the result.

  • max_q_len (Optional[int]) – Maximum query length in the varlen batch. Required with cum_seq_lens_q.

  • enable_pdl (Optional[bool]) – Whether to enable Programmatic Dependent Launch. Used by the TRTLLM-GEN path.

  • swa_topk_lens (Optional[torch.Tensor]) – Active SWA segment lengths, shape [sum_q] INT32. On SM120/SM121 these are sparse-segment lengths. HCA requires them as its per-row visible sliding-window lengths in the range 0 through 128.

  • extra_sparse_indices (Optional[torch.Tensor]) – Optional SM120/SM121 compressed segment indices into compressed_kv_cache.

  • extra_sparse_topk_lens (Optional[torch.Tensor]) – Active compressed segment lengths for SM120/SM121, shape [sum_q] INT32.

  • backend ({"auto", "trtllm-gen", "cute-dsl", "sparse"}) – Backend selection. "auto" preserves the architecture-based default: TRTLLM-GEN on SM100/SM103 and sparse on SM120/SM121. HCA is selected only when "cute-dsl" is requested explicitly.

  • hca_swa_indices (Optional[torch.Tensor]) – Absolute SWA token-row indices, shape [B * Q, 128] INT32. Ring rotation and wraparound are supported. Every entry, including masked padding, must name a legal row of the flattened SWA cache.

  • hca_compressed_block_tables (Optional[torch.Tensor]) – HCA compressed-pool page IDs, shape [B * Q, max_pages] INT32. For each query row, the valid page-ID prefix must cover the compressed footprint implied by hca_seq_lens, rounded up to 128-slot tiles. Those page IDs must be legal even when sparse_topk_lens masks the corresponding slots.

  • hca_seq_lens (Optional[torch.Tensor]) – Per-request backing HCA slot counts [B] INT32 used to schedule TMA loads. Each value counts the 128 window slots plus compressed slots, not original raw KV tokens, and may exceed an effective per-row sparse_topk_lens.

  • hca_is_causal (bool) – Must currently be True. HCA index/page-table rows and valid-length tensors are per query token (B * Q rows).

  • hca_use_persistent (bool) – Select the CuTe DSL persistent tile scheduler. This removes the non-persistent B * Q <= 65535 launch-grid restriction.

  • hca_sparse_indices_format (Optional[Literal["compressed-page-aligned"]]) – Opt-in compatibility mode for legacy TRTLLM-GEN metadata. With "compressed-page-aligned", active SWA entries remain arbitrary absolute token rows, while the active compressed segment must be the canonical expansion page_id * page_size + page_offset. The dispatcher validates and converts them into HCA gather indices, a compressed block table, hca_seq_lens, and swa_topk_lens. This tagged path is a one-shot compatibility path: it allocates, synchronizes the device, immediately launches the decode, and is not CUDA Graph capture safe. Performance-sensitive callers must precompute with convert_compressed_page_aligned_sparse_indices_to_hca_metadata() and reuse the returned metadata through the explicit HCA arguments.

  • remapped_sparse_indices_buffer (Optional[torch.Tensor]) – Optional INT32 output buffer matching sparse_indices. TRTLLM-GEN uses it when either KV pool has a strided page layout. Passing a stable buffer avoids an allocation in serving and CUDA graph paths. Concurrent calls must use distinct buffers. The buffer must not share storage with sparse_indices when remapping is required.

  • sparse_indices_are_storage_offsets (Optional[bool]) – Encoding of sparse_indices for strided TRTLLM-GEN KV pools. Set to False for logical flattened token indices or True for indices already adjusted to storage-row offsets. Strided pools require an explicit value to prevent accidental double remapping.

  • dsv4_inv_rope_cos_sin_cache (Optional[torch.Tensor]) – Enables the TRTLLM-GEN DSv4 RopeQuant epilogue. Must be a contiguous FP32 tensor shaped [max_position, 64] with each row laid out as [cos(32), sin(32)] for interleaved RoPE. The kernel derives each query position as seq_len - q_len + local_query_index; the cache must cover every derived position. This mode requires FP8 E4M3 Q/K/V, 128 query heads, and backend="trtllm-gen".

  • dsv4_output_scale (Optional[torch.Tensor]) – Optional preallocated packed UE8M0 scale output. Shape must be [sum_q, 16, 8] with strides (1, 8 * scale_buf_m, scale_buf_m); scale_buf_m must be a multiple of four and at least sum_q. If omitted in RopeQuant mode, FlashInfer allocates it with zeroed physical padding. A caller-provided buffer must also have its padded token rows zero-initialized before first use; the cubin intentionally leaves those rows unwritten. Supplying this tensor together with out avoids RopeQuant output allocations. CUDA Graph use still requires normal JIT warmup and stable caller-controlled temporary buffers; the existing DSV4 path allocates its internal counter buffer per invocation.

Returns:

  • torch.Tensor or tuple[torch.Tensor, torch.Tensor] – The existing BF16 output when RopeQuant is disabled. With RopeQuant, returns (out_fp8, out_scale). out_fp8 has shape [sum_q, 16, 4096] and group-major strides (4096, sum_q * 4096, 1); out_scale uses the packed UE8M0 layout described above.

  • kv_cache_format ({“fp8”, “nvfp4”}) – SM120/SM121 sparse-cache storage format. "fp8" preserves the existing 584-byte DSv4 cache ABI. "nvfp4" selects the 384-byte group-16 NVFP4 cache ABI and its native prefill/decode kernels. NVFP4 currently supports 16/32/64/128 heads, primary top-k 128 or 512, primary page size 64, and optional extra-cache page size 2 or 64.