From eb07ced0b150adf27944086ab1cd97ed5272bf10 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Thu, 27 Aug 2026 20:11:42 -0700 Subject: [PATCH] [None][feat] Page the DSpark drafter context through the draft KV cache manager The standalone drafter kept a private context arena, dense in max_seq_len and allocated after the KV pool had been carved, so free_gpu_memory_fraction never bounded it: at max_seq_len 997696 with max_batch 8 it wants 21.4 GiB/rank and the worker cannot start. Read the already-funded draft pool through the manager's per-request block tables instead, so the footprint follows the sequences served, and size what remains from max_seq_len rather than the drafter's advertised max_position_embeddings (1048576 for K3, hundreds of GiB on its own). The pool index space differs between the V1 and V2 managers, so the block-table divisor is derived from stride(0) rather than assumed. Three gates kept an external drafter off that path, each written for a mode it does not describe: - attention DP: the bail suits MTP, whose draft layers are target-shaped and appendable to the target pool. An external drafter has its own architecture, so nothing is appended and it stayed on the arena -- which under attention DP is sized with KV heads unsharded (20480 vs 2560 B/token for K3). - disaggregation: nvbugs/5807902 reported an Eagle3 RMSNorm failure and was worked around by disabling the separate draft KV cache for every speculative mode. Keep the workaround where it was reported. - disagg slot allocation: _store_prefill_context was the only place that assigned a drafter slot, so a generation worker -- which receives prompt KV instead of prefilling -- collapsed every concurrent request onto the single dummy slot. A context-only worker also releases the target's IndexMapper slot after prefill; the draft mirror never got that call and saturated. Accuracy cannot detect a broken drafter: speculative decoding is lossless, so one producing garbage scores the same and only runs slower. The added test therefore asserts on acceptance length. That is not hypothetical -- the slot collapse showed up as AL 1.087 vs 3.441 (decode steps 123860 -> 39072 for the same output length) while every accuracy gate passed. Signed-off-by: Zhenhuan Chen --- tensorrt_llm/_torch/models/modeling_dflash.py | 5 +- tensorrt_llm/_torch/pyexecutor/_util.py | 13 +- .../_torch/pyexecutor/kv_cache_manager_v2.py | 67 ++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 13 +- .../_torch/pyexecutor/py_executor_creator.py | 10 +- tensorrt_llm/_torch/speculative/dflash.py | 452 +++++++++++++++--- .../references/acceptance_length.yaml | 7 + .../defs/accuracy/references/gsm8k.yaml | 7 + .../defs/accuracy/test_llm_api_pytorch.py | 77 +++ .../l0_gb300_multi_nodes_node2_gpu8.yml | 3 + 10 files changed, 565 insertions(+), 89 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py index 4ea0222e0c52..d9cc746f9e9d 100644 --- a/tensorrt_llm/_torch/models/modeling_dflash.py +++ b/tensorrt_llm/_torch/models/modeling_dflash.py @@ -873,7 +873,10 @@ def dflash_forward( ) block_tables = ctx_page_table.index_select(0, cache_batch_idx_i32.long()) pages_per_slot = block_tables.size(1) - page_size = ctx_kv_cache.size(-2) + # Index the layer first: ctx_kv_cache is an [L, ...] tensor for the + # private arena but a per-layer list when bound to the draft KV + # cache manager's pool. page_size sits at -2 either way. + page_size = ctx_kv_cache[0].size(-2) kv_indices = block_tables.flatten() kv_indptr = torch.arange( 0, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 084da74711ed..93863f1e5b07 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1435,7 +1435,18 @@ def _should_create_separate_draft_kv_cache(self) -> bool: in the target model and don't produce a separate ModelConfig. We fall back to the target model's config via _get_effective_draft_config(). """ - if self._mapping.enable_attention_dp: + if self._speculative_config is None: + # No drafter at all, so there is nothing to give a manager to. + return False + # Narrower than is_external_drafter(): PARD and DRAFT_TARGET_ONE_MODEL + # never reach the arena this carve-out exists for. + spec_dec_mode = self._speculative_config.spec_dec_mode + is_standalone_drafter = (spec_dec_mode.is_dflash() + or spec_dec_mode.is_dspark()) + if self._mapping.enable_attention_dp and not is_standalone_drafter: + # This bail suits MTP, whose draft layers are target-shaped and + # appendable to the target pool. A standalone drafter has nothing to + # append, so it would be stranded on the private arena instead. logger.info( "Attention DP is enabled, separate draft KV cache is not supported." ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 63cce4af5a3c..afb4525a6306 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -2755,6 +2755,41 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self._prepare_draft_resources(scheduled_batch) return + def _mirror_draft_kv_cache(self, req: LlmRequest): + """The draft manager's entry for ``req``, created on first sight. + + Mirrors what the target manager does in ``_prepare_context_impl``: + create when missing rather than fail, passing ``is_dummy`` through. The + draft mirror never looks up block reuse (``tokens=None``) and never + commits, so this cannot share the target's implementation. + + Reached from both the context and the generation loop, because not + every generation request passed through context on this worker: + attention DP injects its idle placeholder straight as a generation + request when a rank has no real work (ATTENTION_DP_DUMMY_REQUEST_ID, + see py_executor._adp_dummy_is_gen), and a disaggregated generation + worker receives prompt KV rather than prefilling it. + + Returns None when the IndexMapper is saturated. A context request can + retry next iteration; a generation request cannot, and skipping it only + defers the failure to copy_batch_block_offsets(), which asserts in C++ + on the unmapped request ID. + """ + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is not None: + return kv_cache + kv_cache = self._create_kv_cache( + req.py_request_id, + req.lora_task_id, + None, + cache_salt=req.cache_salt, + is_dummy=req.is_dummy, + ) + if kv_cache is None: + return None + kv_cache.stop_committing() + return kv_cache + def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): """Create/resize KV caches in the draft V2 manager for scheduled requests. @@ -2765,23 +2800,16 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): """ with request_context(True, scheduled_batch): for req in scheduled_batch.context_requests: - kv_cache = self.kv_cache_map.get(req.py_request_id) + kv_cache = self._mirror_draft_kv_cache(req) if kv_cache is None: - kv_cache = self._create_kv_cache( - req.py_request_id, - req.lora_task_id, - None, - cache_salt=req.cache_salt, - is_dummy=req.is_dummy, + # Retryable here, unlike the generation loop below: a + # context request has not drafted yet, so the next + # iteration can mirror it once slots free up. + logger.warning( + f"Draft KV cache mirror has no free IndexMapper slot for " + f"context request {req.py_request_id}; retrying next iteration." ) - if kv_cache is None: - # Saturated IndexMapper (e.g. slots held by disagg - # generation transfers in flight): skip mirroring this - # request for now; it is retried next iteration once - # slots free up, before the request runs any spec-dec - # forward that needs the mirror. - continue - kv_cache.stop_committing() + continue if not self._resume_and_restore(req.py_request_id, kv_cache): raise RuntimeError( f"Failed to resume draft KV cache for request {req.py_request_id}" @@ -2800,10 +2828,11 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): ) for req in scheduled_batch.generation_requests: - kv_cache = self.kv_cache_map.get(req.py_request_id) + kv_cache = self._mirror_draft_kv_cache(req) if kv_cache is None: raise RuntimeError( - f"Missing draft KV cache for generation request {req.py_request_id}" + f"Draft KV cache mirror exhausted its IndexMapper on " + f"generation request {req.py_request_id}" ) if not self._resume_and_restore(req.py_request_id, kv_cache): raise RuntimeError( @@ -3650,6 +3679,10 @@ def release_index_slot(self, request_id: int) -> None: the KV cache blocks are still being transferred via NIXL/UCX. """ kv_cache = self.kv_cache_map.get(request_id) + if self.is_draft and (kv_cache is None or request_id in self._early_freed_index_requests): + # The draft mirror only holds a slot for requests it actually + # mirrored, and the target may release the same request twice. + return if kv_cache is not None: for i in range(self.max_beam_width): for pool_idx in range(self.num_pools): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ad23219f7ebf..4f26fb31d62f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -589,6 +589,10 @@ def __init__( # kv cache events self.kv_cache_manager = self.resource_manager.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER) + # A separate draft manager mirrors the target's slot lifecycle, so it + # needs the same early index-slot release on a context-only worker. + self.draft_kv_cache_manager = self.resource_manager.resource_managers.get( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) # V2 owns KV allocation, suspend, resume, and context finalization. # The executor skips the V1 terminate/pause paths and finalizes V2 # context resources before transfer or response handling can terminate @@ -7670,9 +7674,12 @@ def _send_disagg_ctx_kv_async(self, # Forward is done for this request — release the # IndexMapper slot so new requests can reuse it. # KV blocks stay allocated for the upcoming transfer. - if hasattr(self.kv_cache_manager, 'release_index_slot'): - self.kv_cache_manager.release_index_slot( - req.py_request_id) + # getattr: executors built by unit tests skip __init__ + # and so never get the draft binding. + for mgr in (self.kv_cache_manager, + getattr(self, 'draft_kv_cache_manager', None)): + if hasattr(mgr, 'release_index_slot'): + mgr.release_index_slot(req.py_request_id) # Order matters: start_transfer commits the request's blocks to the reuse # tree and pins them, and must run before respond_and_send_async sends the # final KV slice and (for the Python transceiver) transitions the request toward completion. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f58a4b2c0864..65bee8069f98 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -496,10 +496,12 @@ def create_py_executor( if hasattr(spec_config, '_max_batch_size'): spec_config._max_batch_size = max_batch_size - # WAR for https://nvbugs/5807902 - # Disable separate draft KV cache in disaggregated mode - # Enable separate pool for None DI + Non-KVBM and Aggregated + KVBM - if cache_transceiver_config is not None: + # WAR for https://nvbugs/5807902 (Eagle3 disagg RMSNorm crash, closed + # will-not-fix). Keep the blanket disable; carve out only the standalone + # drafters, which it stranded on their private max_seq_len-dense arena. + is_standalone_drafter = (spec_config.spec_dec_mode.is_dflash() + or spec_config.spec_dec_mode.is_dspark()) + if cache_transceiver_config is not None and not is_standalone_drafter: spec_config._allow_separate_draft_kv_cache = False # chunk_unit_size may be changed to 64 when using flash mla diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index b92c1a7c7ac2..0fde44377a0d 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -25,6 +25,7 @@ from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata +from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.resource_manager import BaseResourceManager from .accept_stats import maybe_create_recorder @@ -130,6 +131,24 @@ def prepare(self): worker._ctx_len[slot] = 0 worker._free_slots.append(slot) + # A disagg generation worker receives prompt KV instead of + # prefilling, so _store_prefill_context -- the only other assigner -- + # never runs and every request would land on the dummy slot. + num_contexts = max(0, len(self.request_ids) - self.num_generations) + for rid in self.request_ids[num_contexts:]: + if ( + rid != ATTENTION_DP_DUMMY_REQUEST_ID + and rid < worker._graph_dummy_id_floor + and rid not in worker._req_to_slot + ): + if worker._assign_slot(rid) is None: + # Not fatal -- spec dec is lossless, so the dummy-slot + # fallback only costs acceptance. Warn so it is not silent. + logger.warning( + f"DFlash: no free drafter slot for generation request {rid}; " + "it shares the dummy slot and will draft poorly." + ) + # Route unknown request IDs (cuda graph padding or warmup dummies) # to dummy slot to avoid corrupting real request's context mapping = torch.tensor( @@ -221,8 +240,18 @@ def __init__( self._max_ctx = 0 self._ctx_k_buf = None # [max_batch+1, L, max_ctx+block, nkv, hd] self._ctx_v_buf = None + # [L, pages, K/V, nkv, page, hd] when privately allocated, or the draft + # KV cache manager's per-layer pool views when bound to it. self._ctx_kv_buf = None + # The manager _ctx_kv_buf is borrowed from, so a swap can be detected. + self._ctx_kv_manager = None + # Static per-slot page table for the private arena. When bound to the + # manager it is replaced by _ctx_block_tables, refreshed per iteration. self._ctx_page_table = None + self._ctx_block_tables = None # [max_batch+1, max_blocks_per_seq] int32 + self._ctx_block_indptr = None + self._ctx_pool_idx = 0 + self._ctx_block_divisor = 1 # encoded offset -> raw pool block index self._ctx_kv_indptr = None self._ctx_kv_last_page_len = None self._ctx_page_size = 32 @@ -233,6 +262,7 @@ def __init__( self._req_to_slot = {} # request_id -> slot index self._free_slots = deque() # available slot indices self._dummy_slot = None # for cudagraph padding or warmup dummy requests + self._graph_dummy_id_floor = 1 << 62 # Opt-in acceptance-statistics recorder (None unless # TLLM_DFLASH_ACCEPT_STATS_DIR is set; see accept_stats.py). Only @@ -285,25 +315,224 @@ def set_draft_model(self, draft_model) -> None: super().set_draft_model(draft_model) self._validate_draft_attention_backend(draft_model) - def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): + def _check_ctx_arena_fits(self, capacity, num_slots, L, nkv, hd, dtype): + """Fail with the arithmetic before allocating the drafter context arena. + + The arena is dense in max_seq_len and lands *after* the KV cache manager + has already carved its pool, so free_gpu_memory_fraction does not bound + it. Sized off an unpinned max_seq_len it reaches hundreds of GiB, and the + bare torch.zeros that follows reports only the failed allocation, naming + neither the buffer nor the knob that controls it. + """ + itemsize = torch.tensor([], dtype=dtype).element_size() + if self._dflash_attention_backend == "TRTLLM": + page_size = self._ctx_page_size + pages_per_slot = (capacity + page_size - 1) // page_size + arena = L * num_slots * pages_per_slot * 2 * nkv * page_size * hd * itemsize + else: + arena = 2 * num_slots * L * capacity * nkv * hd * itemsize + + # mem_get_info() excludes the caching allocator's reserved-but-unused + # segments, which torch.zeros() below can serve from. + free = torch.cuda.mem_get_info()[0] + ( + torch.cuda.memory_reserved() - torch.cuda.memory_allocated() + ) + gib = 1 << 30 + logger.info( + f"DFlash: drafter ctx arena needs {arena / gib:.2f} GiB " + f"(slots={num_slots}, layers={L}, capacity={capacity}, " + f"kv_heads={nkv}, head_dim={hd}), {free / gib:.2f} GiB free" + ) + if arena >= free: + # Warn, not raise: this estimate has under-reported real + # headroom, and torch.zeros below is the authority. + logger.warning( + f"DFlash drafter context arena needs {arena / gib:.2f} GiB but only " + f"{free / gib:.2f} GiB is free. It is dense in sequence length: " + f"slots={num_slots} x layers={L} x capacity={capacity} x " + f"kv_heads={nkv} x head_dim={hd} x {itemsize}B, K and V. " + f"capacity comes from max_seq_len (+{self._resolved_block_size} " + f"block slack), so pass an explicit max_seq_len matching the " + f"sequences you actually serve instead of letting it fall back to " + f"the drafter's max_position_embeddings; lowering max_batch_size " + f"scales it down linearly too." + ) + + def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype): + """Per-layer views of the draft KV cache manager's pool, or None. + + The manager already funds a pool sized from the drafter's own config + (K3 at TP8: 5 layers x 2 x 2 kv heads x 64 x 2B = 2560 B/token), and + ``get_buffers(..., "HND")`` hands back exactly the per-layer shape the + private arena uses, so nothing but the ownership of the memory changes. + + A list, not one stacked tensor: only the two consumers that index it by + layer need it, and the two managers disagree on what a stack would mean. + V1 slices one page-major allocation, so the layers are strided views of + a shared storage; V2 wraps each layer's pool base address in its own + tensor. Everything downstream indexes ``[layer_idx]`` first anyway. + + Note the page count is not the drafter's to interpret: a V2 view spans + the whole interleaved pool, not this layer's slice, so only shape[1:] is + checked here and the index space is settled in _init_ctx_block_tables. + """ + if draft_kv_cache_manager is None: + # No separate draft KV cache: attention DP disables it + # (_util.py:_should_create_separate_draft_kv_cache), and the + # two-model paths never build one. + return None + layers = [draft_kv_cache_manager.get_buffers(i, kv_layout="HND") for i in range(L)] + base = layers[0] + expected = (2, nkv, base.size(-2), hd) + for i, layer in enumerate(layers): + if tuple(layer.shape[1:]) != expected or layer.dtype != dtype: + # Fall back rather than fail: the private arena is what every + # DFlash model shipped with, so an unrecognized pool costs the + # old memory footprint, not the run. + logger.warning( + f"DFlash: draft KV pool layer {i} is {tuple(layer.shape)} / {layer.dtype}, " + f"not {('pages',) + expected} / {dtype}; using the private ctx arena." + ) + return None + logger.info( + f"DFlash: binding the drafter ctx cache to the draft KV cache manager pool: " + f"{L} x {tuple(base.shape)} {base.dtype}, page_size={base.size(-2)}" + ) + return layers + + def _init_ctx_block_tables( + self, draft_kv_cache_manager, pool, num_slots, L, nkv, hd, page_size + ) -> bool: + """Size the per-iteration block table and the offset decode constant. + + ``kv_cache_block_offsets`` entries are ``pool_block_index * + num_pool_layers * kv_factor`` plus the K/V field index. What that has to + be divided by to index a ``get_buffers`` view depends on the manager -- + see the stride check below; attention_backend/utils.py hardcodes the V1 + divisor and is not a general contract. + """ + mgr = draft_kv_cache_manager + pool_mapping = mgr.kv_cache_pool_mapping + layer_offset = mgr.layer_offsets[0] + self._ctx_pool_idx = int(pool_mapping[layer_offset, 0]) + num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum()) + if num_pool_layers != L: + logger.warning( + f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but " + f"the drafter has {L}; using the private ctx arena." + ) + return False + # V1 slices a page-major pool, so stride(0) spans every layer; V2 wraps + # the layer base in a dense view over interleaved pages, so stride(0) is + # one layer slot. Guessing wrong is silent: acceptance quietly drops. + layer_block_elems = mgr.kv_factor * nkv * page_size * hd + stride0 = pool[0].stride(0) + if stride0 == layer_block_elems: + self._ctx_block_divisor = mgr.kv_factor + elif stride0 == num_pool_layers * layer_block_elems: + self._ctx_block_divisor = num_pool_layers * mgr.kv_factor + else: + logger.warning( + f"DFlash: draft KV pool layer stride {stride0} matches neither V1's " + f"page-major slice ({num_pool_layers * layer_block_elems}) nor V2's " + f"interleaved view ({layer_block_elems}); using the private ctx arena " + f"rather than guessing block indices." + ) + return False + max_blocks = mgr.max_blocks_per_seq + self._ctx_block_tables = torch.zeros( + (num_slots, max_blocks), dtype=torch.int32, device="cuda" + ) + self._ctx_block_indptr = torch.arange( + 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" + ) + logger.info( + f"DFlash: ctx block tables {tuple(self._ctx_block_tables.shape)}, " + f"pool_idx={self._ctx_pool_idx}, divisor={self._ctx_block_divisor}, " + f"view_pages={pool[0].size(0)}" + ) + return True + + def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: + """Decode this iteration's draft block table into the persistent buffer. + + In place and fixed-shape so a captured graph keeps reading the same + tensor. Returns False when there is no table to fill, i.e. every path + that stayed on the private arena. + """ + if self._ctx_block_tables is None or num_seqs <= 0: + return False + src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None) + if src is None: + # The table is bound but unfillable. Leaving it at zeros would send + # every request's every layer to pool block 0 -- silently. + raise RuntimeError( + "DFlash: bound to the draft KV pool but " + f"{type(attn_metadata).__name__} carries no " + "draft_kv_cache_block_offsets this iteration." + ) + # Field 0 of the K/V pair: both name the same block for a kv_factor=2 + # pool, and the drafter addresses blocks, not planes. + # clone(): .to() is a no-op for an int64 source, and the in-place ops + # below would then edit attn_metadata's own tensor. + encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone() + decoded = encoded.clamp_(min=0).div_(self._ctx_block_divisor, rounding_mode="floor") + self._ctx_block_tables[:num_seqs].copy_(decoded.to(torch.int32)) + return True + + def _lazy_init_ctx_buffers( + self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None + ): self._validate_draft_attention_backend(draft_model) - if self._ctx_buf_inited: - return + # Only TrtllmAttentionMetadata builds draft_kv_cache_block_offsets. + # Without it the table stays all-zero and every request would read and + # write pool block 0 -- silently, at acceptance 1.0. + if draft_kv_cache_manager is not None and not hasattr( + attn_metadata, "draft_kv_cache_block_offsets" + ): + logger.warning( + f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; " + "using the private ctx arena." + ) + draft_kv_cache_manager = None + if self._ctx_buf_inited: + if self._ctx_kv_manager is draft_kv_cache_manager: + return + # KV cache estimation swaps a probe manager for the real one. + # Holding the old pool's tensors across that writes freed memory. + logger.info("DFlash: draft KV cache manager replaced, rebinding the ctx cache.") + self._ctx_buf_inited = False + + # Recorded even when the buffers stay private, so that the estimation + # pass falling back to the arena still rebinds once the real pool lands. + self._ctx_kv_manager = draft_kv_cache_manager + self._ctx_block_tables = None max_batch = spec_metadata.max_num_requests - # Prefer runtime max_seq_len over max_position_embeddings: YaRN - # models advertise 100k+ positions, which would OOM the ctx buffer - # or, if silently capped, corrupt the slot cache on long prompts. - config_max_ctx = getattr(self.spec_config, "max_ctx_len", None) - if config_max_ctx is not None: - self._max_ctx = config_max_ctx - else: - config = getattr(draft_model, "config", None) - max_pos = getattr(config, "max_position_embeddings", None) if config else None - runtime_max = getattr(attn_metadata, "max_seq_len", None) - candidates = [c for c in (runtime_max, max_pos) if c is not None] - self._max_ctx = min(candidates) if candidates else 8192 + # ctx_len is 1:1 with the target's positions, so max_seq_len bounds it. + # max_position_embeddings does not: K3's drafter advertises 1048576, + # which sizes the arena at hundreds of GiB. + config = getattr(draft_model, "config", None) + max_pos = getattr(config, "max_position_embeddings", None) if config else None + runtime_max = getattr(attn_metadata, "max_seq_len", None) + candidates = [c for c in (runtime_max, max_pos) if c is not None] + self._max_ctx = min(candidates) if candidates else 8192 + # py_executor_creator inflates max_seq_len by at most this much for spec + # dec (one term is overlap-only), so compare against what is served. + from .utils import get_num_extra_kv_tokens + + spec_slack = 2 * (self.spec_config.tokens_per_gen_step - 1) + get_num_extra_kv_tokens( + self.spec_config + ) + if runtime_max is not None and self._max_ctx < runtime_max - spec_slack: + # Warn, not raise: the per-request skip in _store_prefill_context + # keeps an oversized request from taking the executor down. + logger.warning( + f"DFlash drafter covers {self._max_ctx} positions but the engine serves " + f"up to {runtime_max - spec_slack}. max_position_embeddings={max_pos} is " + "the binding constraint, so requests past it will draft nothing." + ) dtype = draft_model.fc.weight.dtype if hasattr(draft_model, "fc") else torch.bfloat16 @@ -312,6 +541,12 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): # writes land here and can't corrupt a real request's context. self._dummy_slot = max_batch num_slots = max_batch + 1 + # CUDA-graph padding uses ids counting down from + # CUDA_GRAPH_DUMMY_REQUEST_ID; a floor of one draft block separates them + # from real ids. Imported lazily to avoid an import cycle. + from ..pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + + self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") @@ -345,7 +580,11 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): hd = draft_model._head_dim capacity = self._max_ctx + self._compute_block_size if self._dflash_attention_backend == "TRTLLM": - page_size = self._ctx_page_size + pool = self._managed_ctx_pool(draft_kv_cache_manager, L, nkv, hd, dtype) + # The manager's page size wins when bound to it: its pool is already + # carved, so the drafter adopts the geometry rather than imposing one. + page_size = self._ctx_page_size if pool is None else pool[0].size(-2) + self._ctx_page_size = page_size has_context_attention = any( not draft_model._get_attention_mask_args(layer_idx)[0] for layer_idx in range(L) ) @@ -357,29 +596,46 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): tokens_per_block=page_size, has_context_attention=has_context_attention, ) - self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size - total_pages = num_slots * self._ctx_pages_per_slot - # HND layout consumed by both FlashInfer's paged append and the - # TRTLLM-Gen launcher: [L, pages, K/V, Hkv, page, D]. - self._ctx_kv_buf = torch.zeros( - (L, total_pages, 2, nkv, page_size, hd), - dtype=dtype, - device="cuda", - ) - self._ctx_page_table = torch.arange(total_pages, dtype=torch.int32, device="cuda").view( - num_slots, self._ctx_pages_per_slot - ) - self._ctx_kv_indptr = torch.arange( - 0, - total_pages + 1, - self._ctx_pages_per_slot, - dtype=torch.int32, - device="cuda", - ) + # Settle the block table before committing to the pool: it is the + # last thing that can rule the pool out, and falling back after + # taking the pool branch would leave no buffer allocated at all. + if pool is not None and not self._init_ctx_block_tables( + draft_kv_cache_manager, pool, num_slots, L, nkv, hd, page_size + ): + pool = None + if pool is None: + # Per-slot reservation is an arena-only notion; the pool path + # takes its pages from the manager one iteration at a time. + self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size + total_pages = num_slots * self._ctx_pages_per_slot + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype) + # HND layout consumed by both FlashInfer's paged append and the + # TRTLLM-Gen launcher: [L, pages, K/V, Hkv, page, D]. + self._ctx_kv_buf = torch.zeros( + (L, total_pages, 2, nkv, page_size, hd), + dtype=dtype, + device="cuda", + ) + self._ctx_page_table = torch.arange( + total_pages, dtype=torch.int32, device="cuda" + ).view(num_slots, self._ctx_pages_per_slot) + self._ctx_kv_indptr = torch.arange( + 0, + total_pages + 1, + self._ctx_pages_per_slot, + dtype=torch.int32, + device="cuda", + ) + else: + # Nothing reserved: pages come from the manager's per-request + # block table one iteration at a time, so the footprint follows + # the sequences served rather than max_batch x max_seq_len. + self._ctx_kv_buf = pool self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) else: # VANILLA DFlash backend (FlashAttention) + self._check_ctx_arena_fits(capacity, num_slots, L, nkv, hd, dtype) kv_shape = (num_slots, L, capacity, nkv, hd) self._ctx_k_buf = torch.zeros(kv_shape, dtype=dtype, device="cuda") self._ctx_v_buf = torch.zeros(kv_shape, dtype=dtype, device="cuda") @@ -391,11 +647,46 @@ def _lazy_init_ctx_buffers(self, draft_model, spec_metadata, attn_metadata): f"dflash_attention_backend={self._dflash_attention_backend}" ) + def _ctx_paged_index_args(self): + """(kv_indices, kv_indptr) for the append op, and what ``rows`` means. + + Bound to the manager the page table is per-request and refreshed every + iteration, so rows are batch positions; on the private arena it is the + static per-slot table and rows are slots. Both are dense and uniformly + strided, so kv_indptr is the same arange either way -- only the width + differs. + """ + if self._ctx_block_tables is not None: + table = self._ctx_block_tables + indptr = self._ctx_block_indptr + else: + table = self._ctx_page_table + indptr = self._ctx_kv_indptr + return table.flatten().contiguous(), indptr + + def _assign_slot(self, req_id: int, reset: bool = False): + """Get (or refresh) this request's context slot; None when none is free. + + ``reset`` recycles an existing slot, which is how a reused request id is + told apart from a continuation: a fresh request starts at position 0. + """ + if reset and req_id in self._req_to_slot: + old_slot = self._req_to_slot.pop(req_id) + self._ctx_len[old_slot] = 0 + self._free_slots.append(old_slot) + if req_id not in self._req_to_slot: + if not self._free_slots: + return None + slot = self._free_slots.popleft() + self._req_to_slot[req_id] = slot + self._ctx_len[slot] = 0 + return self._req_to_slot[req_id] + def _store_context_kv_paged( self, k: torch.Tensor, v: torch.Tensor, - slots: torch.Tensor, + rows: torch.Tensor, positions: torch.Tensor, ) -> None: trtllm_gen_ops = get_dflash_trtllm_gen_ops() @@ -403,18 +694,18 @@ def _store_context_kv_paged( # Convert at most once and reuse for every layer. Calling ``.to()`` in # the op call would allocate converted tensors once per layer when the # inputs are not already int32. - slots_i32 = slots.to(torch.int32) + rows_i32 = rows.to(torch.int32) positions_i32 = positions.to(torch.int32) - kv_indices = self._ctx_page_table.flatten().contiguous() + kv_indices, kv_indptr = self._ctx_paged_index_args() for layer_idx in range(k.size(1)): trtllm_gen_ops.append_paged_kv_cache( append_key=k[:, layer_idx].contiguous(), append_value=v[:, layer_idx].contiguous(), - batch_indices=slots_i32, + batch_indices=rows_i32, positions=positions_i32, paged_kv_cache=self._ctx_kv_buf[layer_idx], kv_indices=kv_indices, - kv_indptr=self._ctx_kv_indptr, + kv_indptr=kv_indptr, kv_last_page_len=self._ctx_kv_last_page_len, kv_layout="HND", ) @@ -509,26 +800,32 @@ def _store_prefill_context( first_pos = chunk_pos[0].item() if slen > 0 else 0 # Assign slot for new requests or reset for reused IDs - if req_id not in self._req_to_slot or first_pos == 0: - if req_id in self._req_to_slot: - old_slot = self._req_to_slot[req_id] - self._ctx_len[old_slot] = 0 - self._free_slots.append(old_slot) - if not self._free_slots: - logger.warning("DFlash: no free slots, skipping context store") - offset += slen - continue - slot = self._free_slots.popleft() - self._req_to_slot[req_id] = slot - self._ctx_len[slot] = 0 + if self._assign_slot(req_id, reset=first_pos == 0) is None: + logger.warning("DFlash: no free slots, skipping context store") + offset += slen + continue slot = self._req_to_slot[req_id] cur = int(self._ctx_len[slot].item()) - end = min(cur + slen, self._max_ctx) + if cur + slen > self._max_ctx: + # Request-level, like the no-free-slots path above: truncating + # would silently draft from a stale prefix, but killing the + # forward would take every other in-flight request with it. + logger.warning( + f"DFlash: ctx overflow on slot {slot} " + f"({cur} + {slen} > {self._max_ctx}); skipping its context " + "store, so this request drafts nothing." + ) + self._req_to_slot.pop(req_id, None) + self._ctx_len[slot] = 0 + self._free_slots.append(slot) + offset += slen + continue + end = cur + slen actual = end - cur if actual > 0: cache_dtype = ( - self._ctx_kv_buf.dtype + self._ctx_kv_buf[0].dtype if self._dflash_attention_backend == "TRTLLM" else self._ctx_k_buf.dtype ) @@ -541,10 +838,13 @@ def _store_prefill_context( ) # chunk_k/v: [actual, L, nkv, hd] → [L, actual, nkv, hd] if self._dflash_attention_backend == "TRTLLM": + # Manager block tables are keyed by batch position, the + # private arena's by slot. See _ctx_paged_index_args. + row = i if self._ctx_block_tables is not None else slot self._store_context_kv_paged( chunk_k, chunk_v, - torch.full((actual,), slot, dtype=torch.long, device="cuda"), + torch.full((actual,), row, dtype=torch.long, device="cuda"), torch.arange(cur, end, dtype=torch.long, device="cuda"), ) else: # VANILLA DFlash backend (FlashAttention) @@ -603,8 +903,13 @@ def _forward_impl( ) # Lazy init buffers and attach worker reference for prepare() - self._lazy_init_ctx_buffers(draft_model, spec_metadata, attn_metadata) + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + self._lazy_init_ctx_buffers( + draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager + ) spec_metadata._dflash_worker = self + # Before any store: prefill and decode both address pages through it. + self._refresh_ctx_block_tables(attn_metadata, batch_size) # Save context lengths so both warmup and a failed forward can roll # back the in-place _ctx_len updates made during drafting. @@ -692,8 +997,6 @@ def _forward_impl( total_target_tokens=total_target_tokens, ) - draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - if num_gens > 0: with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): hidden_states_out = draft_model.dflash_forward( @@ -907,6 +1210,9 @@ def prepare_1st_drafter_inputs( # Get slots for gen requests from pre-computed mapping slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] + gen_rows_out = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) K_plus_1 = K + 1 @@ -944,7 +1250,14 @@ def prepare_1st_drafter_inputs( gen_num_accepted_long = gen_num_accepted.long() col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0) write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1) - col_idx = col_idx.clamp(max=self._max_ctx - 1) + if self._ctx_block_tables is not None: + # Bound by the table width. A per-request bound is not + # recoverable here: V2 writes 0 for BAD_PAGE_INDEX and V1 + # leaves the tail stale, so a placeholder reads as block 0. + ctx_capacity = self._ctx_block_tables.size(1) * self._ctx_page_size + col_idx = col_idx.clamp(max=ctx_capacity - 1) + else: + col_idx = col_idx.clamp(max=self._max_ctx - 1) # Fixed-size writes for CUDA graph compatibility: # Write ALL entries but zero out invalid ones. Invalid @@ -959,7 +1272,7 @@ 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. cache_dtype = ( - self._ctx_kv_buf.dtype + self._ctx_kv_buf[0].dtype if self._dflash_attention_backend == "TRTLLM" else self._ctx_k_buf.dtype ) @@ -972,7 +1285,13 @@ def prepare_1st_drafter_inputs( slot_long = slot_flat.long() col_long = col_flat.long() if self._dflash_attention_backend == "TRTLLM": - self._store_context_kv_paged(k_new, v_new, slot_long, col_long) + if self._ctx_block_tables is not None: + # Batch positions of the gen requests, matching the + # per-request block table's row order. + rows_long = gen_rows_out.unsqueeze(1).expand(-1, K + 1).reshape(-1) + else: + rows_long = slot_long + self._store_context_kv_paged(k_new, v_new, rows_long, col_long) else: # VANILLA DFlash backend (FlashAttention) self._ctx_k_buf[slot_long, :, col_long] = k_new self._ctx_v_buf[slot_long, :, col_long] = v_new @@ -995,6 +1314,7 @@ def prepare_1st_drafter_inputs( query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") slots = torch.empty(0, dtype=torch.long, device="cuda") + gen_rows_out = slots bonus = torch.empty(0, dtype=torch.long, device="cuda") return { @@ -1003,10 +1323,16 @@ def prepare_1st_drafter_inputs( "num_ctx_per_req": num_ctx_per_req_t, "ctx_k_cache": self._ctx_k_buf, "ctx_v_cache": self._ctx_v_buf, - "ctx_cache_batch_idx": slots, + # Slots index the private arena; the manager's block table is keyed + # by batch position, so the drafter reads its own rows there. + "ctx_cache_batch_idx": gen_rows_out if self._ctx_block_tables is not None else slots, # Anchor token per gen request (block slot 0): last accepted # token. The dspark Markov chain conditions its first step on it. "first_prev_tokens": bonus, "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": self._ctx_page_table, + "ctx_page_table": ( + self._ctx_block_tables + if self._ctx_block_tables is not None + else self._ctx_page_table + ), } diff --git a/tests/integration/defs/accuracy/references/acceptance_length.yaml b/tests/integration/defs/accuracy/references/acceptance_length.yaml index fe695e2210f6..b882dd834daa 100644 --- a/tests/integration/defs/accuracy/references/acceptance_length.yaml +++ b/tests/integration/defs/accuracy/references/acceptance_length.yaml @@ -42,3 +42,10 @@ TestKimiK3::test_w4a16_mxfp4: TestQwen3_8B::test_dspark: ref_al: 6.259 min_al: 5.946 +# Measured on 8x GB300 with this harness's stock prompt. Higher K3 DSpark +# figures elsewhere are --apply_chat_template runs, worth ~40% acceptance. +# min_al is a collapse tripwire: a drafter whose slots collapse scores 1.09 +# with every accuracy gate still green. +TestKimiK3DSpark::test_gsm8k_tep8: + ref_al: 4.323 + min_al: 4.0 diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 1f21cc2a7cfe..f05cc59eacc3 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -300,6 +300,13 @@ moonshotai/Kimi-K3: # SA spec dec is lossless; scores match the baseline within noise. - spec_dec_algo: SA accuracy: 96.5 + # The NVFP4 requant (Kimi-K3-NVFP4), which ships hf_quant_config.json and so + # reads as MIXED_PRECISION. Same 96.5 as above on purpose: the MEGAMOE_CUTEDSL + # MoE kernel measures ~95.2 where CUTLASS measures 96.209 on these weights, + # and that gap belongs in a kernel fix, not in this threshold. + - quant_algo: MIXED_PRECISION + spec_dec_algo: DSpark + accuracy: 96.5 nvidia/Nemotron-MOE: - accuracy: 88.249 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 126df48a6aa1..cd59a2d3008f 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4741,6 +4741,83 @@ def test_w4a16_mxfp4(self, mode: str, acceptance_length) +# 3600s to match the test-db TIMEOUT (60) budget; measured 25m30s. +@pytest.mark.timeout(3600) +# TP8 leaves ~182 GiB of weights per GPU before any KV cache, past what a +# 184 GiB B200/GB200 part holds. Scheduled on 2x4 GB300 +# (l0_gb300_multi_nodes_node2_gpu8). +@pytest.mark.skip_less_device_memory(250000) +@skip_pre_blackwell +class TestKimiK3DSpark(LlmapiAccuracyTestHarness): + """Kimi K3 driven by a standalone DSpark drafter. + + Split from TestKimiK3 because both the drafter and the target precision + are checkpoint axes here: the RadixArk artifact below is a qwen3-shaped + DSparkDraftModel while Inferact ships a K3DSparkModel, and the target is + the NVFP4 requant rather than TestKimiK3's MXFP4 weights. + """ + + MODEL_NAME = "moonshotai/Kimi-K3" + # The NVFP4 requant, not the MXFP4 weights TestKimiK3 runs: it is the + # Blackwell deployment precision and nothing else covers it end to end. + MODEL_PATH = f"{llm_models_root()}/Kimi-K3-NVFP4" + DSPARK_MODEL_PATH = f"{llm_models_root()}/RadixArk-Kimi-K3-DSpark" + + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k_tep8(self): + """GSM8K plus an acceptance-length gate on the TEP8 DSpark recipe. + + TEP, not TestKimiK3's DEP16 recipe: with attention DP off the drafter gets + its own KV cache pool, which is the path this covers. Accuracy alone + cannot detect a broken drafter -- speculative decoding is lossless, so + a drafter producing garbage scores the same and only runs slower -- + hence the acceptance-length assertion. + """ + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.20, + tokens_per_block=64) + spec_config = DSparkDecodingConfig( + max_draft_len=7, + speculative_model=self.DSPARK_MODEL_PATH, + attention_backend="TRTLLM") + with LLM( + self.MODEL_PATH, + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=False, + max_batch_size=8, + max_num_tokens=4096, + max_seq_len=8192, + trust_remote_code=True, + # The drafter runs between target steps, and its hidden-state + # capture needs a whole-sequence prefill. + disable_overlap_scheduler=True, + enable_chunked_prefill=False, + cuda_graph_config=CudaGraphConfig(max_batch_size=8), + # Not the TRTLLM backend TestKimiK3 uses: trtllm-gen ships + # SiTu cubins only for W4A8_MXFP4_MXFP8, so K3's NVFP4 routed + # experts leave MEGAMOE_CUTEDSL and CUTLASS as the choices. + moe_config=MoeConfig(backend="MEGAMOE_CUTEDSL", + max_num_tokens=33024, + use_low_precision_moe_combine=True), + kv_cache_config=kv_cache_config, + max_stats_len=-1, + enable_iter_perf_stats=True, + speculative_config=spec_config, + ) as llm: + # Unlike TestKimiK3's MXFP4 weights, the requant ships + # hf_quant_config.json, so the matcher sees MIXED_PRECISION -- + # only the experts are 4-bit. + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + acceptance_length = _compute_acceptance_length(llm) + print(f"[AL] TestKimiK3DSpark::test_gsm8k_tep8 " + f"acceptance_length = {acceptance_length:.3f}") + assert_acceptance_length("TestKimiK3DSpark::test_gsm8k_tep8", + acceptance_length) + + class TestQwen3_4B(LlmapiAccuracyTestHarness): MODEL_NAME = "Qwen3/Qwen3-4B" diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml index d1f8be97bc7c..3b2e4c14d5fe 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml @@ -14,3 +14,6 @@ l0_gb300_multi_nodes_node2_gpu8: backend: pytorch tests: - accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_tp8_mtp3_trtllm TIMEOUT (60) + # Kimi K3 GSM8K accuracy plus an acceptance-length gate on the TEP8 DSpark + # recipe. Measured 25m45s on 2x4 GB300. + - accuracy/test_llm_api_pytorch.py::TestKimiK3DSpark::test_gsm8k_tep8 TIMEOUT (60)