diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 6e17258067a1..5f8c27b4c71a 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1218,6 +1218,9 @@ def dflash_forward( ctx_k_cache: torch.Tensor, ctx_v_cache: torch.Tensor, ctx_cache_batch_idx: torch.Tensor, + hybrid_k_bufs: Optional[List[torch.Tensor]] = None, + hybrid_v_bufs: Optional[List[torch.Tensor]] = None, + hybrid_block_idx: Optional[torch.Tensor] = None, ) -> torch.Tensor: """DFlash draft forward with cross-attention over a pooled K/V buffer. @@ -1230,10 +1233,17 @@ def dflash_forward( ctx_k_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] ctx_v_cache: [pool_batch, L, max_ctx+block_size, nkv, hd] ctx_cache_batch_idx: [B] — slot index into the pool per batch entry + hybrid_k_bufs/hybrid_v_bufs: per-layer manager pool views + [pages, tokens_per_block, nkv, hd] (hybrid ctx only) + hybrid_block_idx: [B, W] page ids per request (hybrid ctx only) Returns: [B * block_size, hidden_size] """ - from flash_attn import flash_attn_with_kvcache + if hybrid_k_bufs is not None: + from ..speculative.dflash_hybrid_attn import \ + dflash_ctx_paged_attention + else: + from flash_attn import flash_attn_with_kvcache if self._fused_kv_weight is None: self._build_fused_kv_buffers() @@ -1365,24 +1375,36 @@ def dflash_forward( num_kv_heads_per_rank, head_dim) - # Per-layer view into the pooled ctx cache. - # [pool_batch, max_ctx+block, nkv, hd]; flash_attn dereferences - # each batch via cache_batch_idx, no gather. - layer_k_cache = ctx_k_cache[:, layer_idx] - layer_v_cache = ctx_v_cache[:, layer_idx] - - # flash_attn appends k_noise/v_noise in-place at - # cache_seqlens[i]..+block_size for each batch i. - out = flash_attn_with_kvcache( - q=Q_bshd, - k_cache=layer_k_cache, - v_cache=layer_v_cache, - k=k_noise_bshd, - v=v_noise_bshd, - cache_seqlens=cache_seqlens_i32, - cache_batch_idx=cache_batch_idx_i32, - causal=False, - ) + if hybrid_k_bufs is not None: + # Hybrid: ctx K/V paged in the target manager's pool + out = dflash_ctx_paged_attention( + Q_bshd, + hybrid_k_bufs[layer_idx], + hybrid_v_bufs[layer_idx], + hybrid_block_idx, + cache_seqlens_i32, + k_noise_bshd, + v_noise_bshd, + ) + else: + # Per-layer view into the pooled ctx cache. + # [pool_batch, max_ctx+block, nkv, hd]; flash_attn dereferences + # each batch via cache_batch_idx, no gather. + layer_k_cache = ctx_k_cache[:, layer_idx] + layer_v_cache = ctx_v_cache[:, layer_idx] + + # flash_attn appends k_noise/v_noise in-place at + # cache_seqlens[i]..+block_size for each batch i. + out = flash_attn_with_kvcache( + q=Q_bshd, + k_cache=layer_k_cache, + v_cache=layer_v_cache, + k=k_noise_bshd, + v=v_noise_bshd, + cache_seqlens=cache_seqlens_i32, + cache_batch_idx=cache_batch_idx_i32, + causal=False, + ) attn_output = out.reshape(B * block_size, q_size) # o_proj (flat 2D, handles all-reduce internally) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b2d45dec50aa..2381ddae0703 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -271,6 +271,18 @@ def __init__( KVCacheManagerV2) self._draft_config = draft_config self._skip_est = skip_est + # Hybrid DFlash ctx: resolve the draft layer count before any manager + # is created (get_num_spec_layers reads it via get_pp_layers). + if (self._is_dflash_hybrid_ctx() + and self._speculative_config._num_draft_layers is None + and draft_config is not None): + self._speculative_config._num_draft_layers = ( + draft_config.pretrained_config.num_hidden_layers) + + def _is_dflash_hybrid_ctx(self) -> bool: + sc = self._speculative_config + return (sc is not None and sc.spec_dec_mode.is_dflash() + and getattr(sc, 'use_hybrid_context', False)) def _get_model_kv_cache_manager_cls( self, @@ -423,6 +435,19 @@ def _get_kv_size_per_token(self, effective_draft_config, kv_cache_config, num_layers=self._get_num_draft_layers()) + if self._is_dflash_hybrid_ctx() and self._mapping.is_last_pp_rank(): + # Hybrid ctx: draft layers live in the target manager, so their + # per-token cost adds to the same budget. + effective_draft_config = self._get_effective_draft_config() + draft_kv_cache_manager_cls = get_kv_cache_manager_cls( + effective_draft_config, + kv_cache_config, + is_disagg=self._is_disagg) + total += self._per_manager_cache_cost( + draft_kv_cache_manager_cls, + effective_draft_config, + kv_cache_config, + num_layers=self._get_num_draft_layers()) return total def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -1527,6 +1552,7 @@ def _build_per_layer_num_kv_heads( num_hidden_layers: int, spec_config: Optional[SpeculativeConfig] = None, draft_config: Optional[ModelConfig] = None, + target_head_dim: Optional[int] = None, ) -> Union[int, List[int]]: """ Returns: @@ -1543,6 +1569,22 @@ def _build_per_layer_num_kv_heads( draft_pretrained, 'num_key_value_heads', getattr(draft_pretrained, 'num_attention_heads', None)) + # The manager has a single head_dim shared by all layers in a pool. + # When the draft's head_dim differs from the target's, express the + # draft layers' KV bytes in target-head_dim units. + draft_head_dim = getattr(draft_pretrained, 'head_dim', None) + if (draft_num_kv_heads is not None and target_head_dim is not None + and draft_head_dim is not None + and draft_head_dim != target_head_dim): + draft_kv_elems = draft_num_kv_heads * draft_head_dim + if draft_kv_elems % target_head_dim != 0: + raise ValueError( + f"Draft KV row ({draft_num_kv_heads} heads x " + f"{draft_head_dim}) is not expressible in target head_dim " + f"({target_head_dim}) units; cannot place draft context in " + f"the target KV cache manager.") + draft_num_kv_heads = draft_kv_elems // target_head_dim + if draft_num_kv_heads is None or draft_num_kv_heads == num_key_value_heads: return num_key_value_heads @@ -1693,8 +1735,11 @@ def _create_kv_cache_manager( per_layer_num_kv_heads = num_key_value_heads else: per_layer_num_kv_heads = _build_per_layer_num_kv_heads( - num_key_value_heads, num_hidden_layers, spec_config, - draft_config_for_kv) + num_key_value_heads, + num_hidden_layers, + spec_config, + draft_config_for_kv, + target_head_dim=head_dim if isinstance(head_dim, int) else None) manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index efa31fa2cc1e..628542dde8bc 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -181,8 +181,20 @@ def __init__( self._req_to_slot = {} # request_id -> slot index self._free_slots = deque() # available slot indices + # Hybrid ctx: draft K/V lives in the target manager as spec layers. + self._use_hybrid_context = getattr(spec_config, "use_hybrid_context", False) + self._hybrid_inited = False + self._hybrid_pool_idx = 0 + self._hybrid_blk_divisor = 1 + self._hybrid_tpb = 0 + self._hybrid_k_bufs = None # per draft layer [pages, tpb, nkv, hd] + self._hybrid_v_bufs = None + self._kv_cache_manager = None + self._cur_block_idx = None # decoded page ids, per batch row + logger.info( - f"DFlashWorker initialized with use_separate_draft_kv_cache={use_separate_draft_kv_cache}" + f"DFlashWorker initialized with use_separate_draft_kv_cache={use_separate_draft_kv_cache}, " + f"use_hybrid_context={self._use_hybrid_context}" ) @property @@ -200,7 +212,82 @@ def _draft_tokens_per_req(self) -> int: """ return self.max_draft_len + 1 + def _lazy_init_hybrid_ctx(self, draft_model, spec_metadata, attn_metadata): + mgr = getattr(attn_metadata, "kv_cache_manager", None) + if mgr is None or not hasattr(mgr, "get_buffers"): + raise RuntimeError("DFlash use_hybrid_context requires a block-based KV cache manager.") + if not hasattr(attn_metadata, "kv_cache_block_offsets"): + raise RuntimeError( + "DFlash use_hybrid_context requires the TRTLLM attention backend " + "(kv_cache_block_offsets not found on attn metadata)." + ) + if self._hybrid_inited and self._kv_cache_manager is mgr: + return + + draft_model._build_fused_kv_buffers() + num_layers = draft_model._num_attn_layers + + # Draft spec layers are appended after the target layers. + layer_ids = sorted(mgr.layer_offsets.keys())[-num_layers:] + bufs = [mgr.get_buffers(lid) for lid in layer_ids] # NHD [P, 2, tpb, nkv, hd] + + # The manager sizes spec layers in target-head_dim units, so when the draft head_dim differs + # from the pool's, view each [P, tpb, nkv_pool, hd_pool] buffer back + # to the draft's per-rank [P, tpb, nkv_draft, hd_draft] geometry. + nkv_draft = draft_model._num_kv_heads + hd_draft = draft_model._head_dim + + def _as_draft_geometry(t: torch.Tensor) -> torch.Tensor: + if t.shape[-2] == nkv_draft and t.shape[-1] == hd_draft: + return t + if t.shape[-2] * t.shape[-1] != nkv_draft * hd_draft: + raise RuntimeError( + f"DFlash hybrid ctx: pool row {tuple(t.shape[-2:])} does " + f"not match draft KV geometry ({nkv_draft}, {hd_draft}) " + f"per rank; check per-layer KV head registration." + ) + return t.flatten(-2).unflatten(-1, (nkv_draft, hd_draft)) + + self._hybrid_k_bufs = [_as_draft_geometry(b[:, 0]) for b in bufs] + self._hybrid_v_bufs = [_as_draft_geometry(b[:, 1]) for b in bufs] + + # Pool of the draft layers + its offsets-decode divisor + # (encoded = block_idx * layers_in_pool * kv_factor). + pm = mgr.kv_cache_pool_mapping.view(len(mgr.layer_offsets), -1) + pool_idx = int(pm[mgr.layer_offsets[layer_ids[0]], 0]) + layers_in_pool = int((pm[:, 0] == pool_idx).sum()) + self._hybrid_pool_idx = pool_idx + self._hybrid_blk_divisor = layers_in_pool * mgr.kv_factor + self._hybrid_tpb = mgr.tokens_per_block + self._resolved_block_size = getattr(draft_model, "block_size", None) or ( + self.max_draft_len + 1 + ) + self._max_ctx = mgr.max_blocks_per_seq * mgr.tokens_per_block + self._kv_cache_manager = mgr + + max_batch = spec_metadata.max_num_requests + self._ctx_len = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._free_slots = deque(range(max_batch)) + self._req_to_slot = {} + self._ctx_buf_inited = True + self._hybrid_inited = True + + logger.info( + f"DFlash: hybrid ctx in target manager pool {pool_idx}: " + f"layers={layer_ids}, dtype={bufs[0].dtype}, " + f"tokens_per_block={mgr.tokens_per_block}" + ) + + def _hybrid_block_idx(self, attn_metadata): + """Decode per-batch-row page ids for the draft pool (device, in-graph).""" + enc = attn_metadata.kv_cache_block_offsets[self._hybrid_pool_idx, :, 0, :] + return enc // self._hybrid_blk_divisor + def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): + if self._use_hybrid_context: + self._lazy_init_hybrid_ctx(draft_model, spec_metadata, attn_metadata) + return if self._ctx_buf_inited: return @@ -348,14 +435,19 @@ def _store_prefill_context( continue slot = self._free_slots.popleft() self._req_to_slot[req_id] = slot - self._ctx_len[slot] = 0 + # Hybrid: reused blocks already hold ctx for [0, first_pos). + self._ctx_len[slot] = first_pos if self._use_hybrid_context else 0 slot = self._req_to_slot[req_id] cur = int(self._ctx_len[slot].item()) end = min(cur + slen, self._max_ctx) actual = end - cur if actual > 0: - chunk_proj_cast = chunk_proj[:actual].to(self._ctx_k_buf.dtype) + if self._use_hybrid_context: + # Cast to the pool dtype at write. + chunk_proj_cast = chunk_proj[:actual] + else: + chunk_proj_cast = chunk_proj[:actual].to(self._ctx_k_buf.dtype) self._ctx_len[slot] = end # Precompute post-norm/post-RoPE K,V for this prefill chunk # so decode iters can read without re-projecting. @@ -363,8 +455,20 @@ def _store_prefill_context( chunk_proj_cast, chunk_pos[:actual] ) # chunk_k/v: [actual, L, nkv, hd] → [L, actual, nkv, hd] - self._ctx_k_buf[slot, :, cur:end] = chunk_k.permute(1, 0, 2, 3) - self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) + if self._use_hybrid_context: + # Position-aligned: token p -> manager block p // tpb of + # this batch row. + tpb = self._hybrid_tpb + pos = torch.arange(cur, end, dtype=torch.long, device="cuda") + blk = self._cur_block_idx[i, pos // tpb].long() + off = pos % tpb + dt = self._hybrid_k_bufs[0].dtype + for li, (kb, vb) in enumerate(zip(self._hybrid_k_bufs, self._hybrid_v_bufs)): + kb[blk, off] = chunk_k[:, li].to(dt) + vb[blk, off] = chunk_v[:, li].to(dt) + else: + self._ctx_k_buf[slot, :, cur:end] = chunk_k.permute(1, 0, 2, 3) + self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) offset += slen def forward( @@ -400,6 +504,9 @@ def forward( self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata) spec_metadata._dflash_worker = self + if self._use_hybrid_context: + self._cur_block_idx = self._hybrid_block_idx(attn_metadata) + # Save context lengths before warmup to prevent accumulation is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() if is_warmup: @@ -480,6 +587,9 @@ def forward( ctx_k_cache=inputs["ctx_k_cache"], ctx_v_cache=inputs["ctx_v_cache"], ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], + hybrid_k_bufs=inputs["hybrid_k_bufs"], + hybrid_v_bufs=inputs["hybrid_v_bufs"], + hybrid_block_idx=inputs["hybrid_block_idx"], ) # Gather K logits per gen request from mask positions (1..K). @@ -664,16 +774,33 @@ def prepare_1st_drafter_inputs( # Fast path: store the pre-projected/pre-RoPE'd K/V. # dflash_forward reads these directly via cache_batch_idx. - k_new, v_new = draft_model.precompute_context_kv( - proj_flat.to(self._ctx_k_buf.dtype), pos_flat - ) + if self._use_hybrid_context: + proj_cast = proj_flat # cast to pool dtype at write + else: + proj_cast = proj_flat.to(self._ctx_k_buf.dtype) + k_new, v_new = draft_model.precompute_context_kv(proj_cast, pos_flat) mask_bc = mask_1d.view(-1, 1, 1, 1).to(k_new.dtype) k_new.mul_(mask_bc) v_new.mul_(mask_bc) slot_long = slot_flat.long() col_long = col_flat.long() - self._ctx_k_buf[slot_long, :, col_long] = k_new - self._ctx_v_buf[slot_long, :, col_long] = v_new + if self._use_hybrid_context: + tpb = self._hybrid_tpb + row_flat = ( + torch.arange(num_contexts, num_contexts + num_gens, device="cuda") + .unsqueeze(1) + .expand(-1, K_plus_1) + .reshape(-1) + ) + blk = self._cur_block_idx[row_flat, col_long // tpb].long() + off = col_long % tpb + dt = self._hybrid_k_bufs[0].dtype + for li, (kb, vb) in enumerate(zip(self._hybrid_k_bufs, self._hybrid_v_bufs)): + kb[blk, off] = k_new[:, li].to(dt) + vb[blk, off] = v_new[:, li].to(dt) + else: + self._ctx_k_buf[slot_long, :, col_long] = k_new + 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) @@ -701,4 +828,11 @@ def prepare_1st_drafter_inputs( "ctx_k_cache": self._ctx_k_buf, "ctx_v_cache": self._ctx_v_buf, "ctx_cache_batch_idx": slots, + "hybrid_k_bufs": self._hybrid_k_bufs if self._use_hybrid_context else None, + "hybrid_v_bufs": self._hybrid_v_bufs if self._use_hybrid_context else None, + "hybrid_block_idx": ( + self._cur_block_idx[num_contexts : num_contexts + num_gens] + if self._use_hybrid_context + else None + ), } diff --git a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py new file mode 100644 index 000000000000..ba12e0a9928e --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py @@ -0,0 +1,628 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DFlash hybrid-context cross-attention. + +Reads the draft context K/V from the target KV cache manager's paged pool +(hybrid mode: draft layers registered as spec layers of the target manager). +No existing kernel fits this op: it needs device-resident sequence lengths +(acceptance-dependent, updated inside CUDA graphs), 32-token pages in the +manager's NHD block layout, fp8 KV, and a dense per-step noise-K/V suffix +that is never written to the cache. + +Two execution paths, chosen by batch size at launch (static under CUDA +graph capture, where the batch is padded to a fixed size): + +- Small batch: flash-decoding style split-KV. ``(B, NKV)`` CTAs cannot fill + the GPU, so the context is partitioned across ``S`` CTAs that emit partial + softmax states (acc, m, l); a merge kernel folds the partials plus the + dense noise suffix. Per-split ranges are derived from the device-resident + ctx_len, so the fixed grid stays CUDA-graph safe. +- Large batch: single-pass kernel (one CTA per (request, kv head)) — already + faster than a dense read at this occupancy, and it skips the partial + buffers and merge round-trip. + +Both paths tile ``BLOCK_N`` context tokens (several pages) per iteration +with a per-token page-id gather, rather than one 32-token page per ``tl.dot``. +""" + +import math + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _attend_ctx_tokens( + q_tile, # [R_PAD, D] bf16 + k_cache_ptr, + v_cache_ptr, + blk_row_ptr, # page-id row for this request + kvh, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + start, + end, # token range [start, end) of this CTA + sm_scale, + m_i, + l_i, + acc, + R_PAD: tl.constexpr, + TPB: tl.constexpr, + D: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Accumulate online-softmax over context tokens [start, end).""" + d = tl.arange(0, D) + tok = tl.arange(0, BLOCK_N) + for n0 in range(start, end, BLOCK_N): + idx = n0 + tok + valid = idx < end + # Per-token page gather: page ids come from the block table row. + page = tl.load(blk_row_ptr + idx // TPB, mask=valid, other=0).to(tl.int64) + toff = idx % TPB + k_ptrs = ( + k_cache_ptr + + page[:, None] * stride_kp + + toff[:, None] * stride_kt + + kvh * stride_kh + + d[None, :] + ) + v_ptrs = ( + v_cache_ptr + + page[:, None] * stride_vp + + toff[:, None] * stride_vt + + kvh * stride_vh + + d[None, :] + ) + k = tl.load(k_ptrs, mask=valid[:, None], other=0.0).to(tl.bfloat16) + v = tl.load(v_ptrs, mask=valid[:, None], other=0.0).to(tl.bfloat16) + + s = tl.dot(q_tile, tl.trans(k)) * sm_scale # [R_PAD, BLOCK_N] + s = tl.where(valid[None, :], s, float("-inf")) + + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + # All-(-inf) rows (empty split) keep m == -inf; exp(-inf - -inf) is + # NaN, so guard the rescale factor. + alpha = tl.where(m_new == float("-inf"), 1.0, tl.exp(m_i - m_new)) + p_ij = tl.exp(s - m_new[:, None]) + p_ij = tl.where(valid[None, :], p_ij, 0.0) + l_i = l_i * alpha + tl.sum(p_ij, axis=1) + acc = acc * alpha[:, None] + tl.dot(p_ij.to(tl.bfloat16), v) + m_i = m_new + return m_i, l_i, acc + + +@triton.jit +def _attend_noise_suffix( + q_tile, # [R_PAD, D] + k_noise_ptr, + v_noise_ptr, + b, + kvh, + stride_nb, + stride_nq, + stride_nh, + sm_scale, + m_i, + l_i, + acc, + R_PAD: tl.constexpr, + Q: tl.constexpr, + D: tl.constexpr, + NOISE_PAD: tl.constexpr, +): + """Fold the dense per-step noise K/V suffix into the softmax state.""" + d = tl.arange(0, D) + tn = tl.arange(0, NOISE_PAD) + n_valid = tn < Q + kn_ptrs = k_noise_ptr + b * stride_nb + tn[:, None] * stride_nq + kvh * stride_nh + d[None, :] + vn_ptrs = v_noise_ptr + b * stride_nb + tn[:, None] * stride_nq + kvh * stride_nh + d[None, :] + kn = tl.load(kn_ptrs, mask=n_valid[:, None], other=0.0).to(tl.bfloat16) + vn = tl.load(vn_ptrs, mask=n_valid[:, None], other=0.0).to(tl.bfloat16) + + s = tl.dot(q_tile, tl.trans(kn)) * sm_scale # [R_PAD, NOISE_PAD] + s = tl.where(n_valid[None, :], s, float("-inf")) + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + alpha = tl.where(m_new == float("-inf"), 1.0, tl.exp(m_i - m_new)) + p_n = tl.exp(s - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p_n, axis=1) + acc = acc * alpha[:, None] + tl.dot(p_n.to(tl.bfloat16), vn) + return m_new, l_i, acc + + +@triton.jit +def _load_q_tile( + q_ptr, + b, + kvh, + stride_qb, + stride_qq, + stride_qh, + Q: tl.constexpr, + GROUP: tl.constexpr, + D: tl.constexpr, + R_PAD: tl.constexpr, +): + # Row r -> (head h = kvh*GROUP + r // Q, query qi = r % Q). tl.arange + # needs a power-of-two span, so rows are padded to R_PAD; padded rows + # load zeros and are masked again at the output store. + R: tl.constexpr = Q * GROUP + r = tl.arange(0, R_PAD) + row_ok = r < R + h = kvh * GROUP + r // Q + qi = r % Q + d = tl.arange(0, D) + q_ptrs = q_ptr + b * stride_qb + qi[:, None] * stride_qq + h[:, None] * stride_qh + d[None, :] + return tl.load(q_ptrs, mask=row_ok[:, None], other=0.0).to(tl.bfloat16) # [R_PAD, D] + + +@triton.jit +def _dflash_ctx_attn_kernel( + q_ptr, # [B, Q, NH, D] + k_cache_ptr, # [pages, TPB, NKV, D] + v_cache_ptr, # [pages, TPB, NKV, D] + blk_ptr, # [B, W] page indices per request + ctx_len_ptr, # [B] valid ctx tokens per request (device) + k_noise_ptr, # [B, Q, NKV, D] + v_noise_ptr, # [B, Q, NKV, D] + out_ptr, # [B, Q, NH, D] + stride_qb, + stride_qq, + stride_qh, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + stride_bb, + stride_nb, + stride_nq, + stride_nh, + stride_ob, + stride_oq, + stride_oh, + sm_scale, + Q: tl.constexpr, # queries per request (draft block size) + GROUP: tl.constexpr, # q heads per kv head + TPB: tl.constexpr, # tokens per page + D: tl.constexpr, # head dim + BLOCK_N: tl.constexpr, # ctx tokens per iteration + NOISE_PAD: tl.constexpr, # Q padded to >=16 for tl.dot + R_PAD: tl.constexpr, # Q*GROUP padded to a power of two +): + """Single-pass path: one CTA per (request, kv head).""" + b = tl.program_id(0) + kvh = tl.program_id(1) + R: tl.constexpr = Q * GROUP + + q_tile = _load_q_tile(q_ptr, b, kvh, stride_qb, stride_qq, stride_qh, Q, GROUP, D, R_PAD) + + m_i = tl.full([R_PAD], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([R_PAD], dtype=tl.float32) + acc = tl.zeros([R_PAD, D], dtype=tl.float32) + + ctx_len = tl.load(ctx_len_ptr + b) + m_i, l_i, acc = _attend_ctx_tokens( + q_tile, + k_cache_ptr, + v_cache_ptr, + blk_ptr + b * stride_bb, + kvh, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + 0, + ctx_len, + sm_scale, + m_i, + l_i, + acc, + R_PAD, + TPB, + D, + BLOCK_N, + ) + + m_i, l_i, acc = _attend_noise_suffix( + q_tile, + k_noise_ptr, + v_noise_ptr, + b, + kvh, + stride_nb, + stride_nq, + stride_nh, + sm_scale, + m_i, + l_i, + acc, + R_PAD, + Q, + D, + NOISE_PAD, + ) + + out = acc / l_i[:, None] + r = tl.arange(0, R_PAD) + row_ok = r < R + qi = r % Q + h = kvh * GROUP + r // Q + d = tl.arange(0, D) + out_ptrs = ( + out_ptr + b * stride_ob + qi[:, None] * stride_oq + h[:, None] * stride_oh + d[None, :] + ) + tl.store(out_ptrs, out.to(out_ptr.dtype.element_ty), mask=row_ok[:, None]) + + +@triton.jit +def _dflash_ctx_attn_split_kernel( + q_ptr, + k_cache_ptr, + v_cache_ptr, + blk_ptr, + ctx_len_ptr, + part_acc_ptr, # [B, NKV, S, R, D] fp32 + part_m_ptr, # [B, NKV, S, R] fp32 + part_l_ptr, # [B, NKV, S, R] fp32 + stride_qb, + stride_qq, + stride_qh, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + stride_bb, + sm_scale, + S, # number of context splits (runtime: varies with eager batch size, + # keeping it non-constexpr avoids a Triton recompile per batch size) + NKV: tl.constexpr, + Q: tl.constexpr, + GROUP: tl.constexpr, + TPB: tl.constexpr, + D: tl.constexpr, + BLOCK_N: tl.constexpr, + R_PAD: tl.constexpr, +): + """Split phase: CTA (b, kvh, s) covers a slice of the context.""" + b = tl.program_id(0) + kvh = tl.program_id(1) + s_id = tl.program_id(2) + + ctx_len = tl.load(ctx_len_ptr + b) + # Even token split, rounded to BLOCK_N so slices don't share tiles. + per_split = tl.cdiv(tl.cdiv(ctx_len, BLOCK_N), S) * BLOCK_N + start = s_id * per_split + end = tl.minimum(start + per_split, ctx_len) + + q_tile = _load_q_tile(q_ptr, b, kvh, stride_qb, stride_qq, stride_qh, Q, GROUP, D, R_PAD) + + m_i = tl.full([R_PAD], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([R_PAD], dtype=tl.float32) + acc = tl.zeros([R_PAD, D], dtype=tl.float32) + + m_i, l_i, acc = _attend_ctx_tokens( + q_tile, + k_cache_ptr, + v_cache_ptr, + blk_ptr + b * stride_bb, + kvh, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + start, + end, + sm_scale, + m_i, + l_i, + acc, + R_PAD, + TPB, + D, + BLOCK_N, + ) + + # Padded rows are stored too (the partial buffers are R_PAD-strided); + # the merge kernel discards them at its masked output store. + r = tl.arange(0, R_PAD) + d = tl.arange(0, D) + base = ((b * NKV + kvh) * S + s_id) * R_PAD + tl.store(part_m_ptr + base + r, m_i) + tl.store(part_l_ptr + base + r, l_i) + tl.store(part_acc_ptr + (base + r)[:, None] * D + d[None, :], acc) + + +@triton.jit +def _dflash_ctx_attn_merge_kernel( + q_ptr, + part_acc_ptr, # [B, NKV, S, R, D] + part_m_ptr, # [B, NKV, S, R] + part_l_ptr, # [B, NKV, S, R] + k_noise_ptr, + v_noise_ptr, + out_ptr, + stride_qb, + stride_qq, + stride_qh, + stride_nb, + stride_nq, + stride_nh, + stride_ob, + stride_oq, + stride_oh, + sm_scale, + S, # runtime, see split kernel + NKV: tl.constexpr, + Q: tl.constexpr, + GROUP: tl.constexpr, + D: tl.constexpr, + NOISE_PAD: tl.constexpr, + R_PAD: tl.constexpr, +): + """Merge phase: fold split partials, then the noise suffix.""" + b = tl.program_id(0) + kvh = tl.program_id(1) + R: tl.constexpr = Q * GROUP + r = tl.arange(0, R_PAD) + d = tl.arange(0, D) + + m_i = tl.full([R_PAD], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([R_PAD], dtype=tl.float32) + acc = tl.zeros([R_PAD, D], dtype=tl.float32) + + for s_id in range(0, S): + base = ((b * NKV + kvh) * S + s_id) * R_PAD + m_s = tl.load(part_m_ptr + base + r) + l_s = tl.load(part_l_ptr + base + r) + a_s = tl.load(part_acc_ptr + (base + r)[:, None] * D + d[None, :]) + + m_new = tl.maximum(m_i, m_s) + guard = m_new == float("-inf") + alpha = tl.where(guard, 1.0, tl.exp(m_i - m_new)) + beta = tl.where(guard | (m_s == float("-inf")), 0.0, tl.exp(m_s - m_new)) + # Empty splits carry l == 0, so they contribute nothing even when + # beta is defined. + l_i = l_i * alpha + l_s * beta + acc = acc * alpha[:, None] + a_s * beta[:, None] + m_i = m_new + + q_tile = _load_q_tile(q_ptr, b, kvh, stride_qb, stride_qq, stride_qh, Q, GROUP, D, R_PAD) + m_i, l_i, acc = _attend_noise_suffix( + q_tile, + k_noise_ptr, + v_noise_ptr, + b, + kvh, + stride_nb, + stride_nq, + stride_nh, + sm_scale, + m_i, + l_i, + acc, + R_PAD, + Q, + D, + NOISE_PAD, + ) + + out = acc / l_i[:, None] + row_ok = r < R + qi = r % Q + h = kvh * GROUP + r // Q + out_ptrs = ( + out_ptr + b * stride_ob + qi[:, None] * stride_oq + h[:, None] * stride_oh + d[None, :] + ) + tl.store(out_ptrs, out.to(out_ptr.dtype.element_ty), mask=row_ok[:, None]) + + +# Reused across steps so CUDA graph capture sees stable addresses; keyed by +# (B, NKV, S, R, D, device). Bounded: shapes are fixed per serving config. +_PARTIAL_BUFS = {} + + +def _get_partial_bufs(B, NKV, S, R, D, device): + key = (B, NKV, S, R, D, device) + bufs = _PARTIAL_BUFS.get(key) + if bufs is None: + acc = torch.empty(B * NKV * S * R * D, dtype=torch.float32, device=device) + m = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) + l_buf = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) + bufs = (acc, m, l_buf) + _PARTIAL_BUFS[key] = bufs + return bufs + + +def _num_splits(B: int, NKV: int) -> int: + """Pick the context split count from the (static) launch batch size. + + Aim for enough CTAs to occupy the GPU; cap so per-split work stays + above ~2 tiles. B is padded/static under CUDA graph capture, so the + resulting grid is capture-safe. + """ + sm_count = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + if B * NKV >= sm_count: + return 1 + return min(16, max(1, sm_count // (B * NKV))) + + +def dflash_ctx_paged_attention( + q: torch.Tensor, # [B, Q, NH, D] bf16 + k_cache: torch.Tensor, # [pages, TPB, NKV, D] fp8/bf16 pool view + v_cache: torch.Tensor, + block_idx: torch.Tensor, # [B, W] int32/int64 page ids + ctx_lens: torch.Tensor, # [B] int32/int64, device-resident + k_noise: torch.Tensor, # [B, Q, NKV, D] bf16 + v_noise: torch.Tensor, +) -> torch.Tensor: + B, Q, NH, D = q.shape + _, TPB, NKV, _ = k_cache.shape + assert NH % NKV == 0 + group = NH // NKV + # tl.arange spans must be powers of two. Query rows (R = Q*GROUP) are + # padded below, so any draft length / GQA ratio works; head_dim and the + # manager page size have no such padding path and must be powers of two + # (every known draft uses 64/128; the manager default page is 32). + if D < 16 or (D & (D - 1)) != 0: + raise ValueError( + f"DFlash hybrid ctx kernel requires a power-of-two head_dim >= 16, got {D}." + ) + if TPB & (TPB - 1) != 0: + raise ValueError( + f"DFlash hybrid ctx kernel requires a power-of-two tokens_per_block, got {TPB}." + ) + assert q.stride(-1) == 1 and k_cache.stride(-1) == 1 + + out = torch.empty_like(q) + sm_scale = 1.0 / math.sqrt(D) + noise_pad = max(16, triton.next_power_of_2(Q)) + BLOCK_N = 128 + R = Q * group + # Pad query rows to a power of two (>= 16 for tl.dot); e.g. a draft + # length of 5 gives Q = 6 and R = 24 -> R_PAD = 32. + r_pad = max(16, triton.next_power_of_2(R)) + + common_q = (q.stride(0), q.stride(1), q.stride(2)) + common_n = (k_noise.stride(0), k_noise.stride(1), k_noise.stride(2)) + common_o = (out.stride(0), out.stride(1), out.stride(2)) + + S = _num_splits(B, NKV) + if S == 1: + _dflash_ctx_attn_kernel[(B, NKV)]( + q, + k_cache, + v_cache, + block_idx, + ctx_lens, + k_noise, + v_noise, + out, + *common_q, + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + block_idx.stride(0), + *common_n, + *common_o, + sm_scale, + Q=Q, + GROUP=group, + TPB=TPB, + D=D, + BLOCK_N=BLOCK_N, + NOISE_PAD=noise_pad, + R_PAD=r_pad, + num_warps=4, + ) + return out + + part_acc, part_m, part_l = _get_partial_bufs(B, NKV, S, r_pad, D, q.device) + _dflash_ctx_attn_split_kernel[(B, NKV, S)]( + q, + k_cache, + v_cache, + block_idx, + ctx_lens, + part_acc, + part_m, + part_l, + *common_q, + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + block_idx.stride(0), + sm_scale, + NKV=NKV, + S=S, + Q=Q, + GROUP=group, + TPB=TPB, + D=D, + BLOCK_N=BLOCK_N, + R_PAD=r_pad, + num_warps=4, + ) + _dflash_ctx_attn_merge_kernel[(B, NKV)]( + q, + part_acc, + part_m, + part_l, + k_noise, + v_noise, + out, + *common_q, + *common_n, + *common_o, + sm_scale, + NKV=NKV, + S=S, + Q=Q, + GROUP=group, + D=D, + NOISE_PAD=noise_pad, + R_PAD=r_pad, + num_warps=4, + ) + return out + + +def dflash_ctx_paged_attention_ref( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + block_idx: torch.Tensor, + ctx_lens: torch.Tensor, + k_noise: torch.Tensor, + v_noise: torch.Tensor, +) -> torch.Tensor: + """Plain-torch reference (slow gather path) for parity tests.""" + B, Q, NH, D = q.shape + _, TPB, NKV, _ = k_cache.shape + group = NH // NKV + out = torch.empty_like(q) + lens = ctx_lens.tolist() + for b in range(B): + n = int(lens[b]) + n_pages = (n + TPB - 1) // TPB + pages = block_idx[b, :n_pages].long() + k_ctx = k_cache[pages].to(torch.float32).reshape(-1, NKV, D)[:n] + v_ctx = v_cache[pages].to(torch.float32).reshape(-1, NKV, D)[:n] + k = torch.cat([k_ctx, k_noise[b].to(torch.float32)], dim=0) # [n+Q, NKV, D] + v = torch.cat([v_ctx, v_noise[b].to(torch.float32)], dim=0) + qb = q[b].to(torch.float32) # [Q, NH, D] + for h in range(NH): + g = h // group + s = (qb[:, h] @ k[:, g].T) / math.sqrt(D) # [Q, n+Q] + out[b, :, h] = (torch.softmax(s, dim=-1) @ v[:, g]).to(out.dtype) + return out diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 164d6758ed4a..c949ab8b05a4 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -361,6 +361,9 @@ def get_num_spec_layers(spec_config): if spec_config.spec_dec_mode.is_eagle3_one_model(): num_draft_hidden_layers = spec_config._num_draft_hidden_layers return num_draft_hidden_layers if num_draft_hidden_layers is not None else 1 + if (spec_config.spec_dec_mode.is_dflash() + and getattr(spec_config, 'use_hybrid_context', False)): + return spec_config._num_draft_layers or 0 return 0 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0f0e5f649244..dd1746c61102 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2510,11 +2510,27 @@ class DFlashDecodingConfig(DecodingBaseConfig): "for cross-attention in the draft model. If None, read from the draft " "model config (dflash_config.target_layer_ids).") + use_hybrid_context: bool = Field( + default=False, + status="prototype", + description= + "Store the draft cross-attention context as spec layers of the target " + "KV cache manager instead of a dense per-request buffer. Context " + "memory then scales with the KV pool, and prefix-cache hits restore " + "the draft context.") + decoding_type: Literal["DFlash"] = Field(default="DFlash") + # Draft attention layer count is resolved from the draft model config by + # KvCacheCreator when use_hybrid_context is enabled. + _num_draft_layers: Optional[int] = PrivateAttr(default=None) + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len + if self.use_hybrid_context: + # Draft context lives in the target manager + self._allow_separate_draft_kv_cache = False return self @property