From b9f3145d7895df276150abf229d566145e51753c Mon Sep 17 00:00:00 2001 From: ForFishes <2282912238@qq.com> Date: Wed, 15 Jul 2026 22:32:11 +0800 Subject: [PATCH 1/2] HySparse MQA sparse attention (FA4 block scoring + SWA MQA + block-sparse DSA gather) - FA4-fused block scoring (block_score_fa4) with PyLayer grads - SWA sliding-window MQA (swa_attn) replacing standalone block_sparse_attn_mqa - block-sparse DSA gather via cuDNN (block_sparse_mqa_dsa) with d_v<512 value-region padding to satisfy FlashMLA sparse kernel (d_v==512) and learnable per-head attn_sink - MLA/transformer wiring in multi_latent_attention, transformer_layer, gpt_layer_specs; enable_hy_sparse_attention config plumbing - runtime guards: fusion_layer_utils fp8-scales ImportError fallback for older paddlefleet_ops; transformer_config MTP-must-be-full-attention check --- src/paddlefleet/cudnn_ops/__init__.py | 6 + .../attn/csa_sparse_attn_fwd_cudnn.py | 14 +- .../cudnn_ops/block_sparse_mqa_dsa.py | 388 ++++++++++ src/paddlefleet/models/gpt/gpt_layer_specs.py | 41 ++ .../tilelang_ops/hysparse/__init__.py | 50 +- .../tilelang_ops/hysparse/benchmark.py | 226 ------ .../tilelang_ops/hysparse/block_score_fa4.py | 227 ++++++ .../hysparse/block_sparse_attn_mqa.py | 255 ------- .../hysparse/block_sparse_attn_mqa_bwd.py | 320 -------- .../tilelang_ops/hysparse/pipeline.py | 126 ++-- .../tilelang_ops/hysparse/reference.py | 230 ------ .../tilelang_ops/hysparse/swa_attn.py | 128 ++++ ...ock_score_attn.py => windowed_mqa_attn.py} | 219 +++--- ...e_attn_bwd.py => windowed_mqa_attn_bwd.py} | 182 ++++- .../transformer/moe/fusion_layer_utils.py | 12 +- .../transformer/multi_latent_attention.py | 681 +++++++++++++++++- .../transformer/transformer_config.py | 50 ++ .../transformer/transformer_layer.py | 321 +++++++++ 18 files changed, 2195 insertions(+), 1281 deletions(-) create mode 100644 src/paddlefleet/cudnn_ops/block_sparse_mqa_dsa.py delete mode 100644 src/paddlefleet/tilelang_ops/hysparse/benchmark.py create mode 100644 src/paddlefleet/tilelang_ops/hysparse/block_score_fa4.py delete mode 100644 src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa.py delete mode 100644 src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa_bwd.py delete mode 100644 src/paddlefleet/tilelang_ops/hysparse/reference.py create mode 100644 src/paddlefleet/tilelang_ops/hysparse/swa_attn.py rename src/paddlefleet/tilelang_ops/hysparse/{block_score_attn.py => windowed_mqa_attn.py} (51%) rename src/paddlefleet/tilelang_ops/hysparse/{block_score_attn_bwd.py => windowed_mqa_attn_bwd.py} (63%) diff --git a/src/paddlefleet/cudnn_ops/__init__.py b/src/paddlefleet/cudnn_ops/__init__.py index 8816b7338..519934e82 100644 --- a/src/paddlefleet/cudnn_ops/__init__.py +++ b/src/paddlefleet/cudnn_ops/__init__.py @@ -15,6 +15,10 @@ """cuDNN frontend ops bridged into PaddleFleet.""" from .attn.csa_sparse_attn_bwd_cudnn import csa_sparse_attn_bwd_cudnn +from .block_sparse_mqa_dsa import ( + block_sparse_mqa_attention_dsa, + is_dsa_available, +) from .indexer.csa_indexer_bwd_cudnn import csa_indexer_bwd from .indexer.csa_indexer_fwd_cudnn import ( cudnn_indexer_forward, @@ -23,9 +27,11 @@ ) __all__ = [ + "block_sparse_mqa_attention_dsa", "csa_indexer_bwd", "csa_sparse_attn_bwd_cudnn", "cudnn_indexer_forward", "cudnn_indexer_topk", "cudnn_indexer_topk_fwd", + "is_dsa_available", ] diff --git a/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py b/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py index 305f80a7c..69d34f815 100644 --- a/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py +++ b/src/paddlefleet/cudnn_ops/attn/csa_sparse_attn_fwd_cudnn.py @@ -47,7 +47,7 @@ def _get_topk_alignment() -> int: def flash_mla_sparse_attn( - q, kv, attn_sink, topk_idxs, sm_scale=None, indexer_topk: int = 0 + q, kv, attn_sink, topk_idxs, sm_scale=None, indexer_topk: int = 0, d_v=None ): if _flash_mla_sparse_fwd is None: raise RuntimeError("flash_mla is not available") @@ -55,6 +55,10 @@ def flash_mla_sparse_attn( b, sq, h, d = q.shape _, skv, _ = kv.shape topk = topk_idxs.shape[-1] + # Value dim may be smaller than the query/key dim (absorbed MLA MQA uses + # d_qk=576 / d_v=512). Default to a symmetric d_v=d. + if d_v is None: + d_v = d q_flat = q.reshape([b * sq, h, d]) kv_flat = kv.reshape([b * skv, d]) @@ -70,7 +74,7 @@ def flash_mla_sparse_attn( kv_flat.unsqueeze(1), global_idxs.unsqueeze(1), sm_scale, - d_v=d, + d_v=d_v, attn_sink=attn_sink, topk_length=None, indexer_topk=indexer_topk, @@ -81,4 +85,8 @@ def flash_mla_sparse_attn( else: out_flat, _max_logits, lse = res lse_indexer = None - return out_flat.reshape([b, sq, h, d]), lse.reshape([b, sq, h]), lse_indexer + return ( + out_flat.reshape([b, sq, h, d_v]), + lse.reshape([b, sq, h]), + lse_indexer, + ) diff --git a/src/paddlefleet/cudnn_ops/block_sparse_mqa_dsa.py b/src/paddlefleet/cudnn_ops/block_sparse_mqa_dsa.py new file mode 100644 index 000000000..3d2f1a53e --- /dev/null +++ b/src/paddlefleet/cudnn_ops/block_sparse_mqa_dsa.py @@ -0,0 +1,388 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DSA (FlashMLA sparse fwd + cuDNN DSA bwd) backend for the HySparse +block-sparse MQA gather branch. + +DeepSeek-v4's CSA sparse attention (FlashMLA sparse forward + cuDNN DSA +backward) natively handles the absorbed-MQA D=576 query / D_v=512 value +single-shared-head layout. This module bridges the HySparse *block* selection +onto that *token*-level DSA path: + +1. **Block -> token index expansion.** Each selected document-relative block + ``j`` (spanning key cols ``[bos + j*block_B, bos + (j+1)*block_B)``) is + expanded into its ``block_B`` absolute token columns. ``block_B == 64`` equals + the SM100 TopK alignment, so one block == one DSA tile chunk (no padding). +2. **Causal/doc masking folded into the index list.** Any expanded column with + ``col >= eos`` (or belonging to a ``-1`` padding block) is set to ``-1``, + which DSA treats as invalid -- reproducing the ``valid_range [bos, eos)`` + semantics without a kernel argument. +3. **Sinkless softmax via a very-negative sink (default).** HySparse is sinkless; + DSA always applies an attention sink. Passing ``sink = -1e30`` per head makes + ``exp(sink - m) -> 0``, recovering plain softmax. The same sink tensor is used + in forward and backward and its gradient is discarded. When a learnable + per-head ``attn_sink`` is supplied instead (attention-sink / softmax + off-by-one bias), its logits fill the real heads, padded heads stay ``-1e30``, + and the sink gradient is computed analytically in the backward (the SM100 + cuDNN DSA op allocates ``d_sink`` but never fills it -- it returns zeros -- + so we derive ``d_sink[h] = -sum_{b,s}(p_sink * Delta)`` from the saved LSE / + out / dO instead). +4. **K==V-unified latent.** DSA takes one ``kv_full`` tensor whose value is its + leading ``kv_lora_rank`` slice. ``shared_key_sq [B, S, 576]`` already has + value == leading 512, so it is passed directly and its 576-wide gradient is + the (combined) gradient w.r.t. the shared latent. + +FlashMLA sparse fwd only supports ``h_q == 64`` on SM100, so query heads are +zero-padded up to 64 when ``H < 64`` (padded heads receive zero output gradient +and contribute no KV gradient). +""" + +import paddle + +_DSA_HEADS = 64 # FlashMLA sparse fwd only supports h_q == 64 on SM100. +_NEG_SINK = -1e30 # sink so large-negative that exp(sink - m) underflows to 0. + + +def is_dsa_available() -> bool: + """Whether the FlashMLA sparse fwd + cuDNN DSA bwd path can run here.""" + try: + import paddlefleet_ops + + from paddlefleet.cudnn_ops.attn import csa_sparse_attn_fwd_cudnn + + if ( + not paddlefleet_ops.is_flash_mla_available() + or csa_sparse_attn_fwd_cudnn._flash_mla_sparse_fwd is None + ): + return False + if not paddlefleet_ops.is_cudnn_frontend_available(): + return False + except (ImportError, RuntimeError, AttributeError): + return False + return True + + +def _expand_blocks_to_token_indices(indices, valid_range, block_B): + """Expand doc-relative block ids to per-token key-column indices. + + Args: + indices: [B, S, topk] int, document-relative block ids (-1 padding). + valid_range: [B, S, 2] int, per-query ``[bos, eos)`` valid key columns. + block_B: block size in tokens. + + Returns: + [B, S, topk * block_B] int32 absolute key columns; entries whose column + is ``>= eos`` or that belong to a ``-1`` padding block are set to -1. + + The whole computation is a pure integer index construction and MUST NOT + carry an autograd graph. Under full-layer recompute, ``indices`` / + ``valid_range`` are recomputed with grad tracking enabled, so the trailing + ``paddle.where(...).astype("int32")`` would otherwise build a stray + Where/Cast grad chain. Passing that grad-tracked integer tensor into the + ``_BlockSparseDSA`` PyLayer registers a backward edge to an orphan + ``CastGradNode`` which the engine then schedules with an empty grad holder + (ref_cnt 0) -> ``cast()`` on an undefined tensor -> segfault. Building the + indices under ``no_grad`` (and returning a detached tensor) removes the + stray grad nodes entirely. + """ + with paddle.no_grad(): + b, s, topk = indices.shape + bos = valid_range[..., 0:1].astype("int64") # [B, S, 1] + eos = valid_range[..., 1:2].astype("int64") # [B, S, 1] + + blk = indices.astype("int64") # [B, S, topk] + start = bos + blk * block_B # [B, S, topk] absolute col of block start + offs = paddle.arange(block_B, dtype="int64").reshape([1, 1, 1, block_B]) + cols = start.unsqueeze(-1) + offs # [B, S, topk, block_B] + cols = cols.reshape([b, s, topk * block_B]) + + blk_invalid = ( + (blk < 0).unsqueeze(-1).expand([b, s, topk, block_B]) + ).reshape([b, s, topk * block_B]) + # cols >= bos always holds (blk >= 0 => start >= bos, offs >= 0); only + # the tail past eos needs masking, plus columns from -1 padding blocks. + col_invalid = cols >= eos # eos broadcasts over the last dim + invalid = paddle.logical_or(blk_invalid, col_invalid) + neg = paddle.full_like(cols, -1) + token_indices = paddle.where(invalid, neg, cols).astype("int32") + token_indices.stop_gradient = True + return token_indices + + +class _BlockSparseDSA(paddle.autograd.PyLayer): + """FlashMLA sparse forward + cuDNN DSA backward for absorbed MQA. + + forward inputs: + query: [B, S, H, Dk] (Dk=576) + shared_key_sq: [B, S, Dk] single shared K/V latent (value = leading Dv) + token_indices: [B, S, L] int32 per-batch-local key cols (-1 invalid), + already block-expanded + doc/causal masked. + sm_scale, kv_lora_rank (Dv), num_heads (real H) + outputs: out [B, S, H * Dv] (differentiable in query and shared_key_sq). + """ + + @staticmethod + def forward( + ctx, + query, + shared_key_sq, + token_indices, + sm_scale, + d_v, + num_heads, + attn_sink=None, + ): + from paddlefleet.cudnn_ops.attn.csa_sparse_attn_fwd_cudnn import ( + flash_mla_sparse_attn, + ) + + b, s, h, dk = query.shape + ctx.num_heads = num_heads + ctx.d_v = d_v + ctx.sm_scale = float(sm_scale) + ctx.query_dtype = query.dtype + ctx.kv_dtype = shared_key_sq.dtype + + # Pad query heads up to the DSA-supported h_q == 64. The FlashMLA + # sparse backend fixes h_q at _DSA_HEADS (sink is [_DSA_HEADS]); it can + # only handle h <= _DSA_HEADS by zero-padding the head dim. h > 64 is + # unsupported and must be rejected here rather than failing deep in the + # CUDA op with an opaque shape error. + if h > _DSA_HEADS: + raise ValueError( + f"HySparse DSA sparse attention supports at most {_DSA_HEADS} " + f"query heads per rank, but got {h}. Reduce " + "num_attention_heads / swa_num_attention_heads (per-rank after " + f"TP) to <= {_DSA_HEADS}." + ) + if h < _DSA_HEADS: + pad = paddle.zeros([b, s, _DSA_HEADS - h, dk], dtype=query.dtype) + q_pad = paddle.concat([query, pad], axis=2) + else: + q_pad = query + + # Attention sink over the DSA-fixed 64 heads. When ``attn_sink`` is None + # HySparse is sinkless: a per-head ``-1e30`` makes ``exp(sink - m) -> 0``, + # recovering plain softmax and discarding the (unused) sink gradient. + # When a learnable ``attn_sink [num_heads]`` is supplied, its logits fill + # the real heads and the padded heads keep ``-1e30`` (they contribute no + # output / gradient). The sink gradient is routed back to the parameter. + ctx.learnable_sink = attn_sink is not None + if attn_sink is None: + sink = paddle.full([_DSA_HEADS], _NEG_SINK, dtype="float32") + else: + assert list(attn_sink.shape) == [num_heads], ( + f"attn_sink must be [num_heads={num_heads}]; got " + f"{attn_sink.shape}" + ) + sink_real = attn_sink.cast("float32") + if h < _DSA_HEADS: + sink_pad = paddle.full( + [_DSA_HEADS - h], _NEG_SINK, dtype="float32" + ) + sink = paddle.concat([sink_real, sink_pad], axis=0) + else: + sink = sink_real + sink = sink.contiguous() + kv = shared_key_sq # [B, S, Dk] value = leading d_v + + out, lse, _ = flash_mla_sparse_attn( + q_pad, + kv, + sink, + token_indices, + sm_scale=ctx.sm_scale, + d_v=d_v, + ) # out [B, S, 64, d_v], lse [B, S, 64] + + ctx.save_for_backward(q_pad, kv, out, lse, token_indices, sink) + ctx.needs_grad = ( + not query.stop_gradient, + not shared_key_sq.stop_gradient, + ctx.learnable_sink and not attn_sink.stop_gradient, + ) + out_h = out[:, :, :h, :].contiguous() # drop padded heads + return out_h.reshape([b, s, h * d_v]) + + @staticmethod + def backward(ctx, grad_output): + from paddlefleet.cudnn_ops import csa_sparse_attn_bwd_cudnn + from paddlefleet.fusions.csa_sparse_attn_utils import ( + _local_to_global_flat, + ) + + q_pad, kv, out, lse, token_indices, sink = ctx.saved_tensor() + b, s, hpad, dk = q_pad.shape + d_v = ctx.d_v + h = ctx.num_heads + _, skv, _ = kv.shape + + # Re-pad the incoming grad back to hpad heads (padded heads get 0 grad, + # so they contribute nothing to dq / dkv). + grad_output = grad_output.reshape([b, s, h, d_v]) + if h < hpad: + gpad = paddle.zeros([b, s, hpad - h, d_v], dtype=grad_output.dtype) + do = paddle.concat([grad_output, gpad], axis=2) + else: + do = grad_output + do = do.contiguous() + + q_flat = q_pad.reshape([b * s, hpad, dk]) + o_flat = out.reshape([b * s, hpad, d_v]) + do_flat = do.reshape([b * s, hpad, d_v]) + kv_flat = kv.reshape([b * skv, dk]) + lse_flat = lse.reshape([b * s, hpad]) + gidx_flat = _local_to_global_flat(token_indices, skv) + + dq_flat, dkv_flat, _d_sink_unused = csa_sparse_attn_bwd_cudnn( + q_flat, + kv_flat, + o_flat, + do_flat, + lse_flat, + sink, + gidx_flat, + softmax_scale=ctx.sm_scale, + ) + + gq, gk, gsink = ctx.needs_grad + dq = None + if gq: + dq = dq_flat.reshape([b, s, hpad, dk])[:, :, :h, :].contiguous() + dq = dq.cast(ctx.query_dtype) + dkv = None + if gk: + dkv = dkv_flat.reshape([b, s, dk]).cast(ctx.kv_dtype) + d_attn_sink = None + if gsink: + # The cuDNN DSA backward (SM100) allocates ``d_sink`` but its kernel + # never populates it -- it always returns zeros. So compute the sink + # gradient analytically here from the saved forward tensors. + # + # For a virtual sink logit ``s_h`` competing in the softmax denom + # ``Z = sum_k exp(logit_k) + exp(s_h)``, weight ``p_k = exp(l_k)/Z`` + # and sink mass ``p_sink = exp(s_h)/Z``. Since + # ``d p_k / d s_h = -p_k * p_sink`` and ``out = sum_k p_k v_k``: + # d out / d s_h = -p_sink * out + # d_sink[h] = sum_{b,s}( dO . (d out / d s_h) ) + # = -sum_{b,s}( p_sink * (dO . out) ) + # = -sum_{b,s}( p_sink * Delta ) + # with ``Delta[b,s,h] = sum_dv(out * dO)``. The forward LSE is + # KV-only (excludes the sink), so the full log-denominator is + # ``logaddexp(lse_kv, s_h)`` and ``p_sink = exp(s_h - lse_full)``. + out_h = out[:, :, :h, :].astype("float32") + do_h = do[:, :, :h, :].astype("float32") + delta = (out_h * do_h).sum(axis=-1) # [b, s, h] + sink_real = sink[:h].astype("float32").reshape([1, 1, h]) + lse_h = lse[:, :, :h].astype("float32") + lse_full = paddle.logaddexp(lse_h, sink_real) + p_sink = paddle.exp(sink_real - lse_full) # [b, s, h] + d_attn_sink = (-(delta * p_sink).sum(axis=[0, 1])).contiguous() + d_attn_sink = d_attn_sink.cast("float32") + + # One returned grad per **tensor** input, in order. Non-tensor inputs + # (sm_scale, d_v, num_heads) occupy no slot. ``attn_sink`` occupies a + # slot only when it was passed as a tensor (sinkless -> None -> no slot), + # so the returned count is 3 (sinkless) or 4 (learnable sink). + grads = [dq, dkv, None] # query, shared_key_sq, token_indices + if ctx.learnable_sink: + grads.append(d_attn_sink) + return tuple(grads) + + +def block_sparse_mqa_attention_dsa( + query, + shared_key_sq, + shared_block_indices, + valid_range, + sm_scale=None, + block_B=64, + kv_lora_rank=512, + attn_sink=None, +): + """HySparse block-sparse gather attention over the absorbed MQA latent. + + Args: + query: [B, S, H, Dk] (Dk = kv_lora_rank + rope, e.g. 576). + shared_key_sq: [B, S, Dk] shared K/V latent; value = leading + ``kv_lora_rank`` slice. + shared_block_indices: [B, S, topk] int, document-relative selected block + ids (-1 padding), from ``select_topk_blocks``. + valid_range: [B, S, 2] int, per-query ``[bos, eos)``. + sm_scale: softmax scale (defaults to ``Dk ** -0.5``). + block_B: block size in tokens (must equal the DSA TopK + alignment, 64 on SM100). + kv_lora_rank: value dim ``Dv`` (leading slice of the latent). + attn_sink: [H] fp32 per-head learnable attention-sink logit, + or ``None`` for HySparse's default sinkless softmax + (a ``-1e30`` sink whose gradient is discarded). + + Returns: + ``(out, None)`` where ``out`` is ``[B, S, H * kv_lora_rank]`` and carries + gradient to ``query`` and ``shared_key_sq`` (and ``attn_sink`` when a + learnable sink is supplied). The second element keeps the ``(out, lse)`` + tuple shape of the consumer call site; DSA does not surface a + differentiable lse here. + + ``kv_lora_rank`` < 512 (e.g. ernielite's 448): the FlashMLA sparse kernel + hard-requires ``d_v == 512`` and ``d_qk in {512, 576}``. We map the smaller + latent onto that fixed layout by zero-padding the *value* region up to 512: + the latent ``[value(kv_lora_rank) | rope]`` is re-laid as + ``[value(kv_lora_rank) | zeros(512 - kv_lora_rank) | rope]`` for both query + and key. The dot-product score is unchanged (the inserted zeros contribute + nothing), the kernel value becomes the leading 512 (= real value padded with + zeros), and the output's leading ``kv_lora_rank`` columns are sliced back out. + The pad/slice run outside the PyLayer so autograd routes the gradients back + to the original ``query`` / ``shared_key_sq`` shapes automatically. + """ + if sm_scale is None: + sm_scale = query.shape[-1] ** -0.5 + token_indices = _expand_blocks_to_token_indices( + shared_block_indices, valid_range, block_B + ) + + b, s, num_heads, d_qk = query.shape + pad_v = 512 - kv_lora_rank + if pad_v > 0: + # Re-lay [value | rope] -> [value | zeros | rope] so value == leading 512. + q_val, q_rope = query[..., :kv_lora_rank], query[..., kv_lora_rank:] + k_val = shared_key_sq[..., :kv_lora_rank] + k_rope = shared_key_sq[..., kv_lora_rank:] + zq = paddle.zeros([b, s, num_heads, pad_v], dtype=query.dtype) + zk = paddle.zeros([b, s, pad_v], dtype=shared_key_sq.dtype) + query_p = paddle.concat([q_val, zq, q_rope], axis=-1) + key_p = paddle.concat([k_val, zk, k_rope], axis=-1) + eff_d_v = 512 + elif pad_v < 0: + raise ValueError( + f"HySparse DSA supports kv_lora_rank <= 512, got {kv_lora_rank}." + ) + else: + query_p, key_p, eff_d_v = query, shared_key_sq, kv_lora_rank + + out = _BlockSparseDSA.apply( + query_p, + key_p, + token_indices, + float(sm_scale), + int(eff_d_v), + int(num_heads), + attn_sink, + ) + + if pad_v > 0: + # Drop the padded value columns: keep the real leading kv_lora_rank. + out = out.reshape([b, s, num_heads, eff_d_v])[..., :kv_lora_rank] + out = out.reshape([b, s, num_heads * kv_lora_rank]) + return out, None diff --git a/src/paddlefleet/models/gpt/gpt_layer_specs.py b/src/paddlefleet/models/gpt/gpt_layer_specs.py index e889e7596..d7efdde11 100644 --- a/src/paddlefleet/models/gpt/gpt_layer_specs.py +++ b/src/paddlefleet/models/gpt/gpt_layer_specs.py @@ -86,12 +86,14 @@ from paddlefleet.transformer.multi_latent_attention import ( MLASelfAttention, MLASelfAttentionSublayersSpec, + MQASelfAttention, ) from paddlefleet.transformer.multi_token_prediction import ( get_mtp_layer_spec_for_backend, ) from paddlefleet.transformer.paddle_norm import L2Norm from paddlefleet.transformer.transformer_layer import ( + HySparseTransformerLayer, TransformerLayer, TransformerLayerSublayersSpec, TransformerLayerWithOverlap, @@ -231,6 +233,10 @@ def get_attention_spec( assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." # Decide attention class: always MLASelfAttention (DSA is a pluggable core_attention) attn_cls = MLASelfAttention + + if config is not None and config.enable_hy_sparse_attention: + attn_cls = MQASelfAttention + # Gated attention gated_attention = getattr(config, "gated_attention", False) @@ -423,6 +429,41 @@ def get_gpt_layer_local_spec( transformer_cls = HyperConnectionTransformerLayer + if config is not None and config.enable_hy_sparse_attention: + # HySparse is only implemented for the MLA-absorbed MQA attention path + # (MQASelfAttention). HySparseTransformerLayer passes a ``shared_kv`` + # kwarg into the attention forward; a plain SelfAttention.forward has no + # such parameter and would raise TypeError. Require MLA here so the + # attention class is MQASelfAttention (see get_attention_spec). + is_mla = ( + multi_latent_attention + or attention_layer_type == "multi_latent_attention" + ) + if not is_mla: + raise ValueError( + "enable_hy_sparse_attention requires multi-latent attention " + "(the MLA-absorbed MQA path); set multi_latent_attention=True " + "(or attention_layer_type='multi_latent_attention'). HySparse " + "is not supported with standard self-attention." + ) + # HySparseTransformerLayer does not implement the hyper-connection or + # block-attention-residual dataflows. Overriding transformer_cls here + # would silently drop those layers instead of applying them, so reject + # the unsupported combinations explicitly rather than downgrading. + if config.enable_hyper_connections: + raise ValueError( + "enable_hy_sparse_attention is incompatible with " + "enable_hyper_connections: HySparseTransformerLayer does not " + "implement the hyper-connection dataflow." + ) + if config.block_attention_residuals: + raise ValueError( + "enable_hy_sparse_attention is incompatible with " + "block_attention_residuals: HySparseTransformerLayer does not " + "implement the block-attention-residual dataflow." + ) + transformer_cls = HySparseTransformerLayer + if paddle.distributed.is_initialized(): try: pp_configs = fleet.fleet._user_defined_strategy.hybrid_configs[ diff --git a/src/paddlefleet/tilelang_ops/hysparse/__init__.py b/src/paddlefleet/tilelang_ops/hysparse/__init__.py index 9fb45716c..a7a73c9b3 100644 --- a/src/paddlefleet/tilelang_ops/hysparse/__init__.py +++ b/src/paddlefleet/tilelang_ops/hysparse/__init__.py @@ -14,29 +14,39 @@ """HySparse block-attention TileLang operators (paper arXiv 2602.03560). -MQA/MLA variant: K/V are a single shared head, block selection is aggregated -across the query group by a group-wise maximum, and the sparse branch gathers -only the selected blocks. +MQA/MLA variant: K/V are a single shared head; block selection is aggregated +across the query group by a group-wise maximum; the sparse branch gathers only +the selected blocks. + +Public entry points (enable_hy_sparse_attention=True): +* full layers -> block_score_fa4_attn_fwd (FA4-fused block scoring) + + select_topk_blocks +* SWA layers -> sliding_window_mqa_attention (windowed MQA main path); + the sparse gather branch runs on the cuDNN DSA op + (paddlefleet.cudnn_ops.block_sparse_mqa_attention_dsa). """ -from .block_score_attn import ( - block_score_mqa_attn_fwd, - block_scores_from_logit, -) -from .block_score_attn_bwd import block_score_mqa_bwd_interface -from .block_sparse_attn_mqa import block_sparse_mqa_attn_fwd -from .block_sparse_attn_mqa_bwd import block_sparse_mqa_bwd_interface -from .pipeline import ( - hysparse_forward_mqa, - select_topk_blocks, -) +from .pipeline import select_topk_blocks +from .swa_attn import sliding_window_mqa_attention + + +def block_score_fa4_attn_fwd(*args, **kwargs): + """Lazy wrapper for the FA4-fused block scorer. + + ``block_score_fa4`` imports ``paddlefleet_ops.flash_mask.cute.*``, which is + only available on SM10/Blackwell (guarded by ``is_flash_mask_available()``). + Importing it eagerly at package init would break ``from ... import + sliding_window_mqa_attention`` / ``select_topk_blocks`` on non-Blackwell + hardware (H20/A100), even for callers that never touch the FA4 full path. + Defer the FA4-only import to the actual call site. + """ + from .block_score_fa4 import block_score_fa4_attn_fwd as _impl + + return _impl(*args, **kwargs) + __all__ = [ - "block_score_mqa_attn_fwd", - "block_scores_from_logit", - "block_score_mqa_bwd_interface", - "block_sparse_mqa_attn_fwd", - "block_sparse_mqa_bwd_interface", - "hysparse_forward_mqa", + "block_score_fa4_attn_fwd", "select_topk_blocks", + "sliding_window_mqa_attention", ] diff --git a/src/paddlefleet/tilelang_ops/hysparse/benchmark.py b/src/paddlefleet/tilelang_ops/hysparse/benchmark.py deleted file mode 100644 index c965b98ee..000000000 --- a/src/paddlefleet/tilelang_ops/hysparse/benchmark.py +++ /dev/null @@ -1,226 +0,0 @@ -# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Latency benchmark for the HySparse (MQA) block-attention operators. - -Times each operator (block-score fwd/bwd, block-sparse gather fwd/bwd, and the -full block-score -> TopK -> block-sparse pipeline) and reports wall time plus -peak allocated memory. Includes the target MLA scenario head_dim=576, -num_heads=64, seq_len up to 32k. - -Run standalone:: - - python -m paddlefleet.tilelang_ops.hysparse.benchmark # full sweep - python -m paddlefleet.tilelang_ops.hysparse.benchmark --mla-32k # only 32k - python -m paddlefleet.tilelang_ops.hysparse.benchmark \ - --seqlens 8192 --H 64 --D 576 --block-B 64 --topk 16 -""" - -import argparse -import time - -import paddle - -paddle.enable_compat(scope={"tilelang"}, silent=True) - -from .block_score_attn import block_score_mqa_attn_fwd -from .block_score_attn_bwd import block_score_mqa_bwd_interface -from .block_sparse_attn_mqa import block_sparse_mqa_attn_fwd -from .block_sparse_attn_mqa_bwd import block_sparse_mqa_bwd_interface -from .pipeline import hysparse_forward_mqa -from .reference import make_causal_valid_range - - -def _sync(): - paddle.device.synchronize() - - -def _time_ms(fn, warmup=2, iters=5): - """Median-free mean latency in ms over ``iters`` runs after ``warmup``.""" - for _ in range(warmup): - fn() - _sync() - t0 = time.time() - for _ in range(iters): - fn() - _sync() - return (time.time() - t0) / iters * 1000.0 - - -def _peak_gb(): - return paddle.device.cuda.max_memory_allocated() / 1e9 - - -def _rand_inputs(B, S, H, D, seed=0): - paddle.seed(seed) - q = paddle.randn([B, S, H, D], dtype="bfloat16") - k = paddle.randn([B, S, D], dtype="bfloat16") - v = paddle.randn([B, S, D], dtype="bfloat16") - return q, k, v - - -def bench_config(B, S, H, D, block_B, topk, do_bwd=True, iters=5): - """Benchmark every HySparse MQA op for one shape; returns a result dict.""" - q, k, v = _rand_inputs(B, S, H, D) - vr = make_causal_valid_range(S, batch=B) - sm = D**-0.5 - res = {"B": B, "S": S, "H": H, "D": D, "block_B": block_B, "topk": topk} - - # block-score full-attention forward (emits per-block max logits) - paddle.device.cuda.reset_max_memory_allocated() - res["score_fwd"] = _time_ms( - lambda: block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm, block_B=block_B - ), - iters=iters, - ) - o, lse, _ = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm, block_B=block_B - ) - - # end-to-end pipeline (scoring -> TopK -> gather sparse attention) - paddle.device.cuda.reset_max_memory_allocated() - res["pipeline"] = _time_ms( - lambda: hysparse_forward_mqa( - q, k, v, vr, topk, sm_scale=sm, block_B=block_B - ), - iters=iters, - ) - _, _, idx, _, _ = hysparse_forward_mqa( - q, k, v, vr, topk, sm_scale=sm, block_B=block_B - ) - res["pipeline_peak_gb"] = _peak_gb() - - # block-sparse gather forward, given the pipeline's selected indices - res["sparse_fwd"] = _time_ms( - lambda: block_sparse_mqa_attn_fwd( - q, k, v, idx, vr, sm_scale=sm, block_B=block_B - ), - iters=iters, - ) - so, slse = block_sparse_mqa_attn_fwd( - q, k, v, idx, vr, sm_scale=sm, block_B=block_B - ) - - if do_bwd: - do = paddle.randn([B, S, H, D], dtype="bfloat16") - res["score_bwd"] = _time_ms( - lambda: block_score_mqa_bwd_interface( - q, k, v, o, do, lse, vr, sm_scale=sm, block_B=block_B - ), - iters=iters, - ) - res["sparse_bwd"] = _time_ms( - lambda: block_sparse_mqa_bwd_interface( - q, k, v, so, do, slse, idx, vr, sm_scale=sm, block_B=block_B - ), - iters=iters, - ) - return res - - -_COLS = [ - ("S", "S", 7), - ("H", "H", 4), - ("D", "D", 5), - ("topk", "topk", 5), - ("score_fwd", "scoreFwd", 10), - ("score_bwd", "scoreBwd", 10), - ("sparse_fwd", "sparseFwd", 10), - ("sparse_bwd", "sparseBwd", 10), - ("pipeline", "pipeline", 10), - ("pipeline_peak_gb", "peakGB", 8), -] - - -def _fmt(res): - cells = [] - for key, _, w in _COLS: - v = res.get(key) - if v is None: - s = "-" - elif isinstance(v, float): - s = f"{v:.1f}" - else: - s = str(v) - cells.append(s.rjust(w)) - return " ".join(cells) - - -def _header(): - return " ".join(h.rjust(w) for _, h, w in _COLS) - - -def run_sweep(seqlens, H, D, block_B, topk, do_bwd, iters): - print( - f"HySparse MQA benchmark (B=1, H={H}, D={D}, block_B={block_B}, " - f"units=ms, bf16)" - ) - print(_header()) - for S in seqlens: - res = bench_config( - 1, S, H, D, block_B, topk, do_bwd=do_bwd, iters=iters - ) - print(_fmt(res)) - - -def main(): - p = argparse.ArgumentParser(description="HySparse MQA latency benchmark") - p.add_argument( - "--seqlens", - type=int, - nargs="+", - default=[2048, 4096, 8192, 16384, 32768], - ) - p.add_argument("--H", type=int, default=64) - p.add_argument("--D", type=int, default=576) - p.add_argument("--block-B", type=int, default=64) - p.add_argument("--topk", type=int, default=16) - p.add_argument("--iters", type=int, default=5) - p.add_argument( - "--no-bwd", - action="store_true", - help="skip backward kernels (forward/pipeline only)", - ) - p.add_argument( - "--mla-32k", - action="store_true", - help="only the MLA target: D=576,H=64,S=32768", - ) - args = p.parse_args() - - if args.mla_32k: - run_sweep( - [32768], - 64, - 576, - args.block_B, - args.topk, - not args.no_bwd, - args.iters, - ) - return - run_sweep( - args.seqlens, - args.H, - args.D, - args.block_B, - args.topk, - not args.no_bwd, - args.iters, - ) - - -if __name__ == "__main__": - main() diff --git a/src/paddlefleet/tilelang_ops/hysparse/block_score_fa4.py b/src/paddlefleet/tilelang_ops/hysparse/block_score_fa4.py new file mode 100644 index 000000000..2c6ec189e --- /dev/null +++ b/src/paddlefleet/tilelang_ops/hysparse/block_score_fa4.py @@ -0,0 +1,227 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FA4-fused full block-score attention (HySparse full layers). + +The fused FlashAttention-v4 sm100 forward kernel (``paddlefleet_ops.flash_mask``) +now optionally emits the HySparse per-(query, key-block) max logit alongside the +attention output, computed inside the softmax epilogue (see the ``has_block_logit`` +path in ``flash_fwd_sm100.py``). This wrapper allocates the block-logit buffer, +runs the fast FA4 path, and returns an ``(out, lse, block_logit)`` triple -- at +~FA4 speed and without an extra HBM Q/K re-read for the block-max reduction. + +Semantics: +* ``out`` ``[B, S, H, D_v]`` attention output (softmax-scaled). +* ``lse`` ``[B, S, H]`` natural-log LSE (transposed from FA4's ``[B, H, S]``). +* ``block_logit`` ``[B, H, S, num_blocks]`` max of the **scaled** attention + logit (``softmax_scale * q.k``, the exact value fed into softmax) over each + key-block's valid columns; fully-masked / never-visited blocks stay ``-inf``. + Storing the scaled logit puts every head on one head-independent scale, so a + downstream ``block_logit - LSE`` yields log(max attention weight in the + block). Convert to the eq.(3) probability with + :func:`pipeline.block_scores_from_logit`. + +Masking: ``causal=True`` (single-document causal). Document masks flow through +FA4's flashmask ``startend_row_indices`` (passed straight through); because the +block-max reduce reads the score fragment *after* ``mask_fn`` runs, any mask FA4 +supports is honoured automatically. +""" + +import paddle +from paddlefleet_ops.flash_mask.cute.flashmask_utils import FlashMaskInfoPaddle +from paddlefleet_ops.flash_mask.cute.interface import ( + _flash_attn_bwd, + _flash_attn_fwd, +) + + +class _BlockScoreFA4Attn(paddle.autograd.PyLayer): + """Differentiable FA4-fused full block-score attention. + + ``out`` carries gradient (FA4 fwd + bwd kernels); ``lse`` and the in-place + ``block_logit`` buffer are non-differentiable (they feed the hard TopK block + selection). Runs the fast fused FA4 sm100 path. + """ + + @staticmethod + def forward( + ctx, + q, + k, + v, + block_logit, + block_bos, + causal, + sm_scale, + startend_row_indices, + block_B, + ): + out, lse = _flash_attn_fwd( + q, + k, + v, + softmax_scale=sm_scale, + causal=causal, + return_lse=True, + startend_row_indices=startend_row_indices, + block_logit=block_logit, + block_size=block_B, + block_bos=block_bos, + pack_gqa=False, + ) + # save_for_backward only accepts Tensors; startend_row_indices is + # optional (None disables flashmask), so stash whether it is present + # and only save it when it is a real Tensor. + ctx.has_mask = startend_row_indices is not None + if ctx.has_mask: + ctx.save_for_backward(q, k, v, startend_row_indices, out, lse) + else: + ctx.save_for_backward(q, k, v, out, lse) + ctx.needs_grad = ( + not q.stop_gradient, + not k.stop_gradient, + not v.stop_gradient, + ) + ctx.mark_non_differentiable(lse) + ctx.sm_scale = sm_scale + ctx.causal = causal + return out, lse + + @staticmethod + def backward(ctx, dout, *_): + if ctx.has_mask: + q, k, v, startend_row_indices, out, lse = ctx.saved_tensor() + else: + q, k, v, out, lse = ctx.saved_tensor() + startend_row_indices = None + flashmask_info = None + if startend_row_indices is not None: + flashmask_info = FlashMaskInfoPaddle( + startend_row_indices=startend_row_indices, + is_causal=ctx.causal, + ) + dq, dk, dv, _ = _flash_attn_bwd( + q, + k, + v, + out, + dout.contiguous(), + lse, + flashmask_info, + softmax_scale=ctx.sm_scale, + causal=ctx.causal, + deterministic=paddle.get_flags(["FLAGS_cudnn_deterministic"])[ + "FLAGS_cudnn_deterministic" + ], + ) + # One grad slot per forward Tensor input, in order: + # q, k, v, block_logit, block_bos [, startend_row_indices]. + # block_logit / block_bos are non-differentiable buffers -> None; + # startend_row_indices (only a Tensor input when has_mask) -> None. + gq, gk, gv = ctx.needs_grad + grads = [ + dq if gq else None, + dk if gk else None, + dv if gv else None, + None, # block_logit + None, # block_bos + ] + if ctx.has_mask: + grads.append(None) # startend_row_indices + return tuple(grads) + + +def block_score_fa4_attn_fwd( + q, + k, + v, + valid_range=None, + sm_scale=None, + block_B=64, + causal=True, + startend_row_indices=None, +): + """FA4-fused MHA block-score forward. + + Args: + q: [B, S, H, D] bf16 query. + k: [B, S_kv, H, D] bf16 key. + v: [B, S_kv, H, D_v] bf16 value. + valid_range: [B, S, 2] int32 per-query [bos, eos). The ``bos`` column is + threaded into the kernel as the per-query document start so the fused + block-max buckets key columns by DOCUMENT-relative block + (``floor((col - bos) / block_B)``). This makes packed (bos>0) + block selection bit-identical to running each document alone (bos=0) + -- "pack-equivalence". When ``None`` the kernel uses bos=0 (single + document, absolute == relative). FA4 still does the actual masking + via ``causal`` and/or ``startend_row_indices``. + sm_scale: softmax scale; defaults to D**-0.5. + block_B: key block size (must divide FA4's n_block_size=128). + causal: apply causal masking (single-document). + startend_row_indices: optional flashmask document-mask indices; when + provided FA4 runs its flashmask path and the block-max honours it. + + Returns: + (out [B,S,H,D_v], lse [B,S,H], block_logit [B,H,S,num_blocks]). + num_blocks = ceil(S_kv / block_B). Block coordinates are DOCUMENT-relative + (relative to each query's ``bos`` from ``valid_range``), matching the + HySparse pipeline / TopK selection convention. + """ + assert q.is_contiguous(), "q must be contiguous" + b, s, h, d = q.shape + s_kv = k.shape[1] + assert list(k.shape) == [b, s_kv, h, d], ( + f"k must be [B, S_kv, H, D] matching q; got k {k.shape}, q {q.shape}" + ) + d_v = v.shape[-1] + assert list(v.shape) == [b, s_kv, h, d_v], ( + f"v must be [B, S_kv, H, D_v]; got {v.shape}" + ) + if sm_scale is None: + sm_scale = d**-0.5 + + num_blocks = (s_kv + block_B - 1) // block_B + # Pre-fill -inf: FA4 only writes the key-blocks its (causal/mask) iteration + # actually visits; skipped blocks must read back as -inf so their host-side + # block-score is 0. The buffer is mutated in place by the kernel. + block_logit = paddle.full( + [b, h, s, num_blocks], float("-inf"), dtype="float32" + ) + + # Per-query document start (bos) for document-relative block bucketing. + # valid_range[..., 0] is the bos of each query row; None => single-document + # (bos=0), which makes relative bucketing degenerate to the absolute one. + if valid_range is not None: + assert list(valid_range.shape) == [b, s, 2], ( + f"valid_range must be [B={b}, S={s}, 2]; got {valid_range.shape}" + ) + block_bos = valid_range[..., 0].contiguous().astype("int32") + else: + block_bos = paddle.zeros([b, s], dtype="int32") + + out, lse = _BlockScoreFA4Attn.apply( + q, + k, + v, + block_logit, + block_bos, + causal, + float(sm_scale), + startend_row_indices, + block_B, + ) + + # FA4 LSE is [B, H, S]; the HySparse pipeline expects [B, S, H]. + lse = lse.transpose([0, 2, 1]).contiguous() + return out, lse, block_logit diff --git a/src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa.py b/src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa.py deleted file mode 100644 index ebe9a0a11..000000000 --- a/src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa.py +++ /dev/null @@ -1,255 +0,0 @@ -# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Block-sparse attention (MQA gather): per-query-token block-sparse attention -with a single shared Key/Value head. - -Unlike the dense-mask variant (:mod:`block_sparse_attn`), this kernel actually -**gathers only the selected key blocks** and skips the rest, so its cost scales -with ``nsel`` (selected blocks) instead of the full sequence length. - -Efficiency comes from the MQA/MLA layout: K/V is a single head shared by all -query heads, and the TopK block indices are shared across heads. The kernel -therefore places the ``H`` query heads on the GEMM ``M`` dimension (one query -token per program), so a single gathered key block feeds every head and the -GEMM is wide enough to saturate tensor cores. This mirrors the repo's -``sparse_mqa`` kernel but gathers whole ``block_B``-sized key blocks and applies -the causal + document ``valid_range`` column mask inside the kernel. - -Layout: -* ``Q`` ``[B, S, H, D]`` bf16. -* ``K``, ``V`` ``[B, S_kv, D]`` bf16 (single shared head). -* ``Indices`` ``[B, S, nsel]`` int32 block ids (``-1`` = padding), shared - across heads. -* ``ValidRange`` ``[B, S, 2]`` int32 per-query ``[bos, eos)``. -* ``Output`` ``[B, S, H, D]`` bf16, ``Lse`` ``[B, S, H]`` fp32 natural-log. -""" - -import paddle -import tilelang -from tilelang import language as T - - -@tilelang.jit( - out_idx=[-2, -1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - }, -) -def block_sparse_mqa_fwd( - H, - D, - nsel, - sm_scale, - block_B=64, - num_stages=2, - threads=128, -): - assert D % 16 == 0, ( - f"D must be a multiple of 16 (tensor-core k-tile), got {D}" - ) - assert H <= 128, "this kernel supports up to 128 query heads" - scale_log2 = sm_scale * 1.44269504 - - batch = T.dynamic("batch") - seq_len = T.dynamic("seq_len") - seq_len_kv = T.dynamic("seq_len_kv") - - q_shape = [batch, seq_len, H, D] - kv_shape = [batch, seq_len_kv, D] - o_shape = [batch, seq_len, H, D] - idx_shape = [batch, seq_len, nsel] - vr_shape = [batch, seq_len, 2] - lse_shape = [batch, seq_len, H] - - dtype = T.bfloat16 - accum_dtype = T.float32 - idx_dtype = T.int32 - PH = max(tilelang.math.next_power_of_2(H), 16) # padded heads on M - BB = block_B - - @T.prim_func - def main( - Q: T.Tensor(q_shape, dtype), - K: T.Tensor(kv_shape, dtype), - V: T.Tensor(kv_shape, dtype), - Indices: T.Tensor(idx_shape, idx_dtype), - ValidRange: T.Tensor(vr_shape, idx_dtype), - Output: T.Tensor(o_shape, dtype), - Lse: T.Tensor(lse_shape, accum_dtype), - ): - with T.Kernel(seq_len, batch, threads=threads) as (bs, bb): - Q_shared = T.alloc_shared([PH, D], dtype) - K_shared = T.alloc_shared([BB, D], dtype) - V_shared = T.alloc_shared([BB, D], dtype) - P_shared = T.alloc_shared([PH, BB], dtype) - - acc_o = T.alloc_fragment([PH, D], accum_dtype) - acc_s = T.alloc_fragment([PH, BB], accum_dtype) - row_max = T.alloc_fragment([PH], accum_dtype) - m_i = T.alloc_fragment([PH], accum_dtype) - m_prev = T.alloc_fragment([PH], accum_dtype) - l_i = T.alloc_fragment([PH], accum_dtype) - l_new = T.alloc_fragment([PH], accum_dtype) - alpha = T.alloc_fragment([PH], accum_dtype) - - T.fill(acc_o, 0) - T.fill(m_i, -(2**30)) - T.fill(l_i, 0) - - bos = ValidRange[bb, bs, 0] - eos = ValidRange[bb, bs, 1] - - # load this token's H query heads onto M (pad rows >= H with 0) - for h, d in T.Parallel(PH, D): - Q_shared[h, d] = T.if_then_else( - h < H, Q[bb, bs, h, d], T.cast(0, dtype) - ) - - for i in T.Pipelined(nsel, num_stages=num_stages): - blk = Indices[bb, bs, i] - valid_blk = blk >= 0 - safe_blk = T.if_then_else(valid_blk, blk, 0) - - # gather one selected block. Block ids are document-relative: - # relative block ``blk`` spans absolute columns - # [bos + blk*BB, bos + (blk+1)*BB). Guard the read against the - # padded K/V length (columns >= eos are masked out below, so a - # clamped dummy read is harmless). - for c, d in T.Parallel(BB, D): - col = bos + safe_blk * BB + c - in_bounds = col < seq_len_kv - safe_col = T.if_then_else(in_bounds, col, 0) - K_shared[c, d] = T.if_then_else( - in_bounds, K[bb, safe_col, d], T.cast(0, dtype) - ) - V_shared[c, d] = T.if_then_else( - in_bounds, V[bb, safe_col, d], T.cast(0, dtype) - ) - - T.clear(acc_s) - T.gemm( - Q_shared, - K_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - - # mask: keep column iff block valid and in [bos, eos). The - # relative column bos + blk*BB + c is always >= bos, so only the - # upper bound needs checking. - for h, c in T.Parallel(PH, BB): - col = bos + safe_blk * BB + c - keep = valid_blk and (col < eos) - acc_s[h, c] = T.if_then_else( - keep, acc_s[h, c], -T.infinity(accum_dtype) - ) - - # online softmax (base 2) - T.reduce_max(acc_s, row_max, dim=1, clear=True) - T.copy(m_i, m_prev) - for h in T.Parallel(PH): - m_i[h] = T.max(m_i[h], row_max[h] * sm_scale) - for h in T.Parallel(PH): - alpha[h] = T.exp2((m_prev[h] - m_i[h]) * 1.44269504) - for h, c in T.Parallel(PH, BB): - acc_s[h, c] = T.exp2( - acc_s[h, c] * scale_log2 - m_i[h] * 1.44269504 - ) - T.reduce_sum(acc_s, l_new, dim=1) - for h in T.Parallel(PH): - l_i[h] = l_i[h] * alpha[h] + l_new[h] - for h, d in T.Parallel(PH, D): - acc_o[h, d] = acc_o[h, d] * alpha[h] - T.copy(acc_s, P_shared) - T.gemm( - P_shared, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow - ) - - # normalize; empty rows (no selected valid key) -> 0 out / -inf lse - for h, d in T.Parallel(PH, D): - acc_o[h, d] = T.if_then_else( - l_i[h] > 0, acc_o[h, d] / l_i[h], 0.0 - ) - for h, d in T.Parallel(PH, D): - if h < H: - Output[bb, bs, h, d] = acc_o[h, d] - for h in T.Parallel(PH): - if h < H: - Lse[bb, bs, h] = T.if_then_else( - l_i[h] > 0, - m_i[h] + T.log(l_i[h]), - -T.infinity(accum_dtype), - ) - - return main - - -def block_sparse_mqa_attn_fwd( - q, k, v, indices, valid_range, sm_scale=None, block_B=64 -): - """Forward interface for the MQA gather block-sparse attention. - - Args: - q: [B, S, H, D] bf16. - k, v: [B, S_kv, D] bf16 shared key/value. - indices: [B, S, nsel] int32 block ids (-1 padding), shared heads. - valid_range: [B, S, 2] int32. - sm_scale: softmax scale; defaults to D**-0.5. - block_B: key block size. - - Returns: - out [B,S,H,D] bf16, lse [B,S,H] fp32 (natural log; empty rows -inf). - """ - assert q.is_contiguous() - b, s, h, d = q.shape - s_kv = k.shape[1] - nsel = indices.shape[-1] - # lightweight host-side shape checks (no device sync) so mismatched inputs - # fail early and clearly instead of causing undefined kernel behaviour. - assert len(k.shape) == 3 and list(k.shape) == list(v.shape), ( - f"k/v must be [B, S_kv, D] and match; got {k.shape} vs {v.shape}" - ) - assert k.shape[0] == b and k.shape[2] == d, ( - f"k/v batch/head_dim must match q; got k {k.shape}, q {q.shape}" - ) - assert list(indices.shape[:2]) == [b, s], ( - f"indices must be [B, S, nsel] matching q; got {indices.shape}" - ) - assert list(valid_range.shape) == [b, s, 2], ( - f"valid_range must be [B, S, 2]; got {valid_range.shape}" - ) - if sm_scale is None: - sm_scale = d**-0.5 - if valid_range.dtype != paddle.int32: - valid_range = valid_range.cast("int32") - if indices.dtype != paddle.int32: - indices = indices.cast("int32") - - # kernel gathers whole block_B blocks without bounds clamping, so pad - # K/V so the last block is fully addressable. - pad = (block_B - s_kv % block_B) % block_B - if pad > 0: - k = paddle.nn.functional.pad(k, [0, 0, 0, pad]) - v = paddle.nn.functional.pad(v, [0, 0, 0, pad]) - k = k.contiguous() - v = v.contiguous() - valid_range = valid_range.contiguous() - indices = indices.contiguous() - - kernel = block_sparse_mqa_fwd(h, d, nsel, float(sm_scale), block_B=block_B) - out, lse = kernel(q, k, v, indices, valid_range) - return out, lse diff --git a/src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa_bwd.py b/src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa_bwd.py deleted file mode 100644 index 500cc005c..000000000 --- a/src/paddlefleet/tilelang_ops/hysparse/block_sparse_attn_mqa_bwd.py +++ /dev/null @@ -1,320 +0,0 @@ -# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Backward for the MQA gather block-sparse attention -(:mod:`block_sparse_attn_mqa`). - -Mirrors the forward's layout: one program per query token, with the ``H`` -query heads placed on the GEMM ``M`` dimension and the single shared K/V head -gathered ``block_B`` at a time over the ``nsel`` selected blocks. dQ for the -token is accumulated locally; dK/dV are scattered into the shared K/V head via -``atomic_add`` (many query tokens may select the same block). -""" - -import paddle -import tilelang -from tilelang import language as T - - -@tilelang.jit( - out_idx=[-1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - }, -) -def block_sparse_mqa_bwd( - H, - D, - nsel, - sm_scale, - block_B=64, - block_H=None, - num_stages=1, - threads=128, -): - assert D % 16 == 0, ( - f"D must be a multiple of 16 (tensor-core k-tile), got {D}" - ) - assert H <= 128, "this kernel supports up to 128 query heads" - - batch = T.dynamic("batch") - seq_len = T.dynamic("seq_len") - seq_len_kv = T.dynamic("seq_len_kv") - - q_shape = [batch, seq_len, H, D] - kv_shape = [batch, seq_len_kv, D] - idx_shape = [batch, seq_len, nsel] - vr_shape = [batch, seq_len, 2] - lse_shape = [batch, seq_len, H] - - dtype = T.bfloat16 - accum_dtype = T.float32 - idx_dtype = T.int32 - # Heads are tiled on the GEMM M dimension in groups of ``BH`` so that the - # ``[ยท, D]`` shared buffers fit on-chip for large head dims (e.g. MLA - # D=576); ``BH == H`` keeps the single-program-per-token fast path. - BH = block_H if block_H is not None else H - PH = max(tilelang.math.next_power_of_2(BH), 16) - num_hg = (H + BH - 1) // BH - BB = block_B - - @T.prim_func - def main( - Q: T.Tensor(q_shape, dtype), - K: T.Tensor(kv_shape, dtype), - V: T.Tensor(kv_shape, dtype), - dO: T.Tensor(q_shape, dtype), - Lse: T.Tensor(lse_shape, accum_dtype), - Delta: T.Tensor(lse_shape, accum_dtype), - Indices: T.Tensor(idx_shape, idx_dtype), - ValidRange: T.Tensor(vr_shape, idx_dtype), - dK: T.Tensor(kv_shape, accum_dtype), - dV: T.Tensor(kv_shape, accum_dtype), - dQ: T.Tensor(q_shape, dtype), - ): - with T.Kernel(num_hg, seq_len, batch, threads=threads) as (hg, bs, bb): - h0 = hg * BH - Q_shared = T.alloc_shared([PH, D], dtype) - dO_shared = T.alloc_shared([PH, D], dtype) - K_shared = T.alloc_shared([BB, D], dtype) - V_shared = T.alloc_shared([BB, D], dtype) - P_shared = T.alloc_shared([PH, BB], dtype) - dS_shared = T.alloc_shared([PH, BB], dtype) - dQ_shared = T.alloc_shared([PH, D], dtype) - - acc_s = T.alloc_fragment([PH, BB], accum_dtype) - acc_p = T.alloc_fragment([PH, BB], accum_dtype) - acc_dp = T.alloc_fragment([PH, BB], accum_dtype) - acc_dq = T.alloc_fragment([PH, D], accum_dtype) - acc_dk = T.alloc_fragment([BB, D], accum_dtype) - acc_dv = T.alloc_fragment([BB, D], accum_dtype) - lse_f = T.alloc_fragment([PH], accum_dtype) - delta_f = T.alloc_fragment([PH], accum_dtype) - - bos = ValidRange[bb, bs, 0] - eos = ValidRange[bb, bs, 1] - - for h, d in T.Parallel(PH, D): - gh = h0 + h - use = (h < BH) and (gh < H) - sh = T.if_then_else(use, gh, 0) - Q_shared[h, d] = T.if_then_else( - use, Q[bb, bs, sh, d], T.cast(0, dtype) - ) - dO_shared[h, d] = T.if_then_else( - use, dO[bb, bs, sh, d], T.cast(0, dtype) - ) - for h in T.Parallel(PH): - gh = h0 + h - use = (h < BH) and (gh < H) - sh = T.if_then_else(use, gh, 0) - lse_f[h] = T.if_then_else(use, Lse[bb, bs, sh], 0.0) - delta_f[h] = T.if_then_else(use, Delta[bb, bs, sh], 0.0) - - T.clear(acc_dq) - - for i in T.Pipelined(nsel, num_stages=num_stages): - blk = Indices[bb, bs, i] - valid_blk = blk >= 0 - safe_blk = T.if_then_else(valid_blk, blk, 0) - - # document-relative gather: relative block ``blk`` spans - # absolute columns [bos + blk*BB, bos + (blk+1)*BB). Guard the - # read against the padded K/V length. - for c, d in T.Parallel(BB, D): - col = bos + safe_blk * BB + c - in_bounds = col < seq_len_kv - safe_col = T.if_then_else(in_bounds, col, 0) - K_shared[c, d] = T.if_then_else( - in_bounds, K[bb, safe_col, d], T.cast(0, dtype) - ) - V_shared[c, d] = T.if_then_else( - in_bounds, V[bb, safe_col, d], T.cast(0, dtype) - ) - - # P = softmax prob = exp(raw*sm_scale - lse); masked -> 0. - T.clear(acc_s) - T.gemm( - Q_shared, - K_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - for h, c in T.Parallel(PH, BB): - col = bos + safe_blk * BB + c - keep = valid_blk and (col < eos) - acc_p[h, c] = T.if_then_else( - keep, - T.exp2( - (acc_s[h, c] * sm_scale - lse_f[h]) * 1.44269504 - ), - 0.0, - ) - T.copy(acc_p, P_shared) - - # dP = dO @ V^T - T.clear(acc_dp) - T.gemm( - dO_shared, - V_shared, - acc_dp, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - # dS = P * (dP - Delta) * sm_scale - for h, c in T.Parallel(PH, BB): - acc_dp[h, c] = ( - acc_p[h, c] * (acc_dp[h, c] - delta_f[h]) * sm_scale - ) - T.copy(acc_dp, dS_shared) - - # dQ += dS @ K - T.gemm( - dS_shared, - K_shared, - acc_dq, - policy=T.GemmWarpPolicy.FullRow, - ) - # dV = P^T @ dO ; dK = dS^T @ Q (scattered to shared KV) - T.clear(acc_dv) - T.gemm( - P_shared, - dO_shared, - acc_dv, - transpose_A=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.clear(acc_dk) - T.gemm( - dS_shared, - Q_shared, - acc_dk, - transpose_A=True, - policy=T.GemmWarpPolicy.FullRow, - ) - for c, d in T.Parallel(BB, D): - col = bos + safe_blk * BB + c - if valid_blk and (col < seq_len_kv): - T.atomic_add(dV[bb, col, d], acc_dv[c, d]) - T.atomic_add(dK[bb, col, d], acc_dk[c, d]) - - T.copy(acc_dq, dQ_shared) - for h, d in T.Parallel(PH, D): - gh = h0 + h - if (h < BH) and (gh < H): - dQ[bb, bs, gh, d] = dQ_shared[h, d] - - return main - - -@tilelang.jit(out_idx=[-1]) -def _cast_bf16_kv(D, block_N=64, threads=128): - batch = T.dynamic("batch") - seq_len_kv = T.dynamic("seq_len_kv") - shape = [batch, seq_len_kv, D] - - @T.prim_func - def main( - X: T.Tensor(shape, T.float32), - Out: T.Tensor(shape, T.bfloat16), - ): - with T.Kernel( - T.ceildiv(seq_len_kv, block_N), batch, threads=threads - ) as (bn, bb): - for i, d in T.Parallel(block_N, D): - if bn * block_N + i < seq_len_kv: - Out[bb, bn * block_N + i, d] = X[bb, bn * block_N + i, d] - - return main - - -def _fit_block_h(H, D, block_B, cap_bytes=200000): - """Largest head-group (a divisor-ish of H) whose per-token backward shared - buffers fit. Q/dO/dQ ``[PH, D]`` + K/V ``[BB, D]`` + P/dS ``[PH, BB]`` in - bf16; for large D (MLA 576) all H=64 heads on M overflow, so tile them. - """ - for bh in (H, 64, 32, 16): - if bh > H: - continue - ph = max(tilelang.math.next_power_of_2(bh), 16) - shared = 6 * ph * D + 4 * block_B * D + 4 * ph * block_B - if shared <= cap_bytes: - return bh - return min(H, 16) - - -def block_sparse_mqa_bwd_interface( - q, k, v, o, do, lse, indices, valid_range, sm_scale=None, block_B=64 -): - """Backward interface for the MQA gather block-sparse attention. - - Args: - q: [B, S, H, D] bf16 forward query. - k, v: [B, S_kv, D] bf16 shared key/value (forward inputs). - o: [B, S, H, D] bf16 forward output. - do: [B, S, H, D] bf16 grad of output. - lse: [B, S, H] fp32 natural-log LSE from forward. - indices: [B, S, nsel] int block ids selected per query token. - valid_range: [B, S, 2] int32. - sm_scale: softmax scale; defaults to D**-0.5. - block_B: key block size. - - Returns: - dq [B,S,H,D] bf16, dk/dv [B,S_kv,D] bf16. - """ - assert q.is_contiguous() and do.is_contiguous() and lse.is_contiguous() - b, s, h, d = q.shape - s_kv = k.shape[1] - if sm_scale is None: - sm_scale = d**-0.5 - if valid_range.dtype != paddle.int32: - valid_range = valid_range.cast("int32") - if indices.dtype != paddle.int32: - indices = indices.cast("int32") - - # gather reads whole blocks without clamping -> pad K/V (and dK/dV) so the - # last block is addressable. - pad = (block_B - s_kv % block_B) % block_B - s_kv_pad = s_kv + pad - if pad > 0: - k = paddle.nn.functional.pad(k, [0, 0, 0, pad]) - v = paddle.nn.functional.pad(v, [0, 0, 0, pad]) - k = k.contiguous() - v = v.contiguous() - valid_range = valid_range.contiguous() - indices = indices.contiguous() - - delta = (o.astype("float32") * do.astype("float32")).sum(-1).contiguous() - dk = paddle.zeros([b, s_kv_pad, d], dtype="float32") - dv = paddle.zeros([b, s_kv_pad, d], dtype="float32") - - bwd = block_sparse_mqa_bwd( - h, - d, - indices.shape[-1], - float(sm_scale), - block_B=block_B, - block_H=_fit_block_h(h, d, block_B), - ) - dq = bwd(q, k, v, do, lse, delta, indices, valid_range, dk, dv) - - cast = _cast_bf16_kv(d) - dk_bf = cast(dk) - dv_bf = cast(dv) - if pad > 0: - dk_bf = dk_bf[:, :s_kv, :].contiguous() - dv_bf = dv_bf[:, :s_kv, :].contiguous() - return dq, dk_bf, dv_bf diff --git a/src/paddlefleet/tilelang_ops/hysparse/pipeline.py b/src/paddlefleet/tilelang_ops/hysparse/pipeline.py index 98ee96fd2..4274b061c 100644 --- a/src/paddlefleet/tilelang_ops/hysparse/pipeline.py +++ b/src/paddlefleet/tilelang_ops/hysparse/pipeline.py @@ -12,33 +12,45 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""HySparse forward pipeline: chain block-score attention, block-TopK, and -block-sparse attention. - -All stages share a single K/V head (MQA/MLA sparse branch): - -1. **Block-score attention** (:func:`block_score_mqa_attn_fwd`): full attention - with the shared K/V head that also emits per-(query, key-block) max raw - logits ``BlockLogit``. -2. **Block TopK selection** (:func:`select_topk_blocks`): recover eq.(3) block - scores, aggregate them across heads by a group-wise **maximum** (shared - block selection), mask blocks that hold no causal/document-valid key, and - TopK to per-query block indices. -3. **Block-sparse attention** (:func:`block_sparse_mqa_attn_fwd`): MQA - block-sparse attention that gathers only the selected blocks, so its cost - scales with ``topk`` rather than the sequence length. - -The block scores feed a non-differentiable TopK, so only the two attention -operators carry gradient. +"""HySparse block-TopK selection. + +:func:`select_topk_blocks` recovers eq.(3) block scores from the scaled +per-block max logits emitted by the block-score attention forward, aggregates +them across heads by a group-wise **maximum** (shared block selection), masks +blocks that hold no causal/document-valid key, and TopK-selects per-query block +indices. The block scores feed a non-differentiable TopK, so no autograd graph +is built for the selection. """ import paddle -from .block_score_attn import ( - block_score_mqa_attn_fwd, - block_scores_from_logit, -) -from .block_sparse_attn_mqa import block_sparse_mqa_attn_fwd + +def block_scores_from_logit(block_logit, lse): + """Recover eq.(3) block-max probability scores on the host. + + The block-score attention forward (FA4-fused, :mod:`block_score_fa4`) emits + the **scaled** per-block max logit (``softmax_scale * q.k``, the exact value + fed into softmax); this turns it into the eq.(3) softmax probability + ``exp(block_logit - lse)``. The scale already lives inside ``block_logit``, + so no ``sm_scale`` multiply is applied here. + + Args: + block_logit: [B, H, S, num_blocks] scaled per-block max logit (-inf if + fully masked), as returned by the forward. + lse: [B, S, H] natural-log LSE from the forward. + + Returns: + [B, H, S, num_blocks] block-max softmax probabilities in [0, 1]; + fully-masked blocks are 0. + """ + lse_bhs = lse.transpose([0, 2, 1]).unsqueeze(-1) # [B,H,S,1] + scaled = block_logit.astype("float32") - lse_bhs + scores = paddle.exp(scaled) + # exp(-inf) already 0, but guard any nan from (-inf)-(-inf) style edge. + scores = paddle.where( + paddle.isfinite(scores), scores, paddle.zeros_like(scores) + ) + return scores def _valid_block_mask(valid_range, num_blocks, block_B): @@ -57,15 +69,14 @@ def _valid_block_mask(valid_range, num_blocks, block_B): def select_topk_blocks( - block_logit, lse, valid_range, sm_scale, topk, block_B, head_agg="max" + block_logit, lse, valid_range, topk, block_B, head_agg="max" ): """Select per-query TopK key blocks from block-score attention outputs. Args: - block_logit: [B, H, S, num_blocks] raw per-block max logit. + block_logit: [B, H, S, num_blocks] scaled per-block max logit. lse: [B, S, H] natural-log LSE from block-score attention. valid_range: [B, S, 2] int, per-query [bos, eos) valid key columns. - sm_scale: softmax scale. topk: number of blocks to select per query token. block_B: key block size. head_agg: how to aggregate block scores across heads so the whole @@ -81,7 +92,14 @@ def select_topk_blocks( if topk <= 0: raise ValueError(f"topk must be positive, got {topk}") b, h, s, num_blocks = block_logit.shape - scores = block_scores_from_logit(block_logit, lse, sm_scale) # [B,H,S,nb] + # Block selection is a hard, non-differentiable TopK. Detach the score + # inputs so no autograd graph is built for it: block_logit / lse come out of + # the differentiable attention PyLayer, and without detaching they drag topk + # into backward, where topk_grad dereferences the (int) index gradient and + # segfaults. + block_logit = block_logit.detach() + lse = lse.detach() + scores = block_scores_from_logit(block_logit, lse) # [B,H,S,nb] # aggregate across heads (block selection shared across the query group) if head_agg == "max": agg = scores.max(axis=1) # [B, S, num_blocks] @@ -106,57 +124,3 @@ def select_topk_blocks( pad = paddle.full([b, s, topk - k], -1, dtype="int32") top_idx = paddle.concat([top_idx, pad], axis=-1) return top_idx.contiguous() - - -def hysparse_forward_mqa(q, k, v, valid_range, topk, sm_scale=None, block_B=64): - """HySparse forward with a single shared K/V head (MQA/MLA sparse branch). - - Three stages (block-score -> TopK -> block-sparse) with K/V as one shared - head - ``[B, S_kv, D]``: the block selection is aggregated across the whole query - group by a group-wise **maximum** (paper eq. 3) so all heads share indices, - and the sparse branch uses the MQA gather kernel - (:func:`block_sparse_mqa_attn_fwd`) whose cost scales with ``topk`` rather - than the sequence length. - - Both stages keep K/V as a single shared head: block-score attention uses - the MQA scoring kernel (:func:`block_score_mqa_attn_fwd`, no head broadcast) - and the sparse branch uses the MQA gather kernel. - - Args: - q: [B, S, H, D] bf16 query (H heads). - k, v: [B, S_kv, D] bf16 single shared key/value head. - valid_range: [B, S, 2] int32 causal + document valid key range. - topk: number of blocks selected per query token. - sm_scale: softmax scale; defaults to D**-0.5. - block_B: key block size. - - Returns: - sparse_out: [B, S, H, D] MQA block-sparse attention output. - sparse_lse: [B, S, H] natural-log LSE of the sparse branch. - indices: [B, S, topk] selected block ids (int32, -1 padding). - full_out: [B, S, H, D] full attention output (block-score). - full_lse: [B, S, H] natural-log LSE of block-score attention. - """ - b, s, h, d = q.shape - if sm_scale is None: - sm_scale = d**-0.5 - - # block-score attention with the single shared K/V head -- no broadcast. - full_out, full_lse, block_logit = block_score_mqa_attn_fwd( - q, k, v, valid_range, sm_scale=sm_scale, block_B=block_B - ) - indices = select_topk_blocks( - block_logit, - full_lse, - valid_range, - sm_scale, - topk, - block_B, - head_agg="max", - ) - # sparse branch uses the shared single-head K/V gather kernel - sparse_out, sparse_lse = block_sparse_mqa_attn_fwd( - q, k, v, indices, valid_range, sm_scale=sm_scale, block_B=block_B - ) - return sparse_out, sparse_lse, indices, full_out, full_lse diff --git a/src/paddlefleet/tilelang_ops/hysparse/reference.py b/src/paddlefleet/tilelang_ops/hysparse/reference.py deleted file mode 100644 index 20fa0fdfb..000000000 --- a/src/paddlefleet/tilelang_ops/hysparse/reference.py +++ /dev/null @@ -1,230 +0,0 @@ -# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Naive Paddle reference (ๆ•ฃ็ฎ—ๅญ) for HySparse block attention operators. - -Two operators, mirroring the paper (arXiv 2602.03560), with a single shared -Key/Value head (MQA/MLA): - -* ``ref_block_score_attn_mqa`` โ€” block-score attention, "Algorithm 1": standard - full attention (shared K/V head) that additionally emits block-level max - attention *probability* scores ``S`` (eq. 3), used for TopK block selection. -* ``ref_block_sparse_attn_mqa`` โ€” block-sparse attention: each query token - attends only to a per-token selected set of key blocks gathered from the - shared K/V head. - -Masking (causal + document) is expressed through ``valid_range`` of shape -``[B, S, 2]`` giving, per query token, the half-open valid key column range -``[bos, eos)``. Causal masking sets ``eos = t + 1``; document masking sets -``bos`` to the document start. - -These functions are written for readability/correctness, not speed, and are -the ground truth the TileLang kernels are validated against. -""" - -import paddle -import paddle.nn.functional as F - -NEG_INF = float("-inf") - - -def _range_mask(valid_range, seq_len_kv): - """Build a boolean key mask [B, 1, S, S_kv] from valid_range [B, S, 2].""" - bos = valid_range[..., 0].unsqueeze(1).unsqueeze(-1) # [B, 1, S, 1] - eos = valid_range[..., 1].unsqueeze(1).unsqueeze(-1) # [B, 1, S, 1] - col = paddle.arange(seq_len_kv, dtype=valid_range.dtype) - col = col.reshape([1, 1, 1, seq_len_kv]) # [1, 1, 1, S_kv] - return (col >= bos) & (col < eos) # [B, 1, S, S_kv] - - -def _to_bhsd(x): - """[B, S, H, D] -> [B, H, S, D].""" - return x.transpose([0, 2, 1, 3]) - - -def ref_block_score_attn_mqa(q, k, v, valid_range, sm_scale=None, block_B=64): - """Block-score attention reference (MQA): full attention with a single - shared K/V head plus block-max probability scores (eq. 3). - - Args: - q: [B, S, H, D] query (H heads). - k, v: [B, S_kv, D] shared key/value (single head). - valid_range: [B, S, 2] int, per-query [bos, eos) valid key columns. - sm_scale: softmax scale; defaults to D**-0.5. - block_B: key block size for the emitted block scores. - - Returns: - out: [B, S, H, D] attention output. - lse: [B, S, H] natural-log-sum-exp of the masked logits. - s_block: [B, H, S, num_blocks] block-max softmax probability (eq. 3), - num_blocks = ceil(S_kv / block_B). Fully-masked blocks -> 0. - """ - b, s, h, d = q.shape - s_kv = k.shape[1] - if sm_scale is None: - sm_scale = d**-0.5 - - qb = _to_bhsd(q).astype("float32") # [B, H, S, D] - # broadcast the single shared K/V head across all query heads - kb = k.astype("float32").unsqueeze(1) # [B, 1, S_kv, D] - vb = v.astype("float32").unsqueeze(1) # [B, 1, S_kv, D] - - logits = paddle.matmul(qb, kb, transpose_y=True) * sm_scale # [B,H,S,S_kv] - mask = _range_mask(valid_range, s_kv) # [B,1,S,S_kv] - neg = paddle.full_like(logits, NEG_INF) - logits = paddle.where(mask, logits, neg) - - # Rows with no valid key (bos >= eos) would make softmax/logsumexp nan; - # guard them to 0 output and leave lse as -inf (matches the kernel). - row_has_key = mask.any(axis=-1, keepdim=True) # [B,1,S,1] - lse = paddle.logsumexp(logits, axis=-1) # [B,H,S] - probs = F.softmax(logits, axis=-1) # [B,H,S,S_kv] - probs = paddle.where( - row_has_key.expand_as(probs), probs, paddle.zeros_like(probs) - ) - out = paddle.matmul(probs, vb.expand([b, h, s_kv, d])) # [B,H,S,D] - - # Block-max probability with **document-relative** block coordinates: block j - # of a query spans key columns [bos + j*block_B, bos + (j+1)*block_B), i.e. - # the grid is anchored at that query's document start ``bos`` (not absolute - # sequence columns). This matches processing each document standalone. - num_blocks = (s_kv + block_B - 1) // block_B - col = paddle.arange(s_kv, dtype="int64").reshape([1, 1, s_kv]) # [1,1,S_kv] - bos = valid_range[..., 0:1].astype("int64") # [B, S, 1] - rel = col - bos # [B, S, S_kv] column position relative to doc start - rel_id = paddle.where( # relative block id; -1 for cols before doc start - rel >= 0, rel // block_B, paddle.full_like(rel, -1) - ) # [B, S, S_kv] - rel_id = rel_id.unsqueeze(1) # [B, 1, S, S_kv] - s_block_list = [] - for j in range(num_blocks): - hit = rel_id == j # [B, 1, S, S_kv] - masked = paddle.where(hit, probs, paddle.zeros_like(probs)) - s_block_list.append(masked.max(axis=-1)) # [B, H, S] - s_block = paddle.stack(s_block_list, axis=-1) # [B,H,S,num_blocks] - - out = out.transpose([0, 2, 1, 3]) # back to [B,S,H,D] - lse = lse.transpose([0, 2, 1]) # [B,S,H] - return out.astype(q.dtype), lse, s_block - - -def _selected_key_mask(indices, valid_range, seq_len_kv, block_B): - """Boolean [B, S, S_kv]: key column selected by this query's block ids. - - Block ids are **document-relative**: block ``j`` of a query spans key - columns ``[bos + j*block_B, bos + (j+1)*block_B)`` where ``bos`` is the - query's document start (``valid_range[..., 0]``). A column is selected iff - it lies at or after ``bos`` and its relative block id is among the query's - (valid) selected ids. - - indices: [B, S, nsel] int block ids; -1 marks an invalid/padding slot. - """ - col = paddle.arange(seq_len_kv, dtype="int64").reshape([1, 1, seq_len_kv]) - bos = valid_range[..., 0:1].astype("int64") # [B, S, 1] - rel = col - bos # [B, S, S_kv] - col_block = paddle.where( # relative block id of each key column - rel >= 0, rel // block_B, paddle.full_like(rel, -1) - ) # [B, S, S_kv] - idx = indices.astype("int64").unsqueeze(-2) # [B,S,1,nsel] - col_block_e = col_block.unsqueeze(-1) # [B,S,S_kv,1] - # a column is selected iff its relative block id equals any valid selected id - hit = (col_block_e == idx) & (idx >= 0) # [B,S,S_kv,nsel] - return hit.any(axis=-1) # [B,S,S_kv] - - -def ref_block_sparse_attn_mqa( - q, k, v, indices, valid_range, sm_scale=None, block_B=64 -): - """Block-sparse attention reference (MQA): per-query-token block-sparse - attention with a single Key/Value head shared across all query heads. - - This mirrors the efficient HySparse sparse branch: the block indices are - shared across heads (GQA group-wise max upstream), and K/V are a single - shared head so one gathered block feeds every query head. - - Args: - q: [B, S, H, D] query (H heads). - k, v: [B, S_kv, D] shared key/value (single head). - indices: [B, S, nsel] int block ids selected per query token - (-1 = padding). Shared across heads. - valid_range: [B, S, 2] int, per-query [bos, eos) valid key columns. - sm_scale: softmax scale; defaults to D**-0.5. - block_B: key block size. - - Returns: - out: [B, S, H, D]; lse: [B, S, H]. - """ - b, s, h, d = q.shape - s_kv = k.shape[1] - if sm_scale is None: - sm_scale = d**-0.5 - - qb = _to_bhsd(q).astype("float32") # [B, H, S, D] - # broadcast the single shared KV head across all query heads - kb = k.astype("float32").unsqueeze(1) # [B, 1, S_kv, D] - vb = v.astype("float32").unsqueeze(1) # [B, 1, S_kv, D] - - logits = paddle.matmul(qb, kb, transpose_y=True) * sm_scale # [B,H,S,S_kv] - - range_m = _range_mask(valid_range, s_kv) # [B,1,S,S_kv] - sel_m = _selected_key_mask(indices, valid_range, s_kv, block_B).unsqueeze( - 1 - ) # [B,1,S,S_kv] - mask = range_m & sel_m - neg = paddle.full_like(logits, NEG_INF) - logits = paddle.where(mask, logits, neg) - - # LSE is taken from the fully -inf-masked logits so empty rows (no valid - # key) stay -inf, matching the kernel and the docstring. softmax then uses - # the same logits and its NaN empty rows are zeroed via ``row_has_key``. - row_has_key = mask.any(axis=-1, keepdim=True) # [B,1,S,1] - lse = paddle.logsumexp(logits, axis=-1) # [B,H,S] - probs = F.softmax(logits, axis=-1) - probs = paddle.where( - row_has_key.expand_as(probs), probs, paddle.zeros_like(probs) - ) - out = paddle.matmul(probs, vb.expand([b, h, s_kv, d])) # [B,H,S,D] - - out = out.transpose([0, 2, 1, 3]) - lse = lse.transpose([0, 2, 1]) - return out.astype(q.dtype), lse - - -def make_causal_valid_range(seq_len, batch=1, doc_lengths=None): - """Helper: build valid_range [B, S, 2] for causal (+ optional document) mask. - - Args: - seq_len: total sequence length S (== S_kv). - batch: batch size B. - doc_lengths: optional list of document lengths (packed along S). If - given, sum must equal seq_len; each token's bos is set to - its document start. If None, a single document is assumed. - - Returns: - valid_range: [B, S, 2] int32. - """ - pos = paddle.arange(seq_len, dtype="int32") - eos = pos + 1 # causal: attend up to and including self - if doc_lengths is None: - bos = paddle.zeros([seq_len], dtype="int32") - else: - assert sum(doc_lengths) == seq_len - starts = [] - cur = 0 - for dl in doc_lengths: - starts += [cur] * dl - cur += dl - bos = paddle.to_tensor(starts, dtype="int32") - vr = paddle.stack([bos, eos], axis=-1) # [S, 2] - return vr.unsqueeze(0).expand([batch, seq_len, 2]).contiguous() diff --git a/src/paddlefleet/tilelang_ops/hysparse/swa_attn.py b/src/paddlefleet/tilelang_ops/hysparse/swa_attn.py new file mode 100644 index 000000000..b2380f38a --- /dev/null +++ b/src/paddlefleet/tilelang_ops/hysparse/swa_attn.py @@ -0,0 +1,128 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Causal sliding-window attention (SWA), MQA, on top of the windowed MQA +flash-attention kernels (:mod:`windowed_mqa_attn`). + +A causal sliding window is just a masking pattern: query token ``t`` attends to +keys ``[max(doc_start, t - W + 1), t + 1)``. That half-open range is exactly the +``valid_range [B, S, 2]`` the windowed kernels consume, and their ``eos - bos`` +early-exit naturally bounds the per-token work to the window ``W`` instead of +the full sequence. So SWA needs **no new kernel** -- we just feed a windowed +``valid_range`` to the forward/backward. + +Why this matters for the HySparse / MLA stack: FlashAttention (FA2/3/4) exposes +``window_size`` and covers SWA directly for ``head_dim <= 256``. But the +**absorbed-MLA MQA** shape has Dk=576 / Dv=512 (> 256), which FA4 does not +support (the model falls back to an eager dense attention). This TileLang SWA +MQA path fills that gap: a fused windowed flash attention at Dk=576/Dv=512 that +is far cheaper than the eager O(S^2) fallback. +""" + +import paddle + +from .windowed_mqa_attn import windowed_mqa_attn_fwd +from .windowed_mqa_attn_bwd import windowed_mqa_bwd_interface + + +class _SlidingWindowMQAAttn(paddle.autograd.PyLayer): + """Autograd wrapper: causal SWA with a single shared K/V head (MQA).""" + + @staticmethod + def forward(ctx, q, k, v, valid_range, attn_sink, sm_scale, block_B): + out, lse = windowed_mqa_attn_fwd( + q, + k, + v, + valid_range, + attn_sink=attn_sink, + sm_scale=sm_scale, + block_B=block_B, + ) + # ``attn_sink`` may be None (sinkless). save_for_backward only takes + # tensors, so save it only when present and flag its absence. + ctx.has_sink = attn_sink is not None + if ctx.has_sink: + ctx.save_for_backward(q, k, v, out, lse, valid_range, attn_sink) + else: + ctx.save_for_backward(q, k, v, out, lse, valid_range) + ctx.sm_scale = sm_scale + ctx.block_B = block_B + # PyLayer contract: backward returns None for stop_gradient inputs. + ctx.needs_grad = ( + not q.stop_gradient, + not k.stop_gradient, + not v.stop_gradient, + ctx.has_sink and not attn_sink.stop_gradient, + ) + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + def backward(ctx, dout, dlse): + if ctx.has_sink: + q, k, v, out, lse, valid_range, attn_sink = ctx.saved_tensor() + else: + q, k, v, out, lse, valid_range = ctx.saved_tensor() + attn_sink = None + dq, dk, dv, d_attn_sink = windowed_mqa_bwd_interface( + q, + k, + v, + out, + dout.contiguous(), + lse, + valid_range, + attn_sink=attn_sink, + sm_scale=ctx.sm_scale, + block_B=ctx.block_B, + ) + gq, gk, gv, gsink = ctx.needs_grad + # One returned grad per **tensor** input, in order. sm_scale/block_B are + # non-tensors (no slot). attn_sink only occupies a slot when it was + # passed as a tensor (sinkless -> None -> no slot). + grads = [ + dq if gq else None, + dk if gk else None, + dv if gv else None, + None, # valid_range: int32 tensor input, never needs grad + ] + if ctx.has_sink: + grads.append(d_attn_sink if gsink else None) + return tuple(grads) + + +def sliding_window_mqa_attention( + q, k, v, valid_range, attn_sink=None, sm_scale=None, block_B=64 +): + """Causal sliding-window attention, MQA (single shared K/V head). + + Args: + q: [B, S, H, D] bf16. + k, v: [B, S_kv, D] bf16 shared key/value. + valid_range: [B, S, 2] int32 windowed [bos, eos) range; for a window + ``W`` query ``t`` uses ``[max(doc_start, t-W+1), t+1)``. + attn_sink: [H] fp32 per-head learnable attention-sink logit, or + ``None`` for plain (sinkless) softmax. + sm_scale: softmax scale; defaults to D**-0.5. + block_B: key block size. + + Returns: + out [B, S, H, D_v] bf16, lse [B, S, H] fp32 (non-differentiable). + """ + if sm_scale is None: + sm_scale = q.shape[-1] ** -0.5 + return _SlidingWindowMQAAttn.apply( + q, k, v, valid_range, attn_sink, float(sm_scale), block_B + ) diff --git a/src/paddlefleet/tilelang_ops/hysparse/block_score_attn.py b/src/paddlefleet/tilelang_ops/hysparse/windowed_mqa_attn.py similarity index 51% rename from src/paddlefleet/tilelang_ops/hysparse/block_score_attn.py rename to src/paddlefleet/tilelang_ops/hysparse/windowed_mqa_attn.py index bc2266828..3ce55bd9b 100644 --- a/src/paddlefleet/tilelang_ops/hysparse/block_score_attn.py +++ b/src/paddlefleet/tilelang_ops/hysparse/windowed_mqa_attn.py @@ -12,27 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Block-score attention (paper "Algorithm 1") with a single shared K/V head -(MQA/MLA): full attention that additionally emits per-block max logits used to -derive block-level selection scores. - -Q is ``[B, S, H, D]`` (H query heads); K, V are one shared head -``[B, S_kv, D]``. Masking (causal + document) is expressed through -``valid_range`` ``[B, S, 2]`` giving each query's half-open valid key column -range ``[bos, eos)``. - -The forward returns: -* ``Output`` ``[B, S, H, D]`` attention output. -* ``Lse`` ``[B, S, H]`` natural log-sum-exp of the *scaled* logits. -* ``BlockLogit`` ``[B, H, S, num_blocks]`` per-(query, key-block) max of the - *raw* (unscaled) ``qยทk`` logit over valid columns; fully-masked blocks store - ``-inf``. The block-max softmax *probability* score of eq. (3) is recovered - on the host as ``exp(BlockLogit * sm_scale - Lse)`` (see - :func:`block_scores_from_logit`). - -The backward (:mod:`block_score_attn_bwd`) is a standard flash-attention -backward (dQ, dK, dV); the block scores feed a non-differentiable TopK and -therefore carry no gradient. +"""Causal windowed MQA flash attention (single shared K/V head). + +The fused attention kernel behind the HySparse **SWA layers'** sliding-window +main path. Q is ``[B, S, H, D]`` (H query heads); K, V are one shared head +``[B, S_kv, D]`` / ``[B, S_kv, D_v]`` (MQA/MLA). Causal + document + sliding +window masking is expressed entirely through ``valid_range`` ``[B, S, 2]`` -- +each query's half-open valid key column range ``[bos, eos)`` -- and the +kernel's ``eos - bos`` early-exit bounds the per-token work to that range. + +Why a bespoke kernel: absorbed-MLA MQA has Dk=576 / Dv=512 (the key carries an +extra RoPE slice, so ``D`` exceeds ``D_v``). FlashAttention (FA2/3/4) only +covers ``head_dim <= 256`` and would fall back to eager O(S^2) attention at +these dims; this fused TileLang path is far cheaper. + +Returns ``(out [B,S,H,D_v], lse [B,S,H])``. ``lse`` is the natural-log sum-exp +of the scaled logits, consumed by the backward (:mod:`windowed_mqa_attn_bwd`); +it carries no gradient. """ import paddle @@ -47,42 +43,50 @@ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, }, ) -def block_score_mqa_fwd( +def windowed_mqa_fwd( H, D, sm_scale, + D_v=None, block_B=64, num_stages=2, threads=128, ): - """Block-score kernel with a single shared K/V head (MQA/MLA scoring). + """Windowed MQA flash-attention kernel (single shared K/V head). One program per query **token**, with the ``H`` query heads placed on the - GEMM ``M`` dimension (like the gather kernel). Blocks are - **document-relative**: block ``j`` of a query spans absolute key columns - ``[bos + j*block_B, bos + (j+1)*block_B)`` where ``bos`` is the query's - document start. The full-attention ``Output``/``Lse`` are unchanged (they - still sum over every valid key in ``[bos, eos)``); only the emitted - per-block max logit is binned relative to ``bos`` so downstream TopK - selects blocks exactly as if each document were processed standalone. + GEMM ``M`` dimension. Keys are streamed in ``block_B``-sized tiles over the + token's valid range ``[bos, eos)``; the ``ceildiv(eos - bos, block_B)`` + early-exit skips fully-masked tiles (pure causal halves the work, a sliding + window bounds work to the window, packed documents scan only their own doc). + + ``D`` is the query/key head dim (for the ``qยทk`` logit); ``D_v`` is the + value/output head dim. They are equal for plain attention; for absorbed MLA + the key carries an extra RoPE slice so ``D`` (e.g. 576) exceeds ``D_v`` + (e.g. 512). ``D_v`` defaults to ``D``. """ + if D_v is None: + D_v = D assert D % 16 == 0, ( f"D must be a multiple of 16 (tensor-core k-tile), got {D}" ) + assert D_v % 16 == 0, ( + f"D_v must be a multiple of 16 (tensor-core k-tile), got {D_v}" + ) assert H <= 128, "this kernel supports up to 128 query heads" scale_log2 = sm_scale * 1.44269504 # log2(e), online softmax in base 2 batch = T.dynamic("batch") seq_len = T.dynamic("seq_len") seq_len_kv = T.dynamic("seq_len_kv") - num_blocks = T.dynamic("num_blocks") q_shape = [batch, seq_len, H, D] kv_shape = [batch, seq_len_kv, D] - o_shape = [batch, seq_len, H, D] + v_shape = [batch, seq_len_kv, D_v] + o_shape = [batch, seq_len, H, D_v] lse_shape = [batch, seq_len, H] vr_shape = [batch, seq_len, 2] - blk_shape = [batch, H, seq_len, num_blocks] + sink_shape = [H] dtype = T.bfloat16 accum_dtype = T.float32 @@ -94,21 +98,21 @@ def block_score_mqa_fwd( def main( Q: T.Tensor(q_shape, dtype), K: T.Tensor(kv_shape, dtype), - V: T.Tensor(kv_shape, dtype), + V: T.Tensor(v_shape, dtype), ValidRange: T.Tensor(vr_shape, idx_dtype), - BlockLogit: T.Tensor(blk_shape, accum_dtype), + AttnSink: T.Tensor(sink_shape, accum_dtype), Output: T.Tensor(o_shape, dtype), Lse: T.Tensor(lse_shape, accum_dtype), ): with T.Kernel(seq_len, batch, threads=threads) as (bs, bb): Q_shared = T.alloc_shared([PH, D], dtype) K_shared = T.alloc_shared([BB, D], dtype) - V_shared = T.alloc_shared([BB, D], dtype) + V_shared = T.alloc_shared([BB, D_v], dtype) P_shared = T.alloc_shared([PH, BB], dtype) - acc_o = T.alloc_fragment([PH, D], accum_dtype) + acc_o = T.alloc_fragment([PH, D_v], accum_dtype) acc_s = T.alloc_fragment([PH, BB], accum_dtype) - blk_max = T.alloc_fragment([PH], accum_dtype) + tile_max = T.alloc_fragment([PH], accum_dtype) m_i = T.alloc_fragment([PH], accum_dtype) m_prev = T.alloc_fragment([PH], accum_dtype) l_i = T.alloc_fragment([PH], accum_dtype) @@ -128,18 +132,14 @@ def main( h < H, Q[bb, bs, h, d], T.cast(0, dtype) ) - # causal/document early-exit: only this token's own valid blocks - # (relative block j in [0, ceil((eos-bos)/block_B))) can hold a - # valid key; every later block is fully masked (all cols >= eos) - # and contributes nothing to out/lse. Its per-block max logit would - # be -inf, which the host ``BlockLogit`` buffer is pre-filled with, - # so skipping those blocks is exact. For pure causal this halves the - # work; for packed documents a token only scans its own document. - num_valid_blocks = T.ceildiv(eos - bos, BB) - for j in T.Pipelined(num_valid_blocks, num_stages=num_stages): - # gather relative block j: cols [bos + j*BB, bos + (j+1)*BB). - # Guard the read against the padded K/V length (cols >= eos are - # masked below, so a clamped dummy read is harmless). + # causal / window / document early-exit: only tiles overlapping the + # token's valid range [bos, eos) can hold an unmasked key; later + # tiles are fully masked (all cols >= eos) and add nothing. + num_valid_tiles = T.ceildiv(eos - bos, BB) + for j in T.Pipelined(num_valid_tiles, num_stages=num_stages): + # gather tile j: cols [bos + j*BB, bos + (j+1)*BB). Guard the + # read against the padded K/V length (cols >= eos are masked + # below, so a clamped dummy read is harmless). for c, d in T.Parallel(BB, D): col = bos + j * BB + c in_bounds = col < seq_len_kv @@ -147,6 +147,10 @@ def main( K_shared[c, d] = T.if_then_else( in_bounds, K[bb, safe_col, d], T.cast(0, dtype) ) + for c, d in T.Parallel(BB, D_v): + col = bos + j * BB + c + in_bounds = col < seq_len_kv + safe_col = T.if_then_else(in_bounds, col, 0) V_shared[c, d] = T.if_then_else( in_bounds, V[bb, safe_col, d], T.cast(0, dtype) ) @@ -161,23 +165,18 @@ def main( policy=T.GemmWarpPolicy.FullRow, ) - # causal + document mask (col >= bos automatic for block j) + # causal + document mask (col >= bos automatic for tile j) for h, c in T.Parallel(PH, BB): col = bos + j * BB + c acc_s[h, c] = T.if_then_else( col < eos, acc_s[h, c], -T.infinity(accum_dtype) ) - # per-block max of raw logit over valid cols -> block score src - T.reduce_max(acc_s, blk_max, dim=1, clear=True) - for h in T.Parallel(PH): - if h < H: - BlockLogit[bb, h, bs, j] = blk_max[h] - # online softmax (base 2) over scaled logits + T.reduce_max(acc_s, tile_max, dim=1, clear=True) T.copy(m_i, m_prev) for h in T.Parallel(PH): - m_i[h] = T.max(m_i[h], blk_max[h] * sm_scale) + m_i[h] = T.max(m_i[h], tile_max[h] * sm_scale) for h in T.Parallel(PH): alpha[h] = T.exp2((m_prev[h] - m_i[h]) * 1.44269504) for h, c in T.Parallel(PH, BB): @@ -187,19 +186,40 @@ def main( T.reduce_sum(acc_s, l_new, dim=1) for h in T.Parallel(PH): l_i[h] = l_i[h] * alpha[h] + l_new[h] - for h, d in T.Parallel(PH, D): + for h, d in T.Parallel(PH, D_v): acc_o[h, d] = acc_o[h, d] * alpha[h] T.copy(acc_s, P_shared) T.gemm( P_shared, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow ) + # Fold the (optional) learnable attention sink as a virtual key + # column: a per-head logit ``AttnSink[h]`` competing in the same + # softmax denominator (reduces every weight so they sum < 1). A + # very-negative sink (e.g. -1e30) makes ``exp(sink - m)`` underflow + # to 0, recovering the plain sinkless softmax bit-for-bit. + # AttnSink is a pre-scaled logit (same units as ``m_i``), so it is + # converted to base-2 with log2(e) only (no ``sm_scale``). + for h in T.Parallel(PH): + safe_h = T.if_then_else(h < H, h, 0) + sink_h = T.if_then_else( + h < H, AttnSink[safe_h], -T.infinity(accum_dtype) + ) + m_prev[h] = m_i[h] + m_i[h] = T.max(m_i[h], sink_h) + alpha[h] = T.exp2((m_prev[h] - m_i[h]) * 1.44269504) + l_i[h] = l_i[h] * alpha[h] + T.exp2( + (sink_h - m_i[h]) * 1.44269504 + ) + for h, d in T.Parallel(PH, D_v): + acc_o[h, d] = acc_o[h, d] * alpha[h] + # normalize; empty rows (no valid key) -> 0 out / -inf lse - for h, d in T.Parallel(PH, D): + for h, d in T.Parallel(PH, D_v): acc_o[h, d] = T.if_then_else( l_i[h] > 0, acc_o[h, d] / l_i[h], 0.0 ) - for h, d in T.Parallel(PH, D): + for h, d in T.Parallel(PH, D_v): if h < H: Output[bb, bs, h, d] = acc_o[h, d] for h in T.Parallel(PH): @@ -213,53 +233,36 @@ def main( return main -def block_scores_from_logit(block_logit, lse, sm_scale): - """Recover eq.(3) block-max probability scores on the host. - - Args: - block_logit: [B, H, S, num_blocks] raw per-block max logit (-inf if - fully masked), as returned by the forward kernel. - lse: [B, S, H] natural-log LSE from the forward kernel. - sm_scale: softmax scale. - - Returns: - [B, H, S, num_blocks] block-max softmax probabilities in [0, 1]; - fully-masked blocks are 0. - """ - lse_bhs = lse.transpose([0, 2, 1]).unsqueeze(-1) # [B,H,S,1] - scaled = block_logit.astype("float32") * sm_scale - lse_bhs - scores = paddle.exp(scaled) - # exp(-inf) already 0, but guard any nan from (-inf)-(-inf) style edge. - scores = paddle.where( - paddle.isfinite(scores), scores, paddle.zeros_like(scores) - ) - return scores - - -def block_score_mqa_attn_fwd(q, k, v, valid_range, sm_scale=None, block_B=64): - """Forward interface for the MQA (shared K/V head) block-score attention. +def windowed_mqa_attn_fwd( + q, k, v, valid_range, attn_sink=None, sm_scale=None, block_B=64 +): + """Forward interface for windowed MQA (shared K/V head) flash attention. Args: q: [B, S, H, D] bf16 query (H heads). - k, v: [B, S_kv, D] bf16 single shared key/value head. - valid_range: [B, S, 2] int32 per-query [bos, eos) valid key columns. + k, v: [B, S_kv, D] / [B, S_kv, D_v] bf16 single shared K/V head. + valid_range: [B, S, 2] int32 per-query [bos, eos) valid key columns + (encodes causal + document + sliding-window masking). + attn_sink: [H] fp32 per-head learnable sink logit (attention sink / + softmax off-by-one). ``None`` -> a very-negative sink is used, + recovering the plain sinkless softmax bit-for-bit. sm_scale: softmax scale; defaults to D**-0.5. - block_B: key block size. + block_B: key tile size streamed by the kernel. Returns: - out [B,S,H,D], lse [B,S,H], block_logit [B,H,S,num_blocks] where - num_blocks = ceil(S_kv / block_B). + out [B,S,H,D_v], lse [B,S,H]. """ assert q.is_contiguous() b, s, h, d = q.shape s_kv = k.shape[1] + d_v = v.shape[-1] # lightweight host-side shape checks (no device sync) so mismatched inputs # fail early and clearly instead of causing undefined kernel behaviour. - assert len(k.shape) == 3 and list(k.shape) == list(v.shape), ( - f"k/v must be [B, S_kv, D] and match; got {k.shape} vs {v.shape}" + assert len(k.shape) == 3 and k.shape[0] == b and k.shape[2] == d, ( + f"k must be [B, S_kv, D] matching q; got k {k.shape}, q {q.shape}" ) - assert k.shape[0] == b and k.shape[2] == d, ( - f"k/v batch/head_dim must match q; got k {k.shape}, q {q.shape}" + assert len(v.shape) == 3 and list(v.shape[:2]) == [b, s_kv], ( + f"v must be [B, S_kv, D_v] matching k seqlen; got {v.shape}, k {k.shape}" ) assert list(valid_range.shape) == [b, s, 2], ( f"valid_range must be [B, S, 2]; got {valid_range.shape}" @@ -269,7 +272,18 @@ def block_score_mqa_attn_fwd(q, k, v, valid_range, sm_scale=None, block_B=64): if valid_range.dtype != paddle.int32: valid_range = valid_range.cast("int32") - # kernel reads whole block_B blocks; pad K/V so the last block is in bounds + # No learnable sink -> a very-negative per-head sink makes exp(sink - m) + # underflow to 0, so the kernel produces the plain sinkless softmax. + if attn_sink is None: + attn_sink = paddle.full([h], -1e30, dtype="float32") + else: + assert list(attn_sink.shape) == [h], ( + f"attn_sink must be [H={h}]; got {attn_sink.shape}" + ) + attn_sink = attn_sink.cast("float32") + attn_sink = attn_sink.contiguous() + + # kernel streams whole block_B tiles; pad K/V so the last tile is in bounds pad = (block_B - s_kv % block_B) % block_B if pad > 0: k = paddle.nn.functional.pad(k, [0, 0, 0, pad]) @@ -277,14 +291,7 @@ def block_score_mqa_attn_fwd(q, k, v, valid_range, sm_scale=None, block_B=64): k = k.contiguous() v = v.contiguous() valid_range = valid_range.contiguous() - num_blocks = (s_kv + pad) // block_B - # Blocks past a token's own valid range are skipped by the kernel's - # causal/document early-exit, so their per-block max logit must already read - # as -inf (score 0). Pre-fill instead of leaving the buffer uninitialised. - block_logit = paddle.full( - [b, h, s, num_blocks], float("-inf"), dtype="float32" - ) - kernel = block_score_mqa_fwd(h, d, float(sm_scale), block_B=block_B) - out, lse = kernel(q, k, v, valid_range, block_logit) - return out, lse, block_logit + kernel = windowed_mqa_fwd(h, d, float(sm_scale), D_v=d_v, block_B=block_B) + out, lse = kernel(q, k, v, valid_range, attn_sink) + return out, lse diff --git a/src/paddlefleet/tilelang_ops/hysparse/block_score_attn_bwd.py b/src/paddlefleet/tilelang_ops/hysparse/windowed_mqa_attn_bwd.py similarity index 63% rename from src/paddlefleet/tilelang_ops/hysparse/block_score_attn_bwd.py rename to src/paddlefleet/tilelang_ops/hysparse/windowed_mqa_attn_bwd.py index 68893605b..1cb8b140e 100644 --- a/src/paddlefleet/tilelang_ops/hysparse/block_score_attn_bwd.py +++ b/src/paddlefleet/tilelang_ops/hysparse/windowed_mqa_attn_bwd.py @@ -12,21 +12,44 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Backward for the MQA block-score attention: flash-attention backward -(dQ, dK, dV) for full attention with a single shared K/V head and causal + -document masking. +"""Backward for the windowed MQA flash attention (:mod:`windowed_mqa_attn`): +flash-attention backward (dQ, dK, dV) with a single shared K/V head and +causal + document + sliding-window masking. K/V are one shared head ``[B, S_kv, D]``; dK/dV from every query head are -scattered into that single head via ``atomic_add``. The block-score output of -the forward is consumed by a non-differentiable TopK and therefore contributes -no gradient here. +scattered into that single head via ``atomic_add``. Masking is expressed +through ``valid_range`` exactly as in the forward. """ import paddle import tilelang from tilelang import language as T -from .block_sparse_attn_mqa_bwd import _cast_bf16_kv + +@tilelang.jit(out_idx=[-1]) +def _cast_bf16_kv(D, block_N=64, threads=128): + """Cast a shared-head K/V grad accumulator ``[B, S_kv, D]`` fp32 -> bf16. + + The backward accumulates dK/dV in fp32 (atomic_add scatter); this kernel + down-casts the result to bf16 to match the forward's K/V dtype. + """ + batch = T.dynamic("batch") + seq_len_kv = T.dynamic("seq_len_kv") + shape = [batch, seq_len_kv, D] + + @T.prim_func + def main( + X: T.Tensor(shape, T.float32), + Out: T.Tensor(shape, T.bfloat16), + ): + with T.Kernel( + T.ceildiv(seq_len_kv, block_N), batch, threads=threads + ) as (bn, bb): + for i, d in T.Parallel(block_N, D): + if bn * block_N + i < seq_len_kv: + Out[bb, bn * block_N + i, d] = X[bb, bn * block_N + i, d] + + return main @tilelang.jit( @@ -36,19 +59,25 @@ tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, }, ) -def block_score_mqa_bwd( +def windowed_mqa_bwd( H, D, sm_scale, + D_v=None, block_M=64, block_N=64, block_B=64, num_stages=1, threads=128, ): + if D_v is None: + D_v = D assert D % 16 == 0, ( f"D must be a multiple of 16 (tensor-core k-tile), got {D}" ) + assert D_v % 16 == 0, ( + f"D_v must be a multiple of 16 (tensor-core k-tile), got {D_v}" + ) assert block_B % block_N == 0, "block_B must be a multiple of block_N" batch = T.dynamic("batch") @@ -58,9 +87,12 @@ def block_score_mqa_bwd( q_shape = [batch, seq_len, H, D] kv_shape = [batch, seq_len_kv, D] + v_shape = [batch, seq_len_kv, D_v] + do_shape = [batch, seq_len, H, D_v] lse_shape = [batch, seq_len, H] vr_shape = [batch, seq_len, 2] br_shape = [batch, num_bm, 2] + sink_shape = [H] dtype = T.bfloat16 accum_dtype = T.float32 @@ -75,14 +107,16 @@ def block_score_mqa_bwd( def main( Q: T.Tensor(q_shape, dtype), K: T.Tensor(kv_shape, dtype), - V: T.Tensor(kv_shape, dtype), - dO: T.Tensor(q_shape, dtype), + V: T.Tensor(v_shape, dtype), + dO: T.Tensor(do_shape, dtype), Lse: T.Tensor(lse_shape, accum_dtype), Delta: T.Tensor(lse_shape, accum_dtype), ValidRange: T.Tensor(vr_shape, idx_dtype), BlockRange: T.Tensor(br_shape, idx_dtype), + AttnSink: T.Tensor(sink_shape, accum_dtype), dK: T.Tensor(kv_shape, accum_dtype), - dV: T.Tensor(kv_shape, accum_dtype), + dV: T.Tensor(v_shape, accum_dtype), + dAttnSink: T.Tensor(sink_shape, accum_dtype), dQ: T.Tensor(q_shape, dtype), ): with T.Kernel(T.ceildiv(seq_len, BM), H, batch, threads=threads) as ( @@ -91,9 +125,9 @@ def main( bb, ): Q_shared = T.alloc_shared([BM, D], dtype) - dO_shared = T.alloc_shared([BM, D], dtype) + dO_shared = T.alloc_shared([BM, D_v], dtype) K_shared = T.alloc_shared([BN, D], dtype) - V_shared = T.alloc_shared([BN, D], dtype) + V_shared = T.alloc_shared([BN, D_v], dtype) P_shared = T.alloc_shared([BM, BN], dtype) dS_shared = T.alloc_shared([BM, BN], dtype) @@ -102,9 +136,11 @@ def main( acc_dp = T.alloc_fragment([BM, BN], accum_dtype) acc_dq = T.alloc_fragment([BM, D], accum_dtype) acc_dk = T.alloc_fragment([BN, D], accum_dtype) - acc_dv = T.alloc_fragment([BN, D], accum_dtype) + acc_dv = T.alloc_fragment([BN, D_v], accum_dtype) lse_f = T.alloc_fragment([BM], accum_dtype) delta_f = T.alloc_fragment([BM], accum_dtype) + sink_contrib = T.alloc_fragment([BM], accum_dtype) + sink_acc = T.alloc_fragment([1], accum_dtype) bos = T.alloc_fragment([BM], idx_dtype) eos = T.alloc_fragment([BM], idx_dtype) @@ -116,8 +152,26 @@ def main( lse_f[i] = T.if_then_else(in_range, Lse[bb, row, bh], 0) delta_f[i] = T.if_then_else(in_range, Delta[bb, row, bh], 0) - T.copy(Q[bb, bm * BM : (bm + 1) * BM, bh, :], Q_shared) - T.copy(dO[bb, bm * BM : (bm + 1) * BM, bh, :], dO_shared) + # Guarded per-row load of Q / dO: the kernel grid is + # ceildiv(seq_len, BM), so when seq_len % BM != 0 the last query + # tile spans rows past seq_len. A plain block T.copy would read + # out of bounds (illegal access / garbage feeding the GEMMs); load + # row-by-row, zero-filling the ragged tail (masked rows contribute + # nothing and their dQ write is already row-guarded below). + for i, d in T.Parallel(BM, D): + row = bm * BM + i + in_range = row < seq_len + safe_row = T.if_then_else(in_range, row, 0) + Q_shared[i, d] = T.if_then_else( + in_range, Q[bb, safe_row, bh, d], T.cast(0, dtype) + ) + for i, d in T.Parallel(BM, D_v): + row = bm * BM + i + in_range = row < seq_len + safe_row = T.if_then_else(in_range, row, 0) + dO_shared[i, d] = T.if_then_else( + in_range, dO[bb, safe_row, bh, d], T.cast(0, dtype) + ) T.clear(acc_dq) # document-tight early-exit: the host precomputes, per query tile, @@ -201,8 +255,9 @@ def main( policy=T.GemmWarpPolicy.FullRow, ) for c, d in T.Parallel(BN, D): - T.atomic_add(dV[bb, col0 + c, d], acc_dv[c, d]) T.atomic_add(dK[bb, col0 + c, d], acc_dk[c, d]) + for c, d in T.Parallel(BN, D_v): + T.atomic_add(dV[bb, col0 + c, d], acc_dv[c, d]) # write dQ straight from the accumulator fragment (no shared-memory # staging) -- one fewer [BM, D] bf16 buffer lets the shared budget @@ -211,31 +266,53 @@ def main( if bm * BM + i < seq_len: dQ[bb, bm * BM + i, bh, d] = acc_dq[i, d] + # dAttnSink[h] = -sum_{b,s}( Delta[b,s,h] * p_sink[b,s,h] ), where + # p_sink = exp(sink - lse) is the (virtual) sink token's softmax + # weight. AttnSink is a pre-scaled logit, so convert to base-2 with + # log2(e) only (no sm_scale). Masked/ragged rows (lse_f == 0 from the + # guarded load) contribute exp2(sink*log2e) which is ~0 for the + # very-negative sinkless sentinel; for a real learnable sink the + # padded rows are past seq_len and excluded via the row guard. + for i in T.Parallel(BM): + row = bm * BM + i + in_range = row < seq_len + sink_contrib[i] = T.if_then_else( + in_range, + -delta_f[i] + * T.exp2((AttnSink[bh] - lse_f[i]) * 1.44269504), + 0.0, + ) + T.reduce_sum(sink_contrib, sink_acc, dim=0, clear=True) + T.atomic_add(dAttnSink[bh], sink_acc[0]) + return main -def _fit_block_mn(D, block_B, cap_bytes=230000): +def _fit_block_mn(D, block_B, D_v=None, cap_bytes=230000): """Pick (block_M, block_N) maximising the query tile then the key sub-tile. - The backward holds Q/dO ``[block_M, D]`` + K/V ``[block_N, D]`` + P/dS - ``[block_M, block_N]`` in bf16 shared memory (dQ writes straight from its - accumulator). Prefer the largest ``block_M`` (<=64) -- big M matches the - forward's tensor-core utilisation and K/V amortisation -- and, for that M, - the largest ``block_N`` (dividing ``block_B``) that fits. At MLA D=576 a - full block_M=64 needs block_N=32 to fit Blackwell's ~227 KB; small dims - (D=64) keep block_N=block_B (single-tile fast path). + The backward holds Q ``[block_M, D]`` + dO ``[block_M, D_v]`` + K + ``[block_N, D]`` + V ``[block_N, D_v]`` + P/dS ``[block_M, block_N]`` in + bf16 shared memory (dQ writes straight from its accumulator). Prefer the + largest ``block_M`` (<=64) -- big M matches the forward's tensor-core + utilisation and K/V amortisation -- and, for that M, the largest + ``block_N`` (dividing ``block_B``) that fits. At MLA D=576 a full + block_M=64 needs block_N=32 to fit Blackwell's ~227 KB; small dims (D=64) + keep block_N=block_B (single-tile fast path). """ + if D_v is None: + D_v = D cands_n = [n for n in (block_B, 32, 16) if block_B % n == 0] cands_n = sorted(set(cands_n), reverse=True) for bm in (64, 48, 32, 16): for bn in cands_n: - shared = 4 * (bm * D + bn * D + bm * bn) + shared = 2 * (bm * (D + D_v) + bn * (D + D_v) + 2 * bm * bn) if shared <= cap_bytes: return bm, bn return 16, min(16, block_B) -def block_score_mqa_bwd_interface( +def windowed_mqa_bwd_interface( q, k, v, @@ -243,20 +320,24 @@ def block_score_mqa_bwd_interface( do, lse, valid_range, + attn_sink=None, sm_scale=None, block_B=64, block_M=None, block_N=None, ): - """Backward interface for the MQA block-score attention. + """Backward interface for the windowed MQA flash attention. Args: q: [B, S, H, D] bf16 forward query. k, v: [B, S_kv, D] bf16 single shared key/value head. o: [B, S, H, D] bf16 forward output. do: [B, S, H, D] bf16 grad of output. - lse: [B, S, H] fp32 natural-log LSE from forward. + lse: [B, S, H] fp32 natural-log LSE from forward (already + includes the sink term in its denominator). valid_range: [B, S, 2] int32. + attn_sink: [H] fp32 per-head sink logit matching the forward. ``None`` + -> a very-negative sink (sinkless); its grad is ignored. sm_scale: softmax scale; defaults to D**-0.5. block_B: key block size. block_M: query tile size on the GEMM M dim; ``None`` auto-fits the @@ -266,16 +347,27 @@ def block_score_mqa_bwd_interface( result is mathematically identical for any valid choice. Returns: - dq [B,S,H,D] bf16, dk/dv [B,S_kv,D] bf16. + dq [B,S,H,D] bf16, dk/dv [B,S_kv,D] bf16, d_attn_sink [H] fp32. """ assert q.is_contiguous() and do.is_contiguous() and lse.is_contiguous() b, s, h, d = q.shape s_kv = k.shape[1] + d_v = v.shape[-1] if sm_scale is None: sm_scale = d**-0.5 if valid_range.dtype != paddle.int32: valid_range = valid_range.cast("int32") + # No learnable sink -> very-negative sentinel (sinkless); its grad is unused. + if attn_sink is None: + attn_sink = paddle.full([h], -1e30, dtype="float32") + else: + assert list(attn_sink.shape) == [h], ( + f"attn_sink must be [H={h}]; got {attn_sink.shape}" + ) + attn_sink = attn_sink.cast("float32") + attn_sink = attn_sink.contiguous() + # kernel reads whole block_B blocks; pad K/V (and dK/dV) so the last block # is addressable. pad = (block_B - s_kv % block_B) % block_B @@ -289,9 +381,10 @@ def block_score_mqa_bwd_interface( delta = (o.astype("float32") * do.astype("float32")).sum(-1).contiguous() dk = paddle.zeros([b, s_kv_pad, d], dtype="float32") - dv = paddle.zeros([b, s_kv_pad, d], dtype="float32") + dv = paddle.zeros([b, s_kv_pad, d_v], dtype="float32") + d_attn_sink = paddle.zeros([h], dtype="float32") - fit_m, fit_n = _fit_block_mn(d, block_B) + fit_m, fit_n = _fit_block_mn(d, block_B, d_v) if block_M is None: block_M = fit_m if block_N is None: @@ -316,20 +409,33 @@ def block_score_mqa_bwd_interface( jh = paddle.maximum(jh, jl) block_range = paddle.stack([jl, jh], axis=-1).astype("int32").contiguous() - bwd = block_score_mqa_bwd( + bwd = windowed_mqa_bwd( h, d, float(sm_scale), + D_v=d_v, block_M=block_M, block_N=block_N, block_B=block_B, ) - dq = bwd(q, k, v, do, lse, delta, valid_range, block_range, dk, dv) + dq = bwd( + q, + k, + v, + do, + lse, + delta, + valid_range, + block_range, + attn_sink, + dk, + dv, + d_attn_sink, + ) - cast = _cast_bf16_kv(d) - dk_bf = cast(dk) - dv_bf = cast(dv) + dk_bf = _cast_bf16_kv(d)(dk) + dv_bf = _cast_bf16_kv(d_v)(dv) if pad > 0: dk_bf = dk_bf[:, :s_kv, :].contiguous() dv_bf = dv_bf[:, :s_kv, :].contiguous() - return dq, dk_bf, dv_bf + return dq, dk_bf, dv_bf, d_attn_sink diff --git a/src/paddlefleet/transformer/moe/fusion_layer_utils.py b/src/paddlefleet/transformer/moe/fusion_layer_utils.py index 7ae1f0d9c..55c042eb1 100644 --- a/src/paddlefleet/transformer/moe/fusion_layer_utils.py +++ b/src/paddlefleet/transformer/moe/fusion_layer_utils.py @@ -37,8 +37,18 @@ from paddlefleet_ops.sonicmoe.enums import ActivationType from paddlefleet_ops.sonicmoe.ernie_compat.deepep_metadata import ( deepep_topk_to_sonic_metadata, - deepep_topk_to_sonic_metadata_with_scales, ) + + try: + from paddlefleet_ops.sonicmoe.ernie_compat.deepep_metadata import ( + deepep_topk_to_sonic_metadata_with_scales, + ) + except ImportError: + # Older installed paddlefleet_ops binaries (pre-#1348) do not export + # the fp8-scales variant; the sibling optional imports below use the + # same guard. Only the fp8 + fp8_scale MoE path calls it, which this + # config does not exercise. + deepep_topk_to_sonic_metadata_with_scales = None from paddlefleet_ops.sonicmoe.ernie_compat.mlp_node_v2 import ( _differentiable_router_scores, ) diff --git a/src/paddlefleet/transformer/multi_latent_attention.py b/src/paddlefleet/transformer/multi_latent_attention.py index 628026d40..957b27c7d 100644 --- a/src/paddlefleet/transformer/multi_latent_attention.py +++ b/src/paddlefleet/transformer/multi_latent_attention.py @@ -53,6 +53,89 @@ from paddlefleet.utils import get_pg_rank, get_pg_size +def build_hysparse_valid_range( + attn_mask_startend_row_indices, + seq_len, + batch_size, + window_size=None, +): + """Build ``valid_range`` [B, S, 2] int32 for the HySparse TileLang ops. + + Each query token ``t`` gets a half-open valid key-column range ``[bos, eos)``: + + * ``eos = t + 1`` (causal upper bound). + * ``bos`` = start of the document containing ``t`` (document mask). When + ``window_size`` is given, ``bos`` is additionally clamped up to + ``t - window_size + 1`` (causal sliding window). + + Document boundaries are recovered from the flashmask + ``attn_mask_startend_row_indices`` of shape ``[B, *, S, *]`` whose first + head / first channel holds, per token, the **exclusive end** of the document + that token belongs to (same convention as + ``utils.get_doc_lens`` / ``csa_attention._derive_csa_doc_boundaries``). + When it is ``None`` a single document (``bos`` document part = 0) is assumed. + + The block grid used by the block-score / block-sparse operators is anchored + at ``bos`` (document-relative blocks), so the *document* range (no window + clamp) must be used for block scoring and the block-sparse branch, while the + windowed range is used for the sliding-window main path. + """ + positions = paddle.arange(seq_len, dtype="int64").unsqueeze(0) # [1, S] + if attn_mask_startend_row_indices is not None: + # (C) Convention guard: we read the exclusive document-end from + # channel [:, 0, :, 0]. This only holds for the flashmask layout + # [B, num_masks, S, num_bounds] whose first mask / first bound carries + # the per-token exclusive doc-end (== utils.get_doc_lens / + # csa_attention._derive_csa_doc_boundaries). A silent upstream layout + # change (extra mask channels, bidirectional bounds, transposed axes) + # would make the [:, 0, :, 0] slice mean something else and corrupt + # every downstream bos -> block bucket. Assert the structural shape + # (host-side, free) so such a change fails loudly here instead of + # silently mis-bucketing. + assert attn_mask_startend_row_indices.ndim == 4, ( + "attn_mask_startend_row_indices must be 4-D " + "[B, num_masks, S, num_bounds] so [:, 0, :, 0] is the per-token " + f"exclusive doc-end; got ndim={attn_mask_startend_row_indices.ndim} " + f"shape={attn_mask_startend_row_indices.shape}" + ) + assert attn_mask_startend_row_indices.shape[2] == seq_len, ( + "attn_mask_startend_row_indices axis-2 must be the query length S=" + f"{seq_len}; got shape={attn_mask_startend_row_indices.shape} " + "(layout changed? [:, 0, :, 0] would no longer be the doc-end)" + ) + # [B, *, S, *] -> [B_mask, S] exclusive document end per token. + de = attn_mask_startend_row_indices[:, 0, :, 0].cast("int64") # [Bm, S] + # The flashmask row indices may carry a batch of 1 that broadcasts over + # the data batch (all sequences share one document layout). Expand so + # the produced valid_range matches the query/key/value batch instead of + # the mask's batch. + if de.shape[0] == 1 and batch_size > 1: + de = de.expand([batch_size, seq_len]) + bsz = de.shape[0] + pos_b = positions.expand([bsz, seq_len]) # [B, S] + is_boundary = paddle.zeros([bsz, seq_len], dtype="bool") + is_boundary[:, 0] = True + # a new document starts at t when the previous token's doc-end equals t + # and the doc-end value actually changes. + is_boundary[:, 1:] = (pos_b[:, 1:] == de[:, :-1]) & ( + de[:, 1:] != de[:, :-1] + ) + doc_start = paddle.cummax( + is_boundary.cast("int64") * pos_b, axis=1 + ).values # [B, S] most-recent document start <= t + else: + bsz = batch_size + doc_start = paddle.zeros([bsz, seq_len], dtype="int64") + + pos_b = positions.expand([bsz, seq_len]) + bos = doc_start + if window_size is not None and window_size > 0: + bos = paddle.maximum(doc_start, pos_b - window_size + 1) + eos = pos_b + 1 + valid_range = paddle.stack([bos, eos], axis=-1).cast("int32") # [B, S, 2] + return valid_range.contiguous() + + def _ec_compatible_rope_apply( q_pe, k_pe, @@ -471,6 +554,7 @@ def forward( packed_seq_params=None, in_recompute: bool = False, position_ids=None, + shared_kv: list[Tensor] | None = None, **kwargs, ): """Forward pass for multi-latent attention""" @@ -540,7 +624,28 @@ def forward( else: q_absorbed, wv_b = None, None - if self.recompute_core_attention and self.training: + if self.config.enable_hy_sparse_attention and shared_kv is not None: + # HySparse full-attention layer. The full (dense) attention here is + # computed by the MHA block-score TileLang op, which additionally + # emits per-(query, key-block) max logits. We select the top-k key + # blocks and share both the compressed KV latent and the selected + # block indices with the downstream SWA layers' block-sparse branch. + # + # This branch is checked BEFORE recompute_core_attention: the FA4 + # block-score full path is a distinct computation that must produce + # block_indices for the downstream SWA layer, and running the plain + # recompute(core_attention) branch here would leave block_indices + # undefined (UnboundLocalError at the shared_kv.append below) while + # also failing to emit the top-k blocks. Activation recompute for + # HySparse full layers is handled at the layer level + # (HySparseTransformerLayer.full_recompute). + core_attn_out, block_indices = self._hy_sparse_full_attention( + query, + key, + value, + attn_mask_startend_row_indices, + ) + elif self.recompute_core_attention and self.training: core_attn_out = recompute( self.core_attention, query, @@ -591,6 +696,16 @@ def forward( _log(core_attn_out, "core_attn_out", layer_num) + if self.config.enable_hy_sparse_attention and shared_kv is not None: + # Compressed KV latent shared with block-sparse attention in SWA + # layers (single MQA head): [B, S, 1, kv_lora_rank + qk_rope_head_dim]. + shared_key = paddle.concat( + [kv_compressed.unsqueeze(2), k_pos_emb], axis=-1 + ) + shared_kv.append(shared_key) + # block_indices produced by the MHA block-score path above. + shared_kv.append(block_indices) + # ================= # Output. [b, sq, h] # ================= @@ -655,6 +770,74 @@ def _gate(self, gate_source, core_attn_out): core_attn_out = core_attn_out * paddle.nn.functional.sigmoid(gate) return core_attn_out + def _hy_sparse_full_attention( + self, + query, + key, + value, + attn_mask_startend_row_indices, + ): + """HySparse full-attention layer using the FA4-fused block-score op. + + Runs dense (decompressed) MHA attention through the FA4 sm100 kernel, + whose softmax epilogue additionally emits per-(query, key-block) max + logits at near-zero extra cost (``block_score_fa4_attn_fwd``). From those + we recover block scores and select the top-k key blocks per query token. + The selected block indices (shared across heads, document-relative) are + returned so the downstream SWA layers' block-sparse branch can gather + exactly the same blocks. + + Args: + query: [B, S, H, Dk] decompressed query (H independent heads). + key: [B, S, H, Dk] decompressed key. + value: [B, S, H, Dv] decompressed value. + attn_mask_startend_row_indices: flashmask doc boundaries or ``None``. + + Returns: + core_attn_out: [B, S, H*Dv] dense attention output. + block_indices: [B, S, topk] int32 selected block ids (-1 padding). + """ + from paddlefleet.tilelang_ops.hysparse import ( + block_score_fa4_attn_fwd, + select_topk_blocks, + ) + + b, s, h, _dv = value.shape + block_B = self.config.hy_sparse_block_size + topk = self.config.hy_sparse_topk + sm_scale = self.softmax_scale + + # Document valid_range (no window clamp): full-layer block scoring and + # the SWA block-sparse branch must share the same document-anchored + # blocks so the selected indices are transferable across layers. It also + # anchors select_topk_blocks' per-token block bounds. FA4 itself masks + # via causal + the raw flashmask ``startend_row_indices`` (same doc + # structure valid_range is derived from), so the fused block-max -- + # taken after FA4's mask_fn -- honours the identical document mask. + valid_range = build_hysparse_valid_range( + attn_mask_startend_row_indices, s, b + ) + + out, lse, block_logit = block_score_fa4_attn_fwd( + query, + key, + value, + valid_range=valid_range, + sm_scale=sm_scale, + block_B=block_B, + causal=True, + startend_row_indices=attn_mask_startend_row_indices, + ) + block_indices = select_topk_blocks( + block_logit, + lse, + valid_range, + topk, + block_B, + ) + core_attn_out = out.reshape([b, s, h * _dv]) + return core_attn_out, block_indices + class MLASelfAttention(MultiLatentAttention): """MLA Self-attention layer class @@ -1332,3 +1515,499 @@ def _backward_q_proj(self): def _backward_output_proj(self): """Computes weight gradients of output projection layer""" self.o_proj.backward_dw() + + +class MQASelfAttention(MLASelfAttention): + """Multi-Query Attention.""" + + def __init__( + self, + config: TransformerConfig, + sublayers_spec: MLASelfAttentionSublayersSpec, + layer_number: int, + attn_mask_type=AttnMaskType.padding, + cp_comm_type: str | None = None, + pg_collection: ProcessGroupCollection | None = None, + is_mtp_layer: bool = False, + ): + super().__init__( + config=config, + sublayers_spec=sublayers_spec, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, + ) + + assert not self.config.apply_rope_fusion, ( + "MQA does not support rope fusion." + ) + + # Use MQA only when HySparse is enabled and this is an SWA layer. + # Otherwise, use its parent's forward method (MLA). + self.is_mqa = config.enable_hy_sparse_attention and self.is_swa + + if self.is_mqa: + # Adjust absorbed kv channels for core attention + k_channels = self.config.kv_lora_rank + self.qk_rope_head_dim + v_channels = self.config.kv_lora_rank + + self.core_attention.hidden_size_per_partition = ( + k_channels * self.num_attention_heads_per_partition + ) + self.core_attention.k_channels = k_channels + self.core_attention.v_channels = v_channels + + # Block sparse attention + self.sparse_core_attention = build_spec_layer( + sublayers_spec.core_attention, + config=self.core_attention.config, + layer_number=self.layer_number, + attn_mask_type=self.attn_mask_type, + attention_type=self.attention_type, + is_mtp_layer=self.is_mtp_layer, + is_swa=False, + softmax_scale=self.config.softmax_scale, + k_channels=k_channels, + v_channels=v_channels, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=1, + cp_comm_type=cp_comm_type, + pg_collection=self.pg_collection, + ) + + # Gate for block sparse attention + if self.gated_attention: + self.sparse_gate_proj = build_spec_layer( + sublayers_spec.gate_proj, + self.gate_proj.input_size, + self.gate_proj.output_size, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=self.config.use_bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="mla_gate", + tp_group=self.pg_collection.tp, + ) + + # Learnable attention-sink bias for the SWA MQA path. Gated by + # ``add_swa_attention_sink_bias`` (mirrors the DotProductAttention + # SWA sink promotion). The two HySparse MQA branches are independent + # softmaxes (main sliding-window path + block-sparse DSA path), so + # each gets its own per-head sink logit; both are zero-initialised + # (a zero sink logit == an off-by-one-style sink at logit 0). When + # the switch is off, both stay ``None`` and the kernels run their + # plain sinkless softmax exactly as before. + self.add_swa_attention_sink_bias = getattr( + self.config, "add_swa_attention_sink_bias", False + ) + if self.add_swa_attention_sink_bias: + num_heads = self.num_attention_heads_per_partition + self.swa_attn_sink = self.create_parameter( + shape=[num_heads], + dtype="float32", + default_initializer=paddle.nn.initializer.Constant(0.0), + ) + self.sparse_attn_sink = self.create_parameter( + shape=[num_heads], + dtype="float32", + default_initializer=paddle.nn.initializer.Constant(0.0), + ) + else: + self.swa_attn_sink = None + self.sparse_attn_sink = None + + def forward( + self, + hidden_states, + attention_mask, + attn_mask_startend_row_indices: Tensor | None = None, + key_value_states=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + attention_bias=None, + packed_seq_params=None, + in_recompute: bool = False, + position_ids=None, + shared_kv: list[Tensor] | None = None, + **kwargs, + ): + """Forward pass for multi-latent attention""" + from paddlefleet.transformer.transformer_layer import TransformerLayer + + _log = TransformerLayer._log_md5 + + assert rotary_pos_emb is None, ( + "Rotary position embeddings should not be passed into MQA." + ) + assert attention_bias is None, ( + "Attention bias should not be passed into MQA." + ) + assert rotary_pos_cos is None and rotary_pos_sin is None, ( + "MQA does not support Flash Decoding" + ) + assert not get_context_parallel_world_size() > 1, ( + "MQA does not support context parallel." + ) + assert get_pg_size(self.pg_collection.tp) == 1, ( + "MQA does not support tensor parallel." + ) + + if not self.is_mqa: + return super().forward( + hidden_states, + attention_mask, + attn_mask_startend_row_indices=attn_mask_startend_row_indices, + key_value_states=key_value_states, + packed_seq_params=packed_seq_params, + in_recompute=in_recompute, + position_ids=position_ids, + shared_kv=shared_kv, + **kwargs, + ) + + # ===================== + # Query, Key, and Value + # ===================== + # Get the query, key and value tensors based on the type of attention + # Also get q_compressed for DSA indexer (if enabled) + query, key, value, q_compressed, kv_compressed, k_pos_emb = ( + self.get_query_key_value_tensors( + hidden_states, + key_value_states, + position_ids, + packed_seq_params, + ) + ) + + layer_num = getattr(self, "layer_number", -1) + _log(query, "attn_query", layer_num) + _log(key, "attn_key", layer_num) + if value is not None: + _log(value, "attn_value", layer_num) + + attn_mask_type = self.attn_mask_type + query = query.contiguous() + key = key.contiguous() + + if value is not None: + value = value.contiguous() + + # ================================== + # core attention computation + # ================================== + from paddlefleet.tilelang_ops.hysparse import ( + sliding_window_mqa_attention, + ) + + b, s = query.shape[0], query.shape[1] + block_B = self.config.hy_sparse_block_size + sm_scale = self.softmax_scale + window_size = self.config.sliding_window[0] + + # Absorbed-MLA MQA: a single shared K/V head with Dk=576 (kv_lora_rank + + # qk_rope_head_dim) / Dv=512 (kv_lora_rank). Squeeze the head axis to the + # [B, S_kv, D] layout the TileLang MQA kernels expect. + shared_k = key.squeeze(2).contiguous() # [B, S, 576] + shared_v = value.squeeze(2).contiguous() # [B, S, 512] + + # Windowed valid_range for the sliding-window main path; document-anchored + # valid_range (no window clamp) for the block-sparse branch so its blocks + # match the full layer's block scoring / selected indices. + window_valid_range = build_hysparse_valid_range( + attn_mask_startend_row_indices, s, b, window_size=window_size + ) + doc_valid_range = build_hysparse_valid_range( + attn_mask_startend_row_indices, s, b + ) + + # Sliding-window main path (MQA, Dk=576/Dv=512). + core_attn_out, _ = sliding_window_mqa_attention( + query, + shared_k, + shared_v, + window_valid_range, + attn_sink=getattr(self, "swa_attn_sink", None), + sm_scale=sm_scale, + block_B=block_B, + ) + core_attn_out = core_attn_out.reshape( + [ + b, + s, + self.num_attention_heads_per_partition + * self.config.kv_lora_rank, + ] + ) + + _log(core_attn_out, "core_attn_out", layer_num) + + # ================= + # Absorb value. [b, sq, num_heads * v_head_dim] + # ================= + + kv_lora_rank = self.config.kv_lora_rank + num_heads = self.num_attention_heads_per_partition + + v_absorb_weight = self.kv_b_proj.weight.reshape( + [kv_lora_rank, num_heads, -1] + )[:, :, self.qk_nope_head_dim :] + + def compute_absorbed_v(core_attn_out): + core_attn_out = core_attn_out.view( + *core_attn_out.shape[:-1], num_heads, kv_lora_rank + ) + core_attn_out = paddle.einsum( + "bshl,lhv->bshv", core_attn_out, v_absorb_weight + ) + core_attn_out = core_attn_out.view( + *core_attn_out.shape[:-2], num_heads * self.v_head_dim + ) + return core_attn_out + + core_attn_out = compute_absorbed_v(core_attn_out) + + # ================= + # Sparse attention computation + # ================= + + shared_key, shared_block_indices = shared_kv + # Shared compressed KV latent from the full layer: [B, S, 1, 576]. + # Squeeze to [B, S_kv, D] and split into key (576) / value (512). + shared_key_sq = shared_key.squeeze(2).contiguous() # [B, S, 576] + + # Block-sparse gather branch: DeepSeek-v4's CSA sparse attention + # (FlashMLA sparse fwd + cuDNN DSA bwd). It routes the HySparse block + # selection through DSA's token-level sparse path and natively handles + # the absorbed-MQA Dk=576 / Dv=512 shared-head layout (value == leading + # kv_lora_rank slice of the shared latent). + from paddlefleet.cudnn_ops import ( + block_sparse_mqa_attention_dsa, + is_dsa_available, + ) + + if not is_dsa_available(): + raise RuntimeError( + "HySparse block-sparse attention requires the DSA backend " + "(FlashMLA sparse fwd + cuDNN DSA bwd), unavailable here: it " + "needs SM100 + FlashMLA + the cuDNN frontend." + ) + sparse_core_attn_out, _ = block_sparse_mqa_attention_dsa( + query, + shared_key_sq, + shared_block_indices, + doc_valid_range, + sm_scale=sm_scale, + block_B=block_B, + kv_lora_rank=self.config.kv_lora_rank, + attn_sink=getattr(self, "sparse_attn_sink", None), + ) + sparse_core_attn_out = sparse_core_attn_out.reshape( + [ + b, + s, + self.num_attention_heads_per_partition + * self.config.kv_lora_rank, + ] + ) + + sparse_core_attn_out = compute_absorbed_v(sparse_core_attn_out) + + # ================= + # Output. [b, sq, h] + # ================= + # Apply gated attention + if self.gated_attention: + # Gate input source: q_compressed (post q_a_layernorm, dim=q_lora_rank) when + # gated_attn_use_q_lora is set, otherwise hidden_states. + gate_source = ( + q_compressed if self.gated_attn_use_q_lora else hidden_states + ) + core_attn_out = self._gate(gate_source, core_attn_out) + + # Add sparse attention output + if self.gated_attention: + gate, _ = self.sparse_gate_proj(gate_source) + sparse_core_attn_out = ( + sparse_core_attn_out * paddle.nn.functional.sigmoid(gate) + ) + core_attn_out += sparse_core_attn_out + + output, bias = self.o_proj(core_attn_out) + + _log(output, "attn_o_proj_out", layer_num) + + return output, bias + + def get_query_key_value_tensors( + self, + hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + ): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ + if not self.is_mqa: + return super().get_query_key_value_tensors( + hidden_states, + key_value_states=key_value_states, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + ) + + # b = batch size, s = sequence length, h = hidden size, n = num attention heads + # Attention heads [b, s, n*h] + assert hidden_states.ndim == 3, ( + f"hidden_states should be 3D, [b, s, n*h], got {hidden_states.ndim}D" + ) + + # ========================================= + # Prepare RoPE and seqlen related params + # ========================================= + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + hidden_states, self.config, packed_seq_params + ) + + # rotary_pos_emb: [1, s, 1, pe_dim] + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + + assert self.config.rope_type == "rope", ( + f"MQA only supports rope_type 'rope', got {self.config.rope_type}" + ) + assert packed_seq_params is None, ( + "MQA doesn't support packed_seq_params" + ) + + rotary_pos_emb = self.rotary_pos_emb( + rotary_seq_len, + position_ids=None if self.training else position_ids, + ) + + # ========================================= + # QKV down projection and layernorm + # ========================================= + if self.config.q_lora_rank is not None: + q_compressed, _ = self.q_a_proj(hidden_states) + else: + q_compressed = hidden_states + + # kv_combined: [b, s, (kv_lora_rank + qk_rope_head_dim)] + kv_combined, _ = self.kv_a_proj_with_mqa(hidden_states) + + # kv_compressed: [b, s, kv_lora_rank], k_pos_emb: [b, s, qk_rope_head_dim] + kv_compressed, k_pos_emb = paddle.split( + kv_combined, + [self.config.kv_lora_rank, self.qk_rope_head_dim], + axis=-1, + ) + + # ========================================= + # Apply norm + # ========================================= + + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] + q_compressed = self.q_a_layernorm(q_compressed) + + kv_compressed = self.kv_a_layernorm(kv_compressed) + + # === MD5 probes for MLA intermediate values === + from paddlefleet.transformer.transformer_layer import TransformerLayer + + _log = TransformerLayer._log_md5 + _log(q_compressed, "mla_q_compressed_normed", self.layer_number) + _log(kv_compressed, "mla_kv_compressed_normed", self.layer_number) + _log(k_pos_emb, "mla_k_pos_emb_raw", self.layer_number) + + # ========================================= + # QKV up projection and RoPE apply + # ========================================= + + def qkv_up_proj_and_rope_apply( + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + position_ids=None, + ): + """ + Apply the up projection and RoPE to the query and key. + When sequence packing enabled, the input tensors adopt a packed shape of [t, ...]; + otherwise, they maintain the unpacked shape [b, s, ...]. In subsequent code comments, + we uniformly use [num_tokens, ...] to denote [b, s, ...] or [t, ...] for two cases. + """ + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] + # q: [num_tokens, n * (qk_nope_head_dim + qk_rope_head_dim)] + q, _ = self.q_b_proj(q_compressed) + else: + # q_compressed: [num_tokens, hidden_size] + # q: [num_tokens, n * (qk_nope_head_dim + qk_rope_head_dim)] + q, _ = self.q_proj(q_compressed) + + # q: [num_tokens, n, q_head_dim] + q = q.view( + *q.size()[:-1], + self.num_attention_heads_per_partition, + self.q_head_dim, + ) + + kv_lora_rank = self.config.kv_lora_rank + num_heads = self.num_attention_heads_per_partition + + q_no_pe = q[..., : self.qk_nope_head_dim] + q_pos_emb = q[..., self.qk_nope_head_dim :] + + q_absorb_weight = self.kv_b_proj.weight.reshape( + [kv_lora_rank, num_heads, -1] + )[:, :, : self.qk_nope_head_dim] + q_nope_absorbed = paddle.einsum( + "bshd,lhd->bshl", q_no_pe, q_absorb_weight + ) + + q_pos_emb = apply_rotary_pos_emb( + q_pos_emb, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + config=self.config, + mscale=mscale, + ) + k_pos_emb = apply_rotary_pos_emb( + k_pos_emb.unsqueeze(-2), + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + config=self.config, + mscale=mscale, + ) + + kv_compressed = kv_compressed.unsqueeze(-2) + + query = paddle.concat([q_nope_absorbed, q_pos_emb], axis=-1) + key = paddle.concat([kv_compressed, k_pos_emb], axis=-1) + value = kv_compressed + + return query, key, value, k_pos_emb + + query, key, value, k_pos_emb = qkv_up_proj_and_rope_apply( + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + position_ids, + ) + + return query, key, value, q_compressed, kv_compressed, k_pos_emb diff --git a/src/paddlefleet/transformer/transformer_config.py b/src/paddlefleet/transformer/transformer_config.py index 53943518a..002e1c9f9 100644 --- a/src/paddlefleet/transformer/transformer_config.py +++ b/src/paddlefleet/transformer/transformer_config.py @@ -814,6 +814,34 @@ class TransformerConfig(ModelParallelConfig): with num_chunks=N. Only compatible with tensor_model_parallel_size == 1 (or parallel_output disabled).""" + enable_hy_sparse_attention: bool = False + """Enable the HySparse Attention variant. + + HySparse has the following features: (1) adding a Block Sparse Attention in SWA + layers. (2) KV sharing between full attention and Block Sparse Attention. (3) using + MQA instead of MLA. + """ + + hy_sparse_block_size: int = 64 + """HySparse key block size (``block_B``) used by the TileLang block-score / + block-sparse attention operators. Key columns are grouped into contiguous + blocks of this size (document-relative) for scoring and sparse selection. + + Default 64 follows the HySparse paper (arXiv:2602.03560, Table 1: "Sparse + Attn Block Size = 64" for all 7B/80B configurations).""" + + hy_sparse_topk: int = 16 + """Number of key *blocks* selected per query token in the HySparse block-sparse + branch (the ``topk`` fed to :func:`select_topk_blocks`). The full attention + layer scores all blocks and the top-``hy_sparse_topk`` (shared across the + query group by group-wise max) are attended by the SWA layers' block-sparse + branch. + + Default 16 follows the HySparse paper (arXiv:2602.03560): the paper reports + selection in *tokens* (k = 1024, "Sparse Attn TopK Tokens = 1024"), which maps + to k / block_size = 1024 / 64 = 16 blocks. This field counts blocks, so 16 is + the block-space equivalent of the paper's 1024-token budget.""" + # cache_mla_latents: bool = False #################### @@ -1446,6 +1474,28 @@ def __post_init__(self): f"self.window_attn_skip_freq ({len(self.window_attn_skip_freq)}) " f"must equal num_hidden_layers + num_nextn_predict_layers ({self.num_hidden_layers + self.num_nextn_predict_layers})." ) + # HySparse: MTP layers must be FULL attention layers, never SWA. + # An SWA layer consumes shared_kv (compressed KV latent + block + # indices) produced by an upstream full layer. The MTP boundary + # (MultiTokenPredictionLayer._proj_and_transformer_layer) rebuilds a + # fresh input_dict and does NOT forward shared_key/shared_block_indices + # from the backbone, so an SWA MTP layer would receive shared_kv=[None, + # None] and crash at shared_key.squeeze(2) in the block-sparse branch. + # Fail fast here with a clear message instead. + if self.enable_hy_sparse_attention: + mtp_window_flags = self.window_attn_skip_freq[ + self.num_hidden_layers : + ] + if any(flag != 0 for flag in mtp_window_flags): + raise ValueError( + "When enable_hy_sparse_attention is True, the MTP portion " + f"of window_attn_skip_freq (indices " + f"[{self.num_hidden_layers}:], i.e. {mtp_window_flags}) " + "must be all 0 (full attention layers). MTP layers cannot " + "be sliding-window (SWA) layers because they do not receive " + "the shared KV latent from the backbone across the MTP " + "boundary. Set the MTP entries to 0." + ) if ( self.num_nextn_predict_layers == 0 diff --git a/src/paddlefleet/transformer/transformer_layer.py b/src/paddlefleet/transformer/transformer_layer.py index 85c363ee9..83d61e29f 100644 --- a/src/paddlefleet/transformer/transformer_layer.py +++ b/src/paddlefleet/transformer/transformer_layer.py @@ -949,6 +949,8 @@ def _forward_attention( self.self_attn, DSv4HybridAttention ): extra_kwargs["input_ids"] = input_ids + if "shared_kv" in kwargs: + extra_kwargs["shared_kv"] = kwargs["shared_kv"] if rope_freqs_cis is not None: attention_output_with_bias = self.self_attn( @@ -1423,6 +1425,325 @@ def recompute_handler( return hidden_states +class HySparseTransformerLayer(TransformerLayer): + """Transformer layer with cross-layer KV sharing.""" + + def _mtp_enabled(self, is_mtp): + return ( + self.config.num_nextn_predict_layers is not None + and self.config.num_nextn_predict_layers > 0 + and not is_mtp + and not self.config.mtp_load_weight_only + and not self.config.enable_mtp_magic_send + ) + + def _mtp_split(self, dict_args, is_mtp): + """Split MTP-stacked tensors into main-decoder parts. + + Mirrors ``TransformerLayer.forward`` (the base class does this inline). + MTP concatenates the main hidden state and the ``num_nextn_predict_layers`` + shifted hidden states along the batch dimension; input_ids / position_ids / + masks carry their MTP parts along the seq dimension. Split so the layer + body sees a consistent (main-decoder) batch/seq, mutating ``dict_args`` in + place. Returns a context dict for :meth:`_mtp_restore`, or ``None`` when + MTP is not active. + """ + if not self._mtp_enabled(is_mtp): + return None + n = self.config.num_nextn_predict_layers + ctx = { + "mtp_ids": None, + "mtp_input_ids": None, + "rotary_pos_emb_full": None, + "rotary_pos_cos_full": None, + "rotary_pos_sin_full": None, + "attn_mask_mtp": None, + } + + # hidden_states: split along batch dim -> main + mtp parts + tensor_list = paddle.split(dict_args["hidden_states"], n + 1) + hidden_states = tensor_list[0] + ctx["mtp_input"] = tuple(tensor_list[1:]) + dict_args["hidden_states"] = hidden_states + + # position_ids: split along seq dim + if not self.config.gpt_model_use_experimental_version: + if ( + "position_ids" in dict_args + and dict_args["position_ids"] is not None + ): + position_ids = dict_args["position_ids"] + dict_args["position_ids"] = position_ids[:, :-n] + ctx["mtp_ids"] = position_ids[:, -n:] + + # rotary_pos_emb/cos/sin: trim to main-decoder seq length + if self.config.sequence_parallel: + main_seq_len = ( + hidden_states.shape[0] * self.config.tensor_model_parallel_size + ) + else: + main_seq_len = hidden_states.shape[1] + if ( + "rotary_pos_emb" in dict_args + and dict_args["rotary_pos_emb"] is not None + ): + ctx["rotary_pos_emb_full"] = dict_args["rotary_pos_emb"] + if self.config.sequence_parallel: + dict_args["rotary_pos_emb"] = ctx["rotary_pos_emb_full"][ + :main_seq_len + ] + else: + dict_args["rotary_pos_emb"] = ctx["rotary_pos_emb_full"][ + :, :main_seq_len + ] + if ( + "rotary_pos_cos" in dict_args + and dict_args["rotary_pos_cos"] is not None + ): + ctx["rotary_pos_cos_full"] = dict_args["rotary_pos_cos"] + dict_args["rotary_pos_cos"] = ctx["rotary_pos_cos_full"][ + :, :main_seq_len + ] + if ( + "rotary_pos_sin" in dict_args + and dict_args["rotary_pos_sin"] is not None + ): + ctx["rotary_pos_sin_full"] = dict_args["rotary_pos_sin"] + dict_args["rotary_pos_sin"] = ctx["rotary_pos_sin_full"][ + :, :main_seq_len + ] + + # input_ids: split along seq dim (only when it carries mtp tokens) + if "input_ids" in dict_args and dict_args["input_ids"] is not None: + full_input_ids = dict_args["input_ids"] + seq_lens = hidden_states.shape[ + 0 if self.config.sequence_parallel else 1 + ] + if get_context_parallel_world_size() > 1: + seq_lens *= get_context_parallel_world_size() + if full_input_ids.shape[-1] > seq_lens: + dict_args["input_ids"] = full_input_ids[:, :-n].contiguous() + ctx["mtp_input_ids"] = full_input_ids[:, -n:].contiguous() + + # attn_mask_startend_row_indices: split along seq dim (old dataflow) + if ( + not self.config.experimental_dataflow + and "attn_mask_startend_row_indices" in dict_args + and dict_args["attn_mask_startend_row_indices"] is not None + ): + mask = dict_args["attn_mask_startend_row_indices"] + dict_args["attn_mask_startend_row_indices"] = mask[:, :, :-n, :] + ctx["attn_mask_mtp"] = mask[:, :, -n:, :] + + return ctx + + def _mtp_restore(self, dict_args, output, ctx): + """Re-stack MTP outputs and restore full-length auxiliary tensors. + + Inverse of :meth:`_mtp_split`. Returns the batch-stacked hidden state and + restores ``dict_args`` (input_ids / position_ids / rotary / mask) to full + length for the next layer. + """ + hidden_states_concat = paddle.concat([output, *ctx["mtp_input"]]) + + if not self.config.gpt_model_use_experimental_version: + if "position_ids" in dict_args and ctx["mtp_ids"] is not None: + dict_args["position_ids"] = paddle.concat( + [dict_args["position_ids"], ctx["mtp_ids"]], axis=1 + ) + if ctx["rotary_pos_emb_full"] is not None: + dict_args["rotary_pos_emb"] = ctx["rotary_pos_emb_full"] + if ctx["rotary_pos_cos_full"] is not None: + dict_args["rotary_pos_cos"] = ctx["rotary_pos_cos_full"] + if ctx["rotary_pos_sin_full"] is not None: + dict_args["rotary_pos_sin"] = ctx["rotary_pos_sin_full"] + if ctx["mtp_input_ids"] is not None and "input_ids" in dict_args: + dict_args["input_ids"] = paddle.concat( + [dict_args["input_ids"], ctx["mtp_input_ids"]], axis=1 + ) + if ( + not self.config.experimental_dataflow + and "attn_mask_startend_row_indices" in dict_args + and ctx["attn_mask_mtp"] is not None + ): + dict_args["attn_mask_startend_row_indices"] = paddle.concat( + [ + dict_args["attn_mask_startend_row_indices"], + ctx["attn_mask_mtp"], + ], + axis=2, + ) + return hidden_states_concat + + def forward( + self, + dict_args: dict, + ): + """ + Perform a forward pass through the transformer layer. + + This method calls the core computation of a transformer layer, including + self-attention, cross-attention (if applicable), and feed-forward operations. + """ + # Remove 'dynamic_inference_decode_only' from kwargs if present + # this is only used to uniquely identify decode and non-decode cuda graph + # runners in the cuda graph manager + dict_args.pop("dynamic_inference_decode_only", None) + keys = tuple(dict_args.keys()) + values = tuple(dict_args.values()) + + is_mtp = dict_args.pop("is_mtp", False) + TransformerLayer._skip_mtp_probes = ( + is_mtp # Suppress MD5 probes for MTP passes + ) + + # MTP stacks the main decoder + shifted hidden states along the batch + # dim (see MTP module's paddle.concat(axis=0)); the base + # TransformerLayer.forward splits them off so attention / MoE router see + # a consistent batch, then re-stacks. HySparseTransformerLayer must do + # the same, otherwise the batch-2 stacked hidden reaches the MoE router + # with batch-1 input_ids and trips its shape assertion. Split now, + # restore after the layer body runs. + mtp_ctx = self._mtp_split(dict_args, is_mtp) + + if self.full_recompute or (not has_recovered()): + + def dict_args_get_clone(key): + """Clone is necessary for some args.""" + value = dict_args.get(key, None) + return value.clone() if value is not None else None + + # Mirror the base TransformerLayer recompute path: recompute both + # when full_recompute is set AND inside the RECOVER_STEP recovery + # window (not has_recovered()), so recovering training keeps the + # same reduced activation footprint. Activation-offload settings are + # threaded via _compute_act_offload_kwargs (consumed by recompute). + offload_kwargs = self._compute_act_offload_kwargs() + outputs = recompute( + self._forward_impl, + hidden_states=dict_args["hidden_states"], + attention_mask=dict_args.get("attention_mask", None), + attn_mask_startend_row_indices=dict_args_get_clone( + "attn_mask_startend_row_indices" + ), + context=dict_args.get("context", None), + context_mask=dict_args.get("context_mask", None), + rotary_pos_emb=dict_args_get_clone("rotary_pos_emb"), + rotary_pos_cos=dict_args_get_clone("rotary_pos_cos"), + rotary_pos_sin=dict_args_get_clone("rotary_pos_sin"), + swa_rotary_pos_emb=dict_args_get_clone("swa_rotary_pos_emb"), + swa_rotary_pos_cos=dict_args_get_clone("swa_rotary_pos_cos"), + swa_rotary_pos_sin=dict_args_get_clone("swa_rotary_pos_sin"), + position_ids=dict_args_get_clone("position_ids"), + attention_bias=dict_args.get("attention_bias", None), + packed_seq_params=dict_args.get("packed_seq_params", None), + input_ids=dict_args.get("input_ids", None), + shared_key=dict_args.get("shared_key", None), + shared_block_indices=dict_args.get( + "shared_block_indices", None + ), + **offload_kwargs, + ) + else: + outputs = self._forward_impl(**dict_args) + + if isinstance(outputs, tuple): + output, shared_key, shared_block_indices = outputs + else: + output, shared_key, shared_block_indices = outputs, None, None + + rst = OrderedDict() + rst = {"hidden_states": output} + if mtp_ctx is not None: + # Re-stack main + mtp hidden along batch and restore full-length + # auxiliary tensors (input_ids / position_ids / rotary / mask) into + # dict_args for the next layer. + rst["hidden_states"] = self._mtp_restore(dict_args, output, mtp_ctx) + if shared_key is not None: + rst["shared_key"] = shared_key + rst["shared_block_indices"] = shared_block_indices + rst = {**dict_args, **rst} + return rst + + def _forward_impl( + self, + hidden_states: Tensor, + attention_mask: Tensor | None = None, + attn_mask_startend_row_indices: Tensor | None = None, + context: Tensor | None = None, + context_mask: Tensor | None = None, + rotary_pos_emb: Tensor | None = None, + rotary_pos_cos: Tensor | None = None, + rotary_pos_sin: Tensor | None = None, + swa_rotary_pos_emb: Tensor | None = None, + swa_rotary_pos_cos: Tensor | None = None, + swa_rotary_pos_sin: Tensor | None = None, + position_ids: Tensor | None = None, + attention_bias: Tensor | None = None, + packed_seq_params: PackedSeqParams | None = None, + input_ids: Tensor | None = None, + shared_key: Tensor | None = None, + shared_block_indices: Tensor | None = None, + **kwargs, + ): + timer_name = "moe-mlp" if isinstance(self.mlp, MoELayer) else "mlp" + + # ไฝฟ็”จ็ปŸไธ€็š„ shared_kv ๅ‚ๆ•ฐๅค„็†่พ“ๅ…ฅ่พ“ๅ‡บ: + # 1. ๅฏนไบŽ swa ๅฑ‚ๆ˜ฏ่พ“ๅ…ฅ, ๅชๆถˆ่ดน shared_kv๏ผŒไธ็”Ÿไบง; + # 2. ๅฏนไบŽ full ๅฑ‚ๆ˜ฏ่พ“ๅ‡บ, ๅช็”Ÿไบง shared_kv, ไธๆถˆ่ดน. + if self.self_attn.is_swa: + if shared_key is None or shared_block_indices is None: + raise ValueError( + f"HySparse SWA layer (layer_number={self.layer_number}) " + "requires shared KV latent and top-k block indices from a " + "preceding full-attention layer, but none were provided. " + "Ensure the first backbone attention layer is a full " + "(non-SWA) attention layer so it can produce the shared " + "state -- e.g. set window_attn_skip_freq so that layer 0 " + "is full attention rather than SWA." + ) + shared_kv = [shared_key, shared_block_indices] + else: + shared_kv = [] + + self._log_md5(hidden_states, "input", self.layer_number) + with profile("attn"): + hidden_states, context = self._forward_attention( + hidden_states=hidden_states, + attention_mask=attention_mask, + attn_mask_startend_row_indices=attn_mask_startend_row_indices, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + swa_rotary_pos_emb=swa_rotary_pos_emb, + swa_rotary_pos_cos=swa_rotary_pos_cos, + swa_rotary_pos_sin=swa_rotary_pos_sin, + position_ids=position_ids, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + in_recompute=self.full_recompute, + input_ids=input_ids, + shared_kv=shared_kv, + **kwargs, + ) + assert context is None, ( + "HySparseTransformerLayer doesn't support cross-attention." + ) + self._log_md5(hidden_states, "post_attn_residual", self.layer_number) + with profile(timer_name): + output = self._forward_mlp(hidden_states, input_ids=input_ids) + self._log_md5(output, "layer_output", self.layer_number) + + if (not self.self_attn.is_swa) and shared_kv: + shared_key, shared_block_indices = shared_kv + if self.training and not paddle.is_grad_enabled(): + shared_key.stop_gradient = False + return output, shared_key, shared_block_indices + return output + + class TransformerLayerWithOverlap(TransformerLayer): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) From da832e0436016204467fef94a492c07a055dbd52 Mon Sep 17 00:00:00 2001 From: ForFishes <2282912238@qq.com> Date: Wed, 15 Jul 2026 22:32:19 +0800 Subject: [PATCH 2/2] Add HySparse MQA sparse-attention test coverage - FA4 block-score grad / topk-consistency / multidoc tests - MQA gather DSA, windowed SWA, pipeline, whole-flow precision tests - MQA self-attention module test, build_hysparse_valid_range test - remove obsolete block-attn test superseded by the above --- .../custom_ops/test_hysparse_block_attn.py | 698 ------------------ .../custom_ops/test_hysparse_fa4_grad.py | 131 ++++ .../test_hysparse_fa4_topk_consistency.py | 229 ++++++ .../test_hysparse_mqa_gather_dsa.py | 296 ++++++++ .../test_hysparse_multidoc_block_score.py | 614 +++++++++++++++ .../custom_ops/test_hysparse_pipeline.py | 184 +++++ .../test_hysparse_whole_flow_precision.py | 378 ++++++++++ .../custom_ops/test_hysparse_windowed_swa.py | 429 +++++++++++ .../test_mqa_self_attention.py | 330 +++++++++ .../test_build_hysparse_valid_range.py | 189 +++++ 10 files changed, 2780 insertions(+), 698 deletions(-) delete mode 100644 tests/single_card_tests/custom_ops/test_hysparse_block_attn.py create mode 100644 tests/single_card_tests/custom_ops/test_hysparse_fa4_grad.py create mode 100644 tests/single_card_tests/custom_ops/test_hysparse_fa4_topk_consistency.py create mode 100644 tests/single_card_tests/custom_ops/test_hysparse_mqa_gather_dsa.py create mode 100644 tests/single_card_tests/custom_ops/test_hysparse_multidoc_block_score.py create mode 100644 tests/single_card_tests/custom_ops/test_hysparse_pipeline.py create mode 100644 tests/single_card_tests/custom_ops/test_hysparse_whole_flow_precision.py create mode 100644 tests/single_card_tests/custom_ops/test_hysparse_windowed_swa.py create mode 100644 tests/single_card_tests/test_mqa_self_attention.py create mode 100644 tests/single_card_tests/transformer/test_build_hysparse_valid_range.py diff --git a/tests/single_card_tests/custom_ops/test_hysparse_block_attn.py b/tests/single_card_tests/custom_ops/test_hysparse_block_attn.py deleted file mode 100644 index b5be1061f..000000000 --- a/tests/single_card_tests/custom_ops/test_hysparse_block_attn.py +++ /dev/null @@ -1,698 +0,0 @@ -# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Accuracy tests for the HySparse (MQA) block-attention TileLang operators. - -Validates, against the naive Paddle reference (ๆ•ฃ็ฎ—ๅญ): -* Block-score attention (Algorithm 1, MQA): full attention with a single shared - K/V head + block-max scores, fwd & bwd. -* Block-sparse attention (MQA gather): per-query block-sparse attn over a single - shared K/V head, fwd & bwd. -* The end-to-end pipeline (block-score -> TopK -> block-sparse). - -Both causal and causal+document masking are covered via ``valid_range``. -""" - -import unittest - -import numpy as np -import paddle - -paddle.enable_compat(scope={"tilelang"}, silent=True) - -from paddlefleet.tilelang_ops.hysparse.block_score_attn import ( - block_score_mqa_attn_fwd, - block_scores_from_logit, -) -from paddlefleet.tilelang_ops.hysparse.block_score_attn_bwd import ( - block_score_mqa_bwd_interface, -) -from paddlefleet.tilelang_ops.hysparse.block_sparse_attn_mqa import ( - block_sparse_mqa_attn_fwd, -) -from paddlefleet.tilelang_ops.hysparse.block_sparse_attn_mqa_bwd import ( - block_sparse_mqa_bwd_interface, -) -from paddlefleet.tilelang_ops.hysparse.pipeline import ( - hysparse_forward_mqa, -) -from paddlefleet.tilelang_ops.hysparse.reference import ( - make_causal_valid_range, - ref_block_score_attn_mqa, - ref_block_sparse_attn_mqa, -) - - -def _cuda_or_skip(testcase): - if not paddle.device.is_compiled_with_cuda(): - testcase.skipTest("CUDA build of Paddle is required") - if paddle.device.cuda.device_count() == 0: - testcase.skipTest("No CUDA device available") - - -def _maxerr(a, b): - a = a.astype("float32").numpy() - b = b.astype("float32").numpy() - return float(np.abs(a - b).max()) - - -def _cos(a, b): - a = a.astype("float32").numpy().ravel() - b = b.astype("float32").numpy().ravel() - denom = float(np.linalg.norm(a) * np.linalg.norm(b)) + 1e-12 - return float((a * b).sum() / denom) - - -def _rand_qkv_mqa(b, s, h, d, seed=0): - """Query with H heads; K/V a single shared head [B, S, D].""" - paddle.seed(seed) - q = paddle.randn([b, s, h, d], dtype="bfloat16") - k = paddle.randn([b, s, d], dtype="bfloat16") - v = paddle.randn([b, s, d], dtype="bfloat16") - return q, k, v - - -def _grad_ref_qkv(q, k, v): - """float32 leaf copies of q/k/v for autograd reference gradients.""" - qf = q.astype("float32") - qf.stop_gradient = False - kf = k.astype("float32") - kf.stop_gradient = False - vf = v.astype("float32") - vf.stop_gradient = False - return qf, kf, vf - - -# tolerances: bf16 matmul + fp32 accumulation -OUT_ATOL = 5e-2 -LSE_ATOL = 1e-3 -GRAD_ATOL = 8e-2 - - -class TestBlockScoreMQA(unittest.TestCase): - """Block-score attention (MQA): shared-K/V full attn + block-max scores.""" - - def _run(self, doc_lengths=None, block_B=64, H=8, D=64): - _cuda_or_skip(self) - B, S = 2, 256 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=1) - vr = make_causal_valid_range(S, batch=B, doc_lengths=doc_lengths) - sm_scale = D**-0.5 - - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_out, ref_lse, ref_sblk = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm_scale, block_B=block_B - ) - - out, lse, block_logit = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm_scale, block_B=block_B - ) - sblk = block_scores_from_logit(block_logit, lse, sm_scale) - - self.assertLess(_maxerr(out, ref_out), OUT_ATOL) - self.assertLess(_maxerr(lse, ref_lse), LSE_ATOL) - self.assertLess(_maxerr(sblk, ref_sblk), 1e-3) - - # backward - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_out * do.astype("float32")).sum().backward() - dq, dk, dv = block_score_mqa_bwd_interface( - q, k, v, out, do, lse, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(dq, qf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dk, kf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dv, vf.grad), GRAD_ATOL) - - def test_causal(self): - self._run(doc_lengths=None) - - def test_causal_document(self): - self._run(doc_lengths=[96, 160]) - - def test_block_b_32(self): - # smaller key block size (must still divide the auto-picked block_N). - self._run(doc_lengths=[96, 160], block_B=32) - - def test_block_b_128(self): - # larger key block size exercises a different tiling / mask layout. - self._run(doc_lengths=None, block_B=128) - - def test_single_query_head(self): - # H=1: heads-on-M padding edge (PH = max(pow2(1), 16) = 16). - self._run(doc_lengths=None, H=1) - - def test_backward_block_n_invariance(self): - """The block-score backward sub-tiles the key dim by ``block_N``; any - valid ``block_N`` dividing ``block_B`` must give the SAME gradients - (only the fp-accumulation / per-sub-tile bf16 recast order changes). - Validates the sub-tiling introduced for the D=576 MB=64 backward. - """ - _cuda_or_skip(self) - B, S, H, D, block_B = 2, 256, 8, 64, 64 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=13) - vr = make_causal_valid_range(S, batch=B, doc_lengths=[96, 160]) - sm = D**-0.5 - - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_out, _, _ = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm, block_B=block_B - ) - out, lse, _ = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm, block_B=block_B - ) - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_out * do.astype("float32")).sum().backward() - - # ratio=1 (block_N == block_B) vs ratio=2 (block_N = block_B // 2) - dq1, dk1, dv1 = block_score_mqa_bwd_interface( - q, - k, - v, - out, - do, - lse, - vr, - sm_scale=sm, - block_B=block_B, - block_M=64, - block_N=64, - ) - dq2, dk2, dv2 = block_score_mqa_bwd_interface( - q, - k, - v, - out, - do, - lse, - vr, - sm_scale=sm, - block_B=block_B, - block_M=64, - block_N=32, - ) - # the two tilings must agree tightly (differences only from fp order) - self.assertLess(_maxerr(dq1, dq2), 1e-2) - self.assertLess(_maxerr(dk1, dk2), 1e-2) - self.assertLess(_maxerr(dv1, dv2), 1e-2) - # and both must match the fp32 reference - for dq, dk, dv in ((dq1, dk1, dv1), (dq2, dk2, dv2)): - self.assertLess(_maxerr(dq, qf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dk, kf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dv, vf.grad), GRAD_ATOL) - - def test_backward_packed_deep_docs(self): - """Backward with several packed documents whose later tiles start deep - into the sequence (bos >> 0). Exercises the document-tight key-block - window (leading-block skip) in the block-score backward: skipped blocks - must contribute nothing so dq/dk/dv still match the fp32 reference. - """ - _cuda_or_skip(self) - B, H, D, block_B = 1, 8, 64, 64 - # deep docs: three documents, each many blocks long and not aligned to - # block_B, so per-tile min(bos) lands well past column 0. - doc_lengths = [130, 200, 180] - S = sum(doc_lengths) - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=17) - vr = make_causal_valid_range(S, batch=B, doc_lengths=doc_lengths) - sm_scale = D**-0.5 - - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_out, _, _ = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm_scale, block_B=block_B - ) - out, lse, _ = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm_scale, block_B=block_B - ) - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_out * do.astype("float32")).sum().backward() - dq, dk, dv = block_score_mqa_bwd_interface( - q, k, v, out, do, lse, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(dq, qf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dk, kf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dv, vf.grad), GRAD_ATOL) - - def test_empty_rows(self): - """A query row with no valid key (bos == eos) must give 0 output, - -inf lse, and 0 gradient -- matching the guarded reference. - """ - _cuda_or_skip(self) - B, S, H, D, block_B = 1, 128, 8, 64, 64 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=31) - vr = make_causal_valid_range(S, batch=B).clone() - empty = 40 - vr[:, empty, 0] = empty # bos == eos -> empty half-open range - vr[:, empty, 1] = empty - vr = vr.contiguous() - sm = D**-0.5 - - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_out, ref_lse, _ = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm, block_B=block_B - ) - out, lse, _ = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm, block_B=block_B - ) - # empty row: zero output and -inf lse - self.assertLess(float(out[:, empty].abs().max()), 1e-6) - self.assertTrue(bool((lse[:, empty] == float("-inf")).all().item())) - # non-empty rows still match the reference - self.assertLess(_maxerr(out, ref_out), OUT_ATOL) - - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_out * do.astype("float32")).sum().backward() - dq, dk, dv = block_score_mqa_bwd_interface( - q, k, v, out, do, lse, vr, sm_scale=sm, block_B=block_B - ) - self.assertLess(float(dq[:, empty].abs().max()), 1e-6) - self.assertLess(_maxerr(dq, qf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dk, kf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dv, vf.grad), GRAD_ATOL) - - -class TestBlockSparseMQA(unittest.TestCase): - """Block-sparse attention (MQA gather): gathers only selected blocks.""" - - def _make_indices(self, vr, block_B, nsel, seed=0): - """Random **document-relative** block ids per query token. - - Relative block ``j`` of a token spans key columns - ``[bos + j*block_B, bos + (j+1)*block_B)``; valid blocks are - ``0 .. ceil((eos-bos)/block_B) - 1``. - """ - rng = np.random.default_rng(seed) - vr_np = vr.numpy() - B, S, _ = vr_np.shape - idx = np.full([B, S, nsel], -1, dtype=np.int32) - for b in range(B): - for t in range(S): - bos = int(vr_np[b, t, 0]) - eos = int(vr_np[b, t, 1]) - n = (eos - bos + block_B - 1) // block_B - cand = list(range(n)) - chosen = ( - cand - if len(cand) <= nsel - else list(rng.choice(cand, size=nsel, replace=False)) - ) - for jj, blk in enumerate(chosen): - idx[b, t, jj] = blk - return paddle.to_tensor(idx, dtype="int32") - - def _run(self, S=256, doc_lengths=None, block_B=64): - _cuda_or_skip(self) - B, H, D = 2, 8, 64 - nsel = 3 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=5) - vr = make_causal_valid_range(S, batch=B, doc_lengths=doc_lengths) - idx = self._make_indices(vr, block_B, nsel) - sm_scale = D**-0.5 - - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_out, ref_lse = ref_block_sparse_attn_mqa( - qf, kf, vf, idx, vr, sm_scale=sm_scale, block_B=block_B - ) - out, lse = block_sparse_mqa_attn_fwd( - q, k, v, idx, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(out, ref_out), OUT_ATOL) - finite = np.isfinite(ref_lse.numpy()) - lse_err = np.abs(lse.numpy()[finite] - ref_lse.numpy()[finite]).max() - self.assertLess(float(lse_err), LSE_ATOL) - - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_out * do.astype("float32")).sum().backward() - dq, dk, dv = block_sparse_mqa_bwd_interface( - q, k, v, out, do, lse, idx, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(dq, qf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dk, kf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dv, vf.grad), GRAD_ATOL) - - def test_causal(self): - self._run(doc_lengths=None) - - def test_causal_document(self): - self._run(doc_lengths=[96, 160]) - - def test_unaligned_seqlen(self): - # S_kv not a multiple of block_B exercises the K/V padding path. - self._run(S=200, doc_lengths=None) - - def test_block_b_32(self): - self._run(doc_lengths=[96, 160], block_B=32) - - def test_block_b_128(self): - self._run(S=384, doc_lengths=None, block_B=128) - - def test_all_blocks_equals_full(self): - """Selecting every block must reproduce full-range attention (block-score).""" - _cuda_or_skip(self) - B, S, H, D = 2, 256, 8, 64 - block_B = 64 - num_blocks = (S + block_B - 1) // block_B - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=3) - vr = make_causal_valid_range(S, batch=B) - sm_scale = D**-0.5 - - idx_all = ( - paddle.arange(num_blocks, dtype="int32") - .reshape([1, 1, num_blocks]) - .expand([B, S, num_blocks]) - .contiguous() - ) - out_a, _ = block_sparse_mqa_attn_fwd( - q, k, v, idx_all, vr, sm_scale=sm_scale, block_B=block_B - ) - out_b, _, _ = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(out_a, out_b), 1e-3) - - -class TestPipelineMQA(unittest.TestCase): - """End-to-end MQA block-score -> TopK -> block-sparse pipeline.""" - - def test_pipeline_matches_reference(self): - _cuda_or_skip(self) - B, S, H, D = 2, 256, 8, 64 - block_B = 64 - topk = 3 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=6) - vr = make_causal_valid_range(S, batch=B) - sm_scale = D**-0.5 - - sparse_out, sparse_lse, indices, full_out, full_lse = ( - hysparse_forward_mqa( - q, k, v, vr, topk, sm_scale=sm_scale, block_B=block_B - ) - ) - - # full-attention branch must match the MQA block-score reference - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_full_out, _, _ = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(full_out, ref_full_out), OUT_ATOL) - - # sparse branch must match the MQA reference given the SAME indices - ref_sparse_out, _ = ref_block_sparse_attn_mqa( - qf, kf, vf, indices, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(sparse_out, ref_sparse_out), OUT_ATOL) - - # indices are shared across heads, in range, causal - idx_np = indices.numpy() - qblk = (np.arange(S) // block_B).reshape([1, S, 1]) - picked = idx_np >= 0 - self.assertTrue( - bool( - ( - idx_np[picked] - <= np.broadcast_to(qblk, idx_np.shape)[picked] - ).all() - ) - ) - - def test_topk_exceeds_num_blocks(self): - # topk larger than the number of key blocks must still return a stable - # [B, S, topk] index tensor, padding the surplus slots with -1. - _cuda_or_skip(self) - B, S, H, D = 1, 96, 4, 64 - block_B = 64 # num_blocks = ceil(96/64) = 2 - topk = 5 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=11) - vr = make_causal_valid_range(S, batch=B) - _, _, indices, _, _ = hysparse_forward_mqa( - q, k, v, vr, topk, block_B=block_B - ) - self.assertEqual(list(indices.shape), [B, S, topk]) - # at most num_blocks real ids per row; the rest are -1 padding - idx_np = indices.numpy() - self.assertTrue(bool((idx_np < 2).all())) - self.assertTrue(bool((idx_np[:, -1, :] >= -1).all())) - self.assertTrue(bool((idx_np == -1).any())) - - -class TestDocEquivalenceMQA(unittest.TestCase): - """Documents packed with a document mask must give the SAME sparse result - as running each document standalone -- even when document lengths are not - multiples of block_B. This is the whole point of the document-relative - block coordinates: block selection is anchored at each document's start, so - a packed batch and per-document batches select identical (relative) blocks. - """ - - def _run(self, doc_lengths, topk, block_B=64, H=8, D=64, seed=11): - _cuda_or_skip(self) - S = sum(doc_lengths) - q, k, v = _rand_qkv_mqa(1, S, H, D, seed=seed) - vr = make_causal_valid_range(S, batch=1, doc_lengths=doc_lengths) - sm_scale = D**-0.5 - - packed_out, _, _, _, _ = hysparse_forward_mqa( - q, k, v, vr, topk, sm_scale=sm_scale, block_B=block_B - ) - - start = 0 - for L in doc_lengths: - end = start + L - q_d = q[:, start:end].contiguous() - k_d = k[:, start:end].contiguous() - v_d = v[:, start:end].contiguous() - vr_d = make_causal_valid_range(L, batch=1) - d_out, _, _, _, _ = hysparse_forward_mqa( - q_d, k_d, v_d, vr_d, topk, sm_scale=sm_scale, block_B=block_B - ) - err = _maxerr(packed_out[:, start:end], d_out) - self.assertLess(err, 2e-3, f"doc [{start},{end}) err={err}") - start = end - - def test_two_docs_select_all(self): - # topk >= #blocks/doc -> every valid block selected in both layouts - self._run(doc_lengths=[96, 160], topk=3) - - def test_two_docs_true_sparsity(self): - # topk < #blocks in the long doc -> genuine subset selection must agree - self._run(doc_lengths=[100, 200], topk=2) - - def test_three_docs(self): - self._run(doc_lengths=[64, 96, 130], topk=2) - - -class TestHeadDim576MQA(unittest.TestCase): - """MLA-shaped head_dim=576 (non-power-of-2), num_heads=64. - - Validates the ops run and stay accurate at the large MLA head dim. Output - is checked against the fp32 reference with the usual absolute tolerance; - gradients are checked by cosine similarity because absolute errors grow - with the (large) gradient magnitudes at D=576. - """ - - def test_score_fwd_bwd(self): - _cuda_or_skip(self) - B, S, H, D = 1, 192, 64, 576 - block_B = 64 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=7) - vr = make_causal_valid_range(S, batch=B) - sm_scale = D**-0.5 - - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_out, ref_lse, ref_sblk = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm_scale, block_B=block_B - ) - out, lse, block_logit = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm_scale, block_B=block_B - ) - sblk = block_scores_from_logit(block_logit, lse, sm_scale) - self.assertLess(_maxerr(out, ref_out), OUT_ATOL) - self.assertLess(_maxerr(lse, ref_lse), LSE_ATOL) - self.assertLess(_maxerr(sblk, ref_sblk), 1e-3) - - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_out * do.astype("float32")).sum().backward() - dq, dk, dv = block_score_mqa_bwd_interface( - q, k, v, out, do, lse, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertGreater(_cos(dq, qf.grad), 0.999) - self.assertGreater(_cos(dk, kf.grad), 0.999) - self.assertGreater(_cos(dv, vf.grad), 0.999) - - def test_sparse_fwd_bwd_and_pipeline(self): - _cuda_or_skip(self) - B, S, H, D = 1, 192, 64, 576 - block_B, topk = 64, 2 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=8) - vr = make_causal_valid_range(S, batch=B) - sm_scale = D**-0.5 - - sparse_out, sparse_lse, idx, full_out, _ = hysparse_forward_mqa( - q, k, v, vr, topk, sm_scale=sm_scale, block_B=block_B - ) - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_full, _, _ = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(full_out, ref_full), OUT_ATOL) - ref_sparse, _ = ref_block_sparse_attn_mqa( - qf, kf, vf, idx, vr, sm_scale=sm_scale, block_B=block_B - ) - self.assertLess(_maxerr(sparse_out, ref_sparse), OUT_ATOL) - - # gather backward with head-tiled M (block_H) for the large head dim - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_sparse * do.astype("float32")).sum().backward() - dq, dk, dv = block_sparse_mqa_bwd_interface( - q, - k, - v, - sparse_out, - do, - sparse_lse, - idx, - vr, - sm_scale=sm_scale, - block_B=block_B, - ) - self.assertGreater(_cos(dq, qf.grad), 0.999) - self.assertGreater(_cos(dk, kf.grad), 0.999) - self.assertGreater(_cos(dv, vf.grad), 0.999) - - -class TestBenchmarkSmoke(unittest.TestCase): - """The benchmark harness runs and yields finite latencies for every op.""" - - def test_bench_config_small(self): - _cuda_or_skip(self) - from paddlefleet.tilelang_ops.hysparse.benchmark import bench_config - - res = bench_config( - B=1, - S=512, - H=16, - D=64, - block_B=64, - topk=4, - do_bwd=True, - iters=1, - ) - for key in ( - "score_fwd", - "score_bwd", - "sparse_fwd", - "sparse_bwd", - "pipeline", - "pipeline_peak_gb", - ): - self.assertIn(key, res) - self.assertTrue(np.isfinite(res[key])) - self.assertGreater(res[key], 0.0) - - -class TestBatchEquivalenceMQA(unittest.TestCase): - """Batch size > 1: batch elements are fully independent, and each row may - carry its own document layout. A packed B=2 batch must reproduce (a) each - row run standalone as B=1, and (b) each document run standalone -- even - when the two rows pack completely different document lengths. - """ - - def _vr_multi(self, S, layouts): - """Stack per-row causal+document valid_ranges into [B, S, 2].""" - rows = [ - make_causal_valid_range(S, batch=1, doc_lengths=dl) - for dl in layouts - ] - return paddle.concat(rows, axis=0).contiguous() - - def test_row_isolation(self): - # B=2 packed == two independent B=1 runs (different per-row layouts). - _cuda_or_skip(self) - S, H, D, block_B, topk = 256, 8, 64, 64, 2 - layouts = [[100, 156], [96, 64, 96]] - q, k, v = _rand_qkv_mqa(2, S, H, D, seed=21) - vr = self._vr_multi(S, layouts) - sm = D**-0.5 - - packed, _, _, _, _ = hysparse_forward_mqa( - q, k, v, vr, topk, sm_scale=sm, block_B=block_B - ) - for b, dl in enumerate(layouts): - vr_b = make_causal_valid_range(S, batch=1, doc_lengths=dl) - solo, _, _, _, _ = hysparse_forward_mqa( - q[b : b + 1], - k[b : b + 1], - v[b : b + 1], - vr_b, - topk, - sm_scale=sm, - block_B=block_B, - ) - self.assertLess(_maxerr(packed[b : b + 1], solo), 2e-3) - - def test_per_row_doc_equivalence(self): - # Each (row, doc) slice equals that document run on its own, with the - # two rows using different, non-block-aligned document layouts. - _cuda_or_skip(self) - S, H, D, block_B, topk = 256, 8, 64, 64, 2 - layouts = [[100, 156], [96, 64, 96]] - q, k, v = _rand_qkv_mqa(2, S, H, D, seed=22) - vr = self._vr_multi(S, layouts) - sm = D**-0.5 - - packed, _, _, _, _ = hysparse_forward_mqa( - q, k, v, vr, topk, sm_scale=sm, block_B=block_B - ) - for b, docs in enumerate(layouts): - start = 0 - for L in docs: - end = start + L - d_out, _, _, _, _ = hysparse_forward_mqa( - q[b : b + 1, start:end].contiguous(), - k[b : b + 1, start:end].contiguous(), - v[b : b + 1, start:end].contiguous(), - make_causal_valid_range(L, batch=1), - topk, - sm_scale=sm, - block_B=block_B, - ) - err = _maxerr(packed[b : b + 1, start:end], d_out) - self.assertLess(err, 2e-3, f"row{b} [{start},{end}) err={err}") - start = end - - def test_backward_batched(self): - # dq/dk/dv for a B=2 batch match the fp32 reference (block-score). - _cuda_or_skip(self) - B, S, H, D, block_B = 2, 256, 8, 64, 64 - q, k, v = _rand_qkv_mqa(B, S, H, D, seed=23) - vr = make_causal_valid_range(S, batch=B, doc_lengths=[96, 160]) - sm = D**-0.5 - - qf, kf, vf = _grad_ref_qkv(q, k, v) - ref_out, _, _ = ref_block_score_attn_mqa( - qf, kf, vf, vr, sm_scale=sm, block_B=block_B - ) - out, lse, _ = block_score_mqa_attn_fwd( - q, k, v, vr, sm_scale=sm, block_B=block_B - ) - do = paddle.randn([B, S, H, D], dtype="bfloat16") - (ref_out * do.astype("float32")).sum().backward() - dq, dk, dv = block_score_mqa_bwd_interface( - q, k, v, out, do, lse, vr, sm_scale=sm, block_B=block_B - ) - self.assertLess(_maxerr(dq, qf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dk, kf.grad), GRAD_ATOL) - self.assertLess(_maxerr(dv, vf.grad), GRAD_ATOL) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/single_card_tests/custom_ops/test_hysparse_fa4_grad.py b/tests/single_card_tests/custom_ops/test_hysparse_fa4_grad.py new file mode 100644 index 000000000..8dd78ab11 --- /dev/null +++ b/tests/single_card_tests/custom_ops/test_hysparse_fa4_grad.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Backward / autograd coverage for the FA4-fused full block-score attention +(:mod:`paddlefleet.tilelang_ops.hysparse.block_score_fa4`). + +The consistency test (``test_hysparse_fa4_topk_consistency``) only drives the +forward path. Here we exercise the ``_BlockScoreFA4Attn`` PyLayer *backward* +(FA4 sm100 bwd kernel) and the ``sm_scale=None`` default, verifying the +attention-output gradient against a plain dense-attention reference (the +``block_logit`` / ``lse`` outputs are non-differentiable and carry no grad). + +FA4 block-score fusion runs only on SM 10.x (Blackwell); the test skips +otherwise. +""" + +import math +import unittest + +import paddle + +paddle.enable_compat(scope={"tilelang"}, silent=True) + +_NEG_INF = float("-inf") + + +def _sm100_or_skip(tc): + if not paddle.device.is_compiled_with_cuda(): + tc.skipTest("CUDA build of Paddle required") + if paddle.device.cuda.device_count() == 0: + tc.skipTest("no CUDA device available") + if paddle.device.cuda.get_device_capability()[0] != 10: + tc.skipTest("FA4 block-score fusion requires SM 10.x (Blackwell)") + + +def _ref_causal_attn(q, k, v, sm_scale): + """Differentiable dense causal MHA reference. q/k/v [B,S,H,D] fp32.""" + b, s, h, d = q.shape + sk = k.shape[1] + qf = q.transpose([0, 2, 1, 3]) + kf = k.transpose([0, 2, 1, 3]) + vf = v.transpose([0, 2, 1, 3]) + logits = paddle.matmul(qf, kf, transpose_y=True) * sm_scale # [B,H,S,Sk] + row = paddle.arange(s).reshape([s, 1]) + col = paddle.arange(sk).reshape([1, sk]) + masked = (col > row + (sk - s)).reshape([1, 1, s, sk]) + logits = paddle.where(masked, paddle.full_like(logits, _NEG_INF), logits) + p = paddle.nn.functional.softmax(logits, axis=-1) + out = paddle.matmul(p, vf) # [B,H,S,D] + return out.transpose([0, 2, 1, 3]) # [B,S,H,D] + + +def _cos(a, b): + import numpy as np + + af = a.astype("float32").numpy().reshape(-1) + bf = b.astype("float32").numpy().reshape(-1) + denom = (np.linalg.norm(af) * np.linalg.norm(bf)) + 1e-12 + return float(np.dot(af, bf) / denom) + + +class TestBlockScoreFA4Backward(unittest.TestCase): + def test_backward_matches_dense_reference(self): + _sm100_or_skip(self) + from paddlefleet.tilelang_ops.hysparse import block_score_fa4_attn_fwd + + b, s, h, d = 1, 256, 8, 64 + paddle.seed(2026) + q = paddle.randn([b, s, h, d], dtype="bfloat16") + k = paddle.randn([b, s, h, d], dtype="bfloat16") + v = paddle.randn([b, s, h, d], dtype="bfloat16") + sm_scale = 1.0 / math.sqrt(d) + + qf = q.detach() + kf = k.detach() + vf = v.detach() + qf.stop_gradient = False + kf.stop_gradient = False + vf.stop_gradient = False + out, lse, block_logit = block_score_fa4_attn_fwd( + qf, kf, vf, sm_scale=sm_scale, block_B=64, causal=True + ) + g = paddle.randn(out.shape, dtype="float32").astype("bfloat16") + out.backward(g) + + qr = q.astype("float32").detach() + kr = k.astype("float32").detach() + vr = v.astype("float32").detach() + qr.stop_gradient = False + kr.stop_gradient = False + vr.stop_gradient = False + ref = _ref_causal_attn(qr, kr, vr, sm_scale) + ref.backward(g.astype("float32")) + + self.assertGreater(_cos(qf.grad, qr.grad), 0.99) + self.assertGreater(_cos(kf.grad, kr.grad), 0.99) + self.assertGreater(_cos(vf.grad, vr.grad), 0.99) + + def test_default_sm_scale(self): + _sm100_or_skip(self) + from paddlefleet.tilelang_ops.hysparse import block_score_fa4_attn_fwd + + b, s, h, d = 1, 128, 4, 64 + paddle.seed(7) + q = paddle.randn([b, s, h, d], dtype="bfloat16") + k = paddle.randn([b, s, h, d], dtype="bfloat16") + v = paddle.randn([b, s, h, d], dtype="bfloat16") + # sm_scale=None -> defaults to d ** -0.5. + out_default, _, _ = block_score_fa4_attn_fwd( + q, k, v, block_B=64, causal=True + ) + out_explicit, _, _ = block_score_fa4_attn_fwd( + q, k, v, sm_scale=d**-0.5, block_B=64, causal=True + ) + self.assertEqual(list(out_default.shape), [b, s, h, d]) + self.assertGreater(_cos(out_default, out_explicit), 0.999) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/custom_ops/test_hysparse_fa4_topk_consistency.py b/tests/single_card_tests/custom_ops/test_hysparse_fa4_topk_consistency.py new file mode 100644 index 000000000..38d063a03 --- /dev/null +++ b/tests/single_card_tests/custom_ops/test_hysparse_fa4_topk_consistency.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HySparse: FA4-fused block-score TopK == scatter (eq.3) reference TopK. + +This test guards the *้‡็บฒ (scale) contract* fixed in flash-attention PR #164: +the FA4 sm100 forward now emits the **scaled** per-block max logit +(``softmax_scale * q.k`` -- the exact value fed into softmax), and the host +:func:`select_topk_blocks` therefore recovers the eq.(3) block score as +``exp(block_logit - lse)`` (no host-side ``* sm_scale``). This test verifies the +whole path is self-consistent, from *computing the score* to *deriving TopK +from the score*, against an independent "scatter" reference operator that +implements HySparse paper (arXiv 2602.03560) eq.(3) with plain matmul + softmax: + +* the FA4 fused ``block_logit`` matches the reference scaled per-block max logit + on every finite (unmasked) entry, and the masked pattern matches exactly; +* the eq.(3) block scores ``exp(block_logit - lse)`` match; and +* the per-query TopK block indices selected from the two score sources are + identical. + +The scatter reference is the model: for query ``i`` and key block ``b`` it takes +the max SCALED logit over the block's causally-valid columns, converts it to the +block's max attention weight ``exp(max_logit - lse)`` (eq.3 block importance), +aggregates across the query-group heads by a group-wise max, and TopK-selects. + +FA4 block-score fusion runs only on SM 10.x (Blackwell); the test skips +otherwise. +""" + +import math +import unittest + +import paddle + +_NEG_INF = float("-inf") + + +def _sm100_or_skip(testcase): + if not paddle.device.is_compiled_with_cuda(): + testcase.skipTest("CUDA build of Paddle is required") + if paddle.device.cuda.device_count() == 0: + testcase.skipTest("No CUDA device available") + major = paddle.device.cuda.get_device_capability()[0] + if major != 10: + testcase.skipTest( + f"FA4 block-score fusion requires SM 10.x (Blackwell); got SM {major}.x" + ) + + +def _num_blocks(seqlen_k, block_B): + return (seqlen_k + block_B - 1) // block_B + + +def _causal_valid_range(b, s): + """Single-document causal ``valid_range`` [B, S, 2]: bos=0, eos=t+1.""" + eos = ( + paddle.arange(1, s + 1, dtype="int32") + .reshape([1, s, 1]) + .expand([b, s, 1]) + ) + bos = paddle.zeros([b, s, 1], dtype="int32") + return paddle.concat([bos, eos], axis=-1).contiguous() + + +def _scatter_reference(q, k, sm_scale, block_B): + """Scatter (eq.3) reference: scaled per-block max logit + natural-log LSE. + + q, k: [B, S, H, D] bf16 (single-document causal, S == S_kv). + Returns (block_logit_ref [B,H,S,nb] fp32, lse_ref [B,S,H] fp32) computed with + a full-precision matmul + softmax -- the independent model for the FA4 path. + """ + b, s, h, d = q.shape + sk = k.shape[1] + nb = _num_blocks(sk, block_B) + qf = q.astype("float32").transpose([0, 2, 1, 3]) # [B,H,S,D] + kf = k.astype("float32").transpose([0, 2, 1, 3]) # [B,H,Sk,D] + logits = paddle.matmul(qf, kf, transpose_y=True) * sm_scale # [B,H,S,Sk] + # bottom-right aligned causal (== standard causal when s == sk). + row = paddle.arange(s).reshape([s, 1]) + col = paddle.arange(sk).reshape([1, sk]) + masked = (col > row + (sk - s)).reshape([1, 1, s, sk]) + logits = paddle.where(masked, paddle.full_like(logits, _NEG_INF), logits) + # natural-log LSE over keys -> [B,H,S] -> [B,S,H] (pipeline layout). + lse_bhs = paddle.logsumexp(logits, axis=-1) # [B,H,S] + lse_ref = lse_bhs.transpose([0, 2, 1]).contiguous() # [B,S,H] + # per-block max of the scaled logit -> [B,H,S,nb]. + pad = nb * block_B - sk + if pad > 0: + logits = paddle.concat( + [logits, paddle.full([b, h, s, pad], _NEG_INF, dtype="float32")], + axis=-1, + ) + logits = logits.reshape([b, h, s, nb, block_B]) + block_logit_ref = logits.max(axis=-1) # [B,H,S,nb] + return block_logit_ref, lse_ref + + +def _maxerr_finite(got, ref): + """Masked-pattern match + max abs diff on finite (unmasked) entries.""" + import numpy as np + + got_np = got.astype("float32").numpy() + ref_np = ref.astype("float32").numpy() + # -inf blocks may surface as true -inf or as -FLT_MAX from a reduce over an + # all -inf slice; treat anything below -1e30 as "masked". + thr = -1e30 + got_masked = got_np <= thr + ref_masked = ref_np <= thr + pattern_mismatch = int((got_masked != ref_masked).sum()) + finite = ~ref_masked + diff = ( + np.abs(got_np[finite] - ref_np[finite]) + if finite.any() + else np.array([0.0]) + ) + return pattern_mismatch, float(diff.max()) + + +class TestHySparseFA4TopKConsistency(unittest.TestCase): + # HySparse decompressed-MHA full-attn config knobs (kept small so the test + # is fast; head dim 256 matches the real d_n(192)+d_r(64) MLA head). + BLOCK_B = 64 + TOPK = 4 + + def _run(self, b, s, h, d, dv, seed=2026): + _sm100_or_skip(self) + from paddlefleet.tilelang_ops.hysparse import ( + block_score_fa4_attn_fwd, + select_topk_blocks, + ) + from paddlefleet.tilelang_ops.hysparse.pipeline import ( + block_scores_from_logit, + ) + + paddle.seed(seed) + q = paddle.randn([b, s, h, d], dtype="bfloat16") + k = paddle.randn([b, s, h, d], dtype="bfloat16") + v = paddle.randn([b, s, h, dv], dtype="bfloat16") + sm_scale = 1.0 / math.sqrt(d) + valid_range = _causal_valid_range(b, s) + + # ---- FA4-fused path: compute score, then TopK from score. ---- + out, lse, block_logit = block_score_fa4_attn_fwd( + q, + k, + v, + valid_range=valid_range, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + causal=True, + ) + fa4_idx = select_topk_blocks( + block_logit, lse, valid_range, self.TOPK, self.BLOCK_B + ) + + # ---- Scatter (eq.3) reference: same two stages, plain matmul. ---- + block_logit_ref, lse_ref = _scatter_reference( + q, k, sm_scale, self.BLOCK_B + ) + ref_idx = select_topk_blocks( + block_logit_ref, lse_ref, valid_range, self.TOPK, self.BLOCK_B + ) + + # (1) scaled per-block max logit agrees (masked pattern + finite values). + pattern_mismatch, logit_err = _maxerr_finite( + block_logit, block_logit_ref + ) + self.assertEqual( + pattern_mismatch, + 0, + "block_logit masked pattern mismatch vs scatter ref", + ) + # bf16 tensor-core MMA error, scaled by sm_scale (raw ~0.06*sqrt(d)). + logit_tol = 0.06 * math.sqrt(d) * sm_scale + self.assertLessEqual( + logit_err, + logit_tol, + f"block_logit finite max|diff|={logit_err:.4e} > tol={logit_tol:.4e}", + ) + + # (2) eq.(3) block scores exp(block_logit - lse) agree. + scores_fa4 = block_scores_from_logit(block_logit, lse) + scores_ref = block_scores_from_logit(block_logit_ref, lse_ref) + _, score_err = _maxerr_finite(scores_fa4, scores_ref) + self.assertLessEqual( + score_err, + 5e-2, + f"eq.3 block scores max|diff|={score_err:.4e} too large", + ) + + # (3) end-to-end: TopK block indices from the two score sources match. + # Compare as sets per query row (TopK order is irrelevant for the + # downstream gather; -1 padding slots must line up too). + fa4_sorted = paddle.sort(fa4_idx.astype("int64"), axis=-1) + ref_sorted = paddle.sort(ref_idx.astype("int64"), axis=-1) + mismatch = int((fa4_sorted != ref_sorted).astype("int32").sum().item()) + total = fa4_sorted.numel().item() + # Allow a tiny tie-driven slack: when two blocks have near-equal scores + # bf16 noise can swap which lands in the last TopK slot. + self.assertLessEqual( + mismatch, + max(1, int(0.005 * total)), + f"TopK index mismatch {mismatch}/{total} between FA4 and scatter ref", + ) + + def test_topk_consistency_small(self): + # small heads / head-dim: exercises the score->TopK path end to end. + self._run(b=1, s=256, h=8, d=64, dv=64) + + def test_topk_consistency_h64(self): + # split-D head_dim=256 (real HySparse decompressed-MHA full-attn head), + # H=64 production MLA head count. S=384 (6 blocks) keeps TopK=4 selective + # while trimming the O(H*S^2) fp32 scatter reference vs a larger S. + self._run(b=1, s=384, h=64, d=256, dv=256) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/custom_ops/test_hysparse_mqa_gather_dsa.py b/tests/single_card_tests/custom_ops/test_hysparse_mqa_gather_dsa.py new file mode 100644 index 000000000..a3a7db9af --- /dev/null +++ b/tests/single_card_tests/custom_ops/test_hysparse_mqa_gather_dsa.py @@ -0,0 +1,296 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Correctness tests for the HySparse block-sparse *DSA* backend +(:func:`paddlefleet.cudnn_ops.block_sparse_mqa_attention_dsa`), which routes the +block selection through DeepSeek-v4's FlashMLA sparse forward + cuDNN DSA +backward. + +We validate the absorbed-MQA layout (Dk=576 query/key, Dv=512 value == leading +slice of the shared latent) against a differentiable dense masked-softmax +reference over the exact selected-block column set (forward output and dq/dkv +grads). The comparison requires cosine > 0.99. The suite skips gracefully when +FlashMLA / cuDNN-frontend / a Blackwell (SM100) GPU is unavailable. +""" + +import math +import unittest + +import paddle + +paddle.enable_compat(scope={"tilelang"}, silent=True) + + +def _dsa_unavailable_reason(): + if not paddle.device.is_compiled_with_cuda(): + return "CUDA build of Paddle required" + if paddle.device.cuda.device_count() == 0: + return "no CUDA device available" + try: + from paddlefleet.cudnn_ops import is_dsa_available + + if not is_dsa_available(): + return "FlashMLA sparse fwd + cuDNN DSA bwd not available" + except (ImportError, RuntimeError): + return "hysparse DSA import failed" + cc = paddle.device.cuda.get_device_capability() + if cc[0] < 10: + return f"DSA sparse fwd requires SM100+, got {cc}" + return None + + +_SKIP_REASON = None + + +def _skip_if_no_dsa(tc): + global _SKIP_REASON + if _SKIP_REASON is None: + _SKIP_REASON = _dsa_unavailable_reason() or "" + if _SKIP_REASON: + tc.skipTest(_SKIP_REASON) + + +def _allow_mask(indices, valid_range, s_kv, block_B): + """Bool [B, S, S_kv]: col allowed iff in a selected block โˆฉ [bos, eos).""" + import numpy as np + + idx = indices.numpy() + vr = valid_range.numpy() + b, s, _ = idx.shape + allow = np.zeros([b, s, s_kv], dtype=bool) + for bi in range(b): + for i in range(s): + bos, eos = int(vr[bi, i, 0]), int(vr[bi, i, 1]) + for blk in idx[bi, i]: + if blk < 0: + continue + c0 = bos + int(blk) * block_B + for col in range(c0, min(c0 + block_B, eos)): + if 0 <= col < s_kv: + allow[bi, i, col] = True + return paddle.to_tensor(allow) + + +def _ref_masked_attn(q, k, v, allow, sm_scale, attn_sink=None): + """Differentiable dense masked MQA attention reference (fp32). + + ``attn_sink`` [H] fp32 (optional) adds a virtual sink column to the softmax + denominator (attention sink / off-by-one). + """ + neg_inf = float("-inf") + logits = paddle.einsum("bshd,bkd->bshk", q, k) * sm_scale + neg = paddle.full_like(logits, neg_inf) + mask = allow.unsqueeze(2) # [B,S,1,Skv] + logits = paddle.where(mask, logits, neg) + row_has = allow.any(axis=-1) + m = logits.max(axis=-1, keepdim=True) + if attn_sink is not None: + m = paddle.maximum( + m, attn_sink.reshape([1, 1, -1, 1]).astype("float32") + ) + m = paddle.where(paddle.isfinite(m), m, paddle.zeros_like(m)) + p = paddle.exp(logits - m) + denom = p.sum(axis=-1, keepdim=True) + if attn_sink is not None: + denom = denom + paddle.exp( + attn_sink.reshape([1, 1, -1, 1]).astype("float32") - m + ) + denom = paddle.where(denom > 0, denom, paddle.ones_like(denom)) + p = p / denom + out = paddle.einsum("bshk,bkc->bshc", p, v) + out = out * row_has.astype("float32").unsqueeze(-1).unsqueeze(-1) + return out + + +def _causal_valid_range(b, s): + eos = ( + paddle.arange(1, s + 1, dtype="int32") + .reshape([1, s, 1]) + .expand([b, s, 1]) + ) + bos = paddle.zeros([b, s, 1], dtype="int32") + return paddle.concat([bos, eos], axis=-1).contiguous() + + +def _cos(a, b): + import numpy as np + + af = a.astype("float32").numpy().reshape(-1) + bf = b.astype("float32").numpy().reshape(-1) + denom = (np.linalg.norm(af) * np.linalg.norm(bf)) + 1e-12 + return float(np.dot(af, bf) / denom) + + +def _select_blocks(b, s, block_B, topk): + """Per-token relative block ids: block 0 plus the running block (pos//BB), + padded with -1 to width ``topk``.""" + pos = paddle.arange(s) + b0 = paddle.zeros([b, s, 1], dtype="int32") + b1 = (pos // block_B).cast("int32").reshape([1, s, 1]).expand([b, s, 1]) + idx = paddle.concat([b0, b1], axis=-1) # [b, s, 2] + if topk > 2: + pad = paddle.full([b, s, topk - 2], -1, dtype="int32") + idx = paddle.concat([idx, pad], axis=-1) + return idx.contiguous() + + +class TestBlockSparseDSA(unittest.TestCase): + BLOCK_B = 64 + Dk = 576 + Dv = 512 + + def _make(self, b, s, h, topk, seed=7): + paddle.seed(seed) + q = paddle.randn([b, s, h, self.Dk]).cast("bfloat16") + kf = paddle.randn([b, s, self.Dk]).cast("bfloat16") + vr = _causal_valid_range(b, s) + idx = _select_blocks(b, s, self.BLOCK_B, topk) + sm_scale = 1.0 / math.sqrt(self.Dk) + return q, kf, vr, idx, sm_scale + + def _run_dsa(self, q, kf, idx, vr, sm_scale): + from paddlefleet.cudnn_ops import block_sparse_mqa_attention_dsa + + qd = q.detach().clone() + qd.stop_gradient = False + kd = kf.detach().clone() + kd.stop_gradient = False + out, _ = block_sparse_mqa_attention_dsa( + qd, + kd, + idx, + vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + kv_lora_rank=self.Dv, + ) + out.sum().backward() + return out, qd.grad, kd.grad + + def _run_ref(self, q, kf, idx, vr, sm_scale): + b, s, h = q.shape[0], q.shape[1], q.shape[2] + s_kv = kf.shape[1] + qr = q.detach().cast("float32") + qr.stop_gradient = False + kr = kf.detach().cast("float32") + kr.stop_gradient = False + vr_ = kr[:, :, : self.Dv] + allow = _allow_mask(idx, vr, s_kv, self.BLOCK_B) + out = _ref_masked_attn(qr, kr, vr_, allow, sm_scale) + out = out.reshape([b, s, h * self.Dv]) + out.sum().backward() + return out, qr.grad, kr.grad + + def test_dsa_vs_dense_reference_h4(self): + _skip_if_no_dsa(self) + b, s, h, topk = 1, 192, 4, 2 + q, kf, vr, idx, sm = self._make(b, s, h, topk) + out_dsa, dq_dsa, dkv_dsa = self._run_dsa(q, kf, idx, vr, sm) + out_ref, dq_ref, dkv_ref = self._run_ref(q, kf, idx, vr, sm) + self.assertEqual(list(out_dsa.shape), [b, s, h * self.Dv]) + self.assertGreater(_cos(out_dsa, out_ref), 0.99) + self.assertGreater(_cos(dq_dsa, dq_ref), 0.99) + self.assertGreater(_cos(dkv_dsa, dkv_ref), 0.99) + + def test_dsa_vs_dense_reference_h64(self): + _skip_if_no_dsa(self) + b, s, h, topk = 1, 256, 64, 4 + q, kf, vr, idx, sm = self._make(b, s, h, topk, seed=13) + out_dsa, dq_dsa, dkv_dsa = self._run_dsa(q, kf, idx, vr, sm) + out_ref, dq_ref, dkv_ref = self._run_ref(q, kf, idx, vr, sm) + self.assertGreater(_cos(out_dsa, out_ref), 0.99) + self.assertGreater(_cos(dq_dsa, dq_ref), 0.99) + self.assertGreater(_cos(dkv_dsa, dkv_ref), 0.99) + + def test_learnable_sink_matches_reference(self): + # A finite learnable per-head sink: DSA forward + dq/dkv/d_sink grads + # must match the dense reference that folds the sink into the softmax. + _skip_if_no_dsa(self) + from paddlefleet.cudnn_ops import block_sparse_mqa_attention_dsa + + b, s, h, topk = 1, 192, 4, 2 + q, kf, vr, idx, sm = self._make(b, s, h, topk, seed=21) + paddle.seed(29) + # Moderate sink magnitude: a large random sink shrinks the effective + # softmax mass and amplifies bf16 rounding in the dq GEMM, which would + # make the comparison test bf16 noise rather than the sink math. + sink0 = paddle.randn([h], dtype="float32") * 0.5 + + qd = q.detach().clone() + kd = kf.detach().clone() + sink_d = sink0.detach().clone() + for t in (qd, kd, sink_d): + t.stop_gradient = False + out_dsa, _ = block_sparse_mqa_attention_dsa( + qd, + kd, + idx, + vr, + sm_scale=sm, + block_B=self.BLOCK_B, + kv_lora_rank=self.Dv, + attn_sink=sink_d, + ) + out_dsa.sum().backward() + + s_kv = kf.shape[1] + qr = q.detach().cast("float32") + kr = kf.detach().cast("float32") + sink_r = sink0.detach().clone() + for t in (qr, kr, sink_r): + t.stop_gradient = False + vr_ = kr[:, :, : self.Dv] + allow = _allow_mask(idx, vr, s_kv, self.BLOCK_B) + out_ref = _ref_masked_attn(qr, kr, vr_, allow, sm, attn_sink=sink_r) + out_ref = out_ref.reshape([b, s, h * self.Dv]) + out_ref.sum().backward() + + self.assertGreater(_cos(out_dsa, out_ref), 0.99) + self.assertIsNotNone(sink_d.grad) + self.assertGreater(_cos(qd.grad, qr.grad), 0.99) + self.assertGreater(_cos(sink_d.grad, sink_r.grad), 0.99) + + def test_neg_sink_matches_sinkless(self): + # Sinkless (attn_sink=None) must match a very-negative explicit sink: + # exp(sink - m) -> 0, so the softmax denominator is unchanged. + _skip_if_no_dsa(self) + from paddlefleet.cudnn_ops import block_sparse_mqa_attention_dsa + + b, s, h, topk = 1, 192, 4, 2 + q, kf, vr, idx, sm = self._make(b, s, h, topk, seed=33) + neg_sink = paddle.full([h], -1e30, dtype="float32") + out_none, _ = block_sparse_mqa_attention_dsa( + q, + kf, + idx, + vr, + sm_scale=sm, + block_B=self.BLOCK_B, + kv_lora_rank=self.Dv, + ) + out_neg, _ = block_sparse_mqa_attention_dsa( + q, + kf, + idx, + vr, + sm_scale=sm, + block_B=self.BLOCK_B, + kv_lora_rank=self.Dv, + attn_sink=neg_sink, + ) + self.assertGreater(_cos(out_none, out_neg), 0.999999) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/custom_ops/test_hysparse_multidoc_block_score.py b/tests/single_card_tests/custom_ops/test_hysparse_multidoc_block_score.py new file mode 100644 index 000000000..4f5dc8796 --- /dev/null +++ b/tests/single_card_tests/custom_ops/test_hysparse_multidoc_block_score.py @@ -0,0 +1,614 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HySparse Bug 2: DOCUMENT-relative block bucketing / pack-equivalence. + +The FA4-fused block-score kernel buckets each key column into a block relative +to its query's document start ``bos`` (``rel = floor((col - bos) / block_B)``), +so the whole downstream pipeline (``_valid_block_mask`` / ``select_topk_blocks`` +/ the gather kernels), which is document-RELATIVE, sees consistent coordinates. + +Two properties are exercised: + +* ``test_scatter_ref_multidoc``: a packed sequence of several arbitrary-length + (NOT 64-aligned) documents scored by FA4 (with a flashmask document-causal + mask) matches an independent relative-scatter reference (matmul + softmax) on + the finite block logits, eq.(3) block scores, and TopK indices. + +* ``test_pack_equivalence`` (the core selection-correctness proof): a document D + run alone (``bos=0``) and the SAME document packed behind an arbitrary-length + (unaligned) prefix must select the exact same relative TopK blocks for D's + rows. This holds for the document-relative kernel and FAILS for the old + absolute-coordinate kernel -- so the test itself่จผๆ˜Žs the coordinate fix. + +FA4 block-score fusion runs only on SM 10.x (Blackwell); the test skips +otherwise. +""" + +import math +import unittest + +import paddle + +_NEG_INF = float("-inf") + + +def _sm100_or_skip(testcase): + if not paddle.device.is_compiled_with_cuda(): + testcase.skipTest("CUDA build of Paddle is required") + if paddle.device.cuda.device_count() == 0: + testcase.skipTest("No CUDA device available") + major = paddle.device.cuda.get_device_capability()[0] + if major != 10: + testcase.skipTest( + f"FA4 block-score fusion requires SM 10.x (Blackwell); got SM {major}.x" + ) + + +def _dsa_or_skip(testcase): + """Skip unless the DSA gather backend (FlashMLA + cuDNN) can run here.""" + _sm100_or_skip(testcase) + try: + from paddlefleet.cudnn_ops import is_dsa_available + + if not is_dsa_available(): + testcase.skipTest( + "FlashMLA sparse fwd + cuDNN DSA bwd not available" + ) + except (ImportError, RuntimeError): + testcase.skipTest("hysparse DSA import failed") + + +def _num_blocks(seqlen_k, block_B): + return (seqlen_k + block_B - 1) // block_B + + +def _doc_bounds(doc_lens): + """[(start, end)] cumulative document boundaries for the packed sequence.""" + bounds, off = [], 0 + for L in doc_lens: + bounds.append((off, off + L)) + off += L + return bounds + + +def _multidoc_valid_range(doc_lens): + """valid_range [1, S, 2] int32 for packed docs: per token (bos, eos=pos+1). + + Each token attends causally only within its own document, so bos is the + document start and eos is the token's own position + 1. + """ + s = sum(doc_lens) + bos = paddle.zeros([s], dtype="int32") + eos = paddle.zeros([s], dtype="int32") + for ds, de in _doc_bounds(doc_lens): + idx = paddle.arange(ds, de, dtype="int32") + bos[ds:de] = ds + eos[ds:de] = idx + 1 + return paddle.stack([bos, eos], axis=-1).reshape([1, s, 2]).contiguous() + + +def _multidoc_startend_row_indices(doc_lens, h): + """flashmask LTS [1, H, S, 1] int32: mask query rows >= doc_end per key col. + + Combined with FA4's built-in causal (rows < col masked), each key column c + in document [ds, de) is visible exactly to rows [c, de) -- i.e. a + block-diagonal (per-document) causal mask. + """ + s = sum(doc_lens) + lts = paddle.zeros([s], dtype="int32") + for ds, de in _doc_bounds(doc_lens): + lts[ds:de] = de + return lts.reshape([1, 1, s, 1]).expand([1, h, s, 1]).contiguous() + + +def _causal_valid_range(b, s): + """Single-document causal valid_range [B, S, 2]: bos=0, eos=t+1.""" + eos = ( + paddle.arange(1, s + 1, dtype="int32") + .reshape([1, s, 1]) + .expand([b, s, 1]) + ) + bos = paddle.zeros([b, s, 1], dtype="int32") + return paddle.concat([bos, eos], axis=-1).contiguous() + + +def _relative_scatter_reference(q, k, valid_range, sm_scale, block_B): + """Document-RELATIVE scatter (eq.3) reference: scaled per-block max logit + LSE. + + q, k: [B, S, H, D] bf16 (packed multi-doc, S == S_kv). ``valid_range`` gives + per-row [bos, eos). Block j of row r covers key columns + ``[bos_r + j*block_B, bos_r + (j+1)*block_B)`` and only causally-valid, + same-document keys (``bos_r <= c <= r``) contribute -- mirroring the kernel. + + Returns (block_logit_ref [B,H,S,nb] fp32, lse_ref [B,S,H] fp32). + """ + b, s, h, d = q.shape + sk = k.shape[1] + nb = _num_blocks(sk, block_B) + qf = q.astype("float32").transpose([0, 2, 1, 3]) # [B,H,S,D] + kf = k.astype("float32").transpose([0, 2, 1, 3]) # [B,H,Sk,D] + logits = paddle.matmul(qf, kf, transpose_y=True) * sm_scale # [B,H,S,Sk] + + bos = valid_range[0, :, 0].astype("int64") # [S] + rows = paddle.arange(s, dtype="int64").reshape([s, 1]) + cols = paddle.arange(sk, dtype="int64").reshape([1, sk]) + bos_col = bos.reshape([s, 1]) + # same-document causal validity: bos_r <= c <= r ([S, Sk] bool) + valid = (cols >= bos_col) & (cols <= rows) + # document-relative block id of each (row, col): floor((c - bos_r)/block_B) + rel_block = ( + cols - bos_col + ) // block_B # [S, Sk] (only meaningful where valid) + + lse_mask = valid.reshape([1, 1, s, sk]) + masked_all = paddle.where( + lse_mask, logits, paddle.full_like(logits, _NEG_INF) + ) + lse_bhs = paddle.logsumexp(masked_all, axis=-1) # [B,H,S] + lse_ref = lse_bhs.transpose([0, 2, 1]).contiguous() # [B,S,H] + + block_logit_ref = paddle.full([b, h, s, nb], _NEG_INF, dtype="float32") + for j in range(nb): + colmask = (valid & (rel_block == j)).reshape([1, 1, s, sk]) + masked_j = paddle.where( + colmask, logits, paddle.full_like(logits, _NEG_INF) + ) + block_logit_ref[:, :, :, j] = masked_j.max(axis=-1) + return block_logit_ref, lse_ref + + +def _maxerr_finite(got, ref): + """Masked-pattern match + max abs diff on finite (unmasked) entries.""" + import numpy as np + + got_np = got.astype("float32").numpy() + ref_np = ref.astype("float32").numpy() + thr = -1e30 + got_masked = got_np <= thr + ref_masked = ref_np <= thr + pattern_mismatch = int((got_masked != ref_masked).sum()) + finite = ~ref_masked + diff = ( + np.abs(got_np[finite] - ref_np[finite]) + if finite.any() + else np.array([0.0]) + ) + return pattern_mismatch, float(diff.max()) + + +class TestHySparseMultiDocBlockScore(unittest.TestCase): + BLOCK_B = 64 + TOPK = 4 + + def _topk_set_mismatch(self, idx_a, idx_b): + """#slots that differ after per-row sort (order-independent set compare).""" + a = paddle.sort(idx_a.astype("int64"), axis=-1) + b = paddle.sort(idx_b.astype("int64"), axis=-1) + return int((a != b).astype("int32").sum().item()), a.numel().item() + + def test_scatter_ref_multidoc(self): + # Packed docs of arbitrary lengths, NONE aligned to BLOCK_B=64. + _sm100_or_skip(self) + from paddlefleet.tilelang_ops.hysparse import ( + block_score_fa4_attn_fwd, + select_topk_blocks, + ) + from paddlefleet.tilelang_ops.hysparse.pipeline import ( + block_scores_from_logit, + ) + + doc_lens = [40, 88, 133, 27] # sum = 288, all unaligned to 64 + s = sum(doc_lens) + h, d, dv = 8, 64, 64 + paddle.seed(2026) + q = paddle.randn([1, s, h, d], dtype="bfloat16") + k = paddle.randn([1, s, h, d], dtype="bfloat16") + v = paddle.randn([1, s, h, dv], dtype="bfloat16") + sm_scale = 1.0 / math.sqrt(d) + valid_range = _multidoc_valid_range(doc_lens) + startend = _multidoc_startend_row_indices(doc_lens, h) + + out, lse, block_logit = block_score_fa4_attn_fwd( + q, + k, + v, + valid_range=valid_range, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + causal=True, + startend_row_indices=startend, + ) + block_logit_ref, lse_ref = _relative_scatter_reference( + q, k, valid_range, sm_scale, self.BLOCK_B + ) + + # (1) relative per-block max logit agrees (masked pattern + finite). + pattern_mismatch, logit_err = _maxerr_finite( + block_logit, block_logit_ref + ) + self.assertEqual( + pattern_mismatch, + 0, + "block_logit masked pattern mismatch vs relative scatter ref", + ) + logit_tol = 0.06 * math.sqrt(d) * sm_scale + self.assertLessEqual( + logit_err, + logit_tol, + f"block_logit finite max|diff|={logit_err:.4e} > tol={logit_tol:.4e}", + ) + + # (2) eq.(3) block scores exp(block_logit - lse) agree. + scores_fa4 = block_scores_from_logit(block_logit, lse) + scores_ref = block_scores_from_logit(block_logit_ref, lse_ref) + _, score_err = _maxerr_finite(scores_fa4, scores_ref) + self.assertLessEqual( + score_err, + 5e-2, + f"eq.3 block scores max|diff|={score_err:.4e} too large", + ) + + # (3) per-query TopK relative block indices match. + fa4_idx = select_topk_blocks( + block_logit, lse, valid_range, self.TOPK, self.BLOCK_B + ) + ref_idx = select_topk_blocks( + block_logit_ref, lse_ref, valid_range, self.TOPK, self.BLOCK_B + ) + mismatch, total = self._topk_set_mismatch(fa4_idx, ref_idx) + self.assertLessEqual( + mismatch, + max(1, int(0.005 * total)), + f"TopK index mismatch {mismatch}/{total} between FA4 and scatter ref", + ) + + def _pack_equivalence(self, prefix_len, doc_len, h, d, dv, seed=2026): + """Core pack-equivalence proof. + + Run document D alone (bos=0) and the SAME D packed behind an + arbitrary-length (unaligned) prefix, both CAUSAL-only (no flashmask). + The document-relative kernel must select the identical relative TopK + blocks for D's rows and emit the identical relative block_logit; the old + absolute-coordinate kernel would NOT (its 64-grid shifts by prefix_len). + + Causal-only is sufficient and robust: for D's rows in the packed run the + prefix key columns land at relative block id < 0 (dropped/guarded by the + kernel), and within-D causal visibility is bit-identical to the solo run; + TopK is invariant to the per-row LSE shift the prefix induces because the + block-score ordering within a row only depends on that row's block + logits. So this isolates the coordinate fix, independent of flashmask. + """ + _sm100_or_skip(self) + from paddlefleet.tilelang_ops.hysparse import ( + block_score_fa4_attn_fwd, + select_topk_blocks, + ) + + sm_scale = 1.0 / math.sqrt(d) + + # ---- (a) document D alone: bos=0, single-doc causal. ---- + paddle.seed(seed) + qd = paddle.randn([1, doc_len, h, d], dtype="bfloat16") + kd = paddle.randn([1, doc_len, h, d], dtype="bfloat16") + vd = paddle.randn([1, doc_len, h, dv], dtype="bfloat16") + vr_solo = _causal_valid_range(1, doc_len) + _, lse_solo, bl_solo = block_score_fa4_attn_fwd( + qd, + kd, + vd, + valid_range=vr_solo, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + causal=True, + ) + idx_solo = select_topk_blocks( + bl_solo, lse_solo, vr_solo, self.TOPK, self.BLOCK_B + ) + + # ---- (b) [unaligned prefix | D] packed, single-doc causal over all. ---- + # A plain causal mask over the whole packed sequence lets D's rows also + # attend to the prefix; but those prefix columns map to relative block + # id < 0 for D's rows (bos = prefix_len), which the fixed kernel drops. + # So D's finite relative blocks (>=0) see exactly D's own keys, matching + # the solo run. valid_range gives D's rows bos = prefix_len. + s = prefix_len + doc_len + paddle.seed(seed + 777) + qp = paddle.randn([1, s, h, d], dtype="bfloat16") + kp = paddle.randn([1, s, h, d], dtype="bfloat16") + vp = paddle.randn([1, s, h, dv], dtype="bfloat16") + # Overwrite D's slice with the SAME q/k/v tensors as the solo run so the + # relative-block logits are directly comparable value-for-value. + qp[:, prefix_len:, :, :] = qd + kp[:, prefix_len:, :, :] = kd + vp[:, prefix_len:, :, :] = vd + + # valid_range for the packed run: prefix rows are their own doc (bos=0), + # D's rows have bos=prefix_len, eos=row+1 (causal within packed seq). + eos = paddle.arange(1, s + 1, dtype="int32").reshape([1, s, 1]) + bos = paddle.zeros([1, s, 1], dtype="int32") + bos[:, prefix_len:, :] = prefix_len + vr_pack = paddle.concat([bos, eos], axis=-1).contiguous() + + _, lse_pack, bl_pack = block_score_fa4_attn_fwd( + qp, + kp, + vp, + valid_range=vr_pack, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + causal=True, + ) + idx_pack = select_topk_blocks( + bl_pack, lse_pack, vr_pack, self.TOPK, self.BLOCK_B + ) + + # ---- Compare D's rows. Both are DOCUMENT-relative: no remap needed. ---- + nb_solo = _num_blocks(doc_len, self.BLOCK_B) + # (1) relative block_logit for D's rows over D's relative columns match. + bl_d_pack = bl_pack[ + :, :, prefix_len:, :nb_solo + ] # [1,H,doc_len,nb_solo] + pattern_mismatch, logit_err = _maxerr_finite(bl_d_pack, bl_solo) + self.assertEqual( + pattern_mismatch, + 0, + "packed-vs-solo block_logit masked pattern differs for D's rows " + "-- coordinate systems disagree (absolute kernel bug)", + ) + logit_tol = 0.06 * math.sqrt(d) * sm_scale + self.assertLessEqual( + logit_err, + logit_tol, + f"packed-vs-solo block_logit max|diff|={logit_err:.4e} > " + f"tol={logit_tol:.4e} for D's rows", + ) + + # (2) relative TopK indices for D's rows are IDENTICAL (the hard req). + idx_pack_d = idx_pack[:, prefix_len:, :] # [1,doc_len,TOPK] + mismatch, total = self._topk_set_mismatch(idx_pack_d, idx_solo) + self.assertLessEqual( + mismatch, + max(1, int(0.005 * total)), + f"pack-equivalence VIOLATED: {mismatch}/{total} TopK slots differ " + f"between packed(bos={prefix_len}) and solo(bos=0) for D " + f"(prefix_len={prefix_len}, doc_len={doc_len})", + ) + + def test_pack_equivalence(self): + # Unaligned prefix (40) shifts the absolute 64-grid; the relative kernel + # must still select D's blocks identically to running D alone. This + # causal-only (no-flashmask) path also exercises the relative-block-id<0 + # guard for the prefix columns -- config-independent integer logic, so a + # single d=64 config suffices to cover it. + self._pack_equivalence(prefix_len=40, doc_len=200, h=8, d=64, dv=64) + + def _multidoc_pack_vs_nonpack(self, doc_lens, h, d, dv, seed=2026): + """SEVERAL unaligned-length docs, run NON-PACKED vs PACKED, must match. + + The core precision guarantee for HySparse packed training: a batch of + documents of arbitrary (NOT 64-aligned) lengths must select the exact + same relative TopK blocks -- and emit the same relative block_logit -- + whether each document is scored ALONE (``bos=0``, one sequence per doc) + or all of them are PACKED into a single sequence (each doc's rows carry + ``bos = doc_start``, doc-causal flashmask). This is strictly stronger + than the single-doc ``_pack_equivalence`` above: EVERY document in the + pack (not just one behind a prefix) is checked against its solo run, so + it exercises every non-zero ``bos`` the relative kernel must handle. + """ + _sm100_or_skip(self) + from paddlefleet.tilelang_ops.hysparse import ( + block_score_fa4_attn_fwd, + select_topk_blocks, + ) + + sm_scale = 1.0 / math.sqrt(d) + + # Independent per-doc q/k/v, reused VERBATIM in the packed run so the + # relative-block logits are directly comparable value-for-value. + qs, ks, vs = [], [], [] + for di, L in enumerate(doc_lens): + paddle.seed(seed + di) + qs.append(paddle.randn([1, L, h, d], dtype="bfloat16")) + ks.append(paddle.randn([1, L, h, d], dtype="bfloat16")) + vs.append(paddle.randn([1, L, h, dv], dtype="bfloat16")) + + # ---- NON-PACKED: each doc alone, bos=0, single-doc causal. ---- + solo_bl, solo_idx = [], [] + for di, L in enumerate(doc_lens): + vr = _causal_valid_range(1, L) + _, lse_d, bl_d = block_score_fa4_attn_fwd( + qs[di], + ks[di], + vs[di], + valid_range=vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + causal=True, + ) + solo_bl.append(bl_d) + solo_idx.append( + select_topk_blocks(bl_d, lse_d, vr, self.TOPK, self.BLOCK_B) + ) + + # ---- PACKED: concat docs into one sequence; per-doc bos via + # valid_range and a doc-causal flashmask so each doc's rows attend only + # to their own (causally-earlier) keys. ---- + qp = paddle.concat(qs, axis=1) + kp = paddle.concat(ks, axis=1) + vp = paddle.concat(vs, axis=1) + vr_pack = _multidoc_valid_range(doc_lens) + startend = _multidoc_startend_row_indices(doc_lens, h) + _, lse_pack, bl_pack = block_score_fa4_attn_fwd( + qp, + kp, + vp, + valid_range=vr_pack, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + causal=True, + startend_row_indices=startend, + ) + idx_pack = select_topk_blocks( + bl_pack, lse_pack, vr_pack, self.TOPK, self.BLOCK_B + ) + + # ---- Per-doc compare: packed rows [ds, de) over doc D's relative + # blocks must match D's solo run (both document-RELATIVE, no remap). ---- + logit_tol = 0.06 * math.sqrt(d) * sm_scale + for di, ((ds, de), L) in enumerate( + zip(_doc_bounds(doc_lens), doc_lens) + ): + nb_d = _num_blocks(L, self.BLOCK_B) + bl_d_pack = bl_pack[:, :, ds:de, :nb_d] # [1,H,L,nb_d] + patt, logit_err = _maxerr_finite(bl_d_pack, solo_bl[di]) + self.assertEqual( + patt, + 0, + f"doc{di} (len={L}, bos={ds}) block_logit masked pattern " + f"differs packed-vs-solo -- coordinate systems disagree", + ) + self.assertLessEqual( + logit_err, + logit_tol, + f"doc{di} (len={L}, bos={ds}) block_logit max|diff|=" + f"{logit_err:.4e} > tol={logit_tol:.4e} packed-vs-solo", + ) + idx_pack_d = idx_pack[:, ds:de, :] # [1,L,TOPK] + mismatch, total = self._topk_set_mismatch(idx_pack_d, solo_idx[di]) + self.assertLessEqual( + mismatch, + max(1, int(0.005 * total)), + f"doc{di} (len={L}, bos={ds}) pack-equivalence VIOLATED: " + f"{mismatch}/{total} TopK slots differ packed-vs-solo", + ) + + def test_multidoc_pack_vs_nonpack(self): + # Four docs, all unaligned to BLOCK_B=64; every doc's non-zero bos is + # checked against its solo (bos=0) run. h=8/d=64 fast config. + self._multidoc_pack_vs_nonpack( + doc_lens=[40, 88, 133, 27], h=8, d=64, dv=64 + ) + + def test_multidoc_pack_vs_nonpack_dv256_h64(self): + # Production dims (head_dim=256 split-D, H=64) with unaligned docs. + self._multidoc_pack_vs_nonpack( + doc_lens=[50, 100, 70], h=64, d=256, dv=256 + ) + + def _dense_relative_indices(self, doc_lens, nsel): + """[1, S, nsel] int32 document-relative block ids for a packed seq. + + For each query row in document D (length L), list D's own relative + block ids ``0..nb_D-1`` (nb_D = ceil(L / block_B)) and pad the rest + with ``-1``. This selects EVERY block of the row's own document, so + the block-sparse gather degenerates to full document-causal attention + -- letting us compare the gather kernel packed-vs-solo directly. + """ + s = sum(doc_lens) + idx = paddle.full([s, nsel], -1, dtype="int32") + for (ds, de), L in zip(_doc_bounds(doc_lens), doc_lens): + nb = _num_blocks(L, self.BLOCK_B) + for j in range(min(nb, nsel)): + idx[ds:de, j] = j + return idx.reshape([1, s, nsel]).contiguous() + + def _gather_pack_vs_nonpack(self, doc_lens, h, seed=2026): + """End-to-end GATHER pack-equivalence (the downstream DSA consumer). + + The block-sparse DSA gather op is what actually CONSUMES the + document-relative TopK block indices produced by the (now fixed) + scorer. This test proves the gather itself is pack-equivalent: with + dense relative indices (all of a doc's own blocks selected), running + each document ALONE (bos=0) and all documents PACKED (per-doc bos = + doc_start) must yield the SAME per-token attention output. If the + gather mixed absolute vs relative coordinates, ``col = bos + blk*BB`` + would address the wrong keys for bos>0 docs and the outputs would + diverge. Absorbed-MLA MQA layout: q [1,S,H,576], shared single-head + latent k [1,S_kv,576] whose value is its leading Dv=512 slice. + """ + _dsa_or_skip(self) + from paddlefleet.cudnn_ops import block_sparse_mqa_attention_dsa + + d, dv = 576, 512 + sm_scale = 1.0 / math.sqrt(d) + nsel = max(_num_blocks(L, self.BLOCK_B) for L in doc_lens) + + # Independent per-doc q and shared latent k (value == leading dv slice), + # reused VERBATIM in the packed run for value-for-value comparison. + qs, ks = [], [] + for di, L in enumerate(doc_lens): + paddle.seed(seed + di) + qs.append(paddle.randn([1, L, h, d], dtype="bfloat16")) + ks.append(paddle.randn([1, L, d], dtype="bfloat16")) + + # ---- NON-PACKED: each doc alone, bos=0, dense relative indices. ---- + solo_out = [] + for di, L in enumerate(doc_lens): + vr = _causal_valid_range(1, L) + idx = self._dense_relative_indices([L], nsel) + out_d, _ = block_sparse_mqa_attention_dsa( + qs[di], + ks[di], + idx, + vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + kv_lora_rank=dv, + ) + solo_out.append(out_d) + + # ---- PACKED: concat docs; per-doc bos via valid_range; each doc's + # rows carry its OWN relative block ids (0..nb_D-1). ---- + qp = paddle.concat(qs, axis=1) + kp = paddle.concat(ks, axis=1) + vr_pack = _multidoc_valid_range(doc_lens) + idx_pack = self._dense_relative_indices(doc_lens, nsel) + out_pack, _ = block_sparse_mqa_attention_dsa( + qp, + kp, + idx_pack, + vr_pack, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + kv_lora_rank=dv, + ) + + # ---- Per-doc compare: packed rows [ds, de) must match D's solo run. ---- + import numpy as np + + for di, ((ds, de), L) in enumerate( + zip(_doc_bounds(doc_lens), doc_lens) + ): + got = out_pack[:, ds:de].astype("float32").numpy() + ref = solo_out[di].astype("float32").numpy() + max_diff = float(np.abs(got - ref).max()) + # bf16 gather: same relative math, so differences are only bf16 + # rounding from block-iteration order. Tight tolerance. + self.assertLessEqual( + max_diff, + 2e-2, + f"doc{di} (len={L}, bos={ds}) GATHER output max|diff|=" + f"{max_diff:.4e} packed-vs-solo -- coordinate systems disagree", + ) + + def test_gather_pack_vs_nonpack_h64(self): + # Absorbed-MLA MQA dims: Dk=576, Dv=512, full H=64 query heads. The MQA + # gather is head-count-agnostic (single shared latent), so production + # H=64 fully covers the relative-coordinate gather pack-equivalence. + self._gather_pack_vs_nonpack(doc_lens=[50, 100, 70], h=64) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/custom_ops/test_hysparse_pipeline.py b/tests/single_card_tests/custom_ops/test_hysparse_pipeline.py new file mode 100644 index 000000000..8300e4460 --- /dev/null +++ b/tests/single_card_tests/custom_ops/test_hysparse_pipeline.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the HySparse block-TopK selection host pipeline +(:mod:`paddlefleet.tilelang_ops.hysparse.pipeline`). + +These are pure-tensor tests (no TileLang / FA4 kernels): they pin down the +eq.(3) block-score recovery, the group-wise head aggregation, the +document-relative valid-block masking, and the ``[B, S, topk]`` shape contract +(including ``topk > num_blocks`` padding and invalid-block ``-1`` slots). +""" + +import unittest + +import paddle + +from paddlefleet.tilelang_ops.hysparse.pipeline import ( + block_scores_from_logit, + select_topk_blocks, +) + +_NEG_INF = float("-inf") + + +def _skip_if_no_cuda(tc): + if not paddle.device.is_compiled_with_cuda(): + tc.skipTest("CUDA build of Paddle required") + if paddle.device.cuda.device_count() == 0: + tc.skipTest("no CUDA device available") + + +def _valid_range(bos, eos, b=1, s=1): + vr = paddle.to_tensor([[bos, eos]], dtype="int32").reshape([1, 1, 2]) + return vr.expand([b, s, 2]).contiguous() + + +class TestBlockScoresFromLogit(unittest.TestCase): + def test_matches_manual_exp(self): + _skip_if_no_cuda(self) + # block_logit [B,H,S,nb], lse [B,S,H]. scores = exp(block_logit - lse). + block_logit = paddle.to_tensor( + [1.0, 3.0, 2.0, 0.0], dtype="float32" + ).reshape([1, 1, 1, 4]) + lse = paddle.to_tensor([0.5], dtype="float32").reshape([1, 1, 1]) + scores = block_scores_from_logit(block_logit, lse) + expect = paddle.exp(block_logit.reshape([4]) - 0.5) + self.assertEqual(list(scores.shape), [1, 1, 1, 4]) + self.assertLess(float((scores.reshape([4]) - expect).abs().max()), 1e-5) + + def test_masked_and_nan_guard_to_zero(self): + _skip_if_no_cuda(self) + # -inf block_logit with -inf lse -> (-inf)-(-inf)=nan -> exp=nan; the + # guard must turn any non-finite score into 0. Finite entries survive. + block_logit = paddle.to_tensor( + [_NEG_INF, 2.0], dtype="float32" + ).reshape([1, 1, 1, 2]) + lse = paddle.to_tensor([_NEG_INF], dtype="float32").reshape([1, 1, 1]) + scores = block_scores_from_logit(block_logit, lse).reshape([2]) + self.assertTrue(bool(paddle.isfinite(scores).all())) + self.assertEqual(float(scores[0]), 0.0) # masked -> 0 + + +class TestSelectTopkBlocks(unittest.TestCase): + def _logit(self, values): + # [1, H, 1, nb] from a python list-of-lists (per head). + t = paddle.to_tensor(values, dtype="float32") + h, nb = t.shape + return t.reshape([1, h, 1, nb]) + + def _zero_lse(self, h): + return paddle.zeros([1, 1, h], dtype="float32") + + def test_head_agg_max_selects_largest(self): + _skip_if_no_cuda(self) + # single head, lse=0 -> scores monotonic in logit; top2 of [1,3,2,0]. + block_logit = self._logit([[1.0, 3.0, 2.0, 0.0]]) + idx = select_topk_blocks( + block_logit, + self._zero_lse(1), + _valid_range(0, 4), + 2, + 1, + head_agg="max", + ) + self.assertEqual(list(idx.shape), [1, 1, 2]) + got = sorted(idx.reshape([2]).numpy().tolist()) + self.assertEqual(got, [1, 2]) + + def test_head_agg_sum_across_heads(self): + _skip_if_no_cuda(self) + # A case where sum and max disagree. Per-head scores (lse=0 -> exp(logit)): + # head0: blk0=.6 blk1=.5 blk2=0 head1: blk0=0 blk1=.5 blk2=.6 + # max-agg: [.6, .5, .6] -> top2 = {blk0, blk2} (blk1 excluded) + # sum-agg: [.6, 1., .6] -> top2 must include blk1. + import math + + l6, l5 = math.log(0.6), math.log(0.5) + block_logit = self._logit([[l6, l5, _NEG_INF], [_NEG_INF, l5, l6]]) + idx_sum = select_topk_blocks( + block_logit, + self._zero_lse(2), + _valid_range(0, 3), + 2, + 1, + head_agg="sum", + ) + self.assertIn(1, set(idx_sum.reshape([2]).numpy().tolist())) + idx_max = select_topk_blocks( + block_logit, + self._zero_lse(2), + _valid_range(0, 3), + 2, + 1, + head_agg="max", + ) + self.assertNotIn(1, set(idx_max.reshape([2]).numpy().tolist())) + + def test_invalid_blocks_become_minus_one(self): + _skip_if_no_cuda(self) + # eos=2, block_B=1 -> blocks 0,1 valid; 2,3 invalid. topk=3 -> the third + # slot lands on an invalid block and must be -1. + block_logit = self._logit([[1.0, 3.0, 9.0, 9.0]]) + idx = select_topk_blocks( + block_logit, + self._zero_lse(1), + _valid_range(0, 2), + 3, + 1, + ) + got = sorted(idx.reshape([3]).numpy().tolist()) + self.assertEqual(got, [-1, 0, 1]) + + def test_topk_exceeds_num_blocks_pads_minus_one(self): + _skip_if_no_cuda(self) + # nb=2, topk=4, block_B=1, both blocks valid -> width stays 4 and the two + # surplus slots are -1 padding. + block_logit = self._logit([[1.0, 2.0]]) + idx = select_topk_blocks( + block_logit, + self._zero_lse(1), + _valid_range(0, 2), + 4, + 1, + ) + self.assertEqual(list(idx.shape), [1, 1, 4]) + vals = idx.reshape([4]).numpy().tolist() + self.assertEqual(sorted(v for v in vals if v >= 0), [0, 1]) + self.assertEqual(sum(1 for v in vals if v == -1), 2) + + def test_topk_non_positive_raises(self): + _skip_if_no_cuda(self) + block_logit = self._logit([[1.0, 2.0]]) + with self.assertRaises(ValueError): + select_topk_blocks( + block_logit, self._zero_lse(1), _valid_range(0, 2), 0, 1 + ) + + def test_unknown_head_agg_raises(self): + _skip_if_no_cuda(self) + block_logit = self._logit([[1.0, 2.0]]) + with self.assertRaises(ValueError): + select_topk_blocks( + block_logit, + self._zero_lse(1), + _valid_range(0, 2), + 1, + 1, + head_agg="mean", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/custom_ops/test_hysparse_whole_flow_precision.py b/tests/single_card_tests/custom_ops/test_hysparse_whole_flow_precision.py new file mode 100644 index 000000000..da36e82e8 --- /dev/null +++ b/tests/single_card_tests/custom_ops/test_hysparse_whole_flow_precision.py @@ -0,0 +1,378 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end HySparse precision test over the WHOLE forward flow's three +attention structures, each checked against a pure scattered-op (eager fp32) +reference, plus their combination: + +1. **FULL-layer block scoring** -- FA4-fused ``block_score_fa4_attn_fwd`` + + ``select_topk_blocks`` -> per-query document-relative TopK block indices. +2. **SWA main path** -- ``sliding_window_mqa_attention`` (windowed MQA over the + Dk=576 / Dv=512 absorbed-MLA shared latent). +3. **Block-sparse gather branch** -- ``block_sparse_mqa_attention_dsa`` + (FlashMLA sparse fwd + cuDNN DSA bwd), gathering exactly the blocks selected + in (1). + +The consumer (``MLASelfAttention``) sums the SWA main and sparse-gather branch +outputs (after a shared value-absorb + gating that are identical linear ops on +both paths and so do not affect attention-precision alignment). This test +therefore also checks ``SWA_out + sparse_out`` against the summed eager +references. + +To isolate kernel precision from TopK tie-flips, the KERNEL-selected block +indices are fed into BOTH the kernel gather and the eager gather reference. + +Single-document: tight alignment (cosine > 0.99). Multi-document packed: +acceptable error (cosine > 0.98) -- bf16 rounding + block-iteration order only. + +Runs only where the DSA backend is available (SM100 + FlashMLA + cuDNN +frontend); skips otherwise. +""" + +import math +import unittest + +import paddle + +_NEG_INF = float("-inf") +_DK = 576 # kv_lora_rank(512) + qk_rope_head_dim(64) +_DV = 512 # value == leading kv_lora_rank slice of the shared latent + + +def _dsa_or_skip(testcase): + """Skip unless the DSA gather backend (FlashMLA + cuDNN, SM100) can run.""" + if not paddle.device.is_compiled_with_cuda(): + testcase.skipTest("CUDA build of Paddle is required") + if paddle.device.cuda.device_count() == 0: + testcase.skipTest("No CUDA device available") + major = paddle.device.cuda.get_device_capability()[0] + if major != 10: + testcase.skipTest( + f"HySparse whole-flow requires SM 10.x (Blackwell); got SM {major}.x" + ) + try: + from paddlefleet.cudnn_ops import is_dsa_available + + if not is_dsa_available(): + testcase.skipTest( + "FlashMLA sparse fwd + cuDNN DSA bwd not available" + ) + except (ImportError, RuntimeError): + testcase.skipTest("hysparse DSA import failed") + + +def _num_blocks(seqlen_k, block_B): + return (seqlen_k + block_B - 1) // block_B + + +def _doc_bounds(doc_lens): + bounds, off = [], 0 + for L in doc_lens: + bounds.append((off, off + L)) + off += L + return bounds + + +def _cos(a, b): + import numpy as np + + af = a.astype("float32").numpy().reshape(-1) + bf = b.astype("float32").numpy().reshape(-1) + denom = (np.linalg.norm(af) * np.linalg.norm(bf)) + 1e-12 + return float(np.dot(af, bf) / denom) + + +def _doc_valid_range(doc_lens): + """Document-anchored causal [1, S, 2]: bos=doc_start, eos=pos+1.""" + s = sum(doc_lens) + bos = paddle.zeros([s], dtype="int32") + eos = paddle.zeros([s], dtype="int32") + for ds, de in _doc_bounds(doc_lens): + idx = paddle.arange(ds, de, dtype="int32") + bos[ds:de] = ds + eos[ds:de] = idx + 1 + return paddle.stack([bos, eos], axis=-1).reshape([1, s, 2]).contiguous() + + +def _window_valid_range(doc_lens, window): + """Windowed causal [1, S, 2]: bos=max(doc_start, pos-W+1), eos=pos+1.""" + s = sum(doc_lens) + bos = paddle.zeros([s], dtype="int32") + eos = paddle.zeros([s], dtype="int32") + for ds, de in _doc_bounds(doc_lens): + for pos in range(ds, de): + bos[pos] = max(ds, pos - window + 1) + eos[pos] = pos + 1 + return paddle.stack([bos, eos], axis=-1).reshape([1, s, 2]).contiguous() + + +def _multidoc_startend_row_indices(doc_lens, h): + """flashmask LTS [1, H, S, 1]: mask query rows >= doc_end per key col.""" + s = sum(doc_lens) + lts = paddle.zeros([s], dtype="int32") + for ds, de in _doc_bounds(doc_lens): + lts[ds:de] = de + return lts.reshape([1, 1, s, 1]).expand([1, h, s, 1]).contiguous() + + +def _range_allow_mask(valid_range, s_kv): + """Bool [1, S, S_kv]: col allowed iff bos <= col < eos (windowed/doc range).""" + bos = valid_range[..., 0:1].astype("int64") # [1,S,1] + eos = valid_range[..., 1:2].astype("int64") # [1,S,1] + cols = paddle.arange(s_kv, dtype="int64").reshape([1, 1, s_kv]) + return (cols >= bos) & (cols < eos) # [1,S,S_kv] + + +def _block_allow_mask(indices, valid_range, s_kv, block_B): + """Bool [1, S, S_kv]: col allowed iff in a selected block AND in [bos, eos). + + Mirrors the DSA op's block->token expansion + doc/causal masking exactly. + """ + import numpy as np + + idx = indices.numpy() + vr = valid_range.numpy() + b, s, _ = idx.shape + allow = np.zeros([b, s, s_kv], dtype=bool) + for bi in range(b): + for i in range(s): + bos, eos = int(vr[bi, i, 0]), int(vr[bi, i, 1]) + for blk in idx[bi, i]: + if blk < 0: + continue + c0 = bos + int(blk) * block_B + for col in range(c0, min(c0 + block_B, eos)): + if 0 <= col < s_kv: + allow[bi, i, col] = True + return paddle.to_tensor(allow) + + +def _eager_masked_mqa(q, k, v, allow, sm_scale): + """Differentiable-free dense masked MQA reference (fp32). + + q [1,S,H,Dk], k [1,S_kv,Dk], v [1,S_kv,Dv], allow [1,S,S_kv] bool. + Returns out [1,S,H,Dv] fp32; rows with no allowed key emit 0. + """ + qf = q.astype("float32") + kf = k.astype("float32") + vf = v.astype("float32") + logits = paddle.einsum("bshd,bcd->bshc", qf, kf) * sm_scale # [1,S,H,Skv] + mask = allow.unsqueeze(2) # [1,S,1,Skv] + neg = paddle.full_like(logits, _NEG_INF) + logits = paddle.where(mask, logits, neg) + row_has = allow.any(axis=-1) # [1,S] + m = logits.max(axis=-1, keepdim=True) + m = paddle.where(paddle.isfinite(m), m, paddle.zeros_like(m)) + p = paddle.exp(logits - m) + denom = p.sum(axis=-1, keepdim=True) + denom = paddle.where(denom > 0, denom, paddle.ones_like(denom)) + p = p / denom + out = paddle.einsum("bshc,bcv->bshv", p, vf) # [1,S,H,Dv] + out = out * row_has.astype("float32").unsqueeze(-1).unsqueeze(-1) + return out + + +def _eager_block_logit_lse(q, k, valid_range, sm_scale, block_B): + """Document-RELATIVE scatter (eq.3) reference for the FA4 MHA block scorer. + + q, k [1,S,H,d] decompressed MHA (independent heads, head_dim <= 256 as FA4 + requires). Block j of row r covers relative key columns + ``[bos_r + j*block_B, bos_r + (j+1)*block_B)`` restricted to causally-valid + same-document keys (``bos_r <= c < eos_r``). Returns + (block_logit [1,H,S,nb] fp32 scaled per-block max logit, lse [1,S,H] fp32). + """ + b, s, h, d = q.shape + sk = k.shape[1] + nb = _num_blocks(sk, block_B) + qf = q.astype("float32").transpose([0, 2, 1, 3]) # [1,H,S,d] + kf = k.astype("float32").transpose([0, 2, 1, 3]) # [1,H,Sk,d] + logits = paddle.matmul(qf, kf, transpose_y=True) * sm_scale # [1,H,S,Sk] + + bos = valid_range[0, :, 0].astype("int64").reshape([s, 1]) # [S,1] + eos = valid_range[0, :, 1].astype("int64").reshape([s, 1]) # [S,1] + cols = paddle.arange(sk, dtype="int64").reshape([1, sk]) + valid = (cols >= bos) & (cols < eos) # [S, Sk] + rel_block = (cols - bos) // block_B # [S, Sk] (meaningful where valid) + + vmask = valid.reshape([1, 1, s, sk]) + masked_all = paddle.where(vmask, logits, paddle.full_like(logits, _NEG_INF)) + lse_bhs = paddle.logsumexp(masked_all, axis=-1) # [1,H,S] + lse = lse_bhs.transpose([0, 2, 1]).contiguous() # [1,S,H] + + block_logit = paddle.full([b, h, s, nb], _NEG_INF, dtype="float32") + for j in range(nb): + cm = (valid & (rel_block == j)).reshape([1, 1, s, sk]) + mj = paddle.where(cm, logits, paddle.full_like(logits, _NEG_INF)) + block_logit[:, :, :, j] = mj.max(axis=-1) + return block_logit, lse + + +class TestHySparseWholeFlowPrecision(unittest.TestCase): + BLOCK_B = 64 + TOPK = 4 + WINDOW = 128 + + def _topk_set_mismatch(self, idx_a, idx_b): + a = paddle.sort(idx_a.astype("int64"), axis=-1) + b = paddle.sort(idx_b.astype("int64"), axis=-1) + return int((a != b).astype("int32").sum().item()), a.numel().item() + + def _run_whole_flow(self, doc_lens, h, tol_cos, seed=2026): + """Build the 3-structure flow and check every structure + their sum + against a pure eager fp32 reference. + + The structures live in DIFFERENT layers and use DIFFERENT tensor dims, + sharing only the document-relative block coordinates: + + * Structure 1 (full-layer scoring) runs on DECOMPRESSED MHA with head_dim + ``d_score`` <= 256 (FA4 flash_mask caps head_dim at 256): independent + query/key/value heads. It emits the document-relative TopK block + indices consumed downstream. + * Structures 2/3 (SWA main + block-sparse gather) run on the ABSORBED-MLA + MQA latent: query [1,S,H,576], a single shared K/V head k [1,S,576] + whose value is its leading Dv=512 slice. They consume the SAME block + indices produced by structure 1 (block coordinates are layer-agnostic + document-relative ids). + """ + _dsa_or_skip(self) + from paddlefleet.cudnn_ops import block_sparse_mqa_attention_dsa + from paddlefleet.tilelang_ops.hysparse import ( + block_score_fa4_attn_fwd, + select_topk_blocks, + sliding_window_mqa_attention, + ) + + s = sum(doc_lens) + multi = len(doc_lens) > 1 + paddle.seed(seed) + + doc_vr = _doc_valid_range(doc_lens) + window_vr = _window_valid_range(doc_lens, self.WINDOW) + startend = ( + _multidoc_startend_row_indices(doc_lens, h) if multi else None + ) + + # ---- Structure 1: FA4-fused block scoring + TopK selection. ---- + # Full layers score DECOMPRESSED MHA (independent heads, head_dim<=256). + d_score, dv_score = 192, 128 # qk_nope+qk_rope / v_head_dim, <= 256. + sm_score = 1.0 / math.sqrt(d_score) + q_score = paddle.randn( + [1, s, h, d_score], dtype="bfloat16" + ).contiguous() + k_score = paddle.randn( + [1, s, h, d_score], dtype="bfloat16" + ).contiguous() + v_score = paddle.randn( + [1, s, h, dv_score], dtype="bfloat16" + ).contiguous() + _, lse1, block_logit = block_score_fa4_attn_fwd( + q_score, + k_score, + v_score, + valid_range=doc_vr, + sm_scale=sm_score, + block_B=self.BLOCK_B, + causal=True, + startend_row_indices=startend, + ) + block_indices = select_topk_blocks( + block_logit, lse1, doc_vr, self.TOPK, self.BLOCK_B + ) + # eager block-score selection over the SAME relative coordinates. + bl_ref, lse_ref = _eager_block_logit_lse( + q_score, k_score, doc_vr, sm_score, self.BLOCK_B + ) + idx_ref = select_topk_blocks( + bl_ref, lse_ref, doc_vr, self.TOPK, self.BLOCK_B + ) + mism, total = self._topk_set_mismatch(block_indices, idx_ref) + self.assertLessEqual( + mism, + max(1, int(0.01 * total)), + f"[scoring] TopK indices differ from eager ref: {mism}/{total}", + ) + + # ---- Structures 2/3 run on the ABSORBED-MLA MQA latent. ---- + sm_scale = 1.0 / math.sqrt(_DK) + query = paddle.randn([1, s, h, _DK], dtype="bfloat16").contiguous() + shared_k = paddle.randn([1, s, _DK], dtype="bfloat16").contiguous() + shared_v = shared_k[:, :, :_DV].contiguous() + + # ---- Structure 2: SWA main path (windowed MQA). ---- + swa_out, _ = sliding_window_mqa_attention( + query, + shared_k, + shared_v, + window_vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + ) # [1,S,H,Dv] + swa_allow = _range_allow_mask(window_vr, s) + swa_ref = _eager_masked_mqa( + query, shared_k, shared_v, swa_allow, sm_scale + ) + cos_swa = _cos(swa_out, swa_ref) + self.assertGreater( + cos_swa, tol_cos, f"[SWA] cos={cos_swa:.6f} <= {tol_cos}" + ) + + # ---- Structure 3: block-sparse gather (DSA), KERNEL indices. ---- + sparse_out, _ = block_sparse_mqa_attention_dsa( + query, + shared_k, + block_indices, + doc_vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + kv_lora_rank=_DV, + ) + sparse_out = sparse_out.reshape([1, s, h, _DV]) + # eager gather over the exact same kernel-selected blocks. + sp_allow = _block_allow_mask(block_indices, doc_vr, s, self.BLOCK_B) + sparse_ref = _eager_masked_mqa( + query, shared_k, shared_v, sp_allow, sm_scale + ) + cos_sp = _cos(sparse_out, sparse_ref) + self.assertGreater( + cos_sp, tol_cos, f"[sparse] cos={cos_sp:.6f} <= {tol_cos}" + ) + + # ---- Combined: SWA_out + sparse_out (the consumer's branch sum). ---- + total_kernel = swa_out + sparse_out + total_ref = swa_ref + sparse_ref + cos_tot = _cos(total_kernel, total_ref) + self.assertGreater( + cos_tot, tol_cos, f"[combined] cos={cos_tot:.6f} <= {tol_cos}" + ) + print( + f"\n[whole-flow docs={doc_lens} h={h}] " + f"scoring TopK mism={mism}/{total} | " + f"cos SWA={cos_swa:.6f} sparse={cos_sp:.6f} combined={cos_tot:.6f} " + f"(tol={tol_cos})" + ) + return cos_swa, cos_sp, cos_tot + + def test_whole_flow_single_doc(self): + # One document (bos=0 everywhere): tight kernel-vs-eager alignment. + self._run_whole_flow(doc_lens=[320], h=64, tol_cos=0.99) + + def test_whole_flow_multidoc(self): + # Packed docs of arbitrary (NOT 64-aligned) lengths: acceptable error. + # Every doc's non-zero bos exercises the document-relative coordinates + # across all three structures at once. + self._run_whole_flow(doc_lens=[40, 88, 133, 27], h=64, tol_cos=0.98) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/custom_ops/test_hysparse_windowed_swa.py b/tests/single_card_tests/custom_ops/test_hysparse_windowed_swa.py new file mode 100644 index 000000000..be595ca46 --- /dev/null +++ b/tests/single_card_tests/custom_ops/test_hysparse_windowed_swa.py @@ -0,0 +1,429 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the HySparse windowed MQA flash attention +(:mod:`paddlefleet.tilelang_ops.hysparse.windowed_mqa_attn` fwd/bwd) and the +sliding-window attention wrapper it powers +(:mod:`paddlefleet.tilelang_ops.hysparse.swa_attn.sliding_window_mqa_attention`). + +The windowed kernel is a fused single-shared-K/V-head (MQA) flash attention +whose ``valid_range [B, S, 2]`` half-open ``[bos, eos)`` expresses causal + +document + sliding-window masking. We verify it against a differentiable dense +masked reference (allowed col iff ``bos <= col < eos``) for both the forward +output and the autograd grads, and exercise the ``sliding_window_mqa_attention`` +autograd wrapper (including its ``sm_scale=None`` default). +""" + +import math +import unittest + +import paddle + +paddle.enable_compat(scope={"tilelang"}, silent=True) + +_NEG_INF = float("-inf") + + +def _skip_if_no_cuda(tc): + if not paddle.device.is_compiled_with_cuda(): + tc.skipTest("CUDA build of Paddle required") + if paddle.device.cuda.device_count() == 0: + tc.skipTest("no CUDA device available") + + +def _window_allow_mask(valid_range, s_kv): + """Bool [B, S, S_kv]: col allowed iff ``bos <= col < eos``.""" + import numpy as np + + vr = valid_range.numpy() + b, s, _ = vr.shape + allow = np.zeros([b, s, s_kv], dtype=bool) + for bi in range(b): + for i in range(s): + bos, eos = int(vr[bi, i, 0]), int(vr[bi, i, 1]) + lo, hi = max(0, bos), min(s_kv, eos) + if hi > lo: + allow[bi, i, lo:hi] = True + return paddle.to_tensor(allow) + + +def _ref_masked_attn(q, k, v, allow, sm_scale, attn_sink=None): + """Differentiable dense masked MQA attention reference. + + q [B,S,H,D] fp32 leaf, k [B,Skv,D], v [B,Skv,Dv], allow [B,S,Skv] bool. + ``attn_sink`` [H] fp32, when given, adds a virtual sink column to the + softmax denominator (attention sink / off-by-one): the sink logit competes + in the row max and its ``exp(sink - m)`` mass reduces every real weight so + they sum to < 1. + """ + logits = paddle.einsum("bshd,bkd->bshk", q, k) * sm_scale + neg = paddle.full_like(logits, _NEG_INF) + logits = paddle.where(allow.unsqueeze(2), logits, neg) + row_has = allow.any(axis=-1) + m = logits.max(axis=-1, keepdim=True) # [B,S,H,1] + if attn_sink is not None: + sink = attn_sink.reshape([1, 1, -1, 1]).astype("float32") + m = paddle.maximum(m, sink) + m = paddle.where(paddle.isfinite(m), m, paddle.zeros_like(m)) + p = paddle.exp(logits - m) + denom = p.sum(axis=-1, keepdim=True) + if attn_sink is not None: + denom = denom + paddle.exp( + attn_sink.reshape([1, 1, -1, 1]).astype("float32") - m + ) + denom = paddle.where(denom > 0, denom, paddle.ones_like(denom)) + p = p / denom + out = paddle.einsum("bshk,bkc->bshc", p, v) + return out * row_has.astype("float32").unsqueeze(-1).unsqueeze(-1) + + +def _sliding_window_valid_range(b, s, window): + """Causal sliding-window ``[max(0, t-W+1), t+1)`` valid_range [B,S,2].""" + t = paddle.arange(s, dtype="int32") + bos = ( + paddle.clip(t - window + 1, min=0).reshape([1, s, 1]).expand([b, s, 1]) + ) + eos = (t + 1).reshape([1, s, 1]).expand([b, s, 1]) + return paddle.concat([bos, eos], axis=-1).contiguous() + + +def _rand_inputs(b, s, h, d, dv, s_kv, seed=7): + paddle.seed(seed) + q = paddle.randn([b, s, h, d], dtype="float32") + k = paddle.randn([b, s_kv, d], dtype="float32") + v = paddle.randn([b, s_kv, dv], dtype="float32") + return q, k, v + + +def _cos(a, b): + import numpy as np + + af = a.astype("float32").numpy().reshape(-1) + bf = b.astype("float32").numpy().reshape(-1) + denom = (np.linalg.norm(af) * np.linalg.norm(bf)) + 1e-12 + return float(np.dot(af, bf) / denom) + + +class TestWindowedMQAForward(unittest.TestCase): + BLOCK_B = 64 + + def test_full_causal_matches_dense(self): + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.windowed_mqa_attn import ( + windowed_mqa_attn_fwd, + ) + + b, s, h, d, dv, s_kv = 1, 192, 8, 64, 64, 192 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv) + # full causal window == whole prefix. + vr = _sliding_window_valid_range(b, s, window=s) + sm_scale = 1.0 / math.sqrt(d) + out, lse = windowed_mqa_attn_fwd( + q.astype("bfloat16"), + k.astype("bfloat16"), + v.astype("bfloat16"), + vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + ) + allow = _window_allow_mask(vr, s_kv) + ref = _ref_masked_attn(q, k, v, allow, sm_scale) + self.assertEqual(list(out.shape), [b, s, h, dv]) + self.assertEqual(list(lse.shape), [b, s, h]) + self.assertGreater(_cos(out, ref), 0.995) + + def test_sliding_window_matches_dense(self): + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.windowed_mqa_attn import ( + windowed_mqa_attn_fwd, + ) + + b, s, h, d, dv, s_kv = 1, 192, 4, 64, 64, 192 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=13) + vr = _sliding_window_valid_range(b, s, window=48) + sm_scale = 1.0 / math.sqrt(d) + out, _ = windowed_mqa_attn_fwd( + q.astype("bfloat16"), + k.astype("bfloat16"), + v.astype("bfloat16"), + vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + ) + allow = _window_allow_mask(vr, s_kv) + ref = _ref_masked_attn(q, k, v, allow, sm_scale) + self.assertGreater(_cos(out, ref), 0.995) + + def test_asymmetric_dk_dv(self): + # absorbed-MLA-like head: D (key/query) != D_v (value). D=128, D_v=64. + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.windowed_mqa_attn import ( + windowed_mqa_attn_fwd, + ) + + b, s, h, d, dv, s_kv = 1, 128, 4, 128, 64, 128 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=17) + vr = _sliding_window_valid_range(b, s, window=s) + sm_scale = 1.0 / math.sqrt(d) + out, _ = windowed_mqa_attn_fwd( + q.astype("bfloat16"), + k.astype("bfloat16"), + v.astype("bfloat16"), + vr, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + ) + allow = _window_allow_mask(vr, s_kv) + ref = _ref_masked_attn(q, k, v, allow, sm_scale) + self.assertEqual(list(out.shape), [b, s, h, dv]) + self.assertGreater(_cos(out, ref), 0.995) + + +class TestSlidingWindowWrapper(unittest.TestCase): + BLOCK_B = 64 + + def test_grads_match_reference(self): + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.swa_attn import ( + sliding_window_mqa_attention, + ) + + b, s, h, d, dv, s_kv = 1, 128, 8, 64, 64, 128 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=5) + vr = _sliding_window_valid_range(b, s, window=64) + sm_scale = 1.0 / math.sqrt(d) + + qk = q.astype("bfloat16").detach() + kk = k.astype("bfloat16").detach() + vk = v.astype("bfloat16").detach() + qk.stop_gradient = False + kk.stop_gradient = False + vk.stop_gradient = False + out, _ = sliding_window_mqa_attention( + qk, kk, vk, vr, sm_scale=sm_scale, block_B=self.BLOCK_B + ) + g = paddle.randn(out.shape, dtype="float32").astype("bfloat16") + out.backward(g) + + qr = q.detach() + kr = k.detach() + vr_ = v.detach() + qr.stop_gradient = False + kr.stop_gradient = False + vr_.stop_gradient = False + allow = _window_allow_mask(vr, s_kv) + ref = _ref_masked_attn(qr, kr, vr_, allow, sm_scale) + ref.backward(g.astype("float32")) + + self.assertGreater(_cos(qk.grad, qr.grad), 0.99) + self.assertGreater(_cos(kk.grad, kr.grad), 0.99) + self.assertGreater(_cos(vk.grad, vr_.grad), 0.99) + + def test_grads_ragged_seqlen(self): + # Regression for the windowed-bwd out-of-bounds guard: when + # seq_len % block_M != 0 the last query tile (grid = ceildiv(S, BM)) + # spans rows past seq_len. The backward must zero-fill that ragged + # tail (per-row guarded Q/dO load) instead of reading OOB. Use S=130 + # (block_M auto-fits to 64 -> 130 % 64 = 2, a 62-row ragged tail). + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.swa_attn import ( + sliding_window_mqa_attention, + ) + + b, s, h, d, dv, s_kv = 1, 130, 8, 64, 64, 130 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=23) + vr = _sliding_window_valid_range(b, s, window=48) + sm_scale = 1.0 / math.sqrt(d) + + qk = q.astype("bfloat16").detach() + kk = k.astype("bfloat16").detach() + vk = v.astype("bfloat16").detach() + qk.stop_gradient = False + kk.stop_gradient = False + vk.stop_gradient = False + out, _ = sliding_window_mqa_attention( + qk, kk, vk, vr, sm_scale=sm_scale, block_B=self.BLOCK_B + ) + g = paddle.randn(out.shape, dtype="float32").astype("bfloat16") + out.backward(g) + + qr = q.detach() + kr = k.detach() + vr_ = v.detach() + qr.stop_gradient = False + kr.stop_gradient = False + vr_.stop_gradient = False + allow = _window_allow_mask(vr, s_kv) + ref = _ref_masked_attn(qr, kr, vr_, allow, sm_scale) + ref.backward(g.astype("float32")) + + # grads must be finite (no OOB garbage) and match the dense reference. + import numpy as np + + for name, got in (("dq", qk.grad), ("dk", kk.grad), ("dv", vk.grad)): + self.assertTrue( + np.isfinite(got.astype("float32").numpy()).all(), + f"{name} has non-finite entries (OOB read?)", + ) + self.assertGreater(_cos(qk.grad, qr.grad), 0.99) + self.assertGreater(_cos(kk.grad, kr.grad), 0.99) + self.assertGreater(_cos(vk.grad, vr_.grad), 0.99) + + def test_wrapper_default_sm_scale(self): + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.swa_attn import ( + sliding_window_mqa_attention, + ) + + b, s, h, d, dv, s_kv = 1, 64, 4, 64, 64, 64 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=9) + vr = _sliding_window_valid_range(b, s, window=s) + qk = q.astype("bfloat16") + qk.stop_gradient = False + # sm_scale=None -> defaults to q.shape[-1] ** -0.5 in the wrapper. + out, lse = sliding_window_mqa_attention( + qk, + k.astype("bfloat16"), + v.astype("bfloat16"), + vr, + block_B=self.BLOCK_B, + ) + self.assertEqual(list(out.shape), [b, s, h, dv]) + self.assertTrue(lse.stop_gradient) + out.backward(paddle.ones_like(out)) + self.assertIsNotNone(qk.grad) + + +class TestWindowedMQAAttnSink(unittest.TestCase): + """Attention-sink (softmax off-by-one bias) support on the SWA MQA path.""" + + BLOCK_B = 64 + + def test_neg_sink_matches_sinkless(self): + # A very-negative sink must reproduce the plain sinkless forward + # bit-for-bit: exp(sink - m) underflows to 0, so the denominator is + # unchanged. Guards the "None path == -1e30 path" invariant. + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.windowed_mqa_attn import ( + windowed_mqa_attn_fwd, + ) + + b, s, h, d, dv, s_kv = 1, 96, 8, 64, 64, 96 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=11) + vr = _sliding_window_valid_range(b, s, window=48) + sm_scale = 1.0 / math.sqrt(d) + qb, kb, vb = ( + q.astype("bfloat16"), + k.astype("bfloat16"), + v.astype("bfloat16"), + ) + + out_none, _ = windowed_mqa_attn_fwd( + qb, kb, vb, vr, sm_scale=sm_scale, block_B=self.BLOCK_B + ) + neg_sink = paddle.full([h], -1e30, dtype="float32") + out_neg, _ = windowed_mqa_attn_fwd( + qb, + kb, + vb, + vr, + attn_sink=neg_sink, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + ) + import numpy as np + + np.testing.assert_array_equal( + out_none.astype("float32").numpy(), + out_neg.astype("float32").numpy(), + ) + + def test_forward_matches_sink_reference(self): + # A finite learnable sink must match the dense masked reference that + # folds the sink into the softmax denominator. + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.windowed_mqa_attn import ( + windowed_mqa_attn_fwd, + ) + + b, s, h, d, dv, s_kv = 1, 128, 8, 64, 64, 128 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=13) + vr = _sliding_window_valid_range(b, s, window=64) + sm_scale = 1.0 / math.sqrt(d) + # Mixed-magnitude sinks: some heads sink-heavy, some near-zero. + paddle.seed(31) + attn_sink = paddle.randn([h], dtype="float32") + + out, _ = windowed_mqa_attn_fwd( + q.astype("bfloat16"), + k.astype("bfloat16"), + v.astype("bfloat16"), + vr, + attn_sink=attn_sink, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + ) + allow = _window_allow_mask(vr, s_kv) + ref = _ref_masked_attn(q, k, v, allow, sm_scale, attn_sink=attn_sink) + self.assertGreater(_cos(out, ref), 0.99) + + def test_sink_grad_matches_reference(self): + # d(attn_sink) from the fused backward must match autograd on the dense + # sink reference (and q/k/v grads stay correct with a live sink). + _skip_if_no_cuda(self) + from paddlefleet.tilelang_ops.hysparse.swa_attn import ( + sliding_window_mqa_attention, + ) + + b, s, h, d, dv, s_kv = 1, 128, 8, 64, 64, 128 + q, k, v = _rand_inputs(b, s, h, d, dv, s_kv, seed=17) + vr = _sliding_window_valid_range(b, s, window=64) + sm_scale = 1.0 / math.sqrt(d) + + qk = q.astype("bfloat16").detach() + kk = k.astype("bfloat16").detach() + vk = v.astype("bfloat16").detach() + paddle.seed(41) + sink_k = paddle.randn([h], dtype="float32").detach() + for t in (qk, kk, vk, sink_k): + t.stop_gradient = False + out, _ = sliding_window_mqa_attention( + qk, + kk, + vk, + vr, + attn_sink=sink_k, + sm_scale=sm_scale, + block_B=self.BLOCK_B, + ) + g = paddle.randn(out.shape, dtype="float32").astype("bfloat16") + out.backward(g) + + qr = q.detach() + kr = k.detach() + vr_ = v.detach() + sink_r = sink_k.detach().astype("float32") + for t in (qr, kr, vr_, sink_r): + t.stop_gradient = False + allow = _window_allow_mask(vr, s_kv) + ref = _ref_masked_attn(qr, kr, vr_, allow, sm_scale, attn_sink=sink_r) + ref.backward(g.astype("float32")) + + self.assertIsNotNone(sink_k.grad) + self.assertGreater(_cos(qk.grad, qr.grad), 0.99) + self.assertGreater(_cos(sink_k.grad, sink_r.grad), 0.99) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/test_mqa_self_attention.py b/tests/single_card_tests/test_mqa_self_attention.py new file mode 100644 index 000000000..7a34f7c79 --- /dev/null +++ b/tests/single_card_tests/test_mqa_self_attention.py @@ -0,0 +1,330 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import os +import unittest + +os.environ["FLAGS_cudnn_deterministic"] = "True" + +import numpy as np +import paddle +from paddle.distributed.fleet.meta_parallel import LayerSpec + +from paddlefleet.fusions.fused_bias_dropout import get_bias_dropout_add +from paddlefleet.tensor_parallel.layers import ( + ColumnParallelLinear, + RowParallelLinear, +) +from paddlefleet.tensor_parallel.random import model_parallel_cuda_manual_seed +from paddlefleet.transformer.dot_product_attention import DotProductAttention +from paddlefleet.transformer.enums import AttnMaskType +from paddlefleet.transformer.multi_latent_attention import ( + MLASelfAttention, + MLASelfAttentionSublayersSpec, + MQASelfAttention, +) +from paddlefleet.transformer.paddle_norm import WrappedPaddleNorm +from paddlefleet.transformer.transformer_config import TransformerConfig +from paddlefleet.transformer.transformer_layer import ( + HySparseTransformerLayer, + TransformerLayerSublayersSpec, +) + + +def _hysparse_backend_or_skip(testcase): + """Skip unless the HySparse backends (FA4 FlashMask + cuDNN DSA) can run. + + The full layer scores blocks through FA4 ``block_score_fa4_attn_fwd`` and + the SWA layer consumes the cuDNN DSA gather backend; both require an + SM 10.x (Blackwell) device with FlashMask CUTE and DSA available. On + H20/SM90 or CPU-only builds this test must skip rather than fail at import + or op-launch time. + """ + if not paddle.device.is_compiled_with_cuda(): + testcase.skipTest("CUDA build of Paddle is required") + if paddle.device.cuda.device_count() == 0: + testcase.skipTest("No CUDA device available") + try: + import paddlefleet_ops + + if not paddlefleet_ops.is_flash_mask_available(): + testcase.skipTest("FlashMask (FA4) backend not available") + from paddlefleet.cudnn_ops import is_dsa_available + + if not is_dsa_available(): + testcase.skipTest("cuDNN DSA backend not available") + except (ImportError, RuntimeError): + testcase.skipTest("HySparse FA4/DSA backend import failed") + + +class TestMQASelfAttention(unittest.TestCase): + @classmethod + def setUpClass(cls): + paddle.seed(2026) + model_parallel_cuda_manual_seed(2026) + + cls.batch_size = 2 + cls.seq_len = 4096 + + cls.config = TransformerConfig( + hidden_size=1536, + head_dim=128, + num_attention_heads=4, + num_key_value_heads=4, + gated_attention=True, + qk_rope_head_dim=64, + qk_nope_head_dim=192, + v_head_dim=256, + kv_lora_rank=192, + rope_theta=5000000, + use_qk_norm=True, + multi_latent_attention=True, + rope_type="rope", + add_swa_attention_sink_bias=False, + sliding_window=[128, 128], + window_attn_skip_freq=2, + enable_hy_sparse_attention=True, + ) + + cls.sublayer_spec = MLASelfAttentionSublayersSpec( + core_attention=DotProductAttention, + o_proj=RowParallelLinear, + gate_proj=ColumnParallelLinear, + q_a_proj=ColumnParallelLinear, + q_b_proj=ColumnParallelLinear, + kv_a_proj_with_mqa=ColumnParallelLinear, + kv_b_proj=ColumnParallelLinear, + q_a_layernorm=WrappedPaddleNorm, + kv_a_layernorm=WrappedPaddleNorm, + ) + + def test_forward_backward(self): + # use larger kv_lora_rank to force eager path + config = dataclasses.replace(self.config, kv_lora_rank=512) + + mla = MLASelfAttention( + config, + self.sublayer_spec, + layer_number=0, + attn_mask_type=AttnMaskType.causal, + ) + mla = paddle.amp.decorate(mla, level="O2", dtype="bfloat16") + + mqa = MQASelfAttention( + config, + self.sublayer_spec, + layer_number=0, + attn_mask_type=AttnMaskType.causal, + ) + mqa = paddle.amp.decorate(mqa, level="O2", dtype="bfloat16") + mqa.set_state_dict(mla.state_dict()) + + hidden_states = paddle.randn( + [self.batch_size, self.seq_len, config.hidden_size], + dtype="bfloat16", + ) + hidden_states.stop_gradient = False + + # Run MLA + mla_out, _ = mla( + hidden_states=hidden_states, + attention_mask=None, + ) + output_grad = paddle.randn_like(mla_out) * 1e-2 + mla_out.backward(output_grad) + mla_input_grad = hidden_states.grad + + hidden_states = hidden_states.detach() + hidden_states.stop_gradient = False + + # Run MQA + mqa_out, _ = mqa( + hidden_states=hidden_states, + attention_mask=None, + ) + mqa_out.backward(output_grad) + mqa_input_grad = hidden_states.grad + + # Compare + np.testing.assert_allclose( + mla_out.float(), mqa_out.float(), atol=2e-3, rtol=2e-3 + ) + np.testing.assert_allclose( + mla_input_grad.float(), mqa_input_grad.float(), atol=2e-3, rtol=2e-3 + ) + for mla_param, mqa_param in zip(mla.parameters(), mqa.parameters()): + np.testing.assert_allclose( + mla_param.grad.float(), + mqa_param.grad.float(), + atol=2e-3, + rtol=2e-3, + ) + + def test_kv_sharing(self): + _hysparse_backend_or_skip(self) + layer_spec = TransformerLayerSublayersSpec( + self_attn=LayerSpec( + layer=MQASelfAttention, + sublayers_spec=self.sublayer_spec, + ), + self_attn_bda=get_bias_dropout_add, + ) + full_layer = HySparseTransformerLayer( + self.config, layer_spec, layer_number=0 + ) + full_layer.self_attn.attn_mask_type = AttnMaskType.causal + full_layer = paddle.amp.decorate( + full_layer, level="O2", dtype="bfloat16" + ) + full_layer.full_recompute = True + + swa_layer = HySparseTransformerLayer( + self.config, layer_spec, layer_number=1 + ) + swa_layer.self_attn.attn_mask_type = AttnMaskType.causal + swa_layer = paddle.amp.decorate(swa_layer, level="O2", dtype="bfloat16") + + hidden_states = paddle.randn( + [self.batch_size, self.seq_len, self.config.hidden_size], + dtype="bfloat16", + ) + attn_mask_startend_row_indices = paddle.full( + [self.batch_size, 1, self.seq_len, 1], self.seq_len, dtype="int32" + ) + + out_dict = full_layer( + { + "hidden_states": hidden_states, + "attn_mask_startend_row_indices": attn_mask_startend_row_indices, + } + ) + + self.assertTrue("shared_key" in out_dict) + self.assertTrue("shared_block_indices" in out_dict) + + out_dict = swa_layer(out_dict) + + self.assertTrue("shared_key" in out_dict) + self.assertFalse(out_dict["shared_key"].stop_gradient) + + +class TestMQAGatedAttnUseQLora(unittest.TestCase): + """``gated_attn_use_q_lora`` wiring for the two HySparse MQA gates. + + The MQA layer runs two independent gated softmax branches -- the SWA main + path (``gate_proj``, inherited from the MLA parent) and the block-sparse + DSA path (``sparse_gate_proj``). Both must honour ``gated_attn_use_q_lora``: + when set, each gate projection consumes the low-rank ``q_compressed`` latent + (dim ``q_lora_rank``) instead of the full ``hidden_states`` (dim + ``hidden_size``). ``sparse_gate_proj`` inherits its dims from ``gate_proj``, + so both track the switch together. + + These are construction/dim checks only (no attention kernel launched), so + they run on any device -- no SM100 / FA4 / DSA backend required. + """ + + HIDDEN = 1536 + Q_LORA_RANK = 256 + + @classmethod + def setUpClass(cls): + paddle.seed(2026) + model_parallel_cuda_manual_seed(2026) + cls.sublayer_spec = MLASelfAttentionSublayersSpec( + core_attention=DotProductAttention, + o_proj=RowParallelLinear, + gate_proj=ColumnParallelLinear, + q_a_proj=ColumnParallelLinear, + q_b_proj=ColumnParallelLinear, + kv_a_proj_with_mqa=ColumnParallelLinear, + kv_b_proj=ColumnParallelLinear, + q_a_layernorm=WrappedPaddleNorm, + kv_a_layernorm=WrappedPaddleNorm, + ) + + def _build(self, gated_attn_use_q_lora): + config = TransformerConfig( + hidden_size=self.HIDDEN, + head_dim=128, + num_attention_heads=4, + num_key_value_heads=4, + gated_attention=True, + gated_attn_use_q_lora=gated_attn_use_q_lora, + qk_rope_head_dim=64, + qk_nope_head_dim=192, + v_head_dim=256, + kv_lora_rank=192, + q_lora_rank=self.Q_LORA_RANK, + rope_theta=5000000, + use_qk_norm=True, + multi_latent_attention=True, + rope_type="rope", + add_swa_attention_sink_bias=False, + sliding_window=[128, 128], + window_attn_skip_freq=2, + enable_hy_sparse_attention=True, + ) + # layer_number=1 -> is_swa True (window_attn_skip_freq=2, 1 % 2 != 0) + # so this is a real MQA layer with both gate_proj and sparse_gate_proj. + return MQASelfAttention( + config, + self.sublayer_spec, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + ) + + def test_is_mqa_layer_has_two_gates(self): + mqa = self._build(gated_attn_use_q_lora=False) + self.assertTrue(mqa.is_mqa) + self.assertIsNotNone(mqa.gate_proj) # SWA main path + self.assertIsNotNone(mqa.sparse_gate_proj) # DSA sparse path + + def test_both_gates_use_hidden_size_by_default(self): + # gated_attn_use_q_lora=False -> both gates take full hidden_states. + mqa = self._build(gated_attn_use_q_lora=False) + self.assertFalse(mqa.gated_attn_use_q_lora) + self.assertEqual(mqa.gate_proj.input_size, self.HIDDEN) + self.assertEqual(mqa.sparse_gate_proj.input_size, self.HIDDEN) + + def test_both_gates_use_q_lora_rank_when_enabled(self): + # gated_attn_use_q_lora=True -> both gates take the q_compressed latent. + mqa = self._build(gated_attn_use_q_lora=True) + self.assertTrue(mqa.gated_attn_use_q_lora) + self.assertEqual(mqa.gate_proj.input_size, self.Q_LORA_RANK) + self.assertEqual(mqa.sparse_gate_proj.input_size, self.Q_LORA_RANK) + + def test_sparse_gate_inherits_main_gate_dims(self): + # sparse_gate_proj is built from gate_proj.input_size/output_size, so + # the two gates always agree on both dims, under either flag value. + for use_q_lora in (False, True): + mqa = self._build(gated_attn_use_q_lora=use_q_lora) + self.assertEqual( + mqa.sparse_gate_proj.input_size, mqa.gate_proj.input_size + ) + self.assertEqual( + mqa.sparse_gate_proj.output_size, mqa.gate_proj.output_size + ) + + def test_gate_output_dim_is_num_heads_times_v_head_dim(self): + # Output dim is unaffected by the q_lora switch: num_heads * v_head_dim. + expected_out = 4 * 256 # num_attention_heads * v_head_dim + for use_q_lora in (False, True): + mqa = self._build(gated_attn_use_q_lora=use_q_lora) + self.assertEqual(mqa.gate_proj.output_size, expected_out) + self.assertEqual(mqa.sparse_gate_proj.output_size, expected_out) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/single_card_tests/transformer/test_build_hysparse_valid_range.py b/tests/single_card_tests/transformer/test_build_hysparse_valid_range.py new file mode 100644 index 000000000..48a7f091f --- /dev/null +++ b/tests/single_card_tests/transformer/test_build_hysparse_valid_range.py @@ -0,0 +1,189 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ``build_hysparse_valid_range`` (flashmask -> valid_range). + +This closes the precision-critical gap that the HySparse block-score / gather +op tests leave open: those tests feed ``valid_range`` built DIRECTLY from +``doc_lens`` (``_multidoc_valid_range``), bypassing the production path that +DERIVES ``valid_range`` from the flashmask ``startend_row_indices`` via a +cummax document-boundary scan. A bug in that derivation (wrong ``bos`` per +token) would silently mis-place every block coordinate and hurt precision, yet +go uncaught by the operator tests. + +Here we build a synthetic flashmask (LTS ``[B, H, S, 1]`` whose value is each +token's exclusive document end -- the same convention as ``utils.get_doc_lens``) +from arbitrary-length (NOT 64-aligned) packed documents, and assert +``build_hysparse_valid_range`` reproduces the ground-truth ``[bos, eos)`` built +straight from the document layout -- for the pure document range, the windowed +range, the single-document (``None``) fast path, mask-batch broadcast, and the +multi-head (read head-0) case. Pure host math -- no GPU / SM requirement. +""" + +import unittest + +import paddle + +from paddlefleet.transformer.multi_latent_attention import ( + build_hysparse_valid_range, +) + + +def _doc_bounds(doc_lens): + """[(start, end)] cumulative document boundaries for the packed sequence.""" + bounds, off = [], 0 + for L in doc_lens: + bounds.append((off, off + L)) + off += L + return bounds + + +def _flashmask_from_doc_lens(doc_lens, h=1, batch=1): + """Synthetic flashmask LTS [B, H, S, 1] int32: per token its doc end (excl). + + Mirrors the flashmask document-causal convention consumed by + ``get_doc_lens`` / ``build_hysparse_valid_range``: the value stored at + position ``t`` is the exclusive end of the document containing ``t``. + """ + s = sum(doc_lens) + de = paddle.zeros([s], dtype="int32") + for ds, dee in _doc_bounds(doc_lens): + de[ds:dee] = dee + return de.reshape([1, 1, s, 1]).expand([batch, h, s, 1]).contiguous() + + +def _reference_valid_range(doc_lens, batch=1, window_size=None): + """Ground-truth [B, S, 2] int32 built straight from the document layout. + + ``eos = pos + 1`` (causal); ``bos = doc_start`` (document mask), optionally + clamped up to ``pos - window_size + 1`` for the causal sliding window. + """ + s = sum(doc_lens) + bos = paddle.zeros([s], dtype="int64") + eos = paddle.zeros([s], dtype="int64") + for ds, dee in _doc_bounds(doc_lens): + pos = paddle.arange(ds, dee, dtype="int64") + start = paddle.full([dee - ds], ds, dtype="int64") + if window_size is not None and window_size > 0: + start = paddle.maximum(start, pos - window_size + 1) + bos[ds:dee] = start + eos[ds:dee] = pos + 1 + vr = paddle.stack([bos, eos], axis=-1).cast("int32") # [S, 2] + return vr.reshape([1, s, 2]).expand([batch, s, 2]).contiguous() + + +class TestBuildHySparseValidRange(unittest.TestCase): + def _assert_equal(self, got, ref, msg=""): + got_np = got.astype("int32").numpy() + ref_np = ref.astype("int32").numpy() + self.assertEqual(got_np.shape, ref_np.shape, f"shape mismatch {msg}") + n_diff = int((got_np != ref_np).sum()) + self.assertEqual( + n_diff, 0, f"{n_diff} valid_range entries differ {msg}" + ) + + def test_single_doc_via_flashmask(self): + # One document spanning the whole sequence: bos=0, eos=t+1. + doc_lens = [200] + mask = _flashmask_from_doc_lens(doc_lens) + s = sum(doc_lens) + got = build_hysparse_valid_range(mask, s, 1) + ref = _reference_valid_range(doc_lens) + self._assert_equal(got, ref, "(single doc)") + + def test_none_mask_is_single_doc(self): + # attn_mask_startend_row_indices=None => single document, bos=0. + s = 137 + got = build_hysparse_valid_range(None, s, 1) + ref = _reference_valid_range([s]) + self._assert_equal(got, ref, "(None mask)") + + def test_multidoc_unaligned(self): + # Several docs, all unaligned to 64 -- the packed-training regime. + doc_lens = [40, 88, 133, 27] + mask = _flashmask_from_doc_lens(doc_lens) + s = sum(doc_lens) + got = build_hysparse_valid_range(mask, s, 1) + ref = _reference_valid_range(doc_lens) + self._assert_equal(got, ref, "(multidoc unaligned)") + + def test_multidoc_windowed(self): + # Windowed clamp: bos = max(doc_start, pos - window + 1). + doc_lens = [40, 88, 133, 27] + window = 32 + mask = _flashmask_from_doc_lens(doc_lens) + s = sum(doc_lens) + got = build_hysparse_valid_range(mask, s, 1, window_size=window) + ref = _reference_valid_range(doc_lens, window_size=window) + self._assert_equal(got, ref, "(multidoc windowed)") + + def test_window_larger_than_doc_is_noop(self): + # window >= longest doc => clamp never binds, equals pure doc range. + doc_lens = [40, 88, 27] + window = 1000 + mask = _flashmask_from_doc_lens(doc_lens) + s = sum(doc_lens) + got = build_hysparse_valid_range(mask, s, 1, window_size=window) + ref = _reference_valid_range(doc_lens) # no window clamp binds + self._assert_equal(got, ref, "(window >= doc len)") + + def test_mask_batch_broadcast(self): + # Flashmask batch=1 must broadcast over a data batch > 1. + doc_lens = [50, 100, 70] + mask = _flashmask_from_doc_lens(doc_lens, batch=1) + s = sum(doc_lens) + got = build_hysparse_valid_range(mask, s, 4) + ref = _reference_valid_range(doc_lens, batch=4) + self._assert_equal(got, ref, "(mask batch broadcast)") + + def test_multihead_reads_head0(self): + # Multi-head flashmask (all heads identical doc layout): head-0 read + # must recover the same document valid_range. + doc_lens = [40, 88, 133, 27] + mask = _flashmask_from_doc_lens(doc_lens, h=8) + s = sum(doc_lens) + got = build_hysparse_valid_range(mask, s, 1) + ref = _reference_valid_range(doc_lens) + self._assert_equal(got, ref, "(multihead head-0)") + + def test_matches_get_doc_lens_convention(self): + # Cross-check: bos recovered here must be consistent with the repo's + # get_doc_lens document-length derivation from the same flashmask. + from paddlefleet.transformer.utils import get_doc_lens + + doc_lens = [40, 88, 133, 27] + mask = _flashmask_from_doc_lens(doc_lens) + s = sum(doc_lens) + got = build_hysparse_valid_range(mask, s, 1).astype("int64").numpy() + derived_lens = get_doc_lens(mask).numpy().tolist() + self.assertEqual( + derived_lens, + doc_lens, + f"get_doc_lens {derived_lens} != {doc_lens}", + ) + # eos - bos at each doc's last token equals that doc's length. + for ds, dee in _doc_bounds(doc_lens): + last = dee - 1 + bos_last = int(got[0, last, 0]) + eos_last = int(got[0, last, 1]) + self.assertEqual( + eos_last - bos_last, + dee - ds, + f"doc [{ds},{dee}) last-token span {eos_last - bos_last} " + f"!= len {dee - ds}", + ) + + +if __name__ == "__main__": + unittest.main()