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) → 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 output is BF16. 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 with 584 bytes per token. sparse_indices and swa_topk_lens describe the active SWA segment. To add a compressed segment, pass compressed_kv_cache as another packed uint8 pool 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 packed uint8 head dim 584. 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.

  • 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.

  • 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.