From 12d3b1cc2aa503ec7ba13a8db578d21951a95af1 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Tue, 7 Jul 2026 00:46:49 +0000 Subject: [PATCH 01/10] [None][feat] DFlash: size target KV pool for manager-resident draft context Add DFlashDecodingConfig.use_hybrid_context (prototype). When enabled, the draft cross-attention context layers are registered as spec layers of the target KV cache manager: get_num_spec_layers reports the draft layer count and KvCacheCreator budgets the extra per-token cost on the last PP rank, instead of allocating a dense per-request context buffer. Signed-off-by: chungen04 --- tensorrt_llm/_torch/pyexecutor/_util.py | 20 ++++++++++++++++++++ tensorrt_llm/_torch/speculative/utils.py | 5 +++++ tensorrt_llm/llmapi/llm_args.py | 17 +++++++++++++++++ 3 files changed, 42 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b2d45dec50aa..b394390eef8a 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,14 @@ 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. + total += self._per_manager_cache_cost( + self._kv_cache_manager_cls, + self._get_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, diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 164d6758ed4a..3c72fb62c63c 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -361,6 +361,11 @@ 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)): + # Hybrid ctx: draft cross-attn K/V lives in the target manager as + # extra spec layers (_num_draft_layers resolved by KvCacheCreator). + 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..15bdcb2887fa 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2510,11 +2510,28 @@ 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; 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; a separate one-model + # draft KV cache would be redundant. + self._allow_separate_draft_kv_cache = False return self @property From 2fc997219d69cbbdfb92bd6e04b53ecbf2440d49 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Tue, 7 Jul 2026 00:58:38 +0000 Subject: [PATCH 02/10] [None][feat] DFlash: Triton kernel for paged context cross-attention Read the draft context K/V directly from the target KV cache manager's paged pool. Existing kernels do not fit: sequence lengths are device-resident (acceptance-dependent, updated inside CUDA graphs), pages are target-pool sized rather than 256-token aligned, and the dense noise suffix must be fused into the same softmax. Online-softmax with bf16 tensor-core dots and fp32 accumulation; supports fp8 context storage. Includes a torch reference implementation for testing. Signed-off-by: chungen04 --- .../_torch/speculative/dflash_hybrid_attn.py | 212 ++++++++++++++++++ 1 file changed, 212 insertions(+) create mode 100644 tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py 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..b764b00ca30d --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py @@ -0,0 +1,212 @@ +# 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. +""" + +import math + +import torch +import triton +import triton.language as tl + + +@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 + NOISE_PAD: tl.constexpr, # Q padded to >=16 for tl.dot +): + b = tl.program_id(0) + kvh = tl.program_id(1) + + # Row r -> (head h = kvh*GROUP + r // Q, query qi = r % Q) + R: tl.constexpr = Q * GROUP + r = tl.arange(0, 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, :] + q_tile = tl.load(q_ptrs).to(tl.bfloat16) # [R, D] + + m_i = tl.full([R], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([R], dtype=tl.float32) + acc = tl.zeros([R, D], dtype=tl.float32) + + ctx_len = tl.load(ctx_len_ptr + b) + n_pages = tl.cdiv(ctx_len, TPB) + t = tl.arange(0, TPB) + + for p in range(0, n_pages): + page = tl.load(blk_ptr + b * stride_bb + p).to(tl.int64) + valid = (p * TPB + t) < ctx_len + k_ptrs = ( + k_cache_ptr + page * stride_kp + t[:, None] * stride_kt + kvh * stride_kh + d[None, :] + ) + v_ptrs = ( + v_cache_ptr + page * stride_vp + t[:, 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, TPB] + s = tl.where(valid[None, :], s, float("-inf")) + + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + alpha = tl.exp(m_i - m_new) + p_ij = tl.exp(s - m_new[:, None]) + 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 + + # Dense noise suffix: Q transient mask/bonus K/V, fully visible. + 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, NOISE_PAD] + s = tl.where(n_valid[None, :], s, float("-inf")) + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + alpha = 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) + + out = acc / l_i[:, None] + 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)) + + +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 + assert (Q * group) >= 16 and D >= 16, "tl.dot needs tiles >= 16" + assert q.stride(-1) == 1 and k_cache.stride(-1) == 1 + + out = torch.empty_like(q) + grid = (B, NKV) + _dflash_ctx_attn_kernel[grid]( + q, + k_cache, + v_cache, + block_idx, + ctx_lens, + k_noise, + v_noise, + out, + q.stride(0), + q.stride(1), + q.stride(2), + 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), + k_noise.stride(0), + k_noise.stride(1), + k_noise.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + 1.0 / math.sqrt(D), + Q=Q, + GROUP=group, + TPB=TPB, + D=D, + NOISE_PAD=max(16, triton.next_power_of_2(Q)), + ) + 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 From 7485c67882eea64fb3ecd00ad03873611b22d098 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Tue, 7 Jul 2026 01:00:16 +0000 Subject: [PATCH 03/10] [None][feat] DFlash: store and read draft context via the target KV cache manager When use_hybrid_context is enabled, the DFlash worker scatters the projected context K/V into the draft spec layers of the target manager's paged pool (position-aligned, so prefix-cache hits restore context), and the drafter reads it back with the paged cross-attention kernel instead of flash_attn_with_kvcache over a dense per-request buffer. Dense-path behavior is unchanged when the flag is off. Signed-off-by: chungen04 --- .../_torch/models/modeling_speculative.py | 61 +++++--- tensorrt_llm/_torch/speculative/dflash.py | 139 ++++++++++++++++-- 2 files changed, 171 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 6e17258067a1..3e377ae3e356 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,37 @@ 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; the + # per-step noise K/V is a dense suffix, never cached. + 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/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index efa31fa2cc1e..1a58695d83f4 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,65 @@ 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 by + # get_pp_layers; take the last num_layers ids. + 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] + self._hybrid_k_bufs = [b[:, 0] for b in bufs] + self._hybrid_v_bufs = [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 +418,20 @@ 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); + # start from there instead of losing the prefix. + 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: + # Keep the compute dtype; 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 +439,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 +488,10 @@ def forward( self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata) spec_metadata._dflash_worker = self + if self._use_hybrid_context: + # Page ids straight from the manager's offsets (device, in-graph). + 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 +572,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 +759,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 +813,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 + ), } From 956bb5e8fa993ed7baca03f1e379873e6071e8e3 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Wed, 8 Jul 2026 21:24:01 +0000 Subject: [PATCH 04/10] [None][perf] DFlash hybrid ctx: split-KV paged cross-attention kernel The single-pass kernel launched (batch, num_kv_heads) CTAs and marched one 32-token page per tl.dot iteration; at small batch that is ~4 CTAs on a 148-SM GPU and cost ~+1 ms/step vs the dense-buffer flash_attn path (measured -24% output throughput at concurrency 1, Qwen3-8B TP2). Rework it flash-decoding style: - Small batch: split-KV. The context is partitioned across S CTAs that emit partial softmax states; a merge kernel folds the partials plus the dense noise suffix. S derives from the launch batch size (static under CUDA graph capture); per-split ranges derive from the device-resident ctx_len, so the fixed grid stays graph-safe. Partial buffers are cached per shape for stable capture addresses. - Large batch: keep the single-pass path (no partial round-trip). - Both paths tile BLOCK_N=128 context tokens (4 pages) per tl.dot with a per-token page-id gather instead of one page per iteration. Kernel time at B=1 (ctx 8192, 4 KV heads/rank): 0.284 ms -> 0.063 ms, now faster than dense flash_attn_with_kvcache (0.077 ms). End-to-end aiperf sweep (Qwen3-8B TP2, ISL 8000/OSL 512): hybrid ctx goes from -24%..-7% vs the dense baseline to +2% (c=1), +8% (c=8), +4..5% (c=16..64), with bit-identical greedy outputs. Signed-off-by: chungen04 --- .../_torch/speculative/dflash_hybrid_attn.py | 413 ++++++++++++++---- 1 file changed, 339 insertions(+), 74 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py index b764b00ca30d..127187502bdc 100644 --- a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py +++ b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py @@ -20,6 +20,21 @@ (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 @@ -29,6 +44,118 @@ import triton.language as tl +@triton.jit +def _attend_ctx_tokens( + q_tile, # [R, 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: 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, 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, D] + k_noise_ptr, + v_noise_ptr, + b, + kvh, + stride_nb, + stride_nq, + stride_nh, + sm_scale, + m_i, + l_i, + acc, + R: 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, 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, +): + # Row r -> (head h = kvh*GROUP + r // Q, query qi = r % Q) + R: tl.constexpr = Q * GROUP + r = tl.arange(0, 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).to(tl.bfloat16) # [R, D] + + @triton.jit def _dflash_ctx_attn_kernel( q_ptr, # [B, Q, NH, D] @@ -60,74 +187,198 @@ def _dflash_ctx_attn_kernel( 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 ): + """Single-pass path: one CTA per (request, kv head).""" b = tl.program_id(0) kvh = tl.program_id(1) - - # Row r -> (head h = kvh*GROUP + r // Q, query qi = r % Q) R: tl.constexpr = Q * GROUP - r = tl.arange(0, 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, :] - q_tile = tl.load(q_ptrs).to(tl.bfloat16) # [R, D] + q_tile = _load_q_tile(q_ptr, b, kvh, stride_qb, stride_qq, stride_qh, Q, GROUP, D) m_i = tl.full([R], float("-inf"), dtype=tl.float32) l_i = tl.zeros([R], dtype=tl.float32) acc = tl.zeros([R, D], dtype=tl.float32) ctx_len = tl.load(ctx_len_ptr + b) - n_pages = tl.cdiv(ctx_len, TPB) - t = tl.arange(0, TPB) + 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, TPB, D, BLOCK_N) - for p in range(0, n_pages): - page = tl.load(blk_ptr + b * stride_bb + p).to(tl.int64) - valid = (p * TPB + t) < ctx_len - k_ptrs = ( - k_cache_ptr + page * stride_kp + t[:, None] * stride_kt + kvh * stride_kh + d[None, :] - ) - v_ptrs = ( - v_cache_ptr + page * stride_vp + t[:, 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) + 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, Q, D, NOISE_PAD) - s = tl.dot(q_tile, tl.trans(k)) * sm_scale # [R, TPB] - s = tl.where(valid[None, :], s, float("-inf")) + out = acc / l_i[:, None] + qi = tl.arange(0, R) % Q + h = kvh * GROUP + tl.arange(0, 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)) - m_new = tl.maximum(m_i, tl.max(s, axis=1)) - alpha = tl.exp(m_i - m_new) - p_ij = tl.exp(s - m_new[:, None]) - 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 - # Dense noise suffix: Q transient mask/bonus K/V, fully visible. - 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) +@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, + NKV: tl.constexpr, + S: tl.constexpr, # number of context splits + Q: tl.constexpr, + GROUP: tl.constexpr, + TPB: tl.constexpr, + D: tl.constexpr, + BLOCK_N: 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) + R: tl.constexpr = Q * GROUP - s = tl.dot(q_tile, tl.trans(kn)) * sm_scale # [R, NOISE_PAD] - s = tl.where(n_valid[None, :], s, float("-inf")) - m_new = tl.maximum(m_i, tl.max(s, axis=1)) - alpha = 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) + 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) + + m_i = tl.full([R], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([R], dtype=tl.float32) + acc = tl.zeros([R, 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, TPB, D, BLOCK_N) + + r = tl.arange(0, R) + d = tl.arange(0, D) + base = ((b * NKV + kvh) * S + s_id) * R + 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, + NKV: tl.constexpr, + S: tl.constexpr, + Q: tl.constexpr, + GROUP: tl.constexpr, + D: tl.constexpr, + NOISE_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) + d = tl.arange(0, D) + + m_i = tl.full([R], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([R], dtype=tl.float32) + acc = tl.zeros([R, D], dtype=tl.float32) + + for s_id in range(0, S): + base = ((b * NKV + kvh) * S + s_id) * R + 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) + 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, Q, D, NOISE_PAD) out = acc / l_i[:, None] + 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)) +# 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 = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) + bufs = (acc, m, l) + _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 @@ -145,38 +396,52 @@ def dflash_ctx_paged_attention( assert q.stride(-1) == 1 and k_cache.stride(-1) == 1 out = torch.empty_like(q) - grid = (B, NKV) - _dflash_ctx_attn_kernel[grid]( - q, - k_cache, - v_cache, - block_idx, - ctx_lens, - k_noise, - v_noise, - out, - q.stride(0), - q.stride(1), - q.stride(2), - k_cache.stride(0), - k_cache.stride(1), - k_cache.stride(2), - v_cache.stride(0), - v_cache.stride(1), - v_cache.stride(2), + sm_scale = 1.0 / math.sqrt(D) + noise_pad = max(16, triton.next_power_of_2(Q)) + BLOCK_N = 128 + R = Q * group + + 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, + num_warps=4, + ) + return out + + part_acc, part_m, part_l = _get_partial_bufs(B, NKV, S, R, 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), - k_noise.stride(0), - k_noise.stride(1), - k_noise.stride(2), - out.stride(0), - out.stride(1), - out.stride(2), - 1.0 / math.sqrt(D), - Q=Q, - GROUP=group, - TPB=TPB, - D=D, - NOISE_PAD=max(16, triton.next_power_of_2(Q)), + sm_scale, + NKV=NKV, S=S, Q=Q, GROUP=group, TPB=TPB, D=D, BLOCK_N=BLOCK_N, + 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, + num_warps=4, ) return out From 7c494d5da2480a801c1d54e97ec6b88b99dac77c Mon Sep 17 00:00:00 2001 From: chungen04 Date: Wed, 8 Jul 2026 21:52:05 +0000 Subject: [PATCH 05/10] [None][fix] DFlash hybrid ctx: make split count a runtime kernel arg The split count S was a tl.constexpr derived from the launch batch size. Mixed context+generation iterations run eagerly at their actual (unpadded) batch size, so serving traffic kept discovering new S values and paying a blocking Triton JIT compile mid-request: p90 ITL jumped 3-4x and TTFT tails reached 13-17 s on the first sweep levels after a server start, disappearing once the in-process compile cache warmed (which made it look like measurement noise). S only enters index arithmetic and the merge loop trip count, so pass it as a runtime argument: each kernel compiles exactly once and no JIT ever runs on the serving path. Kernel parity and timings unchanged. Signed-off-by: chungen04 --- tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py index 127187502bdc..351e725d81e4 100644 --- a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py +++ b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py @@ -242,8 +242,9 @@ def _dflash_ctx_attn_split_kernel( 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, - S: tl.constexpr, # number of context splits Q: tl.constexpr, GROUP: tl.constexpr, TPB: tl.constexpr, @@ -300,8 +301,8 @@ def _dflash_ctx_attn_merge_kernel( stride_oq, stride_oh, sm_scale, + S, # runtime, see split kernel NKV: tl.constexpr, - S: tl.constexpr, Q: tl.constexpr, GROUP: tl.constexpr, D: tl.constexpr, From 40a31e37346b52100eda53c2865528c9ef14046c Mon Sep 17 00:00:00 2001 From: chungen04 Date: Wed, 8 Jul 2026 22:16:05 +0000 Subject: [PATCH 06/10] [None][fix] DFlash hybrid ctx: support drafts whose head_dim differs from the target The KV cache manager has a single head_dim per pool, so draft spec layers inherited the target's head_dim and the context store failed with a shape mismatch whenever the two differ (e.g. a 128-head_dim draft on a 256-head_dim target). Register draft spec layers in target-head_dim units (heads scaled so bytes stay exact) and view the pool buffers back to draft geometry in the worker; pure view, no copy. Raises a clear error when the draft KV row is not expressible in target head_dim units. Signed-off-by: chungen04 --- tensorrt_llm/_torch/pyexecutor/_util.py | 26 +++++++++++++++++++++-- tensorrt_llm/_torch/speculative/dflash.py | 23 ++++++++++++++++++-- 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b394390eef8a..f08a21996d9d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1547,6 +1547,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: @@ -1563,6 +1564,24 @@ 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 so pool geometry + # stays exact; consumers view the buffers back to draft geometry + # (see DFlashSpecWorker._lazy_init_hybrid_ctx). + 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 @@ -1713,8 +1732,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 1a58695d83f4..ca5a49a6ffe5 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -231,8 +231,27 @@ def _lazy_init_hybrid_ctx(self, draft_model, spec_metadata, attn_metadata): # get_pp_layers; take the last num_layers ids. 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] - self._hybrid_k_bufs = [b[:, 0] for b in bufs] - self._hybrid_v_bufs = [b[:, 1] for b in bufs] + + # The manager sizes spec layers in target-head_dim units (see + # _build_per_layer_num_kv_heads), 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. + # The last two dims are contiguous per token, so this is a pure view. + 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). From 7328ba75e360af6a68cf5a8f1ce010faf495b898 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Wed, 8 Jul 2026 22:16:26 +0000 Subject: [PATCH 07/10] [None][fix] DFlash hybrid ctx: fix KV size estimation for hybrid-Mamba targets The hybrid-ctx cost path passed the draft's plain-attention config into the target's Mamba-hybrid cache manager class, raising 'Qwen3Config is not a supported hybrid Mamba config' at startup on GDN/Mamba-hybrid targets (e.g. Qwen3.5). Resolve the manager class from the draft config instead, matching the external-drafter path. Signed-off-by: chungen04 --- tensorrt_llm/_torch/pyexecutor/_util.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f08a21996d9d..711a97803f45 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -437,10 +437,17 @@ def _get_kv_size_per_token(self, 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. + # per-token cost adds to the same budget. Resolve the manager + # class from the draft config: the target may be hybrid-Mamba + # while the draft context layers are plain attention. + 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( - self._kv_cache_manager_cls, - self._get_effective_draft_config(), + draft_kv_cache_manager_cls, + effective_draft_config, kv_cache_config, num_layers=self._get_num_draft_layers()) return total From 5df8b518e817acba3951c0890ced91defb0b0727 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Thu, 9 Jul 2026 20:11:13 +0000 Subject: [PATCH 08/10] trim comments Signed-off-by: chungen04 --- .../_torch/models/modeling_speculative.py | 3 +-- tensorrt_llm/_torch/pyexecutor/_util.py | 8 ++------ tensorrt_llm/_torch/speculative/dflash.py | 16 ++++++---------- tensorrt_llm/_torch/speculative/utils.py | 2 -- tensorrt_llm/llmapi/llm_args.py | 5 ++--- 5 files changed, 11 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 3e377ae3e356..5f8c27b4c71a 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1376,8 +1376,7 @@ def dflash_forward( head_dim) if hybrid_k_bufs is not None: - # Hybrid: ctx K/V paged in the target manager's pool; the - # per-step noise K/V is a dense suffix, never cached. + # Hybrid: ctx K/V paged in the target manager's pool out = dflash_ctx_paged_attention( Q_bshd, hybrid_k_bufs[layer_idx], diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 711a97803f45..2381ddae0703 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -437,9 +437,7 @@ def _get_kv_size_per_token(self, 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. Resolve the manager - # class from the draft config: the target may be hybrid-Mamba - # while the draft context layers are plain attention. + # 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, @@ -1573,9 +1571,7 @@ def _build_per_layer_num_kv_heads( # 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 so pool geometry - # stays exact; consumers view the buffers back to draft geometry - # (see DFlashSpecWorker._lazy_init_hybrid_ctx). + # 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 diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index ca5a49a6ffe5..628542dde8bc 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -227,16 +227,13 @@ def _lazy_init_hybrid_ctx(self, draft_model, spec_metadata, attn_metadata): draft_model._build_fused_kv_buffers() num_layers = draft_model._num_attn_layers - # Draft spec layers are appended after the target layers by - # get_pp_layers; take the last num_layers ids. + # 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 (see - # _build_per_layer_num_kv_heads), so when the draft head_dim differs + # 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. - # The last two dims are contiguous per token, so this is a pure view. nkv_draft = draft_model._num_kv_heads hd_draft = draft_model._head_dim @@ -247,7 +244,8 @@ def _as_draft_geometry(t: torch.Tensor) -> torch.Tensor: 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.") + 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] @@ -437,8 +435,7 @@ def _store_prefill_context( continue slot = self._free_slots.popleft() self._req_to_slot[req_id] = slot - # Hybrid: reused blocks already hold ctx for [0, first_pos); - # start from there instead of losing the prefix. + # 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] @@ -447,7 +444,7 @@ def _store_prefill_context( actual = end - cur if actual > 0: if self._use_hybrid_context: - # Keep the compute dtype; cast to the pool dtype at write. + # 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) @@ -508,7 +505,6 @@ def forward( spec_metadata._dflash_worker = self if self._use_hybrid_context: - # Page ids straight from the manager's offsets (device, in-graph). self._cur_block_idx = self._hybrid_block_idx(attn_metadata) # Save context lengths before warmup to prevent accumulation diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 3c72fb62c63c..c949ab8b05a4 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -363,8 +363,6 @@ def get_num_spec_layers(spec_config): 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)): - # Hybrid ctx: draft cross-attn K/V lives in the target manager as - # extra spec layers (_num_draft_layers resolved by KvCacheCreator). 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 15bdcb2887fa..dd1746c61102 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2521,7 +2521,7 @@ class DFlashDecodingConfig(DecodingBaseConfig): decoding_type: Literal["DFlash"] = Field(default="DFlash") - # Draft attention layer count; resolved from the draft model config by + # 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) @@ -2529,8 +2529,7 @@ class DFlashDecodingConfig(DecodingBaseConfig): 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; a separate one-model - # draft KV cache would be redundant. + # Draft context lives in the target manager self._allow_separate_draft_kv_cache = False return self From 623d0a221da85d4a417069136a21e770e458bec5 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Thu, 9 Jul 2026 20:47:51 +0000 Subject: [PATCH 09/10] [None][fix] DFlash hybrid ctx: pad query rows to a power of two in the paged attention kernel tl.arange requires a power-of-two span, but the query-row count R = block_size x GQA group was used directly, so draft lengths not of the form 2^k - 1 (e.g. max_draft_len=5 -> R=24) or non-power-of-two GQA ratios failed Triton compilation at warmup. Every published DFlash draft happens to be power-of-two clean, which is why this never fired in benchmarks. Pad rows to R_PAD = next_power_of_2(R) with masked query loads and masked output stores; padded rows carry zeros through the softmax pipeline and are discarded at the store, so results are exact. The partial buffers of the split path are R_PAD-strided. head_dim and tokens_per_block keep explicit power-of-two checks with actionable errors (no known model needs padding there). Power-of-two shapes compile with the masks folded away, so the existing fast path is unchanged; kernel parity verified for R=24, R=20 (group 5), and R=6 against the torch reference, including CUDA-graph replay. Signed-off-by: chungen04 --- .../_torch/speculative/dflash_hybrid_attn.py | 104 +++++++++++------- 1 file changed, 67 insertions(+), 37 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py index 351e725d81e4..787269ef7496 100644 --- a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py +++ b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py @@ -46,7 +46,7 @@ @triton.jit def _attend_ctx_tokens( - q_tile, # [R, D] bf16 + q_tile, # [R_PAD, D] bf16 k_cache_ptr, v_cache_ptr, blk_row_ptr, # page-id row for this request @@ -63,7 +63,7 @@ def _attend_ctx_tokens( m_i, l_i, acc, - R: tl.constexpr, + R_PAD: tl.constexpr, TPB: tl.constexpr, D: tl.constexpr, BLOCK_N: tl.constexpr, @@ -88,7 +88,7 @@ def _attend_ctx_tokens( 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, BLOCK_N] + 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)) @@ -105,7 +105,7 @@ def _attend_ctx_tokens( @triton.jit def _attend_noise_suffix( - q_tile, # [R, D] + q_tile, # [R_PAD, D] k_noise_ptr, v_noise_ptr, b, @@ -117,7 +117,7 @@ def _attend_noise_suffix( m_i, l_i, acc, - R: tl.constexpr, + R_PAD: tl.constexpr, Q: tl.constexpr, D: tl.constexpr, NOISE_PAD: tl.constexpr, @@ -131,7 +131,7 @@ def _attend_noise_suffix( 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, NOISE_PAD] + 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)) @@ -145,15 +145,19 @@ def _attend_noise_suffix( 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) + # 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) + 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).to(tl.bfloat16) # [R, D] + return tl.load(q_ptrs, mask=row_ok[:, None], other=0.0).to(tl.bfloat16) # [R_PAD, D] @triton.jit @@ -189,36 +193,40 @@ def _dflash_ctx_attn_kernel( 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) + 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], float("-inf"), dtype=tl.float32) - l_i = tl.zeros([R], dtype=tl.float32) - acc = tl.zeros([R, D], dtype=tl.float32) + 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, TPB, D, BLOCK_N) + 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, Q, D, NOISE_PAD) + stride_nh, sm_scale, m_i, l_i, acc, R_PAD, Q, D, NOISE_PAD) out = acc / l_i[:, None] - qi = tl.arange(0, R) % Q - h = kvh * GROUP + tl.arange(0, R) // Q + 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)) + tl.store(out_ptrs, out.to(out_ptr.dtype.element_ty), mask=row_ok[:, None]) @triton.jit @@ -250,12 +258,12 @@ def _dflash_ctx_attn_split_kernel( 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) - R: tl.constexpr = Q * GROUP ctx_len = tl.load(ctx_len_ptr + b) # Even token split, rounded to BLOCK_N so slices don't share tiles. @@ -263,20 +271,23 @@ def _dflash_ctx_attn_split_kernel( 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) + 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], float("-inf"), dtype=tl.float32) - l_i = tl.zeros([R], dtype=tl.float32) - acc = tl.zeros([R, D], dtype=tl.float32) + 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, TPB, D, BLOCK_N) + start, end, sm_scale, m_i, l_i, acc, R_PAD, TPB, D, BLOCK_N) - r = tl.arange(0, R) + # 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 + 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) @@ -307,20 +318,21 @@ def _dflash_ctx_attn_merge_kernel( 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) + r = tl.arange(0, R_PAD) d = tl.arange(0, D) - m_i = tl.full([R], float("-inf"), dtype=tl.float32) - l_i = tl.zeros([R], dtype=tl.float32) - acc = tl.zeros([R, D], dtype=tl.float32) + 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 + 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, :]) @@ -335,18 +347,20 @@ def _dflash_ctx_attn_merge_kernel( 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) + 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, Q, D, NOISE_PAD) + 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)) + 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 @@ -393,7 +407,18 @@ def dflash_ctx_paged_attention( _, TPB, NKV, _ = k_cache.shape assert NH % NKV == 0 group = NH // NKV - assert (Q * group) >= 16 and D >= 16, "tl.dot needs tiles >= 16" + # 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 " + f">= 16, got {D}.") + if TPB & (TPB - 1) != 0: + raise ValueError( + f"DFlash hybrid ctx kernel requires a power-of-two " + f"tokens_per_block, got {TPB}.") assert q.stride(-1) == 1 and k_cache.stride(-1) == 1 out = torch.empty_like(q) @@ -401,6 +426,9 @@ def dflash_ctx_paged_attention( 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)) @@ -418,12 +446,12 @@ def dflash_ctx_paged_attention( *common_o, sm_scale, Q=Q, GROUP=group, TPB=TPB, D=D, BLOCK_N=BLOCK_N, - NOISE_PAD=noise_pad, + 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, D, q.device) + 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, @@ -433,6 +461,7 @@ def dflash_ctx_paged_attention( 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)]( @@ -442,6 +471,7 @@ def dflash_ctx_paged_attention( *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 From 873b11182cb09e9411d870ebcc9496bd2522e9f3 Mon Sep 17 00:00:00 2001 From: chungen04 Date: Thu, 9 Jul 2026 21:31:23 +0000 Subject: [PATCH 10/10] fix lint Signed-off-by: chungen04 --- .../_torch/speculative/dflash_hybrid_attn.py | 204 ++++++++++++++---- 1 file changed, 162 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py index 787269ef7496..ba12e0a9928e 100644 --- a/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py +++ b/tensorrt_llm/_torch/speculative/dflash_hybrid_attn.py @@ -78,12 +78,18 @@ def _attend_ctx_tokens( 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, :] + 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, :] + 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) @@ -143,8 +149,15 @@ def _attend_noise_suffix( @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, + 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 @@ -200,8 +213,7 @@ def _dflash_ctx_attn_kernel( 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) + 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) @@ -209,13 +221,47 @@ def _dflash_ctx_attn_kernel( 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) + 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) + 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) @@ -271,17 +317,35 @@ def _dflash_ctx_attn_split_kernel( 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) + 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) + 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. @@ -347,11 +411,25 @@ def _dflash_ctx_attn_merge_kernel( 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) + 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) + 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 @@ -374,8 +452,8 @@ def _get_partial_bufs(B, NKV, S, R, D, device): 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 = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) - bufs = (acc, m, l) + l_buf = torch.empty(B * NKV * S * R, dtype=torch.float32, device=device) + bufs = (acc, m, l_buf) _PARTIAL_BUFS[key] = bufs return bufs @@ -387,8 +465,7 @@ def _num_splits(B: int, NKV: int) -> int: 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 + 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))) @@ -413,12 +490,12 @@ def dflash_ctx_paged_attention( # (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 " - f">= 16, got {D}.") + 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 " - f"tokens_per_block, got {TPB}.") + 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) @@ -437,40 +514,83 @@ def dflash_ctx_paged_attention( 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, + 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), + 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, + 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, + 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), + 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, + 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, + 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, + NKV=NKV, + S=S, + Q=Q, + GROUP=group, + D=D, + NOISE_PAD=noise_pad, R_PAD=r_pad, num_warps=4, )