diff --git a/3rdparty/patches/msa_strided_paged_kv.patch b/3rdparty/patches/msa_strided_paged_kv.patch index beed6c6b3bdb..a1be1ca5cc82 100644 --- a/3rdparty/patches/msa_strided_paged_kv.patch +++ b/3rdparty/patches/msa_strided_paged_kv.patch @@ -1,5 +1,5 @@ diff --git a/python/fmha_sm100/api.py b/python/fmha_sm100/api.py -index b02f747..b69cd8b 100644 +index b02f747..d9d0a41 100644 --- a/python/fmha_sm100/api.py +++ b/python/fmha_sm100/api.py @@ -532,9 +532,8 @@ def _fmha_sm100_plan( @@ -14,8 +14,227 @@ index b02f747..b69cd8b 100644 return sparse_fmha_plan(qo_segment_lens=qo_segment_lens, kv_segment_lens=kv_segment_lens, num_qo_heads=num_qo_heads, causal=causal, qo_offset=qo_offset, num_kv_splits=num_kv_splits, page_size=page_size, output_maxscore=output_maxscore, +@@ -941,6 +940,62 @@ def _fmha_sm100( + return out, max_score + + ++def _mixed_batch_split(qo_segment_lens, split_threshold, *, enable, sparse): ++ """Row index splitting a mixed prefill/decode batch, or 0 for no split. ++ ++ Returns ``split`` such that rows ``[0, split)`` and ``[split, batch)`` are ++ each uniformly prefill or uniformly decode. Both halves must be contiguous ++ row ranges because that is all the run path can slice. ++ ++ Returns 0 when the batch is uniform, when the two classes interleave, or ++ when splitting is disabled. ++ ++ Both orderings are recognised: ++ ++ * **Decode-first** (prefill rows form a suffix) is this module's original ++ convention and is honoured for every plan flavour, so existing callers ++ keep their exact behaviour. ++ * **Prefill-first** (prefill rows form a prefix), which is what ++ TensorRT-LLM produces, is honoured only for ``sparse`` plans. The dense ++ flavour already handles a mixed batch in one call, whereas the sparse ++ MM-SA-Nv prefill kernel has no mixed path and would otherwise drag every ++ decode row through it. Confining it to sparse plans also leaves the proxy ++ and dense plans byte-identical to their pre-split behaviour, which ++ matters because the proxy's max-score output drives top-k block ++ selection. ++ """ ++ batch_size = int(qo_segment_lens.shape[0]) ++ if not enable or batch_size < 2: ++ return 0 ++ is_prefill = qo_segment_lens > split_threshold ++ n_prefill = int(is_prefill.sum()) ++ if n_prefill == 0 or n_prefill == batch_size: ++ return 0 ++ if bool(is_prefill[batch_size - n_prefill:].all()): ++ return batch_size - n_prefill ++ if sparse and bool(is_prefill[:n_prefill].all()): ++ return n_prefill ++ # Interleaved: neither half would be uniform, so plan the batch as one. ++ return 0 ++ ++ ++def _cache_split_boundaries(sub_lo): ++ """Precompute the boundaries the run path slices a split batch at. ++ ++ Both are properties of the low sub-plan alone, so they are constant for the ++ whole step. Reading them at plan time keeps the device-to-host read off the ++ per-layer path, where a sparse model would otherwise repeat it once per ++ layer. ++ """ ++ offsets = sub_lo.get("qo_segment_offsets") ++ if offsets is not None: ++ sub_lo["_split_nnz"] = int(offsets[-1].item()) // sub_lo.get("pack_factor", 1) ++ if "kv_page_indptr" in sub_lo: ++ sub_lo["_split_kv_pages"] = int(sub_lo["kv_page_indptr"][-1].item()) ++ elif "total_rows" in sub_lo: ++ sub_lo["_split_kv_pages"] = int(sub_lo["total_rows"]) ++ ++ + def fmha_sm100_plan( + qo_segment_lens: torch.Tensor, + kv_segment_lens: torch.Tensor, +@@ -973,9 +1028,11 @@ def fmha_sm100_plan( + ``kv_segment_lens - qo_segment_lens`` for bottom-right causal masking. + A tensor must have shape ``[batch_size]``. + split_prefill_decode : bool, optional +- If True, a mixed batch ordered as decode requests followed by prefill +- requests is split into two sub-plans. The original order must already +- group short decode sequences before long prefill sequences. ++ If True, a mixed batch whose decode and prefill requests form two ++ contiguous groups is split into two sub-plans. Either group may come ++ first; see ``_mixed_batch_split`` for which orderings apply to which ++ plan flavours. A batch that interleaves the two classes is planned as ++ one. + **kwargs + Planner options forwarded to ``_fmha_sm100_plan``. Common options are + ``num_kv_heads``, ``num_kv_splits``, ``page_size``, +@@ -985,8 +1042,13 @@ def fmha_sm100_plan( + Returns + ------- + tuple +- ``(has_mixed_prefill, split, batch_size, decode_plan, prefill_plan)``. +- Pass this tuple unchanged as ``plan_info`` to ``fmha_sm100``. ++ ``(has_mixed_prefill, split, batch_size, sub_plan_lo, sub_plan_hi)``, ++ where ``sub_plan_lo`` describes batch rows ``[0, split)`` and ++ ``sub_plan_hi`` describes rows ``[split, batch_size)``. Which of the ++ two is the prefill half depends on the caller's ordering, so consumers ++ must test a sub-plan's own tags rather than its position. ++ ``sub_plan_hi`` is None when the batch was not split. Pass this tuple ++ unchanged as ``plan_info`` to ``fmha_sm100``. + """ + + # assert qo_segment_lens.device.type == 'cpu' \ +@@ -999,27 +1061,19 @@ def fmha_sm100_plan( + qo_offset = torch.full_like(qo_segment_lens, qo_offset) + + batch_size = qo_segment_lens.shape[0] +- has_mixed_prefill = False +- qmax = qo_segment_lens.max().item() + sparse = kwargs.get("kv_block_num", -1) > 0 + split_threshold = _prefill_qlen_threshold(sparse) +- if split_prefill_decode and qmax > split_threshold: +- split = (qo_segment_lens > split_threshold).nonzero(as_tuple=False)[0, 0].item() +- has_mixed_prefill = split > 0 +- if has_mixed_prefill: +- # print(f"Split into 2 parts at index {split}") +- decode_qo_segment_lens = qo_segment_lens[:split] +- decode_kv_segment_lens = kv_segment_lens[:split] +- decode_qo_offset = qo_offset[:split] +- decode = _fmha_sm100_plan(decode_qo_segment_lens, decode_kv_segment_lens, *args, +- qo_offset=decode_qo_offset, **kwargs) +- decode = {k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in decode.items()} +- prefill_qo_segment_lens = qo_segment_lens[split:] +- prefill_kv_segment_lens = kv_segment_lens[split:] +- prefill_qo_offset = qo_offset[split:] +- prefill = _fmha_sm100_plan(prefill_qo_segment_lens, prefill_kv_segment_lens, *args, +- qo_offset=prefill_qo_offset, **kwargs) +- return (True, split, batch_size, decode, prefill) ++ split = _mixed_batch_split( ++ qo_segment_lens, split_threshold, enable=split_prefill_decode, sparse=sparse ++ ) ++ if split > 0: ++ sub_lo = _fmha_sm100_plan(qo_segment_lens[:split], kv_segment_lens[:split], *args, ++ qo_offset=qo_offset[:split], **kwargs) ++ sub_lo = {k: v.clone() if isinstance(v, torch.Tensor) else v for k, v in sub_lo.items()} ++ sub_hi = _fmha_sm100_plan(qo_segment_lens[split:], kv_segment_lens[split:], *args, ++ qo_offset=qo_offset[split:], **kwargs) ++ _cache_split_boundaries(sub_lo) ++ return (True, split, batch_size, sub_lo, sub_hi) + else: + plan = _fmha_sm100_plan(qo_segment_lens, kv_segment_lens, *args, + qo_offset=qo_offset, **kwargs) +@@ -1081,16 +1135,23 @@ def fmha_sm100( + ------- + tuple[torch.Tensor | None, torch.Tensor | None] + ``(out, max_score)``. Either item may be ``None`` if the corresponding +- output was disabled. When both decode and prefill sub-plans are used, +- outputs are concatenated back into the original batch order. ++ output was disabled. When the plan was split into two sub-plans, each ++ half writes its own rows of ``out`` directly; if no ``out`` was supplied ++ the halves are concatenated back into the original batch order instead. + """ + has_mixed_prefill, split, batch_size, decode, prefill = plan_info + if not has_mixed_prefill: + return _fmha_sm100(q, k, v, decode, out=out, max_score=max_score, kv_indices=kv_indices,kv_block_indexes=kv_block_indexes, q_offset_override=q_offset_override, **kwargs) + else: +- +- decode_pack = decode.get("pack_factor", 1) +- decode_nnz = decode["qo_segment_offsets"][-1].item() // decode_pack ++ # The `decode` / `prefill` names are historical: these are really the ++ # sub-plans for batch rows [0, split) and [split, batch_size), and which ++ # one holds the prefill requests depends on the caller's ordering. None ++ # of the slicing below needs to know, because every boundary is derived ++ # from the low sub-plan's own offsets. ++ decode_nnz = decode.get("_split_nnz") ++ if decode_nnz is None: ++ decode_pack = decode.get("pack_factor", 1) ++ decode_nnz = decode["qo_segment_offsets"][-1].item() // decode_pack + is_paged = kv_indices is not None + nnz_qo = q.shape[0] + num_qo_heads = q.shape[1] +@@ -1101,10 +1162,12 @@ def fmha_sm100( + if is_paged: + k_decode, v_decode = k, v + k_prefill, v_prefill = k, v +- if "kv_page_indptr" in decode: +- kv_page_split = decode["kv_page_indptr"][-1].item() +- else: +- kv_page_split = decode["total_rows"] ++ kv_page_split = decode.get("_split_kv_pages") ++ if kv_page_split is None: ++ if "kv_page_indptr" in decode: ++ kv_page_split = decode["kv_page_indptr"][-1].item() ++ else: ++ kv_page_split = decode["total_rows"] + decode_kv_indices = kv_indices[:kv_page_split] + prefill_kv_indices = kv_indices[kv_page_split:] + else: +@@ -1129,27 +1192,35 @@ def fmha_sm100( + decode_qo_offset = None + prefill_qo_offset = None + ++ # Let each half write straight into its own rows of the caller's buffer. ++ # Row views of a contiguous [nnz, H, D] output are themselves ++ # contiguous, which is what the kernels require, so this removes a ++ # full-size concatenate plus copy from every call. ++ write_in_place = out is not None and out.is_contiguous() ++ + # ---- Run kernels ---- + decode_out, decode_ms = _fmha_sm100( + q_decode, k_decode, v_decode, decode, +- out=None, max_score=None, ++ out=out[:decode_nnz] if write_in_place else None, max_score=None, + kv_indices=decode_kv_indices, kv_block_indexes=decode_block_idx, + q_offset_override=decode_qo_offset, + **kwargs) + prefill_out, prefill_ms = _fmha_sm100( + q_prefill, k_prefill, v_prefill, prefill, +- out=None, max_score=None, ++ out=out[decode_nnz:] if write_in_place else None, max_score=None, + kv_indices=prefill_kv_indices, kv_block_indexes=prefill_block_idx, + q_offset_override=prefill_qo_offset, + **kwargs) + + # ---- Merge out ---- +- if decode_out is not None and prefill_out is not None: ++ if write_in_place: ++ combined_out = out ++ elif decode_out is not None and prefill_out is not None: + combined_out = torch.cat([decode_out, prefill_out], dim=0) + else: + combined_out = None + +- if out is not None and combined_out is not None: ++ if out is not None and combined_out is not None and combined_out is not out: + out.copy_(combined_out) + + # ---- Merge max_score ---- diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py -index d72b17a..e4ad52e 100644 +index d72b17a..98dadae 100644 --- a/python/fmha_sm100/cute/interface.py +++ b/python/fmha_sm100/cute/interface.py @@ -136,6 +136,32 @@ def _prepare_paged_kv_for_tma(k, v, blk_kv: int): @@ -51,7 +270,26 @@ index d72b17a..e4ad52e 100644 def _validate_cu_seqlens( cu_seqlens: torch.Tensor, *, -@@ -736,10 +762,21 @@ def sparse_atten_func( +@@ -622,6 +648,7 @@ def sparse_atten_func( + usable_SM_count: int = -1, + qk_dtype: Optional[torch.dtype] = None, + pv_dtype: Optional[torch.dtype] = None, ++ out: Optional[torch.Tensor] = None, + ): + """Run SM100 CSR block-sparse varlen attention. + +@@ -691,6 +718,10 @@ def sparse_atten_func( + pv_dtype : torch.dtype, optional + Compile-time MMA operand dtype for PV. Defaults to V storage dtype, + except supported FP8 K/V cache staging modes. ++ out : torch.Tensor, optional ++ Destination for the attention output, written in place by the reduction ++ instead of being allocated here. Must be contiguous BF16 of shape ++ ``[total_q, Hq, 128]`` on the same device as ``q``. + + Returns + ------- +@@ -736,10 +767,21 @@ def sparse_atten_func( max_seqlen_q = int(max_seqlen_q) max_seqlen_k = int(max_seqlen_k) @@ -75,6 +313,48 @@ index d72b17a..e4ad52e 100644 k2q_row_ptr.contiguous(), k2q_q_indices.contiguous(), int(topK), +@@ -762,6 +804,7 @@ def sparse_atten_func( + int(max_seqlen_k), + qk_dtype, + pv_dtype, ++ out=out, + ) + + +@@ -1449,6 +1492,7 @@ def _sparse_atten_csr_varlen_forward( + max_seqlen_k: int, + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, ++ out: Optional[torch.Tensor] = None, + ): + total_q, head_q, dim = q.shape + if head_q % head_kv != 0: +@@ -1478,7 +1522,24 @@ def _sparse_atten_csr_varlen_forward( + if kernel_return_temperature_lse + else None + ) +- O_out = torch.empty(total_q, head_q, dim, dtype=torch.bfloat16, device=q.device) ++ if out is None: ++ O_out = torch.empty(total_q, head_q, dim, dtype=torch.bfloat16, device=q.device) ++ else: ++ # `combine` writes its output with a plain contiguous layout, so a ++ # caller that already owns a matching buffer can have it written ++ # directly instead of copying the result afterwards. ++ if tuple(out.shape) != (total_q, head_q, dim): ++ raise ValueError( ++ f"out must have shape {(total_q, head_q, dim)}, " ++ f"got {tuple(out.shape)}" ++ ) ++ if out.dtype != torch.bfloat16: ++ raise TypeError(f"out must be torch.bfloat16, got {out.dtype}") ++ if out.device != q.device: ++ raise ValueError("out must be on the same device as q") ++ if not out.is_contiguous(): ++ raise ValueError("out must be contiguous") ++ O_out = out + LSE_out = torch.empty(total_q, head_q, dtype=torch.float32, device=q.device) + LSE_temperature_out = ( + torch.empty_like(LSE_out) if kernel_return_temperature_lse else None diff --git a/python/fmha_sm100/cute/test_sparse_atten.py b/python/fmha_sm100/cute/test_sparse_atten.py index 21c777e..c22beef 100644 --- a/python/fmha_sm100/cute/test_sparse_atten.py @@ -199,7 +479,7 @@ index 21c777e..c22beef 100644 @pytest.mark.parametrize("causal", [True]) @pytest.mark.parametrize("batch", [3]) diff --git a/python/fmha_sm100/sparse_fmha_adapter.py b/python/fmha_sm100/sparse_fmha_adapter.py -index 306b416..81bf1ff 100644 +index 306b416..a564d47 100644 --- a/python/fmha_sm100/sparse_fmha_adapter.py +++ b/python/fmha_sm100/sparse_fmha_adapter.py @@ -129,11 +129,19 @@ def sparse_fmha_plan( @@ -367,3 +647,33 @@ index 306b416..81bf1ff 100644 ) # build_k2q_csr(return_schedule=True) builds schedule using hardware SM count internally +@@ -382,6 +454,17 @@ def sparse_fmha( + + softmax_scale = sm_scale if sm_scale is not None else q.shape[-1] ** -0.5 + ++ # Hand the caller's buffer to the kernel when its layout already matches ++ # what `combine` produces, so the result does not have to be copied into ++ # place afterwards. Anything else falls back to allocate-then-copy. ++ writable_out = ( ++ out is not None ++ and out.is_contiguous() ++ and out.dtype == torch.bfloat16 ++ and out.device == q.device ++ and tuple(out.shape) == tuple(q.shape) ++ ) ++ + # print(q.shape, k.shape) + result = sparse_atten_func( + q, k, v, +@@ -398,9 +481,10 @@ def sparse_fmha( + seqused_k=seqused_k, + schedule=schedule, + usable_SM_count=usable_SM_count, ++ out=out if writable_out else None, + ) + +- if out is not None: ++ if out is not None and not writable_out: + out.copy_(result) + return out, None + return result, None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index 4c1b7410b008..25bbcb4912f3 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -669,8 +669,10 @@ def on_update_kv_lens(self) -> None: n_valid_buf[:total_q].copy_(n_valid.to(torch.int32)) # Plan length mirrors. A plan is (has_mixed, split, batch, decode_sub, - # prefill_sub); a mixed batch is split at the first long request, so - # each sub-plan mirrors only its own request range. Within a range a + # prefill_sub), whose last two entries cover batch rows [0, split) and + # [split, batch). Which of the two holds the prefill requests depends on + # the batch order, so those names are positional only and nothing here + # may key off them; each sub-plan mirrors just its own range. Within a # sub-plan holds either one row per request or one per query token, # since the planner row-expands dense plans over query tokens, so the row # count selects the source. qo_offset must stay non-negative: negative diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py index fa9b4c8e4f5c..c60acc3e84c8 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py @@ -55,6 +55,7 @@ def msa_package_available() -> bool: # One symbol per module touched by 3rdparty/patches/msa_strided_paged_kv.patch. _MSA_PATCH_MARKERS = ( + ("fmha_sm100.api", "_mixed_batch_split"), ("fmha_sm100.cute.interface", "_prepare_paged_hnd_input"), ("fmha_sm100.sparse_fmha_adapter", "_page_table_for_plan"), )