Repository navigation
[None][feat] Muse Spark Open modeling_v2 TEP target: a per-deployment lever resolver and the Vera Rubin high-concurrency, mid-concurrency and low-latency fast paths - #21377
Conversation
…y chain memo and lazy slot scrub Three independent changes to the KV cache manager v2 admission path. Each is off by default; with its variable unset the manager behaves as before. - TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS=1: page resume and close skip event waits and merges that the stream order already implies, and wait once per distinct event instead of once per page. - TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO=1: the radix tree memoises the block-key chain of a token sequence, so the reuse probe, the first-block probe and the cache creation of one request hash its prompt once. - TLLM_KV_CACHE_MANAGER_V2_LAZY_SLOT_SCRUB=1: a slot allocator scrubs the events of recycled slots once per batch allocation instead of polling them per slot. The nanobind module exposes the switches and the event counters for the unit tests. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ressed layer types; DFlash without a separate draft KV cache ColdPageQuantizationCompressionConfig gains two fields, both defaulting to the previous behaviour: - global_scale: identity (default) or fp8_block_amax. For an FP8 KV cache without a ModelOpt scale, fp8_block_amax uses a global scale of 6, which maps each block's E4M3 absolute maximum onto an exact E4M3 block scale. - uncompressed_layer_types: layer types (sliding_window, full_attention) whose pages stay uncompressed in the cold tier. The KV cache compression check now also admits DFlash when it has no separate draft KV cache: its drafter keeps the context K/V in worker-owned buffers, so the codec only sees target pages. The llm_args golden manifest records the two new fields. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ission and router held placement Generation workers: - TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS=N (0 by default) admits up to N generation-only requests beyond max_num_active_requests, so their KV transfer overlaps the running decode. The forward batch stays capped at max_batch_size. The engine widens its sequence-slot pool and the KV cache manager v2 its lease budget by the same N. It is off under TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN, attention DP and pipeline parallelism. Requests that do not fit go back to the front of the waiting queue, which keeps FCFS order. Router: - LoadBalancingRouter gains max_inflight_per_server and hold_timeout_s. With a cap set, a request is placed only on a server below the cap, and otherwise held until one frees up; a request held longer than the timeout falls back to the least-loaded server. Without a cap, placement is unchanged. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…ap loop on the execution stream RuntimeHooks gains step_loop_on_execution_stream (False by default). When a model declares it, the overlap loop's thread makes the execution stream current, so the sampler and the next step's input preparation queue behind the forward on one stream instead of crossing streams twice per step. Models that do not declare it are unchanged. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…on-excluded MoE layers The checkpoint keeps the experts of its quantization-excluded MoE layers (0 and 61) in BF16. With MUSESPARK_OPEN_L061_EXPERTS_NVFP4=1 (off by default) those experts are built on the global NVFP4 leaf and quantized from BF16 at load. MUSESPARK_OPEN_L061_EXPERTS_NVFP4_LAYERS selects a subset of the layers, and MUSESPARK_OPEN_L061_EXPERTS_NVFP4_FC2_AMAX supplies the FC2 input amax per layer. This is a precision change and is never enabled implicitly. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…projections for large context calls New kernels/context_fp8_dense: at load, FP8 copies of the qkv, o_proj, pool and unpool projection weights; at run time, context calls of at least min_rows rows (1024 and above) quantize their input with a one-pass MXFP8 quant and run the FP8 GEMM with fixed Rubin tactics, plus an optional MXFP8 attention-input gather on sequence-parallel steps. Decode steps never reach the row floor. The install refuses unless the device is sm107, the Rubin CuTe DSL is available and the projections are plain bias-free BF16. This is a precision change. The lever resolver keeps it off unless it is requested explicitly. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…p append and a parallel windowed context commit Decode steps of several requests on sm107 TEP4: - multi_request_prep_append runs the packed-geometry kernel for calls of at least packed_min_rows rows (one program per head over all sequences) instead of the per-sequence grid; the result is bitwise the same. - context_kv_commit_windowed commits the drafter's windowed context arena with one program per request and a tuned compaction for steps of at least parallel_min_requests requests; smaller steps keep the one-program commit. Both write the same bytes. - hot_replicas skips experts that were quantized at load, which have no census table. The thresholds are arguments set by the lever resolver. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…or decode steps of 16 to 32 rows Decode steps of two to four requests (16 to 32 verify rows) on sm107 TEP4: - wide_prologue routes each row with its own one-request bits on the wide router (router shard rows), so a row routes as it would alone; - pool_into_route and unpool_into_boundary run the direct pool and unpool transports with the half expert window at 16, 24 and 32 rows; - attention_output_into_hop publishes o_proj directly into the wide post-attention hop transport at 16, 24 and 32 rows; - finalize_transport serves finalize calls of 17 to 32 tokens. The 8-row layout shares the kernel builders; its arithmetic is unchanged. Which of these paths a deployment serves is decided by the lever resolver. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…et chain and a cooperative head-sampler fold Paths for decode steps of one or a few requests on sm107: - drafter K/V side chain: on a dense full-length drafter arena, the context K/V of the step is projected and stored on a side stream forked after the last capture layer, with the drafter's fixed GEMM tactics; the fc projection can fork early, and the commit kernel also writes the drafter forward's integer preamble (context_kv_store, context_kv_commit, context_kv_fork, verify_fold, drafter_block_decode). - head sampler: a cooperative kernel A computes the per-row states of calls of up to 128 rows with the same per-element cap and combine order as the grid kernel, so states and tokens are bitwise the same; wider calls keep the grid kernel. - the target head can serve without the lm_head inbox, and can gather its logits with a multimem all-gather when every rank builds the symmetric buffer. - runtime: the block-offset stager copies only the runs of page indices that changed since the previous step, keyed by value, and resets on storage, scale or offset changes. Each path declines per step when its arena or shape does not apply. Which of them a deployment serves is decided by the lever resolver. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…cks the fast paths per deployment New levers.py: one table of levers with a rule each. dispatch.compose resolves it once per process from the routing key and the deployment it sees: device, tensor-parallel width, max_batch_size, max_num_tokens, the serving role (aggregated, or the context / generation side of a disaggregated server) and the drafter's context arena (dense, windowed or paged). Installers read their decision and thresholds from the plan; no other module of the target reads the environment. - Lever kinds: composition and bitwise levers follow their rule; numeric levers also need their deployment class; lossy levers are never on unless requested; flagged levers stay off until requested. - Knobs: MUSESPARK_OPEN_LEVER_<NAME>=auto|on|off per lever, declared parameters MUSESPARK_OPEN_LEVER_<NAME>__<PARAM>, MUSESPARK_OPEN_LEVER__ROLE and MUSESPARK_OPEN_LEVER__DEFAULT=auto|baseline. Unset means auto. baseline selects the levers the target served before this change, as a control. Unknown names, bad values, out-of-range parameters, a request a structural rule excludes and disagreement across ranks raise before anything is installed. - Structural rules: the drafter arena selects the commit and side-chain paths; TRTLLM_KV_DEFER_CLOSE=1 excludes the execution-stream loop; UCX_RNDV_THRESH may not be set on a disaggregated role. - The resolved and the installed lever tables are logged once on rank 0, the installed one as JSON after the last install hook, and per-step declines are counted. - The modeling_v2 router records the engine's deployment facts per target class. Host-only tests cover the grammar, the overrides, the structural rules, baseline reproduction and the per-deployment lever tables. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…neration servers wide_router_per_row is a lossless-numeric lever: each verify row routes with its own one-request bits, so near-tie experts can resolve differently from the built-in wide router. Its accuracy gate now holds only on generation servers at TEP4 with max_batch_size 2..4; aggregated servers of the same batch range resolve it off under auto (source "gate"). The lever tests now expect it off at agg mbs 2..4 and check that an aggregated server built from a generation config differs from that generation server only in this lever. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
…election Aggregated sm107 TEP4 servers with max_batch_size of at least HC_MIN_BATCH (48) resolve under auto to the high-concurrency selection: the wide decode transports, the wide o_proj publish, the wide finalize transport and the cooperative head-sampler fold are off; the graph-start QKV prefetch is on; and the parallel windowed context commit starts from one request. A lever Param can now declare a per-deployment default (default_for); a knob still overrides it. Generation servers, TEP8 and smaller batches keep their selection. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #81813 [ run ] triggered by Bot. Commit: |
|
PR_Github #81813 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #81824 [ run ] triggered by Bot. Commit: |
|
PR_Github #81824 [ run ] completed with state
|
feat/musespark-open-agentx's aarch64 build configs and docker CUDA_ARCHS predate main's sm_107 addition, so the SBSA wheel and release image carry no Vera Rubin code, while this PR's TEP fast paths target Vera Rubin. Add 107-real to both aarch64 WHEEL_ARCHS and to the arm64 docker CUDA_ARCHS, so a side-branch release of this branch runs on Vera Rubin as well as GB300. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #82033 [ run ] triggered by Bot. Commit: |
|
PR_Github #82033 [ run ] completed with state
|
… for bandit bandit's B105 reads PRIORITIZE_FIRST_TOKEN_ENV = "TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN" as a hardcoded password (the value ends in TOKEN) and fails the Release-Check and the pre-commit workflow's bandit step. It is an environment variable's name: mark the line # nosec B105, as the other token-named constants are. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
…ters PLCScanningSetup renamed runContainerScanning to runContainerOssScanning and added runContainerMalwareScanning (main: PLC Pipeline update, NVIDIA#19557). This branch's L0_MergeRequest.groovy predates that and still passes runContainerScanning=false, which the job ignores: a merge request that touches an OSS-compliance file (here tensorrt_llm/usage/llm_args_golden_manifest.json) ran the container OSS scan at its default, looked for a post-merge image of the PR commit and failed [OSS-Compliance-Check]. Pass the job's current parameter names, with the values main passes, in both PLC calls. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #82117 [ run ] triggered by Bot. Commit: |
|
PR_Github #82117 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #82156 [ run ] triggered by Bot. Commit: |
|
PR_Github #82156 [ run ] completed with state
|
Description
This PR adds a per-deployment lever resolver to the Muse Spark Open modeling_v2 TEP target, and with it the Vera Rubin (sm107) fast paths for the high-concurrency, mid-concurrency (TEP4 decode steps of 2–4 requests) and low-latency (one request) regions. One build then serves the whole agentic Pareto curve: aggregated TEP8 / TEP4 and the context and generation servers of disaggregated deployments. It follows #21143. The changes are 12 commits:
_experimental/modeling_v2/models/musespark_open/...__tep): the resolver, five groups of fast paths, and the high-concurrency selection;The lever resolver (
<target>/levers.py).dispatch.composeresolves one table of levers once per process from the deployment it sees: device, TP width,max_batch_size,max_num_tokens, the serving role (aggregated, or the context / generation side of a disaggregated server) and the drafter's context arena.MUSESPARK_OPEN_LEVER_<NAME>=auto|on|off, plus declared thresholds…__<PARAM>,MUSESPARK_OPEN_LEVER__ROLEandMUSESPARK_OPEN_LEVER__DEFAULT=baseline.baselineselects exactly what the target served before this change, as a control.Region selections on sm107 under
auto:max_batch_size6–47:max_batch_size≥ 48: the high-concurrency selection. That is packed prep, parallel commit from 1 request and the graph-start QKV prefetch, with the wide paths and the cooperative fold off.Shared paths: opt-in, off by default. With the variables or arguments unset, behaviour is unchanged.
KVCacheManagerV2, each switch independent:TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS=1: skips event waits the stream order implies;…_KEY_CHAIN_MEMO=1: memoises a prompt's block-key chain;…_LAZY_SLOT_SCRUB=1: scrubs recycled slots once per batch allocation.ColdPageQuantizationCompressionConfiggainsglobal_scale(identitydefault,fp8_block_amax) anduncompressed_layer_types. DFlash without a separate draft KV cache now passes the compression check. The llm_args golden manifest records the two fields.TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS=Nadmits up to N generation-only requests beyondmax_num_active_requestsso their KV transfer overlaps the running decode; the forward batch stays capped atmax_batch_size.LoadBalancingRoutergainsmax_inflight_per_server/hold_timeout_s.RuntimeHooks.step_loop_on_execution_stream(False by default): a model can declare it to run the overlap loop's sampler and input preparation on the execution stream.MUSESPARK_OPEN_L061_EXPERTS_NVFP4=1builds the experts of the quantization-excluded MoE layers 0 and 61 as NVFP4 from BF16 at load. This is a precision change, never enabled implicitly.Measured gains. AgentX traces on the 1M-token corpus, Vera Rubin, full protocol, forced draft acceptance 3.05. x = p10 per-user output throughput; y = total throughput per GPU. Each row compares this tree with, in the same wave:
MUSESPARK_OPEN_LEVER__DEFAULT=baseline, i.e. the target's previous selection.max_batch_size4max_batch_size8max_batch_size32max_batch_size48, with the high-concurrency recipe and its context-FP8 opt-in). Against the earlier high-concurrency build of this target, three high-mode runs of this tree give +0.1 % x / +0.3 % y. The resolver's earlier choice at this size (wide paths and cooperative fold on) was −2.2 % y, which is why the high-concurrency selection is the default frommax_batch_size48.Accuracy. Spec decoding on with natural acceptance; long-preamble GSM8K, GPQA-Diamond and a prefix-reuse probe, greedy; pooled paired McNemar.
max_batch_size2–4: −4.0 pp, p 0.002. With the lever off, the same check passes (−0.25 pp, p 1.0). On a generation server it costs nothing (+1.0 pp, p 0.50).[None][fix] … router per row only on generation servers).Validation.
feat/musespark-open-agentxatc55183f445.51daa26acb([None][fix] Muse modeling_v2: half-window early expert ids on Vera Rubin, and context_kv_fork under KV cache estimation #21364). One conflict, inops/comm/pool_direct_route.py, is resolved by keeping both sides: the wide path's CTA-relative shared index and [None][fix] Muse modeling_v2: half-window early expert ids on Vera Rubin, and context_kv_fork under KV cache estimation #21364'splace_half. The resulting tree equals the measured tree plus [None][fix] Muse modeling_v2: half-window early expert ids on Vera Rubin, and context_kv_fork under KV cache estimation #21364's two hunks exactly. The replayed tree was not rebuilt.VANILLAin the recipes; the TRTLLM drafter backend is not measured on Vera Rubin.Test Coverage
tests/unittest/_torch/modeling_v2/models/musespark_open/.../test_modeling_v2_levers.py(new, CPU-only): the knob grammar, overrides, structural rules, baseline reproduction, and every region's lever table, the high-concurrency selection and router-per-row gates included. 180 tests. It is in no CI test list yet; the candidate isl0_cpu..../test_modeling_v2_drafter_context.py: the plan over a paged drafter arena, whose only structural exclusion is the window compaction.tests/unittest/kv_cache_manager_v2_tests/test_stream_ordered_events.py,test_key_chain_memo.py;tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py: the KV cache manager v2 switches.tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py,test_nvfp4_cold_page_fp8_global_scale_gpu.py;tests/unittest/llmapi/test_llm_args.py: the cold-page fields.tests/unittest/_torch/executor/test_disagg_gen_prefetch.py,tests/unittest/disaggregated/test_router_hold.py: generation prefetch admission and router hold.tests/unittest/_torch/executor/test_runtime_hooks.py: the execution-stream step loop hook.tests/unittest/_torch/modeling/test_musespark_open_l061_nvfp4.py,test_musespark_open_context_fp8_dense.py: the NVFP4 layer-0/61 experts and the context FP8 dense projections.accuracy/modeling_v2/test_musespark_open.py(on this branch): GSM8K and mean acceptance length undermodeling_v2="require".PR Checklist
🤖 Generated with Claude Code