From 0c6fd6c08612d92b4bb97b4f4157bcffb1de4a00 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 11:01:33 -0700 Subject: [PATCH] [None][fix] Compile the DSpark non-causal block merge and size its LSE workspace The TRTLLM backend serves a non-causal windowed draft block as the causal windowed decode plus a log-sum-exp merge of the later block keys. Run eagerly, the merge launched dozens of small kernels per layer: 61-148 us on GB300, on top of a 7-94 us decode. The merge is now unrolled over the block keys and runs compiled (maybe_compile, automatic dynamic shapes), in place on the decode's base-2 log-sum-exp: 5-29 us per layer, the same result up to bf16 rounding. Blocks without later keys (one block token, or no right window) skip the log-sum-exp and the merge. The workspace now also holds the softmax stats that flashinfer takes off its front when the decode returns the log-sum-exp. Without them, 64 query heads per rank with a 3-token block outgrew the workspace beyond ~400 requests, and flashinfer refused the launch. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/models/modeling_dflash.py | 151 ++++++++++++++---- .../test_dspark_noncausal_window.py | 90 +++++++++-- 2 files changed, 196 insertions(+), 45 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index aa63d42b4343..5e093953a338 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -28,6 +28,7 @@ get_dflash_trtllm_gen_ops, ) from ..speculative.interface import SpeculativeDecodingMode +from ..utils import maybe_compile from .modeling_utils import FUSED_MODULE_COMPONENTS, get_model_architecture, register_draft_model @@ -82,6 +83,10 @@ def dflash_merge_future_block_attention( backend serves a non-causal sliding-window layer, which the trtllm-gen context kernel rejects while its causal decode kernel windows natively. + The softmax over the later keys, their output and the merge are unrolled + over the (few) block keys, so that under torch.compile they fuse into one + pointwise kernel (see :func:`dflash_merge_future_block_attention_`). + Args: out_causal: [B * blk, Hq, D] causal attention output. lse_causal: [B * blk, Hq] float32 log-sum-exp of the scaled causal logits. @@ -95,28 +100,106 @@ def dflash_merge_future_block_attention( B, blk, num_heads, head_dim = query.shape num_kv_heads = block_key.shape[2] group = num_heads // num_kv_heads - q = query.float().reshape(B, blk, num_kv_heads, group, head_dim) - logits = torch.einsum("bjhgd,bkhd->bhgjk", q, block_key.float()) * scale + q = query.float().reshape(B, blk, num_kv_heads, group, 1, head_dim) + keys = block_key.float().permute(0, 2, 1, 3)[:, None, :, None] # [B, 1, Hkv, 1, blk, D] + logits = (q * keys).sum(-1) * scale # [B, query j, Hkv, g, key k] offsets = torch.arange(blk, device=query.device) distance = offsets.view(1, -1) - offsets.view(-1, 1) # [j, k] = k - j future = distance > 0 if window_right >= 0: future &= distance <= window_right - logits = logits.masked_fill(~future, float("-inf")) - lse_future = torch.logsumexp(logits, dim=-1) # [B, Hkv, g, blk], -inf for no future keys - lse_future_finite = torch.where( - torch.isfinite(lse_future), lse_future, torch.zeros_like(lse_future) + future_logits = [ + logits[..., k].masked_fill(~future[:, k].view(1, blk, 1, 1), float("-inf")) + for k in range(blk) + ] + max_future = future_logits[0] + for logit in future_logits[1:]: + max_future = torch.maximum(max_future, logit) + # A query without later keys has only -inf logits; its weight is 0 below. + max_future = torch.where(torch.isfinite(max_future), max_future, 0.0) + exp_future = [torch.exp(logit - max_future) for logit in future_logits] + sum_future = exp_future[0] + for exp_logit in exp_future[1:]: + sum_future = sum_future + exp_logit + has_future = sum_future > 0 + sum_future = torch.where(has_future, sum_future, 1.0) + lse_future = max_future + torch.log(sum_future) + lse_causal = lse_causal.float().reshape(B, blk, num_kv_heads, group) + # exp(lse_future) / (exp(lse_causal) + exp(lse_future)) + weight_future = torch.where(has_future, torch.sigmoid(lse_future - lse_causal), 0.0) + out = out_causal.float().reshape(B, blk, num_kv_heads, group, head_dim) + merged = out * (1.0 - weight_future).unsqueeze(-1) + value = block_value.float() + prob_scale = weight_future / sum_future + for k in range(blk): + merged = merged + (prob_scale * exp_future[k]).unsqueeze(-1) * value[:, k, None, :, None] + return merged.reshape(B * blk, num_heads, head_dim).to(out_causal.dtype) + + +@maybe_compile +def dflash_merge_future_block_attention_( + out: torch.Tensor, + lse_log2: torch.Tensor, + query: torch.Tensor, + block_key: torch.Tensor, + block_value: torch.Tensor, + scale: float, + window_right: int, +) -> None: + """In-place :func:`dflash_merge_future_block_attention` on a base-2 log-sum-exp. + + Compiled, so that the merge runs as a few fused kernels per layer instead of + dozens of small eager ones. + """ + out.copy_( + dflash_merge_future_block_attention( + out, lse_log2 * math.log(2.0), query, block_key, block_value, scale, window_right + ) ) - probs = torch.exp(logits - lse_future_finite.unsqueeze(-1)) - out_future = torch.einsum("bhgjk,bkhd->bjhgd", probs, block_value.float()) - out_future = out_future.reshape(B * blk, num_heads, head_dim) - lse_future = lse_future.permute(0, 3, 1, 2).reshape(B * blk, num_heads) - lse_total = torch.logaddexp(lse_causal, lse_future) - weight_causal = torch.exp(lse_causal - lse_total).unsqueeze(-1) - weight_future = torch.exp(lse_future - lse_total).unsqueeze(-1) - merged = weight_causal * out_causal.float() + weight_future * out_future - return merged.to(out_causal.dtype) + +def dflash_merges_future_block_keys( + causal: bool, window_size: tuple[int, int], block_size: int +) -> bool: + """Whether the TRTLLM block decode adds later block keys to its causal result. + + True for a non-causal windowed layer whose block has a key after some query + within the right window; such a layer requests the decode's log-sum-exp. + """ + window_left, window_right = window_size + return not causal and window_left >= 0 and block_size > 1 and window_right != 0 + + +def dflash_trtllm_gen_workspace_size( + trtllm_gen_ops, + dtype: torch.dtype, + max_batch_size: int, + block_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + return_lse: bool, +) -> int: + """Bytes of trtllm-gen workspace for DFlash's block attention. + + With ``return_lse``, flashinfer takes the decode's softmax stats + (``round_up(block_size, 256)`` float2 per query head and request) and a + 1 MiB guard off the front of the workspace; the kernel scratch gets the rest. + """ + workspace_bytes = trtllm_gen_ops.get_workspace_size( + dtype=dtype, + num_tokens=max_batch_size * block_size, + num_gen_tokens=max_batch_size * block_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_dim, + max_num_requests=max_batch_size, + rotary_embedding_dim=0, + fp8_context_fmha=False, + ) + if return_lse: + workspace_bytes += 8 * num_heads * max_batch_size * 256 * -(-block_size // 256) + (1 << 20) + return workspace_bytes def dflash_trtllm_windowed_block_attention( @@ -140,7 +223,8 @@ def dflash_trtllm_windowed_block_attention( non-causal unwindowed layer takes the context kernel instead. trtllm-gen windows only causal attention, so the non-causal windowed case runs the causal decode with its log-sum-exp and then adds the block keys after each - query (:func:`dflash_merge_future_block_attention`). + query (:func:`dflash_merge_future_block_attention_`). With no such keys (a + single-token block, or no right window) the causal decode is the result. Args: query: [B, blk, Hq, D] block queries. @@ -152,7 +236,7 @@ def dflash_trtllm_windowed_block_attention( """ B, block_size, num_heads, head_dim = query.shape window_left, window_right = window_size - merge_future = not causal and window_left >= 0 + merge_future = dflash_merges_future_block_keys(causal, window_size, block_size) decoded = trtllm_gen_ops.batch_decode_with_kv_cache( query=query.reshape(-1, num_heads, head_dim), kv_cache=(layer_cache[:, 0], layer_cache[:, 1]), @@ -179,11 +263,8 @@ def dflash_trtllm_windowed_block_attention( ) if merge_future: # trtllm-gen reports the log-sum-exp in base 2. - lse = decoded[1] * math.log(2.0) - out.copy_( - dflash_merge_future_block_attention( - out, lse, query, block_key, block_value, head_dim**-0.5, window_right - ) + dflash_merge_future_block_attention_( + out, decoded[1], query, block_key, block_value, head_dim**-0.5, window_right ) @@ -1490,18 +1571,18 @@ def _prepare_dflash_trtllm_gen_buffers( num_heads: int, num_kv_heads: int, head_dim: int, + return_lse: bool = False, ) -> None: trtllm_gen_ops = self._dflash_trtllm_gen_ops - workspace_bytes = trtllm_gen_ops.get_workspace_size( - dtype=dtype, - num_tokens=max_batch_size * block_size, - num_gen_tokens=max_batch_size * block_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_size=head_dim, - max_num_requests=max_batch_size, - rotary_embedding_dim=0, - fp8_context_fmha=False, + workspace_bytes = dflash_trtllm_gen_workspace_size( + trtllm_gen_ops, + dtype, + max_batch_size, + block_size, + num_heads, + num_kv_heads, + head_dim, + return_lse, ) device = torch.device(device) is_capturing = torch.cuda.is_current_stream_capturing() @@ -1687,6 +1768,12 @@ def dflash_forward( num_heads_per_rank, num_kv_heads_per_rank, head_dim, + return_lse=any( + dflash_merges_future_block_keys( + *self._resolve_block_attention(layer_idx), block_size + ) + for layer_idx in range(len(self.model.layers)) + ), ) else: # FA4 needs no workspace or counter buffers. self._prepare_dflash_index_buffers(hidden_states.device, max_batch_size, block_size) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py index 668266bca8b2..43f634c599bc 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_noncausal_window.py @@ -7,8 +7,10 @@ ``C + j - W + 1`` and every block key. The TRTLLM backend serves that mask as the causal windowed decode plus the block keys after each query, merged by log-sum-exp. These tests check the mask resolution (CPU), the merge against a -dense oracle (CPU), and the full trtllm-gen path against a dense fp32 oracle -over a paged cache (SM100/SM103 only). +dense oracle (CPU), the compiled in-place merge against the eager one (CUDA), +and the full trtllm-gen path against a dense fp32 oracle over a paged cache, +up to a batch whose softmax stats outgrow the plain workspace (SM100/SM103 +only). """ import math @@ -20,6 +22,9 @@ from tensorrt_llm._torch.models.modeling_dflash import ( DFlashForCausalLM, dflash_merge_future_block_attention, + dflash_merge_future_block_attention_, + dflash_merges_future_block_keys, + dflash_trtllm_gen_workspace_size, dflash_trtllm_windowed_block_attention, ) from tensorrt_llm._torch.models.modeling_dspark import GQADSparkForCausalLM @@ -226,17 +231,17 @@ def __init__(self, ops, ctx_lens, block_size, num_heads, num_kv_heads, head_dim, self.block_key = torch.stack([k[-block_size:] for k in self.keys]) self.block_value = torch.stack([v[-block_size:] for v in self.values]) self.seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=dev) + # Sized as the drafter sizes it for a layer that merges later block keys. self.workspace = torch.empty( - ops.get_workspace_size( - dtype=torch.bfloat16, - num_tokens=B * block_size, - num_gen_tokens=B * block_size, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - head_size=head_dim, - max_num_requests=B, - rotary_embedding_dim=0, - fp8_context_fmha=False, + dflash_trtllm_gen_workspace_size( + ops, + torch.bfloat16, + B, + block_size, + num_heads, + num_kv_heads, + head_dim, + return_lse=True, ), dtype=torch.uint8, device=dev, @@ -316,7 +321,7 @@ def _ctx_lens_around(window, block_size, page_size): @pytest.mark.parametrize("page_size", [32, 64]) @pytest.mark.parametrize("block_size", [1, 2, 4, 7]) -@pytest.mark.parametrize("heads", [(16, 1), (16, 2)]) +@pytest.mark.parametrize("heads", [(16, 1), (16, 2), (64, 4)]) @pytest.mark.parametrize("window", [64, 4096]) def test_trtllm_windowed_non_causal_matches_the_oracle(page_size, block_size, heads, window): """Mixed context lengths in one batch, straddling the window and page @@ -386,3 +391,62 @@ def test_trtllm_windowed_non_causal_replays_under_cuda_graph(): graph.replay() expected, _ = _oracle(case.query, case.keys, case.values, new_ctx, 4095, 4095, causal=False) torch.testing.assert_close(out.float(), expected, atol=2e-2, rtol=2e-2) + + +def test_only_blocks_with_later_keys_take_the_merge(): + assert dflash_merges_future_block_keys(False, (4095, 4095), 7) + assert dflash_merges_future_block_keys(False, (4095, -1), 2) + # One block token, or no right window: nothing follows a query. + assert not dflash_merges_future_block_keys(False, (4095, 4095), 1) + assert not dflash_merges_future_block_keys(False, (4095, 0), 7) + # Causal layers, and unwindowed non-causal ones (context kernel). + assert not dflash_merges_future_block_keys(True, (4095, 0), 7) + assert not dflash_merges_future_block_keys(False, (-1, -1), 7) + + +@pytest.mark.parametrize("block_size", [2, 3, 7]) +@pytest.mark.parametrize("heads", [(16, 1), (64, 4)]) +def test_compiled_in_place_merge_matches_the_eager_merge(block_size, heads): + """The layer runs the merge compiled, in place, on the decode's base-2 + log-sum-exp; at every batch size it equals the eager merge.""" + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + # Earlier tests may have used up the recompile budget, after which the + # function would run eagerly; start from a clean compile cache. + torch._dynamo.reset() + torch.manual_seed(block_size) + num_heads, num_kv_heads = heads + for batch in (1, 5, 64): + query = torch.randn(batch, block_size, num_heads, 128, device="cuda").bfloat16() + block_key = torch.randn(batch, block_size, num_kv_heads, 128, device="cuda").bfloat16() + block_value = torch.randn_like(block_key) + out = torch.randn(batch * block_size, num_heads, 128, device="cuda").bfloat16() + lse_log2 = 4 * torch.randn(batch * block_size, num_heads, device="cuda") + expected = dflash_merge_future_block_attention( + out, lse_log2 * math.log(2.0), query, block_key, block_value, 128**-0.5, 4095 + ) + dflash_merge_future_block_attention_( + out, lse_log2, query, block_key, block_value, 128**-0.5, 4095 + ) + torch.testing.assert_close(out, expected) + + +def test_workspace_holds_the_lse_stats_at_large_batch(): + """With its log-sum-exp the decode takes 8 B x heads x requests x 256 + 1 MiB + of softmax stats off the front of the workspace. At 64 heads, 512 requests + and 3 block tokens that alone outgrows the trtllm-gen workspace (flashinfer + refuses the launch); the drafter's size includes the stats.""" + ops = _trtllm_gen_ops_or_skip() + torch.manual_seed(5) + batch, block_size, num_heads, num_kv_heads = 512, 3, 64, 4 + trtllm_gen_bytes = dflash_trtllm_gen_workspace_size( + ops, torch.bfloat16, batch, block_size, num_heads, num_kv_heads, 128, return_lse=False + ) + assert trtllm_gen_bytes < 8 * num_heads * batch * 256 + (1 << 20) + ctx_lens = [1 + (97 * b) % 300 for b in range(batch)] + case = _PagedCase(ops, ctx_lens, block_size, num_heads, num_kv_heads, 128, 64) + case.append_block() + expected, _ = _oracle(case.query, case.keys, case.values, ctx_lens, 4095, 4095, causal=False) + torch.testing.assert_close( + case.run(False, (4095, 4095)).float(), expected, atol=2e-2, rtol=2e-2 + )