From afce4ff36ef3681daf31344e182a17c19e13742b Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 11:09:00 -0700 Subject: [PATCH 1/6] [None][feat] Let the Muse modeling_v2 drafter window attention read a paged context The drafter's window-attention ops, dflash_window_attention and its multi-request form, read the DFlash context from a VANILLA drafter's dense per-slot arena. A TRTLLM drafter keeps that context in the KV cache manager's pages instead, which is what lets disaggregated serving transfer it with the target's KV. Given a page table, both ops now read and append the context through it: the caches are pools of pages and each request's batch index names its table row. The kernels see a pool as one space of rows, so a paged call differs from a dense one only in which row a column maps to, and on the same values the two are bitwise equal. Signed-off-by: Zheyu Fu --- .../ops/spec/dflash_window_attention.py | 107 +++++++++++++++--- .../ops/spec/dflash_window_attention_wide.py | 70 +++++++++--- .../test_modeling_v2_drafter_context.py | 64 +++++++++++ 3 files changed, 211 insertions(+), 30 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention.py index aee028d90aab..0f5fefa943e9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention.py @@ -45,6 +45,14 @@ Nothing reads a dummy slot's block, so only which padding row lands there differs. +The context may instead live in pages, as a KV cache manager keeps it +(``page_table``): the caches are then pools ``[pages, page_size, Hkv, D]`` (the +planes of an HND pool, transposed), and ``cache_batch_idx`` names each request's +table row, which lists its pages in token order, so column ``c`` is row +``c % page_size`` of page ``page_table[row, c // page_size]``. A paged call +differs from a dense one only in which cache row a column is read from or +appended to; on the same values the two are bitwise equal. + Numerics follow FlashAttention-2's split path: fp32 scores, base-2 online softmax, the probabilities rounded to the operand dtype for the P.V product with the row sum taken from the fp32 values, each split normalized and merged @@ -92,6 +100,30 @@ MAX_WINDOW_SPAN = MAX_SPLITS * MAX_TILES * SPLIT_BLOCK_N +@triton.jit +def _context_rows(table, cols, mask, page_rows, PAGE: tl.constexpr): + """The cache rows of context columns ``cols``: the columns of a dense slot, or row + ``c % PAGE`` of page ``table[c // PAGE]`` in a pool whose pages are ``page_rows`` apart.""" + if PAGE: + page = tl.load(table + cols // PAGE, mask=mask, other=0).to(tl.int64) + rows = page * page_rows + cols % PAGE + else: + rows = cols.to(tl.int64) + return rows + + +def _page_arguments(k_cache, v_cache, page_table, cache_batch_idx): + """The split kernel's ``(page_table, PAGE, page_rows, stride_pt, stride_ks, stride_vs)``. + + A dense call reads its slot's columns (``PAGE`` 0, the table unused); a paged + call reads rows of the whole pool through its table row (slot strides 0). + """ + if page_table is None: + return cache_batch_idx, 0, 0, 0, k_cache.stride(0), v_cache.stride(0) + page_rows = k_cache.stride(0) // k_cache.stride(1) + return page_table, k_cache.shape[1], page_rows, page_table.stride(0), 0, 0 + + @triton.jit def _window_attention_split_kernel( q_ptr, @@ -101,6 +133,7 @@ def _window_attention_split_kernel( v_new_ptr, seqlens_ptr, batch_idx_ptr, + page_table_ptr, out_ptr, o_part_ptr, lse_part_ptr, @@ -115,6 +148,8 @@ def _window_attention_split_kernel( stride_vs, stride_vn, stride_vh, + stride_pt, + page_rows, stride_knt, stride_knh, stride_vnt, @@ -135,6 +170,7 @@ def _window_attention_split_kernel( APPEND: tl.constexpr, DIRECT: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, + PAGE: tl.constexpr, ): split = tl.program_id(0) bh = tl.program_id(1) @@ -153,6 +189,8 @@ def _window_attention_split_kernel( return slot = tl.load(batch_idx_ptr + b).to(tl.int64) dims = tl.arange(0, D) + # A paged call's slot names the request's page-table row; its slot strides are 0. + table = page_table_ptr + slot * stride_pt k_base = k_ptr + slot * stride_ks + kvh * stride_kh v_base = v_ptr + slot * stride_vs + kvh * stride_vh block_start = seqlen - M @@ -166,7 +204,7 @@ def _window_attention_split_kernel( j = tl.arange(0, M_P2) j_ok = j < M new_rows = new_row0 + j - cols = (block_start + j).to(tl.int64) + kv_rows = _context_rows(table, block_start + j, j_ok, page_rows, PAGE) k_blk = tl.load( k_new_ptr + new_rows[:, None] * stride_knt + kvh * stride_knh + dims[None, :], mask=j_ok[:, None], @@ -175,8 +213,12 @@ def _window_attention_split_kernel( v_new_ptr + new_rows[:, None] * stride_vnt + kvh * stride_vnh + dims[None, :], mask=j_ok[:, None], ) - tl.store(k_base + cols[:, None] * stride_kn + dims[None, :], k_blk, mask=j_ok[:, None]) - tl.store(v_base + cols[:, None] * stride_vn + dims[None, :], v_blk, mask=j_ok[:, None]) + tl.store( + k_base + kv_rows[:, None] * stride_kn + dims[None, :], k_blk, mask=j_ok[:, None] + ) + tl.store( + v_base + kv_rows[:, None] * stride_vn + dims[None, :], v_blk, mask=j_ok[:, None] + ) # Packed tile row r = g * M + i: query head kvh * GROUP + g, block row i, # which sits at key position seqlen - M + i (bottom-right aligned). @@ -202,7 +244,7 @@ def _window_attention_split_kernel( for t in tl.static_range(TILES): cols = start + t * BLOCK_N + tl.arange(0, BLOCK_N) col_ok = cols < seqlen - cols64 = cols.to(tl.int64) + kv_rows = _context_rows(table, cols, col_ok, page_rows, PAGE) # Both loads are issued before the products: V's latency hides behind # QK^T and the softmax. if APPEND: @@ -210,7 +252,7 @@ def _window_attention_split_kernel( in_arena = col_ok & (cols < block_start) new_rows = new_row0 + (cols - block_start) k = tl.load( - k_base + cols64[:, None] * stride_kn + dims[None, :], + k_base + kv_rows[:, None] * stride_kn + dims[None, :], mask=in_arena[:, None], other=0.0, ) @@ -224,7 +266,7 @@ def _window_attention_split_kernel( k, ) v = tl.load( - v_base + cols64[:, None] * stride_vn + dims[None, :], + v_base + kv_rows[:, None] * stride_vn + dims[None, :], mask=in_arena[:, None], other=0.0, ) @@ -239,12 +281,12 @@ def _window_attention_split_kernel( ) else: k = tl.load( - k_base + cols64[:, None] * stride_kn + dims[None, :], + k_base + kv_rows[:, None] * stride_kn + dims[None, :], mask=col_ok[:, None], other=0.0, ) v = tl.load( - v_base + cols64[:, None] * stride_vn + dims[None, :], + v_base + kv_rows[:, None] * stride_vn + dims[None, :], mask=col_ok[:, None], other=0.0, ) @@ -369,11 +411,13 @@ def window_attention_launch( num_warps: int, merge_warps: int, merge_entry_trigger: bool = False, + page_table: torch.Tensor | None = None, ) -> torch.Tensor: """Launch the split and merge kernels with an explicit split geometry. ``merge_entry_trigger`` (with ``enable_pdl``): the merge kernel releases its dependent before its own wait instead of at its exit (the PDL chain's ``pdl_merge``). + ``page_table``: the caches are pools of pages (see the module docstring). Raises: ValueError: if the geometry cuts the window into more than ``MAX_SPLITS`` splits. @@ -399,6 +443,9 @@ def window_attention_launch( append = k_new is not None if not append: k_new = v_new = k_cache # unused + page_table, page, page_rows, stride_pt, stride_ks, stride_vs = _page_arguments( + k_cache, v_cache, page_table, cache_batch_idx + ) direct = num_splits == 1 if direct: o_part = lse_part = out # unused @@ -417,6 +464,7 @@ def window_attention_launch( v_new, cache_seqlens, cache_batch_idx, + page_table, out, o_part, lse_part, @@ -425,12 +473,14 @@ def window_attention_launch( q.stride(0), q.stride(1), q.stride(2), - k_cache.stride(0), + stride_ks, k_cache.stride(1), k_cache.stride(2), - v_cache.stride(0), + stride_vs, v_cache.stride(1), v_cache.stride(2), + stride_pt, + page_rows, k_new.stride(0), k_new.stride(1), v_new.stride(0), @@ -451,6 +501,7 @@ def window_attention_launch( APPEND=append, DIRECT=direct, LAUNCH_WITH_PDL=enable_pdl, + PAGE=page, num_warps=num_warps, launch_pdl=enable_pdl, ) @@ -491,8 +542,9 @@ class DflashWindowAttention(OpWrapper): """Non-causal sliding-window attention of a DFlash block over its slot's pooled K/V. ``__call__(q, k_cache, v_cache, cache_seqlens, cache_batch_idx, window_size, - softmax_scale=None, enable_pdl=False, k_new=None, v_new=None)`` takes the - arguments the drafter hands ``flash_attn_with_kvcache`` (non-causal): + softmax_scale=None, enable_pdl=False, k_new=None, v_new=None, + merge_entry_trigger=False, page_table=None)`` takes the arguments the drafter + hands ``flash_attn_with_kvcache`` (non-causal): * ``q`` ``[B, M, Hq, D]``: the block's queries, unit last-dim stride; * ``k_cache`` / ``v_cache`` ``[slots, capacity, Hkv, D]``: one layer's view @@ -513,7 +565,11 @@ class DflashWindowAttention(OpWrapper): kernel triggers its dependents before its ``griddepcontrol.wait`` instead of at its exit. It reads only the split's partials and the lengths, all after that wait; its dependent must read the output after - its own wait. The same values. + its own wait. The same values; + * ``page_table`` ``[rows, max_pages]`` int32, unit last stride: the caches + are pools ``[pages, page_size, Hkv, D]`` whose page stride is the same + whole number of rows for both, and ``cache_batch_idx`` names table rows + (see the module docstring). Returns ``[B, M, Hq, D]`` in ``q``'s dtype, contiguous. @@ -544,9 +600,18 @@ def is_valid( k_new: torch.Tensor | None = None, v_new: torch.Tensor | None = None, merge_entry_trigger: bool = False, + page_table: torch.Tensor | None = None, ) -> None: reason = unsupported_reason( - q, k_cache, v_cache, cache_seqlens, cache_batch_idx, window_size, k_new, v_new + q, + k_cache, + v_cache, + cache_seqlens, + cache_batch_idx, + window_size, + k_new, + v_new, + page_table, ) if reason is not None: raise ValueError(f"dflash_window_attention: {reason}") @@ -564,6 +629,7 @@ def __call__( k_new: torch.Tensor | None = None, v_new: torch.Tensor | None = None, merge_entry_trigger: bool = False, + page_table: torch.Tensor | None = None, ) -> torch.Tensor: m, num_q_heads = q.shape[1], q.shape[2] group = num_q_heads // k_cache.shape[2] @@ -581,6 +647,7 @@ def __call__( v_new, **config, merge_entry_trigger=merge_entry_trigger, + page_table=page_table, ) def reference( @@ -705,6 +772,7 @@ def unsupported_reason( window_size: tuple[int, int], k_new: torch.Tensor | None = None, v_new: torch.Tensor | None = None, + page_table: torch.Tensor | None = None, ) -> str | None: """Why the kernels would answer these inputs wrongly, or None if they take them.""" if q.dim() != 4 or k_cache.dim() != 4 or v_cache.dim() != 4: @@ -738,6 +806,17 @@ def unsupported_reason( return "k_new / v_new must have the operands' dtype" if k_new.stride(2) != 1 or v_new.stride(2) != 1: return "k_new / v_new head dimension must have unit stride" + if page_table is not None: + if page_table.dtype != torch.int32 or page_table.dim() != 2 or page_table.stride(1) != 1: + return "page_table must be [rows, max_pages] int32 with unit last stride" + if not page_table.is_cuda: + return "all tensors must be on the GPU" + page_rows = k_cache.stride(0) // k_cache.stride(1) + if (k_cache.stride(0), v_cache.stride(0)) != ( + page_rows * k_cache.stride(1), + page_rows * v_cache.stride(1), + ): + return "a page's stride must be the same whole number of rows for K and V" return None diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention_wide.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention_wide.py index 2892cc71fc37..777e9c7a77e4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention_wide.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/dflash_window_attention_wide.py @@ -32,6 +32,9 @@ slot and write it in no defined order, as ``index_put`` with duplicate indices does. +The context may live in pages (``page_table``), addressed as the single-request +op addresses them; on the same values a paged call is bitwise the dense call. + Numerics follow the single-request op: fp32 scores, base-2 online softmax, the probabilities rounded to the operand dtype for the P.V product with the row sum taken from the fp32 values, each split normalized and merged in fp32 by @@ -58,6 +61,8 @@ from .dflash_window_attention import ( MAX_SPLITS, WINDOW, + _context_rows, + _page_arguments, _window_attention_merge_kernel, attention_reference, unsupported_reason, @@ -78,6 +83,7 @@ def _window_attention_wide_split_kernel( v_new_ptr, seqlens_ptr, batch_idx_ptr, + page_table_ptr, out_ptr, o_part_ptr, lse_part_ptr, @@ -93,6 +99,8 @@ def _window_attention_wide_split_kernel( stride_vs, stride_vn, stride_vh, + stride_pt, + page_rows, stride_knt, stride_knh, stride_vnt, @@ -113,6 +121,7 @@ def _window_attention_wide_split_kernel( APPEND: tl.constexpr, DIRECT: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, + PAGE: tl.constexpr, ): split = tl.program_id(0) bh = tl.program_id(1) @@ -131,6 +140,8 @@ def _window_attention_wide_split_kernel( return slot = tl.load(batch_idx_ptr + b).to(tl.int64) dims = tl.arange(0, D) + # A paged call's slot names the request's page-table row; its slot strides are 0. + table = page_table_ptr + slot * stride_pt k_base = k_ptr + slot * stride_ks + kvh * stride_kh v_base = v_ptr + slot * stride_vs + kvh * stride_vh block_start = seqlen - M @@ -142,7 +153,7 @@ def _window_attention_wide_split_kernel( # The append: the block's K/V into the arena at [L - M, L). No program of # this grid reads those columns (the tail tile reads k_new / v_new). new_rows = new_row0 + jj - cols = (block_start + jj).to(tl.int64) + kv_rows = _context_rows(table, block_start + jj, jj_ok, page_rows, PAGE) k_blk = tl.load( k_new_ptr + new_rows[:, None] * stride_knt + kvh * stride_knh + dims[None, :], mask=jj_ok[:, None], @@ -151,8 +162,12 @@ def _window_attention_wide_split_kernel( v_new_ptr + new_rows[:, None] * stride_vnt + kvh * stride_vnh + dims[None, :], mask=jj_ok[:, None], ) - tl.store(k_base + cols[:, None] * stride_kn + dims[None, :], k_blk, mask=jj_ok[:, None]) - tl.store(v_base + cols[:, None] * stride_vn + dims[None, :], v_blk, mask=jj_ok[:, None]) + tl.store( + k_base + kv_rows[:, None] * stride_kn + dims[None, :], k_blk, mask=jj_ok[:, None] + ) + tl.store( + v_base + kv_rows[:, None] * stride_vn + dims[None, :], v_blk, mask=jj_ok[:, None] + ) # Packed tile row r = g * M + i: query head kvh * GROUP + g, block row i at key # position seqlen - M + i (bottom-right aligned). @@ -181,12 +196,12 @@ def _window_attention_wide_split_kernel( for c0 in tl.range(start, arena_end, BLOCK_N, num_stages=STAGES): cols = c0 + tl.arange(0, BLOCK_N) ok = cols < arena_end - cols64 = cols.to(tl.int64) + kv_rows = _context_rows(table, cols, ok, page_rows, PAGE) k = tl.load( - k_base + cols64[:, None] * stride_kn + dims[None, :], mask=ok[:, None], other=0.0 + k_base + kv_rows[:, None] * stride_kn + dims[None, :], mask=ok[:, None], other=0.0 ) v = tl.load( - v_base + cols64[:, None] * stride_vn + dims[None, :], mask=ok[:, None], other=0.0 + v_base + kv_rows[:, None] * stride_vn + dims[None, :], mask=ok[:, None], other=0.0 ) s = tl.dot(q, tl.trans(k)) keep = ok[None, :] & (cols[None, :] >= first_key[:, None]) @@ -218,14 +233,14 @@ def _window_attention_wide_split_kernel( other=0.0, ) else: - cols64 = cols.to(tl.int64) + kv_rows = _context_rows(table, cols, jj_ok, page_rows, PAGE) k = tl.load( - k_base + cols64[:, None] * stride_kn + dims[None, :], + k_base + kv_rows[:, None] * stride_kn + dims[None, :], mask=jj_ok[:, None], other=0.0, ) v = tl.load( - v_base + cols64[:, None] * stride_vn + dims[None, :], + v_base + kv_rows[:, None] * stride_vn + dims[None, :], mask=jj_ok[:, None], other=0.0, ) @@ -356,9 +371,13 @@ def wide_window_attention_launch( num_warps: int, stages: int, merge_warps: int, + page_table: torch.Tensor | None = None, ) -> torch.Tensor: """Launch the split and merge kernels with an explicit geometry. + ``page_table``: the caches are pools of pages (``dflash_window_attention``'s + module docstring). + Raises: ValueError: if ``split_keys`` is not a whole number of tiles or the geometry cuts the window into more than ``MAX_SPLITS`` splits. @@ -384,6 +403,9 @@ def wide_window_attention_launch( append = k_new is not None if not append: k_new = v_new = k_cache # unused + page_table, page, page_rows, stride_pt, stride_ks, stride_vs = _page_arguments( + k_cache, v_cache, page_table, cache_batch_idx + ) direct = num_splits == 1 if direct: o_part = lse_part = out # unused @@ -402,6 +424,7 @@ def wide_window_attention_launch( v_new, cache_seqlens, cache_batch_idx, + page_table, out, o_part, lse_part, @@ -411,12 +434,14 @@ def wide_window_attention_launch( q.stride(0), q.stride(1), q.stride(2), - k_cache.stride(0), + stride_ks, k_cache.stride(1), k_cache.stride(2), - v_cache.stride(0), + stride_vs, v_cache.stride(1), v_cache.stride(2), + stride_pt, + page_rows, k_new.stride(0), k_new.stride(1), v_new.stride(0), @@ -437,6 +462,7 @@ def wide_window_attention_launch( APPEND=append, DIRECT=direct, LAUNCH_WITH_PDL=enable_pdl, + PAGE=page, num_warps=num_warps, launch_pdl=enable_pdl, ) @@ -483,9 +509,9 @@ class DflashWindowAttentionWide(OpWrapper): """Non-causal sliding-window attention of several DFlash blocks over their slots' pooled K/V. ``__call__(q, k_cache, v_cache, cache_seqlens, cache_batch_idx, window_size, - softmax_scale=None, enable_pdl=False, k_new=None, v_new=None)``: the arguments - and the result of ``dflash_window_attention`` (see there), without its - single-request merge entry trigger. The split geometry follows the request + softmax_scale=None, enable_pdl=False, k_new=None, v_new=None, page_table=None)``: + the arguments and the result of ``dflash_window_attention`` (see there), without + its single-request merge entry trigger. The split geometry follows the request count and the device's SM count (``wide_geometry``, ``wide_launch_config``). Fusion boundary: the output projection (after) stays outside; so does the @@ -515,9 +541,18 @@ def is_valid( enable_pdl: bool = False, k_new: torch.Tensor | None = None, v_new: torch.Tensor | None = None, + page_table: torch.Tensor | None = None, ) -> None: reason = wide_unsupported_reason( - q, k_cache, v_cache, cache_seqlens, cache_batch_idx, window_size, k_new, v_new + q, + k_cache, + v_cache, + cache_seqlens, + cache_batch_idx, + window_size, + k_new, + v_new, + page_table, ) if reason is not None: raise ValueError(f"dflash_window_attention_wide: {reason}") @@ -534,6 +569,7 @@ def __call__( enable_pdl: bool = False, k_new: torch.Tensor | None = None, v_new: torch.Tensor | None = None, + page_table: torch.Tensor | None = None, ) -> torch.Tensor: config = wide_launch_config( q.shape[0], window_size[0] + q.shape[1], wide_geometry(q.device) @@ -550,6 +586,7 @@ def __call__( k_new, v_new, **config, + page_table=page_table, ) def reference( @@ -599,10 +636,11 @@ def wide_unsupported_reason( window_size: tuple[int, int], k_new: torch.Tensor | None = None, v_new: torch.Tensor | None = None, + page_table: torch.Tensor | None = None, ) -> str | None: """Why the kernels would answer these inputs wrongly, or None if they take them.""" reason = unsupported_reason( - q, k_cache, v_cache, cache_seqlens, cache_batch_idx, window_size, k_new, v_new + q, k_cache, v_cache, cache_seqlens, cache_batch_idx, window_size, k_new, v_new, page_table ) if reason is not None: return reason diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py new file mode 100644 index 000000000000..19e6cd83f368 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The musespark_open TEP target's DFlash drafter paths over each layout of the drafter's context. + +* ``dflash_window_attention`` and ``dflash_window_attention_wide`` (at each device's split + geometry) over pages are bitwise their dense calls, the appended block included. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import ( # noqa: E501 + dflash_window_attention_wide as wide_ops, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec.dflash_window_attention import ( # noqa: E501 + WINDOW, + dflash_window_attention, +) + +__extra_import_path__ = ["../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _randn(gen, *shape): + return torch.randn(*shape, device="cuda", generator=gen).bfloat16() + + +@pytest.mark.parametrize("append", [False, True], ids=["read", "append"]) +@pytest.mark.parametrize("sms", [None, *sorted(wide_ops.WIDE_GEOMETRY_BY_SMS)]) +def test_the_window_ops_over_pages_are_their_dense_calls(sms, append, monkeypatch): + """``sms``: the wide op at that device's split geometry, run on the device at hand; None: + the single-request op. Eight requests, a TP4 rank's heads, the served window, PDL.""" + op = dflash_window_attention + if sms is not None: + monkeypatch.setattr(wide_ops, "_device_sms", lambda index: sms) + op = wide_ops.dflash_window_attention_wide + gen = torch.Generator(device="cuda").manual_seed(0) + # Windows starting inside a page; blocks inside one, straddling two or starting one. + ctx = [5_000, 13, 1_020, 2_047, 300, 64, 1, 3_333] + batch, m, hq, hkv, d, page, capacity = len(ctx), 8, 16, 2, 128, 32, 5_024 + slots, per_slot = batch + 1, capacity // page + k_arena, v_arena = _randn(gen, slots, capacity, hkv, d), _randn(gen, slots, capacity, hkv, d) + q, new = _randn(gen, batch, m, hq, d) * 2, {} + if append: + new = dict(k_new=_randn(gen, batch * m, hkv, d), v_new=_randn(gen, batch * m, hkv, d)) + # The slots' columns on scattered pages of a [pages, K/V, Hkv, page, D] pool. + table = torch.randperm(slots * per_slot + 3, device="cuda", generator=gen) + table = table[: slots * per_slot].view(slots, per_slot).int() + pool = _randn(gen, slots * per_slot + 3, 2, hkv, page, d) + k_pages, v_pages = pool[:, 0].transpose(1, 2), pool[:, 1].transpose(1, 2) + cols = torch.arange(capacity, device="cuda") + rows = table[:, cols // page].long(), cols % page + k_pages[rows], v_pages[rows] = k_arena, v_arena + # Request b on slot slots - 1 - b, so its table row is not its batch index. + seqlens = torch.tensor(ctx, dtype=torch.int32, device="cuda") + m + call = (seqlens, (slots - 1 - torch.arange(batch, device="cuda")).int(), WINDOW, None, True) + dense = op(q, k_arena, v_arena, *call, **new) + with validating(op): + paged = op(q, k_pages, v_pages, *call, page_table=table, **new) + assert torch.equal(paged, dense) + # The blocks went to the requests' pages as to their slots, and nothing else moved. + assert torch.equal(k_pages[rows], k_arena) and torch.equal(v_pages[rows], v_arena) From 54051104bbf9289d3b921da11c7701affb317dfd Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 11:11:45 -0700 Subject: [PATCH 2/6] [None][refactor] Split the DFlash worker's drafter-context stages out of its prepare prepare_1st_drafter_inputs stored the step's context K/V, committed it (lengths, dummy slot, window compaction), read the lengths the drafter attends with the allocation clamp, and built the drafter's inputs, all inline. A fast path that schedules part of that elsewhere -- the Muse modeling_v2 composition stores the context on a side stream, for one -- had to copy those rules, and every worker change then had to be made twice. Name each stage as a worker method: _store_step_context, _commit_context (_advance_ctx_len, then _settle_context: the window compaction and _drafter_context_lengths), _empty_dummy_slot and _drafter_inputs. The built-in prepare is those stages in order: the same operations on the same operands, so its results are unchanged. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/speculative/dflash.py | 243 ++++++++++-------- .../hw_agnostic/test_dflash_dummy_slot.py | 1 + 2 files changed, 141 insertions(+), 103 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 59333b9278af..ee0e8fb35024 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -1631,11 +1631,144 @@ def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> N """ self._ctx_len[slots] += num_accepted self._ctx_len.clamp_(max=self._max_ctx) + self._empty_dummy_slot() + + def _empty_dummy_slot(self) -> None: + """Empty the dummy slot, as ``_advance_ctx_len`` does after every add.""" self._ctx_len[self._dummy_slot].zero_() if self._ctx_base is not None: # The windowed arena's base moves with the length; empty is (0, 0). self._ctx_base[self._dummy_slot].zero_() + def _store_step_context( + self, + draft_model, + gen_hs: torch.Tensor, + slots: torch.Tensor, + rows: torch.Tensor, + ctx_len: torch.Tensor, + num_accepted: Optional[torch.Tensor] = None, + ) -> None: + """Store the drafter context K/V of each request's ``K + 1`` captured target rows. + + Row ``r`` of a request goes to position ``ctx_len + r`` of its context: by batch row + (``rows``) in the KV cache manager's pages, else by slot. Rows at or past + ``num_accepted`` are zeroed but still written; without ``num_accepted`` every row is + stored as computed. + """ + offsets = torch.arange(gen_hs.shape[0] // slots.shape[0], dtype=torch.long, device="cuda") + projected_to_store = draft_model.project_target_hidden(gen_hs.reshape(-1, gen_hs.shape[-1])) + if self._ctx_base is None: + col_idx = ctx_len.unsqueeze(1) + offsets.unsqueeze(0) + else: + # Windowed arena: physical columns. + col_idx = (ctx_len - self._ctx_base[slots]).unsqueeze(1) + offsets.unsqueeze(0) + if self._ctx_block_tables is not None: + # Bound by what the manager handed THIS request: the masked + # entries below are still written (fixed-size, for graph + # safety) and a placeholder page resolves to another's. + gen_capacity = self._ctx_block_counts[rows] * self._ctx_page_size + col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) + elif self._ctx_keep: + col_idx = col_idx.clamp(max=self._ctx_k_buf.size(2) - 1) + else: + col_idx = col_idx.clamp(max=self._max_ctx - 1) + col_long = col_idx.reshape(-1) + pos_flat = (ctx_len.unsqueeze(1) + offsets.unsqueeze(0)).reshape(-1) + + # Fast path: store the pre-projected/pre-RoPE'd K/V. + # dflash_forward reads these directly via cache_batch_idx. + cache_dtype = self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + k_new, v_new = draft_model.precompute_context_kv( + projected_to_store.to(cache_dtype), pos_flat + ) + if num_accepted is not None: + # Fixed-size writes for CUDA graph compatibility: + # Write ALL entries but zero out invalid ones. Invalid + # writes land at clamped column indices beyond valid + # _ctx_len range, so they're harmless. + write_mask = offsets.unsqueeze(0) < num_accepted.unsqueeze(1) + mask_bc = write_mask.view(-1, 1, 1, 1).to(k_new.dtype) + k_new.mul_(mask_bc) + if v_new is not None: + v_new.mul_(mask_bc) + if self._ctx_paged: + # Batch positions of the gen requests key the manager's + # per-request block table; slots key the private one. + keys = rows if self._ctx_block_tables is not None else slots + rows_long = keys.unsqueeze(1).expand_as(col_idx).reshape(-1) + self._store_context_kv_paged(k_new, v_new, rows_long, col_long) + else: # VANILLA DFlash backend (FlashAttention) + slot_long = slots.unsqueeze(1).expand_as(col_idx).reshape(-1) + self._ctx_k_buf[slot_long, :, col_long] = k_new + if v_new is not None: + self._ctx_v_buf[slot_long, :, col_long] = v_new + + def _commit_context( + self, slots: torch.Tensor, num_accepted: torch.Tensor, rows: torch.Tensor + ) -> torch.Tensor: + """Commit a step's stored context: advance the lengths, then ``_settle_context``.""" + self._advance_ctx_len(slots, num_accepted) + return self._settle_context(slots, rows) + + def _settle_context(self, slots: torch.Tensor, rows: torch.Tensor) -> torch.Tensor: + """Compact a windowed arena's slots to their window; return the lengths the drafter reads.""" + if self._ctx_keep and not self._ctx_paged: + self._compact_ctx_window(slots) + return self._drafter_context_lengths(slots, rows) + + def _drafter_context_lengths(self, slots: torch.Tensor, rows: torch.Tensor) -> torch.Tensor: + """The context each request's draft block appends after and attends over.""" + if self._ctx_base is None: + num_ctx = self._ctx_len[slots] + else: + # Windowed arena: the physical length the draft forward appends + # after and attends over. + num_ctx = self._ctx_len[slots] - self._ctx_base[slots] + if self._ctx_block_tables is not None: + # The store clamps columns to the request's allocation, so a + # context that outruns it has its tail written over the last + # valid slot. Truncate what is advertised to match, or the read + # attends that clobbered value -- or another request's page. + block_size = self._compute_block_size + allocated = dflash_allocated_ctx_limit( + self._ctx_block_counts, self._ctx_page_size, block_size + )[rows] + if self._ctx_unified_cache: + self._check_block_within_allocation(num_ctx, allocated, block_size) + num_ctx = torch.minimum(num_ctx, allocated) + return num_ctx + + def _drafter_inputs( + self, + noise_embedding: torch.Tensor, + query_positions: torch.Tensor, + num_ctx_per_req: torch.Tensor, + slots: torch.Tensor, + rows: torch.Tensor, + first_prev_tokens: torch.Tensor, + ) -> dict: + """The draft forward's inputs, ``prepare_1st_drafter_inputs``' result.""" + return { + "noise_embedding": noise_embedding, + "query_positions": query_positions, + "num_ctx_per_req": num_ctx_per_req, + "ctx_k_cache": self._ctx_k_buf, + "ctx_v_cache": self._ctx_v_buf, + # Slots index the private arena; the manager's block table is keyed + # by batch position, so the drafter reads its own rows there. + "ctx_cache_batch_idx": rows if self._ctx_block_tables is not None else slots, + # Anchor token per gen request (block slot 0): last accepted + # token. The dspark Markov chain conditions its first step on it. + "first_prev_tokens": first_prev_tokens, + "ctx_kv_cache": self._ctx_kv_buf, + "ctx_page_table": ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ), + } + def _write_ctx_len( self, updates: dict[int, int], base_updates: Optional[dict[int, int]] = None ) -> None: @@ -2422,7 +2555,6 @@ def prepare_1st_drafter_inputs( gen_accepted_tokens = accepted_tokens[num_contexts : num_contexts + num_gens, :] total_tokens_per_req = self._draft_tokens_per_req # K+1 - K = spec_metadata.runtime_draft_len # Get captured multi-layer hidden states from spec_metadata captured_hs = spec_metadata.get_hidden_states(total_target_tokens) @@ -2442,14 +2574,11 @@ def prepare_1st_drafter_inputs( num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" ) - K_plus_1 = K + 1 - bonus_idx = (gen_num_accepted - 1).clamp_min(0).long().unsqueeze(1) bonus = gen_accepted_tokens.gather(1, bonus_idx).squeeze(1).long() ctx_len_gen = self._ctx_len[slots] j_block = torch.arange(query_tokens_per_req, dtype=torch.long, device="cuda") - offsets_kp1 = torch.arange(K_plus_1, dtype=torch.long, device="cuda") # _ctx_len is clamped to _max_ctx only AFTER this step's accepted # tokens are folded in (see the update below), so the running length @@ -2458,7 +2587,6 @@ def prepare_1st_drafter_inputs( # any position the sequence can legitimately reach. ctx_len_now = (ctx_len_gen + gen_num_accepted.long()).clamp_(max=self._max_ctx) query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) - ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) # Go through embed_tokens.forward (NOT .weight[...]) so TP-sharded # vocabs mask out ranks that don't own the token id and all-reduce. @@ -2475,88 +2603,13 @@ def prepare_1st_drafter_inputs( # Target now processes exactly K+1 tokens per gen req, so the # captured slice is already the full set we need to project. gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] - gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) - projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) gen_num_accepted_long = gen_num_accepted.long() - if self._ctx_base is None: - col_idx = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) - else: - # Windowed arena: physical columns. - col_idx = (ctx_len_gen - self._ctx_base[slots]).unsqueeze( - 1 - ) + offsets_kp1.unsqueeze(0) - write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) - if self._ctx_block_tables is not None: - # Bound by what the manager handed THIS request: the masked - # entries below are still written (fixed-size, for graph - # safety) and a placeholder page resolves to another's. - gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size - col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) - elif self._ctx_keep: - col_idx = col_idx.clamp(max=self._ctx_k_buf.size(2) - 1) - else: - col_idx = col_idx.clamp(max=self._max_ctx - 1) - - # Fixed-size writes for CUDA graph compatibility: - # Write ALL entries but zero out invalid ones. Invalid - # writes land at clamped column indices beyond valid - # _ctx_len range, so they're harmless. - slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) - col_flat = col_idx.reshape(-1) - proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] - pos_flat = ctx_position_ids.long().reshape(-1) - mask_1d = write_mask.reshape(-1) - - # Fast path: store the pre-projected/pre-RoPE'd K/V. - # dflash_forward reads these directly via cache_batch_idx. - cache_dtype = ( - self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype + self._store_step_context( + draft_model, gen_hs, slots, gen_rows_out, ctx_len_gen, gen_num_accepted_long ) - k_new, v_new = draft_model.precompute_context_kv( - proj_flat.to(cache_dtype), pos_flat - ) - mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) - k_new.mul_(mask_bc) - if v_new is not None: - v_new.mul_(mask_bc) - slot_long = slot_flat.long() - col_long = col_flat.long() - if self._ctx_paged: - if self._ctx_block_tables is not None: - # Batch positions of the gen requests, matching the - # per-request block table's row order. - rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) - else: - rows_long = slot_long - self._store_context_kv_paged(k_new, v_new, rows_long, col_long) - else: # VANILLA DFlash backend (FlashAttention) - self._ctx_k_buf[slot_long, :, col_long] = k_new - if v_new is not None: - self._ctx_v_buf[slot_long, :, col_long] = v_new - - self._advance_ctx_len(slots, gen_num_accepted_long) - if self._ctx_keep and not self._ctx_paged: - self._compact_ctx_window(slots) - - if self._ctx_base is None: - num_ctx_per_req_t = self._ctx_len[slots] + num_ctx_per_req_t = self._commit_context(slots, gen_num_accepted_long, gen_rows_out) else: - # Windowed arena: the physical length the draft forward appends - # after and attends over. - num_ctx_per_req_t = self._ctx_len[slots] - self._ctx_base[slots] - if self._ctx_block_tables is not None: - # The write above clamps columns to the request's allocation, so - # a context that outruns it has its tail written over the last - # valid slot. Truncate what is advertised to match, or the read - # attends that clobbered value -- or another request's page. - allocated = dflash_allocated_ctx_limit( - self._ctx_block_counts, self._ctx_page_size, block_size - ) - if self._ctx_unified_cache: - self._check_block_within_allocation( - num_ctx_per_req_t, allocated[gen_rows_out], block_size - ) - num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + num_ctx_per_req_t = self._drafter_context_lengths(slots, gen_rows_out) noise_embedding = noise_embed_2d query_positions = query_position_ids.long() @@ -2574,22 +2627,6 @@ def prepare_1st_drafter_inputs( gen_rows_out = slots bonus = torch.empty(0, dtype=torch.long, device="cuda") - return { - "noise_embedding": noise_embedding, - "query_positions": query_positions, - "num_ctx_per_req": num_ctx_per_req_t, - "ctx_k_cache": self._ctx_k_buf, - "ctx_v_cache": self._ctx_v_buf, - # Slots index the private arena; the manager's block table is keyed - # by batch position, so the drafter reads its own rows there. - "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, - # Anchor token per gen request (block slot 0): last accepted - # token. The dspark Markov chain conditions its first step on it. - "first_prev_tokens": bonus, - "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": ( - self._ctx_block_tables - if self._ctx_block_tables is not None - else self._ctx_page_table - ), - } + return self._drafter_inputs( + noise_embedding, query_positions, num_ctx_per_req_t, slots, gen_rows_out, bonus + ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py index fe29e3395073..b1a05ce2d415 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -39,6 +39,7 @@ def _worker(lengths, bases=None): _ctx_base=None if bases is None else torch.tensor(bases, dtype=torch.long, device="cuda"), ) w._advance_ctx_len = partial(DFlashWorker._advance_ctx_len, w) + w._empty_dummy_slot = partial(DFlashWorker._empty_dummy_slot, w) return w From 1b5ebf5015278146bf9da11ec8f095857a661d51 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 11:17:01 -0700 Subject: [PATCH 3/6] [None][fix] Commit the Muse modeling_v2 drafter context through the DFlash worker The verify fold and the context-K/V fork carry their own copies of the DFlash worker's context commit, and the copies have drifted from the worker: - After every length update the worker empties the dummy slot, which CUDA-graph padding rows share: they attend no context, and their next context K/V starts inside the padding request's own pages. The fork's commit kernel and the fold's store did not, so the dummy slot's length grew by the padding rows' accepted counts every step. Over the KV cache manager's pages a padding row then stores past the padding request's pages, onto page 0 of another request. - On a unified target/draft KV cache the worker checks the drafter's block against the request's allocation before clamping the lengths it advertises; the fold clamped without the check. Rather than patch the copies again, both now run the worker's stages: the fold stores and commits with _store_step_context and _commit_context, the fork's side chain stores with _store_step_context, and the fork's commit kernel keeps only the part it fuses, the zeroing of the masked rows, before the worker's _commit_context. Over a windowed arena that drafter_window_compact compacts, the side stream advances the lengths and the prepare settles the context after the join, where the compaction ran before; the sm107 windowed commit still advances and compacts in its one launch, then applies the worker's dummy-slot rule and reads its lengths. The drafter's inputs are the worker's _drafter_inputs. Signed-off-by: Zheyu Fu --- .../kernels/context_kv_fork/fork.py | 162 +++++--------- .../kernels/verify_fold/fold.py | 210 +++--------------- .../ops/spec/context_kv_commit.py | 125 +++-------- .../test_modeling_v2_drafter_context.py | 29 +++ 4 files changed, 153 insertions(+), 373 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py index 77090ff2cb08..27e84adc3a80 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py @@ -20,23 +20,23 @@ forked from the main stream at a fork point, it stores all rows unmasked; the main stream meanwhile runs the head, the sampler and the acceptance. Once the prepare has read the lengths it needs (the query positions), the side stream -waits for the main stream and runs one kernel (``ops/spec/context_kv_commit.py``) -that zeroes the masked rows in place, advances the lengths and writes the new -lengths the drafter reads, while the main stream runs the bonus-token -embedding; the main stream joins the side stream at the end of -``prepare_1st_drafter_inputs``. The arena, the lengths and everything the -drafter reads are the built-in's bytes. The fork point: right after the -target's LM-head GEMM (``llm.lm_head``'s ``apply_linear``, bound on the -instance), before its column gather. +waits for the main stream, zeroes the masked rows in place in one kernel +(``ops/spec/context_kv_commit.py``) and commits the context with the worker's +``_commit_context``, while the main stream runs the bonus-token embedding; the +main stream joins the side stream at the end of ``prepare_1st_drafter_inputs``. +The arena, the lengths and everything the drafter reads are the built-in's +bytes. The fork point: right after the target's LM-head GEMM +(``llm.lm_head``'s ``apply_linear``, bound on the instance), before its column +gather. Over a windowed arena (``TRTLLM_DFLASH_WINDOWED_CTX_ARENA``: the worker keeps the last ``_ctx_keep`` columns of each slot, physical column 0 holding logical -token ``_ctx_base``) the side chain stores at the physical columns the -built-in uses (``ctx_len - base + r``, clamped to the arena), with the +token ``_ctx_base``) the worker's store writes the physical columns, with the logical RoPE positions. Without ``drafter_window_compact`` (``compacts_after_join`` -False) the commit is ``ops/spec/context_kv_commit_windowed.py``: the zeroing, the -length update, then the built-in's compaction of each slot to its window and the -physical lengths the drafter reads, all on the side stream. +False) the zeroing, the length update and the compaction of each slot to its +window are one launch (``ops/spec/context_kv_commit_windowed.py``), followed by +the worker's dummy-slot rule and its read of the physical lengths, all on the +side stream. ``DecodeTarget.forward`` arms the fork with the step's metadata, so only decode steps (no context rows) fork; the step must also be one the side chain @@ -61,10 +61,11 @@ physical column 0 holds its logical token ``_ctx_base[slot]``, and the built-in stores row ``r`` at column ``min(_ctx_len - _ctx_base + r, capacity - 1)``, then compacts the slot's window (``_compact_ctx_window``) and hands the drafter ``_ctx_len - _ctx_base``. -The side chain and the commit address those columns (``arena_columns``), and -the prepare runs the compaction and that read on the main stream after the -join: the compaction moves the columns the side chain wrote, so it must follow -them, and nothing else in the step writes ``_ctx_base`` before it. +The zeroing addresses those columns (``arena_columns``) and the side stream +advances the lengths; the prepare settles the context (``_settle_context``: the +compaction and that read) on the main stream after the join: the compaction +moves the columns the side chain wrote, so it must follow them, and nothing else +in the step writes ``_ctx_base`` before it. """ from __future__ import annotations @@ -163,35 +164,16 @@ def compacts_after_join(worker) -> bool: True when ``drafter_window_compact`` re-classed the worker (``WindowCompactDFlashWorker``, installed after this fork): the commit then - only zeroes and advances the lengths (``context_kv_commit`` with the bases), - and the prepare runs that compaction and reads the physical lengths on the main - stream after the join (``drafter_lengths``). Otherwise the commit is - ``context_kv_commit_windowed``, which compacts and writes the physical lengths - itself on the side stream. + zeroes the masked rows (``context_kv_commit`` with the bases) and advances the + lengths, and the prepare settles the context on the main stream after the join + (``_settle_context``). Otherwise the commit is ``context_kv_commit_windowed``, + which advances the lengths and compacts in its own launch on the side stream. """ from ..drafter_window_compact.compact import WindowCompactDFlashWorker return isinstance(worker, WindowCompactDFlashWorker) -def drafter_lengths(worker, slots: torch.Tensor, new_lengths: torch.Tensor) -> torch.Tensor: - """The built-in prepare's end once this step's rows are stored: ``num_ctx_per_req``. - - ``new_lengths`` are what the commit returned: ``_ctx_len[slots]`` after the - update, which is what the drafter reads over the dense arena, or the physical - lengths after the compaction when the windowed commit compacted - (``compacts_after_join`` False). When the worker's compaction follows the join, - over the windowed arena the built-in compacts the slots' windows first - (``_compact_ctx_window``) and reads the physical lengths - ``_ctx_len[slots] - _ctx_base[slots]``. - """ - if worker._ctx_base is None or not compacts_after_join(worker): - return new_lengths - if worker._ctx_keep and not worker._ctx_paged: - worker._compact_ctx_window(slots) - return worker._ctx_len[slots] - worker._ctx_base[slots] - - # -------------------------------------------------------------------------- # The side chain and its controller # -------------------------------------------------------------------------- @@ -222,6 +204,8 @@ def __init__(self, worker: DFlashWorker, draft_model: nn.Module, point: str): self.draft_model = draft_model self.point = point self.stream = dedicated_stream() + #: A decode step's batch rows (``[:requests]``), as the worker's stages take them. + self.rows = torch.arange(MAX_REQUESTS, device="cuda") self.launches = {"head": 0, "verify": 0} self._armed: Optional[tuple] = None self._launched: Optional[Launched] = None @@ -291,14 +275,16 @@ def join(self, launched: Launched) -> None: main.wait_stream(self.stream) self._launched = None - def commit(self, launched: Launched, slots, num_accepted) -> torch.Tensor: - """Zero the masked rows and advance the lengths on the side stream; return the new lengths. + def commit(self, launched: Launched, slots, num_accepted) -> Optional[torch.Tensor]: + """Zero the masked rows and commit the context on the side stream; return the drafter's lengths. The side stream waits for the main stream's work so far -- the acceptance and the prepare's last read of ``_ctx_len`` -- and runs the - commit kernel behind the stored rows; the main stream does not wait - until ``join``. The returned ``[gens]`` lengths are allocated on the - main stream and written by the side, so they are valid after the join. + zeroing and the worker's commit behind the stored rows; the main stream + does not wait until ``join``. The returned ``[gens]`` lengths are + allocated on the main stream and written by the side, so they are valid + after the join. None over a windowed arena the worker compacts + (``compacts_after_join``): the prepare settles it after the join. """ worker = self.worker main = torch.cuda.current_stream() @@ -308,33 +294,25 @@ def commit(self, launched: Launched, slots, num_accepted) -> torch.Tensor: "capture state than the one committing it" ) num_ctx = torch.empty(launched.num_gens, dtype=torch.int64, device=worker._ctx_len.device) + rows = self.rows[: launched.num_gens] base, last = arena_columns(worker) self.stream.wait_stream(main) with torch.cuda.stream(self.stream): - if base is None: - context_kv_commit( - worker._ctx_k_buf, - worker._ctx_v_buf, - worker._ctx_len, - slots, - num_accepted, - num_ctx, - launched.k + 1, - worker._max_ctx, - ) - elif compacts_after_join(worker): + if base is None or compacts_after_join(worker): context_kv_commit( worker._ctx_k_buf, worker._ctx_v_buf, worker._ctx_len, slots, num_accepted, - num_ctx, launched.k + 1, - worker._max_ctx, - base, last, + base, ) + if base is not None: + worker._advance_ctx_len(slots, num_accepted.long()) + return None + num_ctx.copy_(worker._commit_context(slots, num_accepted.long(), rows)) else: context_kv_commit_windowed( worker._ctx_k_buf, @@ -349,6 +327,9 @@ def commit(self, launched: Launched, slots, num_accepted) -> torch.Tensor: column_limit(worker), worker._ctx_keep, ) + # The launch advanced and compacted; the worker's rules for the rest. + worker._empty_dummy_slot() + num_ctx.copy_(worker._drafter_context_lengths(slots, rows)) return num_ctx def settle(self) -> None: @@ -403,38 +384,19 @@ def plan(self, spec_metadata, attn_metadata) -> Optional[tuple[int, int]]: def side_chain(self, spec_metadata, num_gens: int, k: int) -> torch.Tensor: """The built-in store of the step's context K/V, unmasked; runs on the side stream. - The same ops on the same operands as the built-in, in the built-in's - order, minus the mask multiply. Its first launch is a torch index - kernel, so the fork is a plain dependency on the main stream's last - kernel. + The worker's ``_store_step_context`` without its mask. Its first launch is a + torch index kernel, so the fork is a plain dependency on the main stream's + last kernel. """ - worker, drafter = self.worker, self.draft_model - rows = num_gens * (k + 1) + worker = self.worker slots = worker._batch_to_slot[:num_gens] - ctx_len = worker._ctx_len[slots] - offsets = torch.arange(k + 1, dtype=torch.long, device="cuda") - positions = ctx_len.unsqueeze(1) + offsets.unsqueeze(0) - base, last = arena_columns(worker) - if base is None: - col = positions.clamp(max=last) - else: - # Windowed arena: physical columns, clamped as the built-in clamps them. - col = ((ctx_len - base[slots]).unsqueeze(1) + offsets.unsqueeze(0)).clamp(max=last) # The built-in's get_hidden_states(num_tokens)[0 : rows]: the same view. gen_hs = spec_metadata.captured_hidden_states[ - :rows, : spec_metadata.hidden_size * spec_metadata.num_capture_layers + : num_gens * (k + 1), : spec_metadata.hidden_size * spec_metadata.num_capture_layers ] - projected = drafter.project_target_hidden(gen_hs.reshape(-1, gen_hs.shape[-1])) - slot_flat = slots.unsqueeze(1).expand(-1, k + 1).reshape(-1) - col_flat = col.reshape(-1) - pos_flat = positions.long().reshape(-1) - k_new, v_new = drafter.precompute_context_kv( - projected.to(worker._ctx_k_buf.dtype), pos_flat + worker._store_step_context( + self.draft_model, gen_hs, slots, self.rows[:num_gens], worker._ctx_len[slots] ) - slot_long, col_long = slot_flat.long(), col_flat.long() - worker._ctx_k_buf[slot_long, :, col_long] = k_new - if v_new is not None: - worker._ctx_v_buf[slot_long, :, col_long] = v_new return gen_hs @@ -529,10 +491,10 @@ def _fork_prepare( for its one path here -- no context rows, captured target features, the dense arena (full-length or windowed) -- with the store moved to the side chain: the bonus token and the query positions as the built-in computes - them, then the commit (zeroing, lengths and the new lengths the drafter - reads) launched on the side stream, the bonus-token embedding on the main - stream meanwhile, and the join; over a windowed arena then the built-in's - window compaction and its read of the physical lengths (``drafter_lengths``). + them, then the commit launched on the side stream, the bonus-token + embedding on the main stream meanwhile, and the join; over a windowed arena + the worker compacts, then the prepare settles the context + (``_settle_context``). """ num_gens = launched.num_gens if ( @@ -589,23 +551,15 @@ def _fork_prepare( noise_embed_2d[:, 0, :] = embed_bonus fork.join(launched) - num_ctx_per_req_t = drafter_lengths(self, slots, num_ctx_per_req_t) + rows = fork.rows[:num_gens] + if num_ctx_per_req_t is None: + num_ctx_per_req_t = self._settle_context(slots, rows) - # The dense arena (plan()) has no block table: the built-in's - # allocation clamp of the advertised length does not apply. attn_metadata._seq_lens_cuda[0:num_gens] = total_tokens_per_req attn_metadata._seq_lens[0:num_gens] = total_tokens_per_req - return { - "noise_embedding": noise_embed_2d, - "query_positions": query_position_ids.long(), - "num_ctx_per_req": num_ctx_per_req_t, - "ctx_k_cache": self._ctx_k_buf, - "ctx_v_cache": self._ctx_v_buf, - "ctx_cache_batch_idx": slots, - "first_prev_tokens": bonus, - "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": self._ctx_page_table, - } + return self._drafter_inputs( + noise_embed_2d, query_position_ids.long(), num_ctx_per_req_t, slots, rows, bonus + ) #: ``ContextKvForkDFlashWorker`` combined with another feature's worker class, by that class. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py index cbb858ad7d0a..23f121dd4d63 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py @@ -48,22 +48,14 @@ follows the call order of ``DFlashWorker._forward_impl``. ``prepare_1st_drafter_inputs`` is written out because its bonus part is -inline; the part that stores this step's accepted context into the drafter's -arena (verify-independent work, plus the count-dependent write mask and -``_ctx_len`` update) is ``_fold_store_accepted_context``, the built-in lines -unchanged, so a later fast path can override it on its own. With the -context-K/V fork (``context_kv_fork/fork.py``, installed before this module, so its -class sits behind this one) a folded step whose store the fork's side chain -ran skips that part: ``_fold_forked_inputs`` launches the fork's commit (the -zeroing and the length update) right away -- the fold kernel was the main -stream's last read of ``_ctx_len`` -- builds the noise block meanwhile and -joins the side stream before returning; a step this class does not fold -reaches the fork's own prepare through ``super()``. - -Over the windowed drafter context arena (``TRTLLM_DFLASH_WINDOWED_CTX_ARENA=1``) -the store and the drafter's lengths follow the built-in's windowed lines: rows -at the physical columns ``_ctx_len - _ctx_base + r``, the window compaction -after the length update, and ``_ctx_len - _ctx_base`` for the drafter. +inline; it stores and commits this step's context with the worker's own stages +(``_store_step_context``, ``_commit_context``). With the context-K/V fork +(``context_kv_fork/fork.py``, installed before this module, so its class sits +behind this one) a folded step whose store the fork's side chain ran skips the +store: ``_fold_forked_inputs`` launches the fork's commit right away -- the fold +kernel was the main stream's last read of ``_ctx_len`` -- builds the noise block +meanwhile and joins the side stream before returning; a step this class does +not fold reaches the fork's own prepare through ``super()``. Once the fold has run in a step it has applied ``kv_lens_cuda += 1`` and owes the rewind, so the methods that consume it (``FOLD_CONSUMERS``) never hand the @@ -82,7 +74,7 @@ class sits behind this one) a folded step whose store the fork's side chain import torch from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata -from tensorrt_llm._torch.speculative.dflash import DFlashWorker, dflash_allocated_ctx_limit +from tensorrt_llm._torch.speculative.dflash import DFlashWorker from tensorrt_llm._torch.speculative.interface import SpecWorkerBase from tensorrt_llm.logger import logger @@ -455,7 +447,6 @@ def prepare_1st_drafter_inputs( gen_num_accepted = num_accepted_tokens[num_contexts : num_contexts + num_gens] total_tokens_per_req = self._draft_tokens_per_req # K+1 - K = spec_metadata.runtime_draft_len captured_hs = spec_metadata.get_hidden_states(total_target_tokens) has_target_features = ( captured_hs is not None @@ -477,53 +468,22 @@ def prepare_1st_drafter_inputs( noise_embed_2d = self._fold_noise_block(embed_tokens, fold, num_gens, block_size) if has_target_features: - self._fold_store_accepted_context( - draft_model, - captured_hs, - attn_metadata, - num_gens, - total_tokens_per_req, - K, - slots, - gen_rows_out, - fold.offsets, - fold.ctx_positions, - gen_num_accepted, - fold.num_accepted_long, + # Target now processes exactly K+1 tokens per gen req, so the + # captured slice is already the full set we need to project. + gen_hs = captured_hs[: num_gens * total_tokens_per_req] + num_accepted_long = fold.num_accepted_long + self._store_step_context( + draft_model, gen_hs, slots, gen_rows_out, self._ctx_len[slots], num_accepted_long ) - - if self._ctx_base is None: - num_ctx_per_req_t = self._ctx_len[slots] + num_ctx_per_req_t = self._commit_context(slots, num_accepted_long, gen_rows_out) else: - # Windowed arena: the physical length the draft forward appends - # after and attends over. - num_ctx_per_req_t = self._ctx_len[slots] - self._ctx_base[slots] - if self._ctx_block_tables is not None: - allocated = dflash_allocated_ctx_limit( - self._ctx_block_counts, self._ctx_page_size, block_size - ) - num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) - noise_embedding = noise_embed_2d - query_positions = fold.query_positions + num_ctx_per_req_t = self._drafter_context_lengths(slots, gen_rows_out) attn_metadata._seq_lens_cuda[num_contexts : num_contexts + num_gens] = total_tokens_per_req attn_metadata._seq_lens[num_contexts : num_contexts + num_gens] = total_tokens_per_req - - return { - "noise_embedding": noise_embedding, - "query_positions": query_positions, - "num_ctx_per_req": num_ctx_per_req_t, - "ctx_k_cache": self._ctx_k_buf, - "ctx_v_cache": self._ctx_v_buf, - "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, - "first_prev_tokens": bonus, - "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": ( - self._ctx_block_tables - if self._ctx_block_tables is not None - else self._ctx_page_table - ), - } + return self._drafter_inputs( + noise_embed_2d, fold.query_positions, num_ctx_per_req_t, slots, gen_rows_out, bonus + ) def _fold_forked_inputs( self, @@ -541,22 +501,13 @@ def _fold_forked_inputs( ): """The drafter inputs of a folded step whose context store the context-K/V fork's side chain ran. - ``context_kv_fork`` forked the store of - this step's ``K + 1`` captured rows onto its side stream, unmasked, from - the context lengths before this step; it is ``_fold_store_accepted_context`` - without the mask multiply and the length update. The fold kernel was the - main stream's last read of ``_ctx_len`` in this step (the query and - context positions are its outputs), so the side stream's commit -- the - zeroing of the rows past the accepted count and ``_ctx_len += accepted``, - clamped -- is launched first, the noise block runs on the main stream - meanwhile, and the main stream joins the side stream before returning. - The commit returns the new lengths, which are ``_ctx_len[slots]`` after - the built-in store (less ``_ctx_base[slots]`` after the compaction when the - windowed commit compacts); the side chain's plan admits only the dense - (unpaged) arena, full-length or windowed, where the built-in neither clamps - them to an allocation nor indexes by batch row. When ``drafter_window_compact`` - compacts the windowed arena, the built-in's window compaction and its read of - the physical lengths follow the join (``context_kv_fork.drafter_lengths``). + ``context_kv_fork`` forked the store of this step's ``K + 1`` captured rows + onto its side stream, unmasked, from the context lengths before this step. + The fold kernel was the main stream's last read of ``_ctx_len`` in this step + (the query and context positions are its outputs), so the fork's commit is + launched first, the noise block runs on the main stream meanwhile, and the + main stream joins the side stream before returning, then settles a windowed + arena that ``drafter_window_compact`` compacts (``_settle_context``). """ num_gens = fold.accepted.shape[0] if ( @@ -572,111 +523,16 @@ def _fold_forked_inputs( + ("; the arena is paged" if self._ctx_block_tables is not None else "") + ("; the step has no target features" if not has_target_features else "") ) - from ..context_kv_fork.fork import drafter_lengths - - num_ctx_per_req_t = fork.commit(launched, slots, gen_num_accepted) + num_ctx = fork.commit(launched, slots, gen_num_accepted) noise_embed_2d = self._fold_noise_block(embed_tokens, fold, num_gens, block_size) fork.join(launched) - num_ctx_per_req_t = drafter_lengths(self, slots, num_ctx_per_req_t) + if num_ctx is None: + num_ctx = self._settle_context(slots, fold.gen_rows) attn_metadata._seq_lens_cuda[0:num_gens] = total_tokens_per_req attn_metadata._seq_lens[0:num_gens] = total_tokens_per_req - return { - "noise_embedding": noise_embed_2d, - "query_positions": fold.query_positions, - "num_ctx_per_req": num_ctx_per_req_t, - "ctx_k_cache": self._ctx_k_buf, - "ctx_v_cache": self._ctx_v_buf, - "ctx_cache_batch_idx": slots, - "first_prev_tokens": fold.bonus, - "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": self._ctx_page_table, - } - - def _fold_store_accepted_context( - self, - draft_model, - captured_hs, - attn_metadata, - num_gens, - total_tokens_per_req, - K, - slots, - gen_rows_out, - offsets_kp1, - ctx_position_ids, - gen_num_accepted, - gen_num_accepted_long, - ) -> None: - """The built-in ``if has_target_features:`` block of ``prepare_1st_drafter_inputs``, unchanged. - - Projects this step's captured target rows, writes their context K/V - into the drafter's arena (rows past the accepted count zeroed; at the - physical columns over a windowed arena) and advances ``_ctx_len``, then - compacts a windowed arena's slots to their window as the built-in does. - ``gen_num_accepted_long`` is the fold's copy of ``gen_num_accepted.long()``. - """ - gen_start = attn_metadata.num_ctx_tokens - # Target now processes exactly K+1 tokens per gen req, so the - # captured slice is already the full set we need to project. - gen_hs = captured_hs[gen_start : gen_start + num_gens * total_tokens_per_req] - gen_hs_to_project = gen_hs.reshape(-1, gen_hs.shape[-1]) - projected_to_store = draft_model.project_target_hidden(gen_hs_to_project) - if self._ctx_base is None: - col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) - else: - # Windowed arena: physical columns. - col_idx = (self._ctx_len[slots] - self._ctx_base[slots]).unsqueeze( - 1 - ) + offsets_kp1.unsqueeze(0) - write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) - if self._ctx_block_tables is not None: - # Bound by what the manager handed THIS request: the masked - # entries below are still written (fixed-size, for graph - # safety) and a placeholder page resolves to another's. - gen_capacity = self._ctx_block_counts[gen_rows_out] * self._ctx_page_size - col_idx = torch.minimum(col_idx, (gen_capacity - 1).unsqueeze(1)).clamp_(min=0) - elif self._ctx_keep: - col_idx = col_idx.clamp(max=self._ctx_k_buf.size(2) - 1) - else: - col_idx = col_idx.clamp(max=self._max_ctx - 1) - - # Fixed-size writes for CUDA graph compatibility: - # Write ALL entries but zero out invalid ones. Invalid - # writes land at clamped column indices beyond valid - # _ctx_len range, so they're harmless. - slot_flat = slots.unsqueeze(1).expand(-1, K + 1).reshape(-1) - col_flat = col_idx.reshape(-1) - proj_flat = projected_to_store # already [num_gens*(K+1), proj_dim] - pos_flat = ctx_position_ids.long().reshape(-1) - mask_1d = write_mask.reshape(-1) - - # Fast path: store the pre-projected/pre-RoPE'd K/V. - # dflash_forward reads these directly via cache_batch_idx. - cache_dtype = self._ctx_kv_buf[0].dtype if self._ctx_paged else self._ctx_k_buf.dtype - k_new, v_new = draft_model.precompute_context_kv(proj_flat.to(cache_dtype), pos_flat) - mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) - k_new.mul_(mask_bc) - if v_new is not None: - v_new.mul_(mask_bc) - slot_long = slot_flat.long() - col_long = col_flat.long() - if self._ctx_paged: - if self._ctx_block_tables is not None: - # Batch positions of the gen requests, matching the - # per-request block table's row order. - rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) - else: - rows_long = slot_long - self._store_context_kv_paged(k_new, v_new, rows_long, col_long) - else: # VANILLA DFlash backend (FlashAttention) - self._ctx_k_buf[slot_long, :, col_long] = k_new - if v_new is not None: - self._ctx_v_buf[slot_long, :, col_long] = v_new - - self._ctx_len[slots] += gen_num_accepted_long - self._ctx_len.clamp_(max=self._max_ctx) - if self._ctx_keep and not self._ctx_paged: - self._compact_ctx_window(slots) + return self._drafter_inputs( + noise_embed_2d, fold.query_positions, num_ctx, slots, fold.gen_rows, fold.bonus + ) # ------------------------------------------------------------------ # path_walk_tail: the step's tail in the path walk diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py index 8d0df1f2a5c5..f33f8e1db849 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py @@ -1,33 +1,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""DFlash prepare: the verify-dependent end of the context-K/V store, in one launch. +"""DFlash prepare: the verify-dependent K/V end of the context-K/V store, in one launch. The built-in ``DFlashWorker.prepare_1st_drafter_inputs`` (``_torch/speculative/dflash.py``) projects the ``K + 1`` captured target rows of every generation request into the drafter's context K/V, multiplies the rows at and past the request's accepted count by zero, stores all ``K + 1`` rows into the dense context arena -at columns ``min(ctx_len[slot] + r, max_ctx - 1)``, then advances -``ctx_len[slot]`` by the accepted count and clamps every slot's length to -``max_ctx``. The ``context_kv_fork`` fast path stores the rows unmasked earlier, on a side stream -that does not wait for the verify; this entry is what is left once the -accepted counts exist: the zeroing of the masked rows, in place in the arena, -the length update, and the built-in's read of the new lengths -(``_ctx_len[slots]``, the drafter's ``num_ctx_per_req``) into an output. +at columns ``min(ctx_len[slot] + r, max_ctx - 1)``, then commits the context +(``DFlashWorker._commit_context``). The ``context_kv_fork`` fast path stores the +rows unmasked earlier, on a side stream that does not wait for the verify; this +entry is the K/V part of what is left once the accepted counts exist: the zeroing +of the masked rows, in place in the arena, before the worker's commit. The zeroing multiplies each stored BF16 element by 0.0 in fp32 and rounds the product back to BF16, as ``k_new.mul_(mask)`` does on a BF16 tensor, so the arena ends up holding the built-in's bytes. Columns are computed from each -slot's length before any length is written, as the built-in's ``col_idx`` is -taken before its update: a CTA barrier separates every thread's reads of the -lengths from the first length write. +slot's length before the commit, as the built-in's ``col_idx`` is. Over the windowed drafter context arena (``TRTLLM_DFLASH_WINDOWED_CTX_ARENA=1``) a slot's physical column 0 holds its logical token ``ctx_base[slot]``, and the built-in stores row ``r`` at column ``min(ctx_len[slot] - ctx_base[slot] + r, capacity - 1)``. Given ``ctx_base`` and that last column, the entry zeroes those -columns instead; the lengths are updated as before (the built-in's window -compaction and its read of ``ctx_len - ctx_base`` follow the call, on the caller's -side). Without ``ctx_base`` the arithmetic is the dense arena's, unchanged. +columns instead. Without ``ctx_base`` the arithmetic is the dense arena's, unchanged. """ from __future__ import annotations @@ -91,11 +85,7 @@ def _context_kv_commit_kernel( base_ptr, slots_ptr, accepted_ptr, - new_len_ptr, - num_ctx_ptr, num_gens, - num_slots, - max_ctx, last_col, k_slot_stride, k_layer_stride, @@ -111,16 +101,12 @@ def _context_kv_commit_kernel( HAS_V: tl.constexpr, HAS_BASE: tl.constexpr, ): - # Every request's new length is computed from the lengths as they were - # before this call (the built-in gathers ``ctx_len[slots]`` once, adds, and - # scatters), so the zeroing and the new lengths read no length this call wrote. for g in range(num_gens): slot = tl.load(slots_ptr + g).to(tl.int64) - ctx = tl.load(ctx_len_ptr + slot) # The physical column of the slot's next context token. - origin = ctx + origin = tl.load(ctx_len_ptr + slot) if HAS_BASE: - origin = ctx - tl.load(base_ptr + slot) + origin = origin - tl.load(base_ptr + slot) accepted = tl.load(accepted_ptr + g).to(tl.int64) # A tile of K columns and its V columns are loaded before either is # stored, so their loads are in flight together. @@ -141,21 +127,6 @@ def _context_kv_commit_kernel( else: k = tl.load(k_cols, mask=zeroed, other=0.0) tl.store(k_cols, _times_zero(k), mask=zeroed) - tl.store(new_len_ptr + g, ctx + accepted) - # The lengths are written by one thread below; every thread's reads of - # them above (the zeroed columns' addresses) come first. - tl.debug_barrier() - # Scatter in request order, then clamp every slot, as the built-in's - # ``ctx_len[slots] += n`` followed by ``ctx_len.clamp_(max=max_ctx)``. - for g in range(num_gens): - slot = tl.load(slots_ptr + g).to(tl.int64) - tl.store(ctx_len_ptr + slot, tl.load(new_len_ptr + g)) - for s in range(num_slots): - tl.store(ctx_len_ptr + s, tl.minimum(tl.load(ctx_len_ptr + s), max_ctx)) - # The built-in's read of the new lengths after its update (``_ctx_len[slots]``). - for g in range(num_gens): - slot = tl.load(slots_ptr + g).to(tl.int64) - tl.store(num_ctx_ptr + g, tl.load(ctx_len_ptr + slot)) def tile_pairs(rows: int, layers: int, width: int) -> int: @@ -165,27 +136,22 @@ def tile_pairs(rows: int, layers: int, width: int) -> int: class ContextKvCommit(OpWrapper): - """Zero the context-K/V rows at and past each request's accepted count, then advance its length. + """Zero the context-K/V rows at and past each request's accepted count. ``k_arena`` / ``v_arena`` are the dense context arenas ``[slots, layers, capacity, kv_heads, head_dim]`` (BF16; ``v_arena`` may be None); the ``rows`` - columns ``min(ctx_len[slot] + r, max_ctx - 1)``, ``r < rows``, of every + columns ``min(ctx_len[slot] + r, last_column)``, ``r < rows``, of every request's slot hold that step's unmasked context K/V. ``ctx_len`` is int64 ``[slots]``, ``slots`` int64 ``[gens]`` (one per generation request), - ``num_accepted`` int32 or int64 ``[gens]``, ``num_ctx_out`` int64 - ``[gens]``. In place: - - * for each request ``g`` and row ``r >= num_accepted[g]``, in every layer of - both arenas: the column's elements become ``bf16(fp32(x) * 0.0)``; - * ``ctx_len[slots[g]] = ctx_len[slots[g]] + num_accepted[g]`` (gathered - before any write, scattered in request order), then every slot's length - is clamped to ``max_ctx``; - * ``num_ctx_out[g] = ctx_len[slots[g]]`` (the new lengths). + ``num_accepted`` int32 or int64 ``[gens]``. In place, for each request ``g`` + and row ``r >= num_accepted[g]``, in every layer of both arenas: the column's + elements become ``bf16(fp32(x) * 0.0)``. ``last_column`` must be a column of + the arenas. ``ctx_base`` (int64 ``[slots]``, the windowed arena's logical index of each slot's column 0) moves the rows to columns ``min(ctx_len[slot] - - ctx_base[slot] + r, last_column)``; ``last_column`` defaults to ``max_ctx - - 1`` and must be a column of the arenas. ``ctx_base`` is read, never written. + ctx_base[slot] + r, last_column)``. ``ctx_len`` and ``ctx_base`` are read, + never written. """ ARCHS = frozenset({Arch.SM_103}) @@ -197,26 +163,23 @@ class ContextKvCommit(OpWrapper): served shape), loading a tile of both arenas before storing either, so the loads of a request's columns are in flight together rather than one column after another. Elements of a (row, layer) column are contiguous - (``kv_heads * head_dim``, one block). A CTA barrier separates the - zeroing, whose column addresses every thread computes from the lengths, - from the length writes, which one thread issues. + (``kv_heads * head_dim``, one block). Call site: ``context_kv_fork.ContextKvFork.commit`` (from ``ContextKvForkDFlashWorker.prepare_1st_drafter_inputs``) under the ``context_kv_fork`` fast path, on the side stream that stored the unmasked - rows, behind the main stream's last read of ``ctx_len`` in the prepare; - the main stream joins before the drafter reads the arena and - ``num_ctx_out``. + rows, behind the main stream's last read of ``ctx_len`` in the prepare and + before the worker's commit; the main stream joins before the drafter reads + the arena. Numerics: the zeroing is the BF16 product with 0.0 (finite x -> +-0 with x's sign, inf / NaN -> the canonical NaN), as the built-in's ``mul_`` by a - BF16 mask; the lengths are integer arithmetic. Bitwise the built-in's - arena and lengths when the arena held the unmasked rows. The built-in also - multiplies the kept rows by 1.0, which leaves every value but a + BF16 mask. Bitwise the built-in's arena when the arena held the unmasked + rows. The built-in also multiplies the kept rows by 1.0, which leaves every value but a non-canonical NaN unchanged; the GPU kernels that compute the stored rows produce only the canonical NaN. - The arena's end: rows whose column clamps share column ``max_ctx - 1``, + The arena's end: rows whose column clamps share column ``last_column``, and the built-in writes them with one ``index_put``, whose winner among duplicate indices is unspecified (a deterministic ``index_put`` keeps the last). The entry zeroes that column when any masked row maps to it; the @@ -244,18 +207,14 @@ def __call__( ctx_len: torch.Tensor, slots: torch.Tensor, num_accepted: torch.Tensor, - num_ctx_out: torch.Tensor, rows: int, - max_ctx: int, + last_column: int, ctx_base=None, - last_column=None, ) -> None: num_gens = slots.numel() if num_gens == 0: return - last = max_ctx - 1 if last_column is None else last_column layers, width = k_arena.shape[1], k_arena.shape[3] * k_arena.shape[4] - new_len = torch.empty(num_gens, dtype=torch.int64, device=ctx_len.device) has_v = v_arena is not None values = v_arena if has_v else k_arena _context_kv_commit_kernel[(1,)]( @@ -265,12 +224,8 @@ def __call__( ctx_len if ctx_base is None else ctx_base, slots, num_accepted, - new_len, - num_ctx_out, num_gens, - ctx_len.numel(), - max_ctx, - last, + last_column, k_arena.stride(0), k_arena.stride(1), k_arena.stride(2), @@ -294,13 +249,11 @@ def reference( ctx_len: torch.Tensor, slots: torch.Tensor, num_accepted: torch.Tensor, - num_ctx_out: torch.Tensor, rows: int, - max_ctx: int, + last_column: int, ctx_base=None, - last_column=None, ) -> None: - """The built-in torch ops, in place: mask the stored rows in row order, then update the lengths. + """The built-in torch ops, in place: mask the stored rows in row order. One row at a time (slots are distinct, so no ``index_put`` below sees a duplicate index): where rows share the arena's last column, the masks @@ -310,8 +263,7 @@ def reference( """ offsets = torch.arange(rows, dtype=torch.long, device=ctx_len.device) origin = ctx_len[slots] if ctx_base is None else ctx_len[slots] - ctx_base[slots] - last = max_ctx - 1 if last_column is None else last_column - col = (origin.unsqueeze(1) + offsets.unsqueeze(0)).clamp(max=last) + col = (origin.unsqueeze(1) + offsets.unsqueeze(0)).clamp(max=last_column) mask = offsets.unsqueeze(0) < num_accepted.long().unsqueeze(1) for arena in (k_arena, v_arena): if arena is None: @@ -320,9 +272,6 @@ def reference( stored = arena[slots, :, col[:, r]] stored.mul_(mask[:, r].view(-1, 1, 1, 1).to(stored.dtype)) arena[slots, :, col[:, r]] = stored - ctx_len[slots] += num_accepted.long() - ctx_len.clamp_(max=max_ctx) - num_ctx_out.copy_(ctx_len[slots]) def is_valid( self, @@ -331,13 +280,10 @@ def is_valid( ctx_len: torch.Tensor, slots: torch.Tensor, num_accepted: torch.Tensor, - num_ctx_out: torch.Tensor, rows: int, - max_ctx: int, + last_column: int, ctx_base=None, - last_column=None, ) -> None: - last = max_ctx - 1 if last_column is None else last_column for name, arena in (("k_arena", k_arena), ("v_arena", v_arena)): if arena is None and name == "v_arena": continue @@ -347,8 +293,8 @@ def is_valid( assert arena.stride(4) == 1 and arena.stride(3) == arena.shape[4], ( f"{name}'s (kv_heads, head_dim) elements of a column must be contiguous" ) - assert 0 <= last < arena.shape[2], ( - f"{name} has {arena.shape[2]} columns; the last column {last} is not one of them" + assert 0 <= last_column < arena.shape[2], ( + f"{name} has {arena.shape[2]} columns; the last column {last_column} is not one of them" ) if v_arena is not None: assert v_arena.shape[:2] == k_arena.shape[:2], ( @@ -367,11 +313,6 @@ def is_valid( "num_accepted must hold one count per request" ) assert num_accepted.dtype in (torch.int32, torch.int64), "num_accepted must be int32/int64" - assert ( - num_ctx_out.dtype == torch.int64 - and num_ctx_out.dim() == 1 - and num_ctx_out.numel() == slots.numel() - ), "num_ctx_out must be int64 [gens]" assert slots.numel() <= MAX_REQUESTS, f"at most {MAX_REQUESTS} requests" assert 1 <= rows <= 64, f"rows must be in [1, 64], got {rows}" diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py index 19e6cd83f368..e506895580b0 100644 --- a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py @@ -4,11 +4,15 @@ * ``dflash_window_attention`` and ``dflash_window_attention_wide`` (at each device's split geometry) over pages are bitwise their dense calls, the appended block included. +* ``context_kv_fork``'s commit empties the dummy slot, as the worker's own commit does. """ import pytest import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.context_kv_fork import ( # noqa: E501 + fork as context_kv_fork, +) from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import ( # noqa: E501 dflash_window_attention_wide as wide_ops, ) @@ -16,6 +20,7 @@ WINDOW, dflash_window_attention, ) +from tensorrt_llm._torch.speculative.dflash import DFlashWorker __extra_import_path__ = ["../../.."] from _validating import validating # noqa: E402 — needs the path declared above @@ -62,3 +67,27 @@ def test_the_window_ops_over_pages_are_their_dense_calls(sms, append, monkeypatc assert torch.equal(paged, dense) # The blocks went to the requests' pages as to their slots, and nothing else moved. assert torch.equal(k_pages[rows], k_arena) and torch.equal(v_pages[rows], v_arena) + + +@pytest.mark.parametrize("windowed", [False, True], ids=["dense", "sm107_windowed_commit"]) +def test_the_fork_commit_empties_the_dummy_slot(windowed): + """The dummy slot, which CUDA-graph padding rows share, is empty after the commit, the DFlash + worker's rule: over the dense arena, and after the sm107 windowed commit's own launch.""" + worker = object.__new__(DFlashWorker) + worker._ctx_k_buf = torch.zeros(4, 5, 64, 1, 128, dtype=torch.bfloat16, device="cuda") + worker._ctx_v_buf = torch.zeros_like(worker._ctx_k_buf) + worker._ctx_len = torch.tensor([5, 0, 0, 6], device="cuda") + # A windowed arena without drafter_window_compact: each slot starts at a base. + worker._ctx_base = torch.tensor([2, 0, 0, 3], device="cuda") if windowed else None + worker._ctx_keep, worker._max_ctx, worker._dummy_slot = (48 if windowed else 0), 64, 3 + worker._ctx_paged, worker._ctx_block_tables = False, None + fork = context_kv_fork.ContextKvFork(worker, None, context_kv_fork.POINT) + main = torch.cuda.current_stream() + launched = context_kv_fork.Launched(main, None, 2, 7, context_kv_fork.POINT, None) + # A request on slot 0 accepts 3 rows, a padding row on the dummy slot 8. + accepted = torch.tensor([3, 8], dtype=torch.int32, device="cuda") + num_ctx = fork.commit(launched, torch.tensor([0, 3], device="cuda"), accepted) + fork.join(launched) + assert worker._ctx_len.tolist() == [8, 0, 0, 0] + # The lengths the drafter attends: physical over the windowed arena. + assert num_ctx.tolist() == ([6, 0] if windowed else [8, 0]) From baef3a4c3015114fb1c4c28afd8dcd83a8d6f247 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 11:24:06 -0700 Subject: [PATCH 4/6] [None][feat] Run the Muse modeling_v2 DFlash fast paths over a paged drafter context The musespark_open TEP composition required the VANILLA DFlash attention backend, so a TRTLLM drafter, whose context lives in the KV cache manager's pages, dropped every fast path, the target's included, and `modeling_v2: require` refused to start. Disaggregated serving needs that drafter: only then is the drafter's context transferred with the target's KV. The plan now admits either layout. A paged context is never windowed, so over pages the plan keeps every path but the window compaction (47 of the 48 sm103 TEP4 paths) and says so in its log line. Every path that touches the drafter's context reads it where the worker keeps it: - the block decode hands each layer's view of the manager's pool and the worker's one page table to the window ops; - the context-K/V fork's side chain decodes this step's page table at the fork point, which the worker then skips, and stores with the worker's own store; its commit is the worker's alone, since the drafter's block append overwrites the rows past the accepted count before anything reads them; - the verify fold takes the fork's paged steps. Signed-off-by: Zheyu Fu --- .../dispatch.py | 18 +- .../kernels/context_kv_fork/fork.py | 65 ++++--- .../drafter_block_decode/block_decode.py | 76 +++++---- .../kernels/verify_fold/fold.py | 2 - .../attention.py | 2 + .../modeling.py | 2 +- .../test_modeling_v2_drafter_context.py | 160 ++++++++++++++++++ 7 files changed, 253 insertions(+), 72 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py index 3151233d545b..1f04bd6622c3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py @@ -545,8 +545,8 @@ def dflash_linear(spec_config: Any) -> bool: """Whether ``spec_config`` is DFlash drafting a linear chain of ``DRAFT_LEN`` tokens. That is the drafter every composition's drafter, worker and head paths are - written for: one draft per position (no tree), a fixed draft length (no - schedule), and the drafter's VANILLA pooled-context attention. + written for: one draft per position (no tree) and a fixed draft length (no + schedule). """ from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig @@ -559,7 +559,6 @@ def dflash_linear(spec_config: Any) -> bool: and getattr(spec_config, "max_draft_len", None) == DRAFT_LEN and getattr(spec_config, "max_total_draft_tokens", None) == DRAFT_LEN and getattr(spec_config, "draft_len_schedule", None) is None - and getattr(spec_config, "attention_backend", None) == "VANILLA" ) @@ -572,8 +571,7 @@ def describe_speculation(spec_config: Any) -> str: f"{getattr(mode, 'name', type(spec_config).__name__)} " f"max_draft_len={getattr(spec_config, 'max_draft_len', None)} " f"max_total_draft_tokens={getattr(spec_config, 'max_total_draft_tokens', None)} " - f"draft_len_schedule={getattr(spec_config, 'draft_len_schedule', None)} " - f"attention_backend={getattr(spec_config, 'attention_backend', None)}" + f"draft_len_schedule={getattr(spec_config, 'draft_len_schedule', None)}" ) @@ -604,7 +602,7 @@ def served_compositions(compositions: Mapping[CompositionKey, TargetPaths]) -> s return ( f"{keys}, each with DFlash drafting a linear chain of {DRAFT_LEN} tokens " "(max_draft_len = max_total_draft_tokens = " - f"{DRAFT_LEN}, no draft_len_schedule, VANILLA attention_backend) and an NVFP4 checkpoint" + f"{DRAFT_LEN}, no draft_len_schedule) and an NVFP4 checkpoint" ) @@ -1143,7 +1141,11 @@ def compose( f"{TAG} no fast path for this deployment ({reason}); running the built-in forward" ) return NO_PATHS - windowed = windowed_context_arena(spec_worker) + from tensorrt_llm._torch.speculative.dflash import _PAGED_ATTENTION_BACKENDS + + # The worker pages the drafter's context on these backends, and windows only a dense arena. + paged = getattr(spec_worker, "_dflash_attention_backend", None) in _PAGED_ATTENTION_BACKENDS + windowed = windowed_context_arena(spec_worker) and not paged paths, gated = capacity_gated(frozenset(served.served), max_batch_size) paths, yields = yielded(paths) if windowed: @@ -1169,6 +1171,8 @@ def compose( plan = Plan(paths) if windowed: composed += f" with {_WINDOWED_ARENA}" + if paged: + composed += " with a paged drafter context" for gate in gated: composed += ( f" without {gate.path} (max_batch_size {max_batch_size} above {gate.most_requests};" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py index 27e84adc3a80..76869aac21f4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py @@ -7,7 +7,7 @@ ``hidden_norm``: the drafter fc shard's sharded GEMM and fabric gather), turns them into the drafter's context K/V for all its layers (``precompute_context_kv``: an fp32 RMSNorm, the fused K/V GEMM, the K norm and RoPE), and stores all ``K + 1`` -rows into the request's slot of the dense context arena -- after the target's +rows into the request's context (its dense arena slot or pages) -- after the target's LM head, the sampler and the acceptance, and before the drafter. None of that reads the verify result: the captured rows are final once the last tapped target layer has run, the columns and RoPE positions come from the request's @@ -40,7 +40,7 @@ ``DecodeTarget.forward`` arms the fork with the step's metadata, so only decode steps (no context rows) fork; the step must also be one the side chain -reproduces (``ContextKvFork.plan``: the dense arena, the captured features, +reproduces (``ContextKvFork.plan``: an allocated context, the captured features, ``K + 1`` rows per request), otherwise the worker runs the built-in method. If an armed fork point did not fire (another feature replaced the LM head's GEMM call), the worker forks at the verify and logs that once. @@ -66,6 +66,12 @@ compaction and that read) on the main stream after the join: the compaction moves the columns the side chain wrote, so it must follow them, and nothing else in the step writes ``_ctx_base`` before it. + +Over the KV cache manager's pages the side chain first decodes this step's page +table (``_refresh_ctx_block_tables``, which the worker then skips), and the commit +is the worker's alone: the rows past the accepted count are not zeroed, since +their columns are the drafter block append's, which overwrites them in the same +step before anything reads them. """ from __future__ import annotations @@ -116,6 +122,7 @@ def column_limit(worker: DFlashWorker) -> int: "sample_and_accept_draft_tokens", "_lazy_init_ctx_buffers", "_ensure_spec_dec_state_restored", + "_refresh_ctx_block_tables", "_forward_impl", ) @@ -250,7 +257,7 @@ def fire(self, point: str) -> None: capture = current_capture_id() self.stream.wait_stream(main) with torch.cuda.stream(self.stream): - gen_hs = self.side_chain(spec_metadata, num_gens, k) + gen_hs = self.side_chain(spec_metadata, attn_metadata, num_gens, k) self._launched = Launched(main, capture, num_gens, k, point, gen_hs) self.launches[point] += 1 @@ -276,7 +283,7 @@ def join(self, launched: Launched) -> None: self._launched = None def commit(self, launched: Launched, slots, num_accepted) -> Optional[torch.Tensor]: - """Zero the masked rows and commit the context on the side stream; return the drafter's lengths. + """Zero a dense arena's masked rows and commit the context on the side stream. The side stream waits for the main stream's work so far -- the acceptance and the prepare's last read of ``_ctx_len`` -- and runs the @@ -298,7 +305,9 @@ def commit(self, launched: Launched, slots, num_accepted) -> Optional[torch.Tens base, last = arena_columns(worker) self.stream.wait_stream(main) with torch.cuda.stream(self.stream): - if base is None or compacts_after_join(worker): + if worker._ctx_paged: + num_ctx.copy_(worker._commit_context(slots, num_accepted.long(), rows)) + elif base is None or compacts_after_join(worker): context_kv_commit( worker._ctx_k_buf, worker._ctx_v_buf, @@ -350,17 +359,16 @@ def plan(self, spec_metadata, attn_metadata) -> Optional[tuple[int, int]]: """``(generation requests, K)`` when the side chain reproduces this step's store, else None. Host-side reads only (per-capture constants under a CUDA graph): the - dense arena (full-length or windowed) is allocated, the step has no context rows and at most - ``MAX_REQUESTS`` requests of ``K + 1`` target rows each, the target - features were captured and the drafter has ``fc`` + ``hidden_norm`` - (the built-in's ``has_target_features``). + context is allocated (pages, or a BF16 dense arena, full-length or + windowed), the step has no context rows and at most ``MAX_REQUESTS`` + requests of ``K + 1`` target rows each, the target features were captured + and the drafter has ``fc`` + ``hidden_norm`` (the built-in's + ``has_target_features``). """ worker, drafter = self.worker, self.draft_model if ( not worker._ctx_buf_inited - or worker._ctx_paged - or worker._ctx_k_buf is None - or worker._ctx_k_buf.dtype != torch.bfloat16 + or (not worker._ctx_paged and worker._ctx_k_buf.dtype != torch.bfloat16) or worker._ctx_len is None or worker._batch_to_slot is None ): @@ -381,14 +389,15 @@ def plan(self, spec_metadata, attn_metadata) -> Optional[tuple[int, int]]: return None return num_gens, k - def side_chain(self, spec_metadata, num_gens: int, k: int) -> torch.Tensor: + def side_chain(self, spec_metadata, attn_metadata, num_gens: int, k: int) -> torch.Tensor: """The built-in store of the step's context K/V, unmasked; runs on the side stream. - The worker's ``_store_step_context`` without its mask. Its first launch is a - torch index kernel, so the fork is a plain dependency on the main stream's - last kernel. + The worker's ``_store_step_context`` without its mask, after this step's page + table decode over the KV cache manager's pages. Its first launch is a torch + kernel, so the fork is a plain dependency on the main stream's last kernel. """ worker = self.worker + DFlashWorker._refresh_ctx_block_tables(worker, attn_metadata, num_gens) slots = worker._batch_to_slot[:num_gens] # The built-in's get_hidden_states(num_tokens)[0 : rows]: the same view. gen_hs = spec_metadata.captured_hidden_states[ @@ -438,6 +447,13 @@ def _lazy_init_ctx_buffers( "side chain was writing the previous one" ) + def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: + fork: Optional[ContextKvFork] = self.__dict__.get("_kv_fork") + if fork is not None and fork.launched: + # The side chain decoded this step's table at the fork point. + return self._ctx_block_tables is not None + return super()._refresh_ctx_block_tables(attn_metadata, num_seqs) + def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): fork: Optional[ContextKvFork] = self.__dict__.get("_kv_fork") if fork is not None: @@ -488,12 +504,11 @@ def _fork_prepare( """The built-in ``prepare_1st_drafter_inputs`` for a step whose store the side chain ran. Written out from the built-in ``DFlashWorker.prepare_1st_drafter_inputs`` - for its one path here -- no context rows, captured target features, the - dense arena (full-length or windowed) -- with the store moved to the side - chain: the bonus token and the query positions as the built-in computes - them, then the commit launched on the side stream, the bonus-token - embedding on the main stream meanwhile, and the join; over a windowed arena - the worker compacts, then the prepare settles the context + for its one path here -- no context rows, captured target features -- with + the store moved to the side chain: the bonus token and the query positions + as the built-in computes them, then the commit launched on the side stream, + the bonus-token embedding on the main stream meanwhile, and the join; over + a windowed arena the worker compacts, then the prepare settles the context (``_settle_context``). """ num_gens = launched.num_gens @@ -604,12 +619,6 @@ def worker_unsupported_reason(spec_worker, draft_model) -> Optional[str]: return f"the spec worker is {type(spec_worker).__name__}, not a DFlashWorker" if not isinstance(draft_model, DFlashForCausalLM): return f"the drafter is {type(draft_model).__name__}, not a DFlashForCausalLM" - backend = getattr(draft_model, "dflash_attention_backend", None) - if backend != "VANILLA" or getattr(draft_model, "_paged_ctx_cache", False): - return ( - f"the drafter's context cache is not the dense arena (attention backend {backend}); " - "the fork reproduces only the dense arena's store" - ) cls = type(spec_worker) for name in FORK_METHODS: mine = getattr(ContextKvForkDFlashWorker, name, None) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py index 292e17633680..1ad691603f43 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py @@ -25,14 +25,15 @@ ``drafter_block_decode``, second part: the block decode itself is this module's ``windowed_block_decode`` -- the built-in ``dflash_forward`` written out for -the one path this drafter takes (VANILLA arena, fused QK-norm + RoPE, DFlash 2 -block convolutions, non-causal windowed layers, no output gate) over the -built-in modules -- with the block's K/V append done inside the attention op -instead of two ``index_put`` launches per layer, and without the built-in's -unused RoPE table gather. This part is installed only when every layer's -attention geometry is one the op takes (it calls the op directly; selected for -a drafter outside that path, the install raises); a step outside that path -runs the built-in forward (with the windowed attention). +the one path this drafter takes (fused QK-norm + RoPE, DFlash 2 block +convolutions, non-causal windowed layers, no output gate) over the built-in +modules -- with the block's K/V append done inside the attention op instead of +the built-in's per-layer append, and without the built-in's unused RoPE table +gather. The op reads the context where the DFlash worker keeps it: the dense +arena, or the worker's pages through its page table. This part is installed +only when every layer's attention geometry is one the op takes (it calls the op +directly; selected for a drafter outside that path, the install raises); a step +outside that path runs the built-in forward. ``wide_drafter_window_attention`` (tep4 on sm_103 only) binds a ``WideWindowAttention`` to the drafter (``kernels/wide_drafter_window_attention``); a block decode of the @@ -205,11 +206,11 @@ def install_windowed_block_attention(draft_model: nn.Module, enable_pdl: bool) - def block_decode_unsupported_reason(draft_model: nn.Module) -> str | None: """Why ``windowed_block_decode`` cannot stand in for this drafter's forward, or None. - Besides the class and backend checks, every layer's attention geometry -- - a block of ``draft_model.block_size`` rows, its query / KV head counts and - head width, its window, bf16 operands -- must be one the attention op - takes: ``windowed_block_decode`` calls the op directly, with no per-call - fallback to the built-in attention. + Besides the class checks, every layer's attention geometry -- a block of + ``draft_model.block_size`` rows, its query / KV head counts and head width, + its window, bf16 operands -- must be one the attention op takes: + ``windowed_block_decode`` calls the op directly, with no per-call fallback + to the built-in attention. """ if not isinstance(draft_model, DFlashForCausalLM): return f"not a DFlashForCausalLM ({type(draft_model).__name__})" @@ -219,8 +220,6 @@ def block_decode_unsupported_reason(draft_model: nn.Module) -> str | None: return f"{type(draft_model).__name__} gates the attention output" if not getattr(draft_model, "_uses_worker_attention_backend", False): return "the drafter brings its own attention" - if getattr(draft_model, "dflash_attention_backend", None) != "VANILLA": - return "the drafter's attention backend is not VANILLA" block_size = getattr(draft_model, "block_size", None) if not block_size: return "the drafter declares no block size" @@ -253,12 +252,13 @@ def windowed_block_decode( ctx_kv_cache: torch.Tensor | None = None, ctx_page_table: torch.Tensor | None = None, ) -> torch.Tensor: - """The built-in ``DFlashForCausalLM.dflash_forward``, VANILLA / fused-RoPE path, append fused. + """The built-in ``DFlashForCausalLM.dflash_forward``, fused-RoPE path, append fused. Bound to the drafter as ``dflash_forward`` by ``install_block_decode``. - Same modules, same order, same arena state afterwards; the per-layer - ``index_put`` of the block K/V and the block attention are one - ``dflash_window_attention`` call with ``k_new`` / ``v_new``. Steps the + Same modules, same order, same context state afterwards; the per-layer + append of the block K/V and the block attention are one + ``dflash_window_attention`` call with ``k_new`` / ``v_new``, over the dense + arena or, given ``ctx_page_table``, the pages of ``ctx_kv_cache``. Steps the written-out path does not cover (fp16/fp32 activations, a drafter without the fused QK-norm + RoPE, a block wider than the one ``install_block_decode`` checked the attention geometry for) run the built-in forward. @@ -266,12 +266,17 @@ def windowed_block_decode( if self._fused_kv_weight is None: self._build_fused_kv_buffers() attn0 = self.model.layers[0].self_attn + if ctx_page_table is None: + caches = [(ctx_k_cache[:, i], ctx_v_cache[:, i]) for i in range(ctx_k_cache.shape[1])] + else: + # Each layer's [pages, K/V, Hkv, page, D] pool as two pools of [page, Hkv, D] rows. + caches = [(pool[:, 0].transpose(1, 2), pool[:, 1].transpose(1, 2)) for pool in ctx_kv_cache] + k0, v0 = caches[0] if not ( self._use_fused_qk_norm_rope - and noise_embedding.dtype == ctx_k_cache.dtype == ctx_v_cache.dtype == torch.bfloat16 - and ctx_k_cache.dim() == 5 - and ctx_k_cache.stride(-1) == 1 - and ctx_v_cache.stride(-1) == 1 + and noise_embedding.dtype == k0.dtype == v0.dtype == torch.bfloat16 + and k0.stride(-1) == 1 + and v0.stride(-1) == 1 and noise_embedding.shape[1] <= self.block_size ): return DFlashForCausalLM.dflash_forward( @@ -355,23 +360,25 @@ def windowed_block_decode( v_new = qkv_query[:, q_size + kv_size :].reshape(-1, num_kv_heads_per_rank, head_dim) _, window_size = self._resolve_block_attention(layer_idx) + k_cache, v_cache = caches[layer_idx] if wide is not None: out = wide( q_bshd, - ctx_k_cache[:, layer_idx], - ctx_v_cache[:, layer_idx], + k_cache, + v_cache, seq_lens_after, cache_batch_idx_i32, window_size, enable_pdl, k_new, v_new, + ctx_page_table, ) else: out = dflash_window_attention( q_bshd, - ctx_k_cache[:, layer_idx], - ctx_v_cache[:, layer_idx], + k_cache, + v_cache, seq_lens_after, cache_batch_idx_i32, window_size, @@ -380,6 +387,7 @@ def windowed_block_decode( k_new, v_new, merge_entry, + ctx_page_table, ) attn_output = out.reshape(B * block_size, q_size) @@ -474,8 +482,8 @@ def install(draft_model: nn.Module | None, block_conv: bool, block_decode: bool) Raises: ValueError: if one asked for cannot be installed on this drafter - (``block_conv``: no DFlash 2 block convolutions; ``block_decode``: no - VANILLA block attention, or ``block_decode_unsupported_reason``). + (``block_conv``: no DFlash 2 block convolutions; ``block_decode``: + ``block_decode_unsupported_reason``). TypeError: from ``install_fused_block_conv``. """ if block_conv: @@ -514,11 +522,11 @@ def _install_windowed_attention(draft_model: nn.Module | None, block_decode: boo wrapped = isinstance( getattr(draft_model, "_dflash_flash_attention", None), WindowedBlockAttention ) - if draft_model is None or not install_windowed_block_attention(draft_model, enable_pdl): - raise ValueError( - "[musespark_open] drafter_block_decode cannot run here: the drafter has no VANILLA block attention" - ) - if not wrapped: + if ( + draft_model is not None + and install_windowed_block_attention(draft_model, enable_pdl) + and not wrapped + ): logger.info( f"[musespark_open] drafter_block_decode: windowed DFlash block attention on the drafter (pdl={enable_pdl})" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py index 23f121dd4d63..b6b7982a42a3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py @@ -513,14 +513,12 @@ def _fold_forked_inputs( if ( launched.num_gens != num_gens or launched.k != spec_metadata.runtime_draft_len - or self._ctx_block_tables is not None or not has_target_features ): raise FoldStateError( "[musespark_open] verify_fold + context_kv_fork: the context-K/V side chain was launched for " f"{launched.num_gens} requests and K {launched.k}, the folded step has " f"{num_gens} and K {spec_metadata.runtime_draft_len}" - + ("; the arena is paged" if self._ctx_block_tables is not None else "") + ("; the step has no target features" if not has_target_features else "") ) num_ctx = fork.commit(launched, slots, gen_num_accepted) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_drafter_window_attention/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_drafter_window_attention/attention.py index 4bbe748d1c79..d96ad64c5ea1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_drafter_window_attention/attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_drafter_window_attention/attention.py @@ -61,6 +61,7 @@ def __call__( enable_pdl: bool, k_new: torch.Tensor, v_new: torch.Tensor, + page_table: Optional[torch.Tensor], ) -> torch.Tensor: """``dflash_window_attention_wide`` with the block decode's arguments (default scale).""" return dflash_window_attention_wide( @@ -74,6 +75,7 @@ def __call__( enable_pdl, k_new, v_new, + page_table, ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py index 63bcc3295d14..e691447cab5d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py @@ -43,7 +43,7 @@ or none (``dispatch.NO_PATHS``) when the deployment has no composition, and the target then runs only the built-in modules as described above. A missing public precondition of a planned fast path is an error, whatever the mode. Over a -full-length drafter context arena the fast paths written for a windowed one +full-length or paged drafter context the fast paths written for a windowed arena are left out of the plan (``dispatch.WINDOWED_CONTEXT_ARENA_PATHS``); over a windowed arena (``TRTLLM_DFLASH_WINDOWED_CTX_ARENA``) the paths that write the arena by column (``dispatch.FULL_CONTEXT_ARENA_PATHS``: the context-K/V fork and diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py index e506895580b0..28a3782e8e1c 100644 --- a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py @@ -5,14 +5,32 @@ * ``dflash_window_attention`` and ``dflash_window_attention_wide`` (at each device's split geometry) over pages are bitwise their dense calls, the appended block included. * ``context_kv_fork``'s commit empties the dummy slot, as the worker's own commit does. +* The block decode over a unified target/draft ``KVCacheManagerV2`` is bitwise the dense one. +* Over pages ``compose`` plans every sm103 TEP4 fast path but the window compaction. """ +import importlib.util +import types +from dataclasses import replace +from pathlib import Path + import pytest import torch +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ModelingV2Mode +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep import ( # noqa: E501 + dispatch, + modeling, +) from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.context_kv_fork import ( # noqa: E501 fork as context_kv_fork, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.drafter_block_decode import ( # noqa: E501 + block_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.wide_drafter_window_attention import ( # noqa: E501 + attention as wide_attention, +) from tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open.muse_spark_open_816b_a42b_nvfp4__sm_103__tep.ops.spec import ( # noqa: E501 dflash_window_attention_wide as wide_ops, ) @@ -20,7 +38,21 @@ WINDOW, dflash_window_attention, ) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_dflash import DFlashForCausalLM +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, + KVCacheManagerV2, + Role, +) +from tensorrt_llm._torch.speculative import dflash_attention from tensorrt_llm._torch.speculative.dflash import DFlashWorker +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.quantization.mode import QuantAlgo __extra_import_path__ = ["../../.."] from _validating import validating # noqa: E402 — needs the path declared above @@ -91,3 +123,131 @@ def test_the_fork_commit_empties_the_dummy_slot(windowed): assert worker._ctx_len.tolist() == [8, 0, 0, 0] # The lengths the drafter attends: physical over the windowed arena. assert num_ctx.tolist() == ([6, 0] if windowed else [8, 0]) + + +PAGE = 16 +#: Six requests, so a block decode of them takes the wide op where it is installed: one past the +#: target's window, and request 33 a CUDA-graph padding dummy. +REQUEST_IDS = range(31, 37) +CAPACITIES, HISTORIES = (200, 40, 16, 64, 120, 32), (180, 24, 0, 40, 100, 9) + + +def _unified_manager(num_draft: int, nkv: int, head_dim: int) -> KVCacheManagerV2: + """The requests in one cache of six FP8 target layers and BF16 draft ones, all windowed.""" + target = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.FP8, + num_kv_heads=nkv, + head_dim=head_dim, + buffer_roles=(Role.KEY, Role.VALUE), + sliding_window_size=4 * PAGE, + ) + specs = [replace(target, global_layer_id=i) for i in range(6)] + draft = replace(target, owner=KVCacheLayerOwner.DRAFT, dtype=DataType.BF16) + specs += [replace(draft, global_layer_id=6 + i) for i in range(num_draft)] + manager = KVCacheManagerV2( + KvCacheConfig(max_gpu_total_bytes=64 << 20, host_cache_size=0, enable_block_reuse=False), + CacheType.SELF, + num_layers=6, + num_kv_heads=nkv, + head_dim=head_dim, + tokens_per_block=PAGE, + max_seq_len=256, + max_batch_size=len(REQUEST_IDS), + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.FP8, + layer_specs=specs, + ) + for request_id, capacity, history in zip(REQUEST_IDS, CAPACITIES, HISTORIES): + cache = manager._create_kv_cache(request_id, None, None, is_dummy=request_id == 33) + assert cache.resume(manager._stream.cuda_stream) + assert cache.resize(capacity, history) + return manager + + +@pytest.mark.parametrize("wide", [False, True], ids=["single", "wide"]) +def test_the_block_decode_over_a_unified_kv_cache_is_the_dense_arena_one(wide): + """Over each draft layer's view of the manager's pool and the worker's one page table, the + block decode is bitwise the VANILLA drafter's over dense arenas of the same context.""" + if (reason := dflash_attention.dflash_trtllm_gen_unavailability_reason()) is not None: + pytest.skip(reason) + # The speculative tests' tiny DFlash 2 drafter, over pages and over dense arenas. + path = Path(__file__).parents[4] / "speculative" / "hw_agnostic" / "test_dflash2_semantics.py" + spec = importlib.util.spec_from_file_location("_dflash2_semantics", path) + sem = importlib.util.module_from_spec(spec) + spec.loader.exec_module(sem) + drafters = [] + for backend in ("TRTLLM", "VANILLA"): + config = ModelConfig(pretrained_config=sem._tiny_config(True), attn_backend="TRTLLM") + drafter = DFlashForCausalLM(config, dflash_attention_backend=backend).to("cuda") + drafter.load_weights(dict(sem._tiny_weights(True))) + drafter._build_fused_kv_buffers() + block_decode.install(drafter, True, True) + if wide: + wide_attention.install(drafter) + drafters.append(drafter) + paged, dense = drafters + layers, nkv, head_dim = paged._num_attn_layers, paged._num_kv_heads, paged._head_dim + manager = _unified_manager(layers, nkv, head_dim) + try: + num_seqs = len(REQUEST_IDS) + shape = (manager.num_attention_op_pools, num_seqs, 2, manager.max_blocks_per_seq) + offsets = torch.full(shape, -7, dtype=torch.int32, device="cuda") + manager.copy_batch_block_offsets(offsets, list(REQUEST_IDS), 1, 0, num_seqs) + worker = object.__new__(DFlashWorker) + worker._ctx_unified_cache = True + pool = worker._managed_ctx_pool(manager, layers, nkv, head_dim, torch.bfloat16, 2) + assert worker._init_ctx_block_tables(manager, pool, num_seqs, layers, nkv, head_dim, PAGE) + metadata = types.SimpleNamespace(kv_cache_block_offsets=offsets) + assert worker._refresh_ctx_block_tables(metadata, num_seqs) + table = worker._ctx_block_tables[:num_seqs] + # Each layer's view spans the whole interleaved pool: fill every one before reading any. + gen = torch.Generator(device="cuda").manual_seed(0) + for view in pool: + view.normal_(generator=gen) + ctx = torch.tensor(HISTORIES, device="cuda") + cols = torch.arange(int(ctx.max()) + sem.BLOCK, device="cuda") + # Each layer's K and V planes as [pages, page, Hkv, D] rows; the arenas [B, L, C, Hkv, D]. + planes = [[view[:, half].transpose(1, 2) for view in pool] for half in (0, 1)] + at = table.long()[:, cols // PAGE], cols % PAGE + arenas = [torch.stack([plane[at] for plane in half], dim=1) for half in planes] + block = ctx[:, None] + torch.arange(sem.BLOCK, device="cuda") + step = dict( + noise_embedding=_randn(gen, num_seqs, sem.BLOCK, sem.TINY["hidden_size"]), + query_positions=block, + num_ctx_per_req=ctx, + ctx_cache_batch_idx=torch.arange(num_seqs, device="cuda"), + ) + got = paged.dflash_forward( + **step, ctx_k_cache=None, ctx_v_cache=None, ctx_kv_cache=pool, ctx_page_table=table + ) + want = dense.dflash_forward(**step, ctx_k_cache=arenas[0], ctx_v_cache=arenas[1]) + assert torch.equal(got, want), (got.float() - want.float()).abs().max().item() + # The block's K/V are in the requests' pages as in their arena slots. + landed = table.long().gather(1, block // PAGE), block % PAGE + rows = torch.arange(num_seqs, device="cuda")[:, None] + for half, arena in zip(planes, arenas): + for layer, plane in enumerate(half): + assert torch.equal(plane[landed], arena[rows, layer, block]) + finally: + manager.shutdown() + + +def test_a_paged_context_plans_every_fast_path_but_the_window_compaction(monkeypatch): + """sm103 TEP4 with a TRTLLM drafter, TRTLLM_DFLASH_WINDOWED_CTX_ARENA set: a paged context is + never windowed, so the window compaction is the one served path the plan leaves out.""" + monkeypatch.setattr(dispatch, "missing_preconditions", lambda *args: []) + monkeypatch.setattr(dispatch, "agreed_union", lambda mapping, local: list(dict.fromkeys(local))) + model_config = types.SimpleNamespace( + mapping=types.SimpleNamespace(tp_size=4, gpus_per_node=4, tp_group=list(range(4))), + spec_config=DFlashDecodingConfig(max_draft_len=7, speculative_model="/drafter"), + quant_config=types.SimpleNamespace(quant_algo=QuantAlgo.NVFP4), + ) + worker = types.SimpleNamespace(_dflash_attention_backend="TRTLLM", windowed_ctx_arena=True) + plan = dispatch.compose( + None, model_config, ModelingV2Mode.REQUIRE, modeling.TARGET_PATHS, (10, 3), worker, 64 + ) + served, _ = dispatch.capacity_gated(frozenset(modeling.TARGET_PATHS[103, 4].served), 64) + served, _ = dispatch.yielded(served) + assert served - plan.paths == {"drafter_window_compact"} From 09716ed634701faa4d64e6219c9f37bbf291b360 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 11:25:01 -0700 Subject: [PATCH 5/6] [None][feat] Let DFlash AUTO resolve to TRTLLM under modeling_v2 The Muse Spark Open modeling_v2 composition now runs over a paged drafter context, so AUTO no longer pins VANILLA under modeling_v2: it follows the checkpoint rule as on any other deployment, and a disaggregated deployment then gets the drafter whose context is transferred with the target's KV. Only a disaggregated worker with TRTLLM_DISAGG_SHARED_DRAFT_KV=0 still gets VANILLA. Signed-off-by: Zheyu Fu --- docs/source/features/speculative-decoding.md | 2 +- tensorrt_llm/llmapi/llm_args.py | 9 ++++----- .../speculative/hw_agnostic/test_dflash_draft_kv.py | 6 +++--- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/docs/source/features/speculative-decoding.md b/docs/source/features/speculative-decoding.md index 401ae37f571d..b5f585450493 100644 --- a/docs/source/features/speculative-decoding.md +++ b/docs/source/features/speculative-decoding.md @@ -152,7 +152,7 @@ Reference: [DFlash: Distilled Flash Speculative Decoding](https://arxiv.org/pdf/ * `speculative_model`: Path or HuggingFace model ID for the DFlash draft model. * `mask_token_id`: Token ID used as the mask token for parallel prediction. If not set, it is read from the draft model config. * `target_layer_ids`: List of target model layer indices whose hidden states are captured for cross-attention in the draft model. If not set, read from the draft model config. -* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker with `TRTLLM_DISAGG_SHARED_DRAFT_KV=0` and under `modeling_v2`. `"VANILLA"` uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only. +* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker with `TRTLLM_DISAGG_SHARED_DRAFT_KV=0`. `"VANILLA"` uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only. ```python from tensorrt_llm.llmapi import DFlashDecodingConfig diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0f25df9d07e0..5a6c913bdfe1 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3316,7 +3316,7 @@ class DFlashDecodingConfig(DecodingBaseConfig): "on the same paged cache and supports SM90 only. AUTO (the default) " "resolves to TRTLLM where it can run and every drafter layer is " "windowed, and to VANILLA otherwise, including on a disaggregated worker " - "with TRTLLM_DISAGG_SHARED_DRAFT_KV=0 and under modeling_v2. In " + "with TRTLLM_DISAGG_SHARED_DRAFT_KV=0. In " "disaggregated serving a TRTLLM drafter keeps its KV with the target's, " "so the context server's drafter context reaches the generation server; " "both servers need the same speculative_config.") @@ -7086,14 +7086,13 @@ def validate_speculative_config(self): and self.cache_transceiver_config.backend is not None) shared_draft_kv = disagg_shares_draft_kv( self.speculative_config) - # An unshared disaggregated worker refuses a windowed TRTLLM draft pool; modeling_v2 drafts on VANILLA. + # An unshared disaggregated worker refuses a windowed TRTLLM draft pool. if (self.speculative_config.attention_backend == "AUTO" - and ((disagg and not shared_draft_kv) - or self.modeling_v2 != "off")): + and disagg and not shared_draft_kv): self.speculative_config.attention_backend = "VANILLA" logger.info( "DFlash attention_backend AUTO resolved to VANILLA " - "(modeling_v2, or a disaggregated worker with " + "(a disaggregated worker with " "TRTLLM_DISAGG_SHARED_DRAFT_KV=0).") assert self.speculative_config.max_draft_len > 0, "DFlash max_draft_len must be > 0" # From the disaggregated-serving-benchmarking overlay (TRT-LLM aa7db814): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index 774ed238af78..a4a3228ffbaf 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -61,14 +61,14 @@ def test_default_dflash_resolves_auto_before_choosing_its_draft_cache( ("NIXL", "0", "off", "VANILLA"), ("NIXL", "1", "off", "AUTO"), ("NIXL", None, "off", "AUTO"), - (None, "0", "auto", "VANILLA"), + (None, "0", "auto", "AUTO"), (None, "0", "off", "AUTO"), ], ) -def test_auto_keeps_vanilla_where_trtllm_is_refused_or_bypassed( +def test_auto_keeps_vanilla_where_trtllm_is_refused( tmp_path, monkeypatch, transceiver, shared, modeling_v2, resolved ) -> None: - """Workers keeping a private draft KV refuse TRTLLM; modeling_v2 DFlash compositions need VANILLA.""" + """Workers keeping a private draft KV refuse TRTLLM; modeling_v2 resolves AUTO like any deployment.""" from tensorrt_llm.llmapi import CacheTransceiverConfig, TorchLlmArgs, llm_args # No GPU: whatever stays AUTO is left to the ranks' windowed-and-probe rule. From 8af51a768eb7a7552f8a922a01a77972c4e6883b Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 27 Oct 2026 11:25:50 -0700 Subject: [PATCH 6/6] [None][test] Run the Muse modeling_v2 drafter context tests in the MuseSpark pre-merge stage Pre-merge CI on feat/musespark-open-agentx runs only the MuseSpark stages, so the tests of the composition over a TRTLLM drafter's paged context did not run on this branch. They need a GPU and none is cpu_only, so they go to l0_gb300_musespark. Signed-off-by: Zheyu Fu --- tests/integration/test_lists/test-db/l0_gb300_musespark.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_gb300_musespark.yml b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml index d71a45b69db4..96846251b211 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_musespark.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_musespark.yml @@ -50,3 +50,5 @@ l0_gb300_musespark: - unittest/disaggregated/test_kv_transfer.py - unittest/disaggregated/test_bounce.py - accuracy/test_disaggregated_serving.py::TestQwen3_8BDFlash::test_gsm8k_1p1d_unified_draft_kv + # The Muse modeling_v2 composition over a TRTLLM drafter's paged context. + - unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py