Repository navigation
[None][feat] Qwen3.8-Flash-Next NVFP4 modeling_v2 unified target on 4×GB300 (TP4/TEP4/ADP4), AgentX-tuned - #21372
Draft
Wanli-Jiang wants to merge 218 commits into
Draft
Wanli-Jiang wants to merge 218 commits into
Wanli-Jiang wants to merge 218 commits into
Conversation
…nd L2 hints in the finalize-fused grouped GEMM The finalize-fused contiguous grouped GEMM gains three opt-in knobs; the defaults keep every existing caller's kernel unchanged: - mma_inst_tile_k: the MMA instruction K tile count (4 before). 2 gives a K tile of 128, which leaves room for five A/B stages beside the C staging buffer instead of two, so a GEMM with a small K (an FC2 of K = 640) pipelines. - row_owned_epilogue: each epilogue thread waits on its own row's previous bulk reduce instead of a per-tile bulk wait plus a CTA barrier. - l2_hints: the A/B operands are loaded EVICT_FIRST and the bf16 scatter-add stores EVICT_LAST. The Qwen3.8-Flash-Next target's prefill MoE FC2 uses them: GB300, per layer 622 -> 555 us at 16384 rows and 427 -> 373 us at 8192. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…tadata Metadata work of a generation-only MTP decode step: - QSA sparse metadata: the decode row mapping of a batch without context rows (up to 256 requests) is one Triton launch instead of ~12 ATen ops, eager and in the CUDA graph. - GDN replay work items take the single-launch Triton builder for every decode count up to 256 whose bookkeeping is contiguous (one compile for all batch sizes), with or without the all-layer commit. - The spec metadata stages its RNG uploads in one pinned tensor and caches the sampling-parameter normalization. The replay-mode metadata test hands the dispatch real bookkeeping tensors and mocks both work-item builders. Exact: the QSA generation mapping matches the general cumulative-sum and search mapping bitwise in unit tests, and every replaced component was compared in a batched acceptance-length eval. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ed MTP decode steps In graph-captured generation-only MTP decode steps (eager and mixed steps keep the original ops): - the greedy argmax of 1-16 rows is the Triton row argmax (greedy_argmax, 64 chunks a row, bf16 or fp32 logits) instead of a memset, a 31-block reduce and a cast per call; - strict and forced acceptance run as one Triton launch (fused_accept); every force mode is launched once at warmup; - the first draft step's row selection (gather ids, hidden rows, positions, KV and sequence lengths) is one launch (draft_loop_ops.draft_row_select), dropping two cats and a dead cumsum, and the next step's tokens take two launches instead of six (draft_loop_ops.next_new_tokens). GLM-5.2's MTP draft worker override passes the fused flag through. Bitwise against the replaced ops (unit tests incl. CUDA graph replay, run on B200; in-graph shadow checks over ~185K serving comparisons). MTP3 con1 on GB300, Qwen3.8-Flash-Next (with its target's fused rollback and step setup): captured bookkeeping 121 nodes / 267 us -> about 42 nodes / 79 us. That is the measured 27 nodes / 53.6 us plus two changes this commit leaves out: a duplicate QSA row mapping (1 node, 4.3 us) and the V2 Mamba replay bookkeeping folded into the conv promotion (14 nodes, 20.9 us). con32 neutral. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
… own the checkpoint commit A model whose cached-replay kernel commits each wrapping checkpoint itself needs the manager's all-layer commit off: kept on, the manager commits those checkpoints a second time after every step of 16 or more decode requests. GDNReplayState gains all_layer_commit_enabled (on by default), and the Qwen4Exp hybrid cache manager gains disable_gdn_cached_replay_all_layer_commit(). Its use_gdn_cached_replay_all_layer_commit holds only while the commit is enabled, so the Mamba2 metadata then builds the replay work items for every decode row instead of only for partitioned steps. Nothing changes for a model that does not call it. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…timizations for ISL 8K / OSL 1K Performance work on the qwen3_8_flash_next_nvfp4__sm_103__tp1 modeling_v2 target over the ISL 8192 / OSL 1024 Pareto curve on one GB300 (TP1, NVFP4 KV cache, bf16 SSM state), for STP and MTP3: target-local ops (CuTe DSL, Triton), each with a GPU test, wired into the target. Decode at small batch - Hyper-connection sites: faster 1-row kernels; 2..16 rows on a tensor-core launch pair with the MoE finalize folded in; tc_first up to 256 rows. - FP8 W8A16 GEMV (fp8_gemv) for the dense projections under dense FP8, the vocab head, and the 1-row output and shared-expert projections. - GDN: the in_proj GEMV with a fused conv1d epilogue (2..16 verify rows, up to 8 decode rows); decode rows that stage the state by TMA; opt-in NVFP4 in_proj up to 8 rows. - NVFP4 sparse GQA decode on a CuTe split-K + merge kernel, behind a fused decode prologue over NVFP4 KV. - MoE: pinned trtllm-gen tactics for the decode buckets; a tactic this build's runner does not list as valid falls back to the AutoTuner with a warning. MTP - The draft step runs on the trunk's fused machinery and carries only the draft loop's rows; an FP8 lm_head copy serves the draft head; opt-in NVFP4 draft head and draft-layer experts. - Verify: request-grouped NVFP4 sparse GQA (coalesced gated epilogue, split last wave, per-row merge, every reachable launch plan compiled up front); the GDN cached-replay verify on a CuTe tensor-core kernel that commits wrapping checkpoints in-kernel, so the target switches the manager's all-layer commit off when it binds a bf16 pool. - One-launch speculative rollback (spec_rollback_commit) and decode-step setup (verify_step_setup); batched rollback of the QSA index planes; fused PLE on verify and mixed steps. Prefill and mixed steps - Faster prefill CuteDSL MoE (one-launch routing and sort, pinned tactics, the finalize-fused GEMM's new knobs); opt-in per-tensor FP8 dense projections; tcgen05 dense projections for 17..256 rows (mid_rows_mm). - HC sites without a stored normed stream; a leaner GDN prefill block. - The sparse attention row-tile pass is persistent and free of per-step Triton recompiles; NVFP4 KV context pages are dequantized into a bf16 scratch. - The PLE host table lives on the GPU's NUMA node. Opt-in switches: TRTLLM_MODELING_V2_DENSE_FP8, TRTLLM_MODELING_V2_GDN_INPROJ_NVFP4, TRTLLM_MODELING_V2_MTP_MOE_NVFP4 and TRTLLM_MODELING_V2_MTP_HEAD_NVFP4. GB300, ISL 8K / OSL 1K, from the NVFP4 KV + bf16 state starting point (con1 tok/s/user; con32 and con256 tok/s/GPU): STP 349 -> 451, 2,907 -> 3,378, 5,986 -> 7,448; MTP3 at a forced acceptance of 3.3: 425 -> 857, 3,570 -> 5,552, 4,957 -> 10,148. These were measured with a few kernels since left out (the small-row NVFP4 MoE at STP con1, the FP8 activation maximum folded into its producers, request-tiled verify index scores, a CuTe plan kernel for the row-tile pass, a lean one-split decode kernel), each within measurement noise or below 2% at one point. Accuracy on this revision: gsm8k STP 96.32, MTP1 96.06 (acceptance 1.969, MMMU 58.11), MTP3 over CuteDSL experts, fp8 KV and a bf16 state 96.25 (acceptance 3.508); the MTP3 and MTP7 long-context pass-key cases pass (acceptance 3.998, 7.117). Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…et's fastest MTP3 configuration The modeling_v2 MTP3 case now runs with an NVFP4 KV cache and a bf16 recurrent state, and with the target's opt-in fast paths: FP8 dense projections, the NVFP4 GDN input projection for decode rows, and NVFP4 copies of the MTP draft layer's routed experts and of the draft head. Concurrency 1 and 256 as before; acceptance stays unpinned. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…GB300 Pareto optimizations, head 1edb516) as the onboarding base Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…block_indices at 128 / 1024 / 8192 rows (r21) Local variant of the tp1 target's attention/qsa_expand_block_indices for the unified tp4/tep4/dep4 target, with cells at the 128-request decode batch, the MTP7 verify step of 128 requests (1024 rows) and an 8192-row context chunk ending a 262144-token prompt, and CUDA-graph replay at 1, 128, 256 and 1024 rows. Certified on sm_103. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…r__r30 variant, copied from the tp1 target The local variant this target's dispatch 30 verifies, started as a verbatim copy of qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/norm/rms_norm_gated_token_major and its test. Class and instance names are kept for the later promotion. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ken_major__r30 at the tp4/tep4 and dep4 GDN epilogue Adds the per-rank cells of this target's GDN gated-norm epilogue on sm_103: 12 value heads with z the column view of a tp4/tep4 rank's 4120-wide input projection at 1, 4, 128, 1024 and 8192 tokens (4-row kernel with and without PDL, and the 8-row prefill launch at 98304 rows), and 48 heads at token stride 16480 for a dep4 rank's 32 tokens. The test derives z's column offset from the [q k v | z | a | b] layout per cell, and checks the token stride only where there is more than one token. The note records the measured error (at most one bf16 ulp, cosine >= 0.9999985 against fp32) per cell. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…r__r30 notes the 12-head 8-row tail as uncertified At 12 heads a token count from 2731 up launches the 8-row kernel, and an odd one leaves its last program half full; no certified cell reaches either. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
… per-rank projection shapes Local variant of catalog/gemm/cublas_mm for the qwen4 unified target (tp4/tep4/dep4, sm_103). Its test adds the target's 33 projection cells to TARGET_CELLS: attention and GDN in/out projections at tp4/tep4 and dep4 widths, shared expert, router up to dep4's 32768 gathered rows, the fp32-out HC first projection, PLE and lm_head, at decode, MTP-7 verify and 8192-row prefill counts. fp32 outputs are now gated in fp32 ulp (1024 / 512, measured at most 192 / 97), so an output rounded through bf16 fails; the bf16-ulp band admitted it. Negative controls cover a product missing one K block and the bf16-rounded fp32 output. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…request decode and MTP-7 verify rows Local variant ops/embedding/ngram_hash__r13 of tp1's ngram_hash for the hf_fc694b5_orig / sm_103 / unified target, with its GPU test. Adds cells at 128 and 1024 call-site windows and a 128-request x 8-query chunk tail, graph capture at 128 and 1024 windows, and a guard test for num_query_tokens above the history length. The note records the measured cost and the built-in qwen4_exp_ple_ngram_hash Triton op's equivalence on the [T, 3] form. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
@routes keeps every route registered for an architecture, in registration order (the imports in models/__init__.py). modeling_v2_resolve tries them in turn: the first that does not raise NoTarget decides, and when all decline the NoTarget carries each family's reason. A routing module registering the same architecture twice still asserts. The claims suite reads routed targets per registered routing module instead of per architecture, and the routing suite pins the order, the first-claim rule, the joined reasons and the registration guard. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
… and dep4 to a unified target shell A second family, qwen4_exp_for_conditional_generation, routes Qwen4ExpForConditionalGeneration on sm_103 at four ranks in one node: tp4 (attention TP4, experts split by intermediate), tep4 (attention TP4, experts EP4) and dep4 (attention DP, experts EP4), with no speculation or one-model MTP-Eagle at max_draft_len 1..7, bf16/FP8/NVFP4 KV, a bf16 recurrent state and the TRTLLM MoE backend. Everything else raises NoTarget, and the one-rank deployment still routes to the qwen4_exp tp1 target. The target hf_fc694b5_orig__sm_103__unified is a shell: its class raises at construction, PrefillTarget and DecodeTarget are empty, and weights.py holds only its docstring. The constants routing admits are declared in modeling.py and the CPU routing test holds the two together. Adds the accuracy gate class TestModelingV2HfFc694b5OrigSm103Unified with test_gsm8k at tep4, and lists the target's unit-test directory on the CPU stage. The directory is a package because its test_routing.py repeats the glm_moe_dsa family's basename. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…pool layout pool.py derives the pools one rank binds at tp4, tep4 and dep4 from the text config and the Mapping, the way the stock Qwen4ExpHybridCacheManager sizes them: KV heads over attention TP (1 per rank at tp4/tep4, 2 at dep4), GDN conv and recurrent state over mamba_effective_tp_size (4 key / 12 value heads per rank, or all 16 / 48), the replicated QSA index and position planes and PLE state. It ports the tp1 target's pool-contract checks and adds the conv [q | k | v] section check and take_pools, which takes and checks every trunk view. test_pool_layout.py pins the per-rank table and builds the stock manager through the engine's own class choice and factory at ranks 0 and 3 of each member, for bf16, fp8 and NVFP4 K/V, and with MTP3 under the cached replay. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…4 local variant r38, copied from the tp1 target Start the unified target's local variant of the tp1-owned NVFP4 sparse GQA entry and its GPU test as they stand; the test imports the variant module. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…gqa_nvfp4 r38 at dep4's kv_nvfp4 context steps Two cells on the prefill kernel at one dep4 rank's operands: an 8192-row chunk reaching max_seq_len 9216 (table column 287 of a [32, 292] table) and a 512-row four-request step, both in the trunk layout (24 role pages, layer 11's planes) over pools whose request pages straddle 2^31 bytes of the scale plane, at page counts in both of Triton's specialization classes. The bf16 twin dequantizes only the mapped pages so such a pool stays inside the cell budget. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…rm_rope at the per-rank rope shapes (r3) Local variant of the catalog's attention/fused_qk_norm_rope for the unified tp4/tep4/dep4 target, with cells for the main rope on the fused input-projection row at attention TP4 (6 q / 1 k heads, 4352-wide row) at 1, 128, 1024 and 8192 tokens, and dep4's 24 / 2 row (14080 wide) at 32 tokens. The mRoPE section split and plain-reduction checks also run at the TP4 geometry. Certified on sm_103; positions past 32768 are recorded as uncertified, with measurements. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…variant, copied from the tp1 target The local variant this target's dispatch 14 verifies, started as a verbatim copy of qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/ple/ngram_token_history and its test (imports pointed at the variant). Class and instance names are kept for the later promotion. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ry__r14 at the per-rank slot pools Adds the four cells vocabulary row 14 names: a 128 x 1 decode, the MTP7 128 x 8 verify and a fresh 8192-token context (max_num_tokens) at a tep4 rank's 1874-slot pool, and a dep4 rank's 32 x 1 decode at its 850-slot pool, each with slots spread over the whole pool. The pool sizes are the stock engine's per-rank PLE n-gram pools at core's config with the n-gram table host-resident; the note records them, the contiguous [slots, 2] int64 view the manager hands over, and the pools left uncertified. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…lain sum at hidden 2560, tp4 Local variant of catalog/comm/allreduce for the qwen4 unified target (sm_103), copied with its 4-rank op matrix. The target sums its attention / GDN out-projection and MoE partials with op NONE at hidden 2560, twice per layer at tp4 and tep4; the catalog entry certifies only RESIDUAL_RMS_NORM at hidden 2880 / 7168. Adds four world-4 NONE cells (1, 128, 1024 and 8192 rows) and matrix checks for NONE: inputs untouched on both kernels AUTO picks (one-shot up to 1024 rows, NCCL from 2048), bitwise agreement across ranks, CUDA-graph replay at the target's captured row counts (1..128 decode, 512 / 1024 MTP verify), and a missing-rank negative control on each kernel. ARCHS is sm_103, where the variant's cells were driven. Measured on GB300: the one-shot cells are bitwise the fp32-sum reference; 8192 rows (NCCL) reads 2.59 ulp per element and 0.76 ulp RMS inside the 5 / 1.25 band; a sum missing one rank reads 189-230 / 148 ulp. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…shot bucket as 1 to 2047 rows AUTO's lookup truncates log2(rows), so every row count below 2048 takes the one-shot kernel at hidden 2560 / world 4, and NCCL starts at 2048. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…_r8 variant, copied from the tp1 target The local variant this target's dispatch 8 verifies, started as a verbatim copy of qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/hc/qwen4_exp_hc_inject_norm and its test, plus the new ops/hc package. Class and instance names are kept for the later promotion. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ct_norm__r8 at every HC site form from 1 row to max_num_tokens Adds the unified target's cells on sm_103: the attention-site injection at 1 and 1024 rows, the MoE-site combine over the routed rows already reduced across ranks (no fold) at 1, 32, 128 and 8192 rows, the stack-final site at 1 row under that mixer's own norm weight (RMS 3.1, up to 12.4, nearly all positive), and norm-only at 1 and 8192 rows (the stack's first site). The test draws the stack-final weight from the checkpoint's statistics and adds a guard test: shared/gate_logits with no injection is silently the norm alone, which is_valid refuses. The note records the measured error per cell (one ulp, or inside the combine's band; cosine 1.0, relative L2 <= 1.8e-5). Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…local variant r9, copied from the tp1 target The tp1 target's ops/hc/qwen4_exp_hc_lowrank_act entry and its GPU test, copied as they stand into the unified tp4/tep4/dep4 target's ops/hc as the r9 variant, class and instance names kept, the test importing the variant. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…nk_act at the target's HC mixer rows (r9) Cells for the unified tp4/tep4/dep4 target's HC mixers on the 352-wide fp32 first-projection logits (hc_first_cols): one decode token, a dep4 rank's largest captured decode batch (32 rows), tep4's MTP-7 verify (1024 rows) and an 8192-row prefill chunk with the injection weights, and one token at the stack-final mixer without them. The narrow-logits guard is now driven: an unguarded view narrower than rank + hc_count is read past its end. The note records the call sites, the JIT source, index limits and measured ulp. Certified on sm_103. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…__r31 variant, copied from the tp1 target A verbatim copy of the tp1 target's ops/ssm/gdn_decode_gated_norm_pdl.py and its test, with one change each so the copy imports here: the gated-norm reference comes from this target's rms_norm_gated_token_major__r30 variant, and the test imports this variant. Class GdnDecodeGatedNormPdl and instance gdn_decode_gated_norm_pdl are kept. ops/ssm/__init__.py is the SPDX header. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…norm_pdl__r31 at 4 / 12 heads and the engines' state pages Four cells beside tp1's three: a tp4/tep4 rank's 4 k / 12 v heads (proj width 4120) at 1 row, 2 rows (row 1 a pad slot) and 4 rows (MAX_ROWS) over tep4's measured 1874-slot state page, and dep4's 16 / 48 heads at 1 row over its 850-slot page. A page is the stock manager's coalesced layout (slot stride 36 GDN layers' states, the view the last layer's sub-page), so the top slots sit 1.3e10 and 2.4e10 elements past the base. In fp32 such a page spans 53 / 96 GB, so the test reserves its virtual range and maps memory under the touched slots only, runs the reference on a dense copy of them, and checks every other byte of their pages unchanged. The graph-replay check also runs at 4 / 12 heads. The note records what this target does with the kernel: its routing serves a bf16 state cache only and the kernel takes fp32 pools only, so it is not launched here; gdn_decode_rows_gated_norm reads its reference. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…r__r4 variant, copied from the catalog The local variant this target's dispatch 4 verifies, started as a verbatim copy of catalog/moe/fp4_block_scale_moe_runner and its test (imports repointed at the variant). Class and instance names are kept for the later promotion. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…romoted from the qwen4_exp targets Moves the unified target's local variant of ssm/causal_conv1d_triton_update into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
… from the qwen4_exp targets Moves the unified target's local variant of ssm/causal_conv1d_update into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…oted from the qwen4_exp targets Moves the unified target's local variant of ssm/cu_seqlens_from_seq_lens into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ule, promoted from the qwen4_exp targets Moves the unified target's local variant of ssm/flashinfer_chunk_gated_delta_rule into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…moted from the qwen4_exp targets Moves the unified target's local variant of ssm/gdn_decode_gated_norm_pdl into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…omoted from the qwen4_exp targets Moves the unified target's local variant of ssm/gdn_decode_rows_gated_norm into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…d from the qwen4_exp targets Moves the unified target's local variant of ssm/gdn_prefill_conv_post into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…d from the qwen4_exp targets Moves the unified target's local variant of vision/patch_embed_conv3d into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…e, promoted from the qwen4_exp targets Moves the unified target's local variant of vision/triton_pos_embed_interpolate into the catalog with its test, removes the tp1 target's copy and test, and points both targets and every test that imported either copy at the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…le_cached_replay_update, promoted from the qwen4_exp targets Moves the unified target's local variant into the catalog with its test, removes the tp1 target's copy and test, and points both targets at the catalog entry. The tp1 target's gdn_cached_replay_cute test borrows that test's harness by path; it now loads it from the catalog tree. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…, the qwen4_exp variant r44 merged in Merges the unified target's variant r44 into the catalog entry: its measured note, the qwen4_exp shared-expert cells (2 x 640 at 1, 128, 1024 and 8192 rows), the CUDA-graph replay, saturated-gate and wrong-activation tests; the variant and its test are removed. The unified target does not call this entry (its shared experts run silu_and_mul_rows), so the cells are keyed under the tp1 target, which calls it at that width. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…4_exp variant r2 merged in Merges the unified target's variant r2 into the catalog entry: its measured note and test, which generalizes the shared test's helpers to a tensor-parallel rank's 6/1 heads and NVFP4 pools and adds the unified target's cells under its TARGET_CELLS key. ARCHS stays sm_100 + sm_103; the unified cells were driven on sm_103 only. The variant is removed and the unified target's modeling and vision tower import the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…e qwen4_exp variant r4 merged in Merges the unified target's variant r4 into the catalog entry: the guard on a window call handed the whole stack's weights (measured ~360 bf16 ulp RMS off, no error), the notes on the 512-expert windows and tp4's zero-padded 160-wide intermediate, and the test's unified layouts, negative controls and autotuner-cache check. ARCHS stays sm_100 + sm_103; the unified cells were driven on sm_103 only. The variant is removed and the unified target imports the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…from its package reach the catalog Four tests import the module itself (`from <target>.ops.<category> import <entry> as m`) rather than names from it; the promotion's rewrite matched module paths only, so they still named the removed target copies. They now import the catalog modules. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ests no longer call themselves a target's variant Docstrings only: the tests' and four entries' opening sentences still described a local variant of one target; they now name the catalog entry and the targets that share it. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ce, promoted from the qwen4_exp targets Merges the tp1 target's copy and the unified target's variants r17 (the trunk's fused projection row at tp4/tep4 and dep4) and mtp (the MTP layer's [T, 6, 512] q projection) into one catalog entry: 13 cells (tp1's three once), one note over the three call sites, one test whose _build takes every source's spec keys, and every source's graph-capture and gate-discrimination tests. The copies are removed and both targets import the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…qwen4_exp variants r3 and mtp merged in Merges the unified target's variants into the catalog entry: their measured notes, and in the test the unified target's TARGET_CELLS (r3's tp4/tep4 fused row and dep4 cells, mtp's MTP-layer cells: 10 in all), the mRoPE tests parametrized over the four head geometries the sources drove, and the MTP rope's CUDA-graph replay. ARCHS stays sm_100 + sm_103; the unified cells were driven on sm_103 only. The variants are removed and the unified target imports the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ants r1, mtp and vision merged in Merges the unified target's three variants into the catalog entry: one note with their measurements (the corrected DeepSeek o_proj worst case, the misaligned-bias kernel switch, the fp32 bands, the gate sliced from the fused shared-expert weight), and in the test the unified target's 63 cells (trunk 33, MTP layer 16, vision tower 14) under its TARGET_CELLS key, one _compare (fp32 outputs at 1024/512 fp32 ulp, one-column outputs scored over the column) and every source's negative controls. ARCHS stays sm_100 + sm_103; the unified cells were driven on sm_103 only. The variants are removed and the unified target's modeling and vision tower import the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ant r6 merged in Merges the unified target's allgather variant into the catalog entry. The entry gains the variant's 13 cells (hidden 2560, worlds 4 and 2, sizes None) and its note: the two symbols (`allgather` over the MPI session, `allgather_pg` over a ProcessGroup), the per-rank row counts sampled, the int-accumulator limit of the ragged form, and where each cell was driven. The catalog matrix becomes the variant's: cells filtered by the job's world size (a cell without `world` runs at every size), the CUDA-graph sizes the union of the default list and the qwen4 capture list, 48 gathers per decode graph, and the served path's checks run again through the MPI-session symbol. The collected test runs worlds 2 and 4. ARCHS stays sm_100 + sm_103; the qwen4 cells, world 2 and the MPI symbol were driven on sm_103 only. The variant and its matrix are removed; the unified target imports the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ant r5 merged in Merges the unified target's allreduce variant into the catalog entry: the plain sum (`NONE`, exported beside `RESIDUAL_RMS_NORM`) and its four cells at world 4, hidden 2560 (1, 128, 1024 and 8192 rows), with the variant's measurements in the note (which kernel AUTO picks per bucket, bitwise one-shot results, the NCCL band, rank agreement, graph replay and the missing-rank negative control). The catalog matrix gains the variant's `check_none_*` checks and keeps every RESIDUAL_RMS_NORM check at worlds 2, 4 and 8. ARCHS stays sm_100 + sm_103; the NONE cells were driven on sm_103 only. The variant and its matrix are removed; the unified target imports the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…variant r7 merged in Merges the unified target's reducescatter variant into the catalog entry: its cells at hidden 2560, worlds 4 and 2, and its note. The catalog matrix becomes the variant's: a cell runs at its `world`, one without it (the deepseek-r1 / gpt-oss cells, hidden 7168) at world 4, the dep4 layout it was certified at; the CUDA-graph checks hold 48 reduce-scatters per decode graph; the separately capped wedge job is kept. The collected test runs worlds 2 and 4. ARCHS stays sm_100 + sm_103; the qwen4 cells and world 2 were driven on sm_103 only. The variant and its matrix are removed; the unified target imports the catalog entry. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…on left empty Every entry of the unified target is now in the catalog, so its `ops/` package held only `__init__.py` files; the tp1 target's `activation`, `embedding` and `norm` categories are empty the same way. Nothing imports them. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…te CELLS as literal Cell calls in the class body flashinfer_top_k_ragged_transform, qsa_rollback_planes (both halves), ple_gate_conv_inject_norm, ple_ngram_embed, ngram_token_history and ngram_hash built their cells through a helper or a module-level tuple. Each class body now lists Cell(why=..., spec=...) itself; the helpers build only the spec dict. Every cell's why and spec is unchanged. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…t the entry's own CELLS are driven flashinfer_silu_and_mul, fused_qk_norm_rope, thop_attention, cublas_mm, fp4_block_scale_moe_runner and fp4_quantize parametrize test_certified_cells over certified_cells(op, TARGET_CELLS). Each test now also reads the entry's CELLS itself and asserts every one of its specs is among the parametrized cells. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…for the 262144-token score tables At max_seq_len 262144 the unified target's QSA score table is 65536 columns (65568 with MTP's extra pages), past the narrow kernel's MAX_WIDTH, so decode fell back to the qsa_topk_select + qsa_expand_block_indices composition. A wide kernel now serves these tables: a lane-maximum prefilter narrows each row's candidates, compacts them, and selects the block top-k and its token indices in one launch. supports() routes tables up to MAX_WIDE_WIDTH to it; the narrow kernel keeps the core and MTP widths. New cells certify the wide tables bitwise against the reference: relu-sparse, summed-ReLU indexer scores, all-equal ties, and the decode CUDA-graph shapes. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
… QSA index-score kernels; context sparse GQA on row tiles Decode and MTP verify: - The HC mixer's first half runs as one tcgen05 launch up to 256 rows (qwen4_exp_hc_decode_tc_first__decode, a 10-way K split up to 4 rows) in place of cuBLAS split-K, its reduce and the low-rank activation. - The GDN in_proj at up to 8 rows runs on FlashInfer's low-M kernels (low_m_bf16_mm__decode); the verify replay runs on the CuTe kernel at a rank's heads (gdn_cached_replay_cute__verify), launched with PDL after the conv, which signals its dependents (causal_conv1d_triton_update's launch_dependent_kernels). - QSA index scores at decode run on a visible-split grid sized by the device (qsa_paged_index_scores). Context: - Sparse GQA shares row tiles across the rows of a group (qsa_sparse_gqa_rowtile__prefill, qsa_rowtile_attention_cute__prefill), an FP8 pool read through a widened bf16 scratch (qsa_fp8_pages_to_bf16__prefill). Each kernel's cells certify it against its reference. On the AgentX serving configs, against the tree before these changes on the same GB300 node: ll_tp4 +9.3% and ll_tep4 +16.5% tok/s/user, ht_dep4 +10.6% output tok/s/gpu. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…ill host overhead, MTP draft HC on the fused entries - MoE front end: the stack's MoE-side mixer writes the routed experts' NVFP4 activation in its epilogue (`qwen4_exp_hc_up_gate_mean_fp4__moe`), the runner routes from the router logits itself (`fp4_block_scale_moe_renorm_routed__moe`), and the decode router runs on a low-M bf16 GEMM (`low_m_bf16_mm__router`). - Prefill host overhead: `catalog/moe/fp4_block_scale_moe_runner` memoizes the AutoTuner's runner and tactic per static arguments and row bucket (bitwise the op's result); `catalog/ssm/flashinfer_chunk_gated_delta_rule` runs carried calls of at most 64 context tokens in place over the pool (the 2-token snapshot step of chunked prefill). - MTP draft layer: its hyper-connection sites run on the trunk's fused entries (`qwen4_exp_hc_inject_norm`, `qwen4_exp_hc_decode_tc_first`), certified at the draft step's 1 and 32 rows; `test_mtp_draft_hc.py` checks them against the HuggingFace math. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…rologue in one launch; decode projections on a PDL GEMV - Attention decode prologue: the tp1 target's `qwen4_exp_attn_decode_prologue` at one rank's heads (6/1 at tp4/tep4, 24/2 at dep4) over bf16 and e4m3 pools (`ops/attention/qwen4_exp_attn_decode_prologue__decode`): the main and index planes' norm + rope, the packed-QKV copy, the K/V append, the index key and coordinate store and the group compression of a sparse decode attention layer run as one PDL launch. NVFP4 pools keep the composition. - Decode projections: the attention input, attention output and GDN output projections run up to 4 rows at K=1536 on a CuTe PDL GEMV (`ops/gemm/bf16_gemv_pdl__decode`), else on `low_m_bf16_mm__decode` (now also at [4352, 2560] and [2560, 1536]); every dep4 shape and every row count past 8 stays on cuBLAS. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
Wanli-Jiang
changed the title
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR to use temporal
@coderabbitai summary
Description
Adds the modeling_v2 target
hf_fc694b5_orig__sm_103__unifiedfor Qwen3.8-Flash-Next NVFP4 (Qwen4ExpForConditionalGeneration) on GB300 (sm_103), with kernels tuned for the AgentX workload. One target serves a 4-GPU node in three layouts, selected withmodeling_v2: require:moe_expert_parallel_size: 1)enable_attention_dp)Depends on #20937 (the single-GPU target and its ops). The branch carries #20937 (head
1edb5163) merged with main at4f5e469378(e3129fcc75); this PR's changes are the 211 commits after that merge:725c55178d(207)catalog/0edbaf0d78,0c3c81ff4a,410abc76f7,80f149b54fCode:
models/qwen4_exp_for_conditional_generation/hf_fc694b5_orig__sm_103__unified/:modeling.py,weights.py,pool.py,vision.py, andops/, the kernels written for this target.models/qwen4_exp_for_conditional_generation/routing.py: routes this architecture to the target at two and four ranks;models/qwen4_exp/routing.pykeeps routing one rank to the [TRTLLM-16460][perf] Qwen3.8-Flash-Next NVFP4 modeling_v2 target: GB300 Pareto optimizations for ISL 8K / OSL 1K #20937 target.catalog/: the ops the qwen4_exp targets share (attention, comm, gemm, hc, moe, norm, ple, ssm, activation, vision), with their tests. The tp1 target's local copies of the shared ops (13 files) are removed.Features
Each feature below is a contract cell: the engine boots, serves, and scores GSM8K within 3 SE of the in-tree reference; MTP also matches the stock engine's acceptance length.
max_num_turns: 3, GDN snapshot 2 tokens before each prompt end)TestQwen38FlashNextNvfp4ModelingV2::test_auto_dtypeAll 14 cells hold on this branch's model code at
725c55178d; scores are each check's latest passing run. GSM8K protocol for all rows: chat template, few-shot as multi-turn,enable_thinking: false, 512 output tokens. Acceptance runs 30 pass-key prompts of 3k–8k tokens.Performance kernels
Decode and MTP verify:
ops/hc/qwen4_exp_hc_decode_tc_first__decode), in place of cuBLAS split-K, its reduce and the low-rank activation. The MTP draft layer's hyper-connection sites run on the trunk's fused entries (qwen4_exp_hc_inject_norm,qwen4_exp_hc_decode_tc_first).ops/attention/qwen4_exp_attn_decode_prologue__decode). NVFP4 pools keep the composition.max_seq_len: 262144(65568 with MTP's extra pages) run on a wide kernel incatalog/attention/qsa_topk_expand_decode. A lane-maximum prefilter narrows each row's candidates, and the block top-k and its token indices are selected in one launch.catalog/attention/qsa_paged_index_scores).ops/ssm/gdn_cached_replay_cute__verify), launched with PDL after the conv;catalog/ssm/causal_conv1d_triton_updatesignals its dependents.ops/gemm/bf16_gemv_pdl__decode).ops/gemm/low_m_bf16_mm__decode).ops/hc/qwen4_exp_hc_up_gate_mean_fp4__moe).ops/moe/fp4_block_scale_moe_renorm_routed__moe).ops/gemm/low_m_bf16_mm__router).Context:
ops/attention/qsa_sparse_gqa_rowtile__prefill,qsa_rowtile_attention_cute__prefill). An FP8 pool is read through a widened bf16 scratch (qsa_fp8_pages_to_bf16__prefill).catalog/moe/fp4_block_scale_moe_runnermemoizes the AutoTuner's runner and tactic per static arguments and row bucket, bitwise the op's result.catalog/ssm/flashinfer_chunk_gated_delta_ruleruns carried calls of at most 64 context tokens in place over the pool. This covers the 2-token snapshot step of chunked prefill.Each kernel's cells certify it against its reference. The
catalog/changes also serve the #20937 tp1 target, whose suite passes on them.Serving configs for AgentX
Measured with AgentX trace replay through aiperf (
semianalysis_cc_traces_weka_062126_256k, seed 42, 900 s per point, 10 warm-up requests per lane,--connection-reuse-strategy never). The trace's prompts average 95k tokens at concurrency 1 and 71k at 128.ll_tp4ll_tep4ht_dep4timeout_iters: 25,batching_wait_iters: 5), FP8 KV, conversation-affinity routingThroughput is measured on one GB300 node (nvl72d197-T16) at the tuning run's tree. That tree has the same kernels as this branch, with the shared ops still as target-local copies; moving them into
catalog/changes imports only. The same config varies by up to ~7% across nodes. GSM8K is scored at this branch's head on each config.The trace's prompt-prefix reuse (aiperf
prefix_hit_pct) is 92.6–93.6%; the engine's KV-cache hit rate is about 91% at both concurrencies.Option choices were made by same-node A/B on this target before its performance kernels:
timeout_iters: 10,batching_wait_iters: 2keeps TTFT at 3.4 s for +13.7%.tokens_per_block: 64(does not boot),avg_seq_len, branch snapshots,stream_interval: 20,max_num_tokens: 16384(−4.4% at 128), and the PDL / expandable-segments / no-GC environment bundle (−4.9% at 1).Configs (extra LLM API options)
ll_tp4(ll_tep4is the same withmoe_expert_parallel_size: 4):ht_dep4:Test Coverage
tests/unittest/_torch/modeling_v2/models/qwen4_exp_for_conditional_generation/hf_fc694b5_orig__sm_103__unified/: the target's tests, eachops/kernel's cells, andtest_mtp_draft_hc.py(the MTP draft layer's hyper-connection sites against the HuggingFace math). The CPU-only tests are inl0_cpu.yml; the GPU tests are under the modeling_v2 entry.tests/unittest/_torch/modeling_v2/<category>/for the catalog entries.l0_gb300_multi_gpus.yml:accuracy/modeling_v2/test_qwen4_exp_for_conditional_generation.py::TestModelingV2HfFc694b5OrigSm103Unified::test_gsm8k[tp4|tep4|dep4]...::test_mtp3_gsm8k...::test_mtp3_long_context_passkeyaccuracy/test_disaggregated_serving.py::TestQwen38FlashNextNvfp4ModelingV2::test_auto_dtypeValidation at the head (
80f149b54f):The tp1 target's suite and the tests of the changed
catalog/entries last ran at410abc76f7, the head's last change tocatalog/: 1302 passed, 0 failed.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.