Skip to content

[None][feat] Muse Spark Open modeling_v2 TEP target: GB300 and Vera Rubin fast paths, and the DFlash 2 drafter transform on the sharded draft head - #21143

Open
xwang233 wants to merge 15 commits into
NVIDIA:feat/musespark-open-agentxfrom
xwang233:mso-hcv3midc/agentx-tep
Open

xwang233 wants to merge 15 commits into
NVIDIA:feat/musespark-open-agentxfrom
xwang233:mso-hcv3midc/agentx-tep

Conversation

@xwang233

Copy link
Copy Markdown
Collaborator

Description

This PR extends the Muse Spark Open modeling_v2 TEP target on GB300 (sm103) and Vera Rubin (sm107) with a further set of decode, prefill and drafter fast paths, and makes the vocab-sharded draft head serve the DFlash 2 drafter's declared logit transform. It follows #20770 (GB300) and #20844 (Vera Rubin). The changes are 15 commits:

  • 10 inside the Staircase target (_experimental/modeling_v2/models/musespark_open/...__tep);
  • 1 opt-in extension of the Rubin CuTe DSL MoE GEMMs;
  • 2 opt-in Muse Spark Open model hooks;
  • 2 opt-in changes to the shared KV cache manager v2 and its scheduler.

The modeling_v2 target. These paths are always on inside the target and selected by GPU architecture and tensor-parallel size; each path owns its row band.

  • Attention: mixed steps with 9..64 generation requests run their generation rows on the decode paths (mixed_step_gen_rows on sm107, the decode KV-split rule on sm103); a packed per-row-count launch geometry for the multi-request decode prep on sm103; the split-KV page ring depth per SM version.
  • Experts: fixed TRTLLM-gen decode expert tactics per (expert window, rows), bitwise as the autotuned ones; on sm107 at TP4, prefill expert calls of 4096..8192 rows on the Rubin CuTe DSL FC1/FC2 (rounding differs from TRTLLM-gen; expert selection is unchanged); the half expert window at sm103 TEP4, as at sm107, behind its capacity gate (max_batch_size ≤ 8).
  • Sequence-parallel context steps: routing on each rank's own rows with the top-8 ids and weights gathered, and a copy-engine gather of the attention input behind the QKV projection on sm103 (bitwise as the built-in path).
  • Model boundary: the replicated embedding gather, embed norm, residual copy and layer 0's pre-QKV norm in one kernel, and a BF16 last-layer join (bitwise).
  • Verify head: on sm103 TEP4, steps of two or more K = 7 verify blocks sample on each rank's vocabulary shard and agree on the token over the multicast exchange instead of all-gathering the logits.
  • Drafter: the fused drafter all-reduce hops on blocks of several requests; drafter split-K GEMM tables for 2..4-request blocks; per-GPU PDL chain widths.

DFlash 2 drafter logit transform on the sharded draft head (25525dc68c, fc6d7acc84). The drafter declares an output multiplier and softcap, which the DFlash worker applies in fp32. The vocab-sharded draft head took only bf16 shard logits, so with the transform on every draft step fell back to the built-in selector (full-vocab fp32 all-gather and torch.topk). The transform never decreases, so the draft head now keeps the bf16 shard and applies it after its top-k, from tables built at install by the drafter's own method and checked exhaustively over the 65536 bf16 values, ties included. The draft tokens equal the built-in selector's over the transformed logits; target verification and the accept rule are unchanged. Any step that has to take the built-in path warns once per reason, and the install's refusals are agreed over the TP group.

Shared paths: opt-in, off by default. With the variables or arguments unset, behaviour is unchanged.

  • KVCacheManagerV2: TRTLLM_KV_DEFER_CLOSE=1 defers closing a finished request's cache until after the next forward step is launched (target-model manager only).
  • KVCacheV2Scheduler: TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES=N backs off first-chunk context admission for N generation passes after the KV pool refuses one; TRTLLM_SCHED_ADMISSION_DIAG=1 logs the admission host time.
  • MuseSparkOpenOmniInputProcessor supports enable_tokenization_cache for media-free prompts.
  • MuseSparkOpenDecoderLayer._sp_attn_input_gather: a hook a model-specific path can set to replace the sequence-parallel attention-input all-gather.
  • Rubin CuTe DSL MoE GEMMs: TRTLLM-gen gated FC1 layout, permuted FC2 store, TRTLLM-gen N shuffle and the 64-deep MMA, all constructor options defaulting to the previous kernels.

Measured gains. AgentX traces on the 1M-token corpus, full protocol, forced draft acceptance (3.05; GB300 TEP4 c1 3.09). Each row compares this PR's tree with the previous target code with the same drafter transform, run in the same wave. x = p10 per-user output throughput; y = total throughput per GPU.

GPU Point Δx Δy
GB300 TEP8 c1 +3.6 % +0.7 %
GB300 TEP4 c1 +3.4 % +1.5 %
GB300 TEP4 c10, max_batch_size 3 +2.2 % +0.0 %
GB300 TEP4 c48 +0.5 % +0.6 %
GB300 TEP4 c72, max_num_tokens 16k (two runs) +3.0 / +4.5 % +1.2 / +1.7 %
Vera Rubin TEP8 c1 +5.7 % +1.0 %
Vera Rubin TEP4 c10, max_batch_size 3 +3.0 % −0.4 %
Vera Rubin TEP4 c32, max_batch_size 8 +0.5 % +0.5 %
Vera Rubin TEP4 c64 +2.9 % +1.9 %
Vera Rubin TEP4 c72 +2.0 % +1.5 %

On Vera Rubin each value is the mean of two runs of this PR. Every run of the sharded draft head logged no built-in-path warning.

Accuracy. Spec decoding on with natural acceptance; GSM8K (1319 questions behind a ≥ 22k-character preamble) and GPQA-Diamond, greedy; pooled paired McNemar against the previous target code. Every probe passes (p ≥ 0.05): accuracy deltas from −1.0 to +1.0 pp at Vera Rubin c72 / c32 and GB300 c72 / c10. Natural acceptance length rises with the correct drafter transform, e.g. 4.09 → 4.51 on GSM8K at Vera Rubin c72.

Validation. The tree was validated as built (wheel for sm103 and sm107) on feat/musespark-open-agentx at af9c7ccd94. It merges and rebases onto the current branch tip without conflicts; the only files both sides change are kv_cache_manager_v2.py and scheduler_v2.py, in disjoint regions. The unified DFlash KV cache (#20806) is not reached under modeling_v2, where DFlash AUTO still resolves to VANILLA. The per-launch decode KV split (#20934) defers to this target's decode stand-ins on the calls they intercept.

Test Coverage

  • tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_deferred_close.py: deferred close.
  • tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py: context page backoff and admission diagnostics.
  • tests/unittest/_torch/modeling/test_musespark_open_tokenization_cache.py: the omni processor's prefix-tokenization cache.
  • tests/unittest/_torch/modeling_v2/models/musespark_open/.../kernels/draft_head/ and .../ops/spec/: the transform tables on every bf16 value, the draft top-k ops on their cells, end-to-end draft tokens against torch.topk and the built-in selector (1 to 24 requests, TP4 and TP8, tie-heavy and saturating logits), the fast path under the fold tail and CUDA-graph replay, every fallback and its warning, install atomicity, and a source pin on DFlashWorker._forward_impl.
  • accuracy/modeling_v2/test_musespark_open.py (on this branch): GSM8K and mean acceptance length under modeling_v2="require".

PR Checklist

  • Please check this after reviewing the above items as appropriate for this PR.

🤖 Generated with Claude Code

…ests' caches

With TRTLLM_KV_DEFER_CLOSE=1 the overlap executor arms deferred closing on
a target-model KVCacheManagerV2. free_resources then releases a finished
request's index slot at once and queues its cache; the close, which
commits the last block and releases the pages, runs right after the next
forward step is launched, before the executor waits on an empty queue, and
at shutdown. Dummy requests always close at once.

Off by default (TRTLLM_KV_DEFER_CLOSE unset or 0): the unset path is
unchanged; the index-slot release is factored into _detach_index_slot.

Tests: test_kv_cache_v2_deferred_close.py.
Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ssion diagnostics

TRTLLM_SCHED_CTX_PAGE_BACKOFF_PASSES=N: once a pass that schedules
generation sees a first context chunk refused for KV pages, the rest of
that pass and the next N generation passes try no more first chunks; one
pass in N + 1 probes again. Chunk continuations and passes without
generation are never held. The state depends only on the scheduling
decisions, so ranks stay in lockstep.

TRTLLM_SCHED_ADMISSION_DIAG=1 (window TRTLLM_SCHED_ADMISSION_DIAG_PASSES)
logs the host time of context admission per phase.

Both are off by default (backoff 0, diagnostics unset): no candidate is
held back and no scheduler method is wrapped.

Tests: test_kv_cache_v2_scheduler.py.
Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…put processor

MuseSparkOpenOmniInputProcessor declares supports_tokenization_cache and
accepts enable_tokenization_cache. A media-free prompt is encoded through
the shared prefix-tokenization cache, whose spliced ids equal what the
processor's own encode would produce. Prompts with media, and processors
built without a tokenizer, keep the existing path.

Off by default: enabled only by the enable_tokenization_cache LLM argument.

Tests: test_musespark_open_tokenization_cache.py.
Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…parallel layers

MuseSparkOpenDecoderLayer gains _sp_attn_input_gather. When a model-specific
path sets it on a layer (with a QKV projection that consumes its result),
the sequence-parallel context step gathers the attention input through it
instead of the NCCL all-gather.

Off by default: the class attribute is None, and the layer is unchanged
unless a path sets it.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…muted store and 64-deep MMA

- The Rubin gather + grouped GEMM + SwiGLU kernel (FC1) and its custom op
  take trtllm_gen_gated_layout: the weights are TRTLLM-gen's gated FC1
  tensors, alpha holds the gate scales followed by the output scales, and
  the epilogue follows TRTLLM-gen's SwiGLU and NVFP4 quantization
  arithmetic.
- The Rubin grouped GEMM finalize kernel (FC2) and its custom op take
  store_permuted (one output row per permuted row, alpha * acc, no token
  scale and no scatter), trtllm_gen_n_shuffle (B and its block scales in
  TRTLLM-gen's shuffled row order; output columns in logical order) and
  mma_k64 (the 64-deep SM100-family MMA, whose fp32 accumulators equal
  those of SM100-family NVFP4 GEMMs).
- Kernel and tuning-config caches are keyed on the new options; the tiled
  MMA construction is factored into _make_tiled_mma with the same
  arguments.

All four options are constructor arguments, False by default; with every
option False the kernels and ops build and run as before.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ion, per-SM prep and split-KV launches

- mixed_step_gen_rows (Rubin, sm107 TEP4): on the full-attention layers, a
  mixed step with 9..64 generation requests runs its context rows on the
  thop op (context phase only) and its generation rows on
  multi_request_split_kv's decode stand-in, as decode steps do.
  Sliding-window layers, pure context steps and mixed steps outside the
  band keep the thop attention.
- decode_kv_splits on mixed steps (GB300, sm103 TEP4): a mixed step with
  9..64 generation requests arms the decode KV-split rule for the
  generation phase's trtllm-gen decode FMHA call, as decode steps of
  72..512 rows already do. The split stand-in counts the calls it split
  and the calls it kept.
- multi_request_prep_append: on sm103 a packed launch geometry per row
  count (items per program, warps, vector lane loads and stores) and an
  L2 prefetch of the next projection's weight slice; sm107 keeps the grid
  of one program per (row, head). Both write the same bytes.
- multi_request_split_kv: the ring depth of KV pages a CTA holds in flight
  is chosen per SM version and number of Q terms (on sm103 a deeper ring
  for a one-term Q; sm107 keeps its depth).

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ics and Rubin CuTe DSL prefill experts

- decode_expert_tactics: the stacked NVFP4 expert call of tep4 decode
  steps runs a fixed TRTLLM-gen [tile_n, config] tactic per (expert
  window, rows) from a per-GPU table instead of the autotuner's
  power-of-two bucket choice. GB300 (sm103): the 128-token tile for
  256..504 rows; every pinned tactic computes the expert rows bitwise as
  the autotuned ones. Rubin (sm107): the 32-token tile for 256..384 rows,
  bitwise as the autotuned tactics with the same FC2 MMA K. Unlisted row
  counts and windows, tuning passes and configs not valid for the call's
  shape keep the autotuner; a GPU without a table is refused.
- prefill_cutedsl_experts (Rubin, sm107 at tp 4 only, selected by
  serves(compute_capability, tp_size)): the unfinalized expert calls of
  4096..8192 NVFP4 rows (prefill and mixed steps) of every NVFP4 MoE layer
  on the rank's contiguous expert block run moe_sort, the Rubin CuTe DSL
  FC1 (TRTLLM-gen gated layout) and FC2 (permuted store, TRTLLM-gen N
  shuffle, 64-deep MMA) on the served TRTLLM-gen tensors, and return the
  TRTLLM-gen call's triple to the model's per-expert post-norm finalize.
  Kernels are compiled at install with fixed tactics per row bucket; the
  install is agreed over the TP group. Rounding differs from TRTLLM-gen;
  expert selection is unchanged. Every other GPU, tp size and call keeps
  the TRTLLM-gen expert call.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ngine gather on sequence-parallel context steps

The sequence-parallel context path wires a per-SM set of parts
(SHARD_PARTS_BY_SM) and refuses any other SM:

- GB300 (sm103): the MoE routes each rank's own rows with the fused
  routing kernel and gathers the top-8 ids and FP32 weights instead of the
  router logits; the attention-side residual join also writes the
  router's FP32 input (router_input); and on steps of at least MIN_ROWS
  rows the copy engines gather the attention input into the peers' receive
  buffers while the QKV projection runs in row blocks behind per-peer
  flags (ce_gather). The gather is agreed on by every rank and falls back
  to the NCCL all-gather when its buffers cannot be built or a projection
  cannot run in row blocks; its host flag words use prefer_pinned().
- Rubin (sm107): the MoE re-class only; the router logits are gathered,
  every row of the step is routed, and the attention input is
  all-gathered.

Routing ids and weights, the router input and the gathered rows are the
built-in path's, bit for bit.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
- Input: a decode step of at most replicated_embedding.MAX_ROWS ids runs
  the replicated embedding gather, embed_norm, the FP32 residual copy and
  layer 0's pre-QKV norm as one CuTe DSL kernel (one 128-thread CTA per
  row, flashinfer's RMSNorm body), and hands its two outputs to layer 0 as
  the residual stream and the precomputed attention input. PDL: wait
  before the first global read, trigger at exit.
- Output: while armed for a decode forward, the last layer's post-FFN
  join stores BF16 directly, so the cast before the final norm is a no-op.
- install checks the model it acts on (replicated table, m_emb == 1, the
  pre-QKV ordering, flashinfer's RMSNorm geometry, the last layer not a
  DFlash capture layer) and is agreed over the TP group. inputs_embeds
  calls, other row counts and prefill keep the built-in chain.

Bitwise with the built-in chain.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
- vocab_parallel_head (GB300, sm103 TEP4): decode steps of VP_MIN_ROWS
  (two K=7 verify blocks) up to the head's widest rows sample and argmax
  on each rank's own vocabulary shard; the ranks exchange per-thread
  states, greedy words and candidates over the tray's multicast and agree
  on the token in the head_sampler JIT CUDA kernels, instead of
  all-gathering the logits. Greedy rows are the argmax of the soft-capped
  row exactly; sampled rows keep the shipped CDF order; filtered rows run
  the shipped body on the row gathered from the ranks.
- target_head.install(llm, mapping, gather, vocab_parallel=False) builds
  the exchange buffers (rebuilt by widen for the wider steps); 8-row steps
  keep the lm_head_inbox gather.
- Split-sampler routing for the vocab-parallel pending rows.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…er GEMM tactics and per-GPU PDL chain widths

- wide_drafter_hops: the fused drafter all-reduce hops (one-shot
  all-reduce, finish convolution and add-RMSNorm per site) also on block
  decodes above 8 rows, up to the device's served rows (GB300: 9..56,
  Rubin: 9..64; the blocks of several requests at K = 7), on a second,
  wide one-shot workspace. Forwards of at most 8 rows keep the first.
- Drafter GEMMs: the fixed low-M split-K tables at TP4 on sm103 and sm107
  list the blocks of 2..4 requests (16, 24 and 32 rows) next to the
  one-request block.
- drafter_pdl_chain: the conv / merge / silu entry triggers act on every
  block width of 1..4 requests at which the drafter's tactic table covers
  every successor; the merge and silu widths are per GPU (sm103: 1..4
  requests for both; sm107: merge on one request, silu on one and four).
  The block decode reads the widths the chain's install recorded.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…hs on GB300 and Rubin

- Dispatch registers mixed_step_gen_rows, model_boundary,
  wide_drafter_hops and vocab_parallel_head with their dependencies.
- GB300 (sm103) TEP4 composes vocab_parallel_head and wide_drafter_hops,
  and decode_kv_splits installs its mixed-step mixin.
- Rubin (sm107) TEP4 composes mixed_step_gen_rows and wide_drafter_hops;
  load_weights passes the device's compute capability, so the CuTe DSL
  prefill experts install on sm107 at tp 4 only.
- Both GPUs, TEP4 and TEP8, compose model_boundary (bitwise).
- DecodeTarget runs the model_boundary input kernel and arms its BF16
  last-layer join; adopt_model installs the mixed-step attention mixins
  after every other attention re-class; the drafter shell installs the
  wide hops and passes vocab_parallel to the target head.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…103 TEP4 and the wide projection table by device

- half_window_experts joins the GB300 (sm103) TEP4 composition, as at
  Rubin (sm107) TEP4: the NVFP4 expert layers also hold their ring
  predecessor's upper expert half. Decode steps of at most 8 rows are
  placed by the prologue's routed kernel, those of 9..32 rows (two to four
  requests) by ring_buddy_assign's half flow, both balancing the selected
  experts over the four ranks. The window costs 45 GiB of KV cache per
  rank, so its capacity gate carries it only up to max_batch_size 8, and a
  plan carrying it leaves hot_replicas out (dispatch.YIELDS). Compositions
  above max_batch_size 8 are unchanged.
- wide_projection_gathers takes its 9..32-row kernel table from
  wide_kernel_steps on the dense GEMM plan's compute capability: the sm103
  table, and a refusal on sm107 (whose TEP4 composition runs those
  projections on wide_dense_gemms' own table) and on any other device. No
  served composition changes: only sm103 TEP4 serves the path.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ransform after the sharded draft top-k

The Muse Spark Open DFlash 2 drafter declares a logit transform (output
multiplier and softcap), which the DFlash worker applies to the draft logits
right after the drafter's logits processor and which returns fp32. The
vocab-sharded draft head takes only bf16 shard logits, so with the transform
on every draft step fell back to the built-in selector: a full-vocab fp32
all-gather, torch.topk over the full vocabulary and the full-vocab fill,
scatter and argmax, and the path-walk tail and the drafter PDL chain were
lost.

The transform never decreases, so each position's top-k is the same before
and after it, except where it maps distinct bf16 logits to the same fp32
value (tanh saturates), which torch.topk then orders by vocab id. The draft
head now keeps the bf16 shard and applies the transform after its top-k:

- kernels/draft_head/logit_transform.py: per drafter, a canonical-key table
  (bf16 radix key -> key of its tie group's representative) and a values
  table (bf16 radix key -> the drafter's own transform of that value),
  built on the device by DFlashForCausalLM.apply_draft_logit_transform
  itself and checked exhaustively over the 65536 bf16 values (fp32 result,
  NaN and signed zeros kept, never decreasing, groups contiguous, canonical
  values comparing as the transformed ones, layout independence on a
  host-seeded probe of the runtime shape).
- draft_shard_topk / draft_rank_topk take the canonical table and its
  identity band (only keys outside the band load it); draft_topk_merge takes
  the values table for the unary values. Without the tables every kernel is
  the previous one. Each op's reference and is_valid follow the new
  arguments, and the ops declare cells for the transform-on configurations.
- draft_head.install puts a deferring transform on the drafter instance: it
  hands back the very shard tensor ShardDraftLogitsProcessor just returned
  and records it; ShardedDraftHeadDFlashWorker._refine_block_logits keeps
  those logits bf16 only when the step takes the sharded selector (which then
  owes the transform and must apply it before sampling), and applies the
  drafter's transform to everything else, so sampled drafts, d2t, other
  layouts and any other tensor keep the built-in computation bit for bit.
  Deferral also stops when a class composed above overrides the refinement
  or the selector, or when the drafter's attributes no longer match the
  tables. selector_draft_tokens takes the transform as a required keyword;
  the verify-fold tail walk passes it too.
- The install builds and checks everything before changing anything, runs
  the tie-order probe also on logits whose transformed values tie (from the
  transform's own tie groups, and at many requests), and agrees its refusal
  and a digest of the tables over the TP group. wide_draft_topk's probe
  runs with the transform as well.

No change outside the target: an identity transform installs nothing new
and runs the same kernels as before.

Tests (tests/unittest/_torch/modeling_v2/models/musespark_open/...): the
tables against the drafter's method on every bf16 value and in the runtime
layout; the ops on their cells; the merge and the end-to-end draft tokens
against torch.topk and the built-in selector over the transformed logits
(1 to 24 requests, TP4 and TP8, all three key layouts, tie-heavy and
saturating logit families); the fast path engaged with the transform on
(no vocabulary all-gather, no full-vocab torch.topk), under the fold tail
and in CUDA-graph replay; every fallback; install atomicity, idempotence
and the probe's teeth; and a source pin on DFlashWorker._forward_impl's
statements between the logits processor and _refine_block_logits.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…es follow the drafter, warn on built-in steps

The vocab-sharded draft head applies the drafter's logit transform after its
top-k from tables built at install from the drafter's logit_multiplier and
logit_softcap. Only the deferring transform checked that the drafter still
declares those attributes; ShardedDraftHeadDFlashWorker._refine_block_logits
and the selector's _draft_head_transform_for took bf16 logits as owing the
installed tables. A drafter whose attributes changed to the identity (1.0 /
None) after the install hands bf16 logits through, and the sharded selector
then ranked and scored them with the stale tables.

- _refine_block_logits keeps a step's logits untransformed only when the
  drafter still matches the tables, and otherwise applies the drafter's class
  method with its current attributes, the built-in computation.
- A new _draft_head_takes decides the sharded route for both selectors
  (ShardedDraftHeadDFlashWorker and the verify-fold tail walk): the existing
  shard conditions plus, with a transform installed, tables that still match
  the drafter. A step it turns down takes the built-in path with the
  drafter's current transform. _draft_head_transform_for takes the drafter
  and raises DraftTransformStateError on tables that do not hold, and
  _draft_head_transformed applies the drafter's current transform.
- Every way a step with a transform installed takes the built-in path warns
  once per reason through tensorrt_llm.logger.warning_once with the
  "[musespark_open] draft_head:" prefix: the deferring transform declining
  (logits processor replaced, another tensor, non-bf16 logits, changed
  attributes, a worker class that overrides the refinement or selector), and
  the refinement or selector turning a step down (full-vocab logits, layout,
  sampled drafts, d2t, changed attributes, logits the transform did not
  defer). The deferring transform records its reason on the drafter so a
  step warns once, not once per stage.
- The install agrees the key exchange's refusal over the TP group too: a
  rank whose make_allreduce returns None and the ranks that built theirs all
  raise the same ValueError, before anything changes.

Tests (test_modeling_v2_draft_head_transform.py): the drafter's multiplier
and softcap changed after the install (to the identity, another softcap, a
multiplier only, a softcap only), before the logits processor, after the
deferring transform and after the refinement, at 1 and 3 requests and under
the fold tail: the built-in transform of this rank's shard reaches the
vocabulary all-gather bit for bit, the sharded selector never runs, the
draft tokens equal the built-in ones of the new transform, and the tables
serve again once the drafter declares their transform; the selector refuses
tables that do not hold; each fallback reason warns exactly once over
repeated steps, the fast path and a drafter without a transform never warn;
the key-exchange refusal is agreed and leaves the worker and drafter as they
were.

Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>

@JunyiXu-nv JunyiXu-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Runtime part: with the new env vars unset, the KV cache manager, scheduler and executor paths are unchanged, but the opt-in TRTLLM_KV_DEFER_CLOSE=1 path also defers preemption releases (repro in the inline comment) and can leave a cache open across a drained control action.

Comment on lines +5733 to +5740
if self._defer_close_armed and not request.is_dummy:
if request.py_request_id in self._early_freed_index_requests:
self._early_freed_index_requests.discard(request.py_request_id)
else:
self._detach_index_slot(kv_cache)
self.index_mapper.remove_sequence(request.py_request_id)
self._deferred_closes = (*self._deferred_closes, (request.py_request_id, kv_cache))
return

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With TRTLLM_KV_DEFER_CLOSE=1 this also defers preemption: preempt_request -> _release_preempted -> free_resources returns before close(), so its "pages were released" result is no longer true. _try_recompute_pause_for_gen (scheduler_v2.py L2437-2449) retries the allocation right after each victim, still fails, and takes the next one. Running the head's own scheduler and manager code against a stub pool, one pool-full generation step recompute-paused 3 victims instead of 1 and was still not scheduled. Could _release_preempted (and the admission rollbacks) close at once, so only finished requests are deferred?

idle = (total_num_live_requests == 0 and len(waiting_queue) == 0
and not self.is_shutdown
and not self._has_pending_connector_transfers())
if idle and self._kv_defer_close_armed:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This flush never runs while a control request is parked: the control_pending return above (L6105-6108) comes first. _handle_control_request then fires a drained update_weights/reset_prefix_cache/sleep once active_requests is empty, while the last finished request's cache is still open; it only closes after the action, at L5651. reset_prefix_cache drains so that no request holds or commits blocks across the reset, and a request cancelled mid-prefill still commits its partial block in close(). Could _handle_control_request flush the deferred closes before firing?

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants