[xguannv]

Description

Kimi K3 NVFP4 on Rubin (SM107) as one Modeling V2 target, kimi_k3_nvfp4__sm_107__default, plus the runtime
changes it needs, so that it builds and runs on main by itself.

Each layer below is self-contained in its commits; review per layer.

1. Runtime (5 commits)

Commit	Change
76637bdd60	Parallelize long KV cache block-offset copies (kvCacheManagerV2Utils.cu): bounded chunks per row, several CTAs per long table; layout unchanged.
a1b8638c2d	Reserve 1 GiB outside the PyTorch caching-allocator cap for driver / KV-transfer allocations made after executor setup.
6b20017b37	MLA context attention accepts caller-quantized FP8 Q / K / V (attentionOp.cpp, mlaKernels.cu, TRTLLM backend); BF16 inputs keep the existing path.
1921268140	Share hybrid KV cache pools with speculative drafters: an opting-in target cache manager appends the drafter's layers to its own pools (KV cache manager V2, mamba V2, FP4 MLA, disaggregated transfer). Opt-in for disaggregated serving via TRTLLM_DISAGG_SHARED_DRAFT_KV=1.
f0cbfb8be3	The shared draft KV pre-check of disaggregated_mpi_worker checks a one-sided config (a DWDP context MPI world lists no generation servers) for its own side; the peer handshake rejects mismatched draft layouts.

2. DWDP test discovery fix (1 commit)

1fc5fa945d: DWDP worker URLs are serialized as YAML strings and every serving role publishes to the same discovery
directory when the launcher strips Slurm variables, so the disaggregated frontend finds the context workers.

3. The Kimi K3 SM107 target (7 commits)

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/; the family router
returns it for KimiK3ForConditionalGeneration on SM107. One target serves the four deployments features.py
knows: tp1 (one-GPU DWDP context server), helix8, helix16_moetp4ep4 (ULL) and dep16; the profile is selected
before the target's modules are imported.

Under the default modeling_v2=auto the route selects this target on its own: it claims only SM107, Kimi K3's
exact text shape, the NVFP4 checkpoint, DSpark or no speculative decoding, an NVFP4 or FP8 KV cache, and one of the
four layouts above; anything else raises NoTarget and the built-in Kimi K3 implementation runs. Like the two Kimi
K3 SM100 targets it is not @disable_v2_by_default. The shipped deployment configs set modeling_v2: require.

Commit	Module
e175294fc2	features, routing, core layers (linear, MXFP8 GEMM, norm, MLP, embedding, all-reduce, streams), Lamport collectives, attention-residual norm
cd0433ba94	FP4 / FP8 MLA attention with Helix: context FMHA and split-K context, FP4 MLA decode and packed reduce, fused q path, Helix tail
0b7ec090d1	KDA linear attention: mixer, prefill, MTP decode, fused projection
94063aaf0b	MoE (router, route tail, small-batch MoE, ULL front, tails) and the low-M / mid-M FP8 and MXFP8 GEMMs with push variants
74672aa51d	DSpark speculative decoding: drafter, vocab-parallel head, Markov chain, acceptance, draft KV
270b61f79c	the model: modeling.py, weights.py and the runtime modules (moe_runtime, fp8_weight_read, attn_res, text_config)
533e3177f4	GSM8K accuracy test of the disaggregated deployments, perf-sanity configs (DSpark drafter attention_backend: CUTEDSL explicit, which is what AUTO resolves to for this MLA drafter), CI / QA test lists

The deployments opt into FP8 storage of the BF16 shared and latent MLP weights (KIMI_K3_FP8_WEIGHT_READ_MOE_MLP=1,
the same opt-in as the built-in Kimi K3 implementation): it is a lossy conversion of checkpoint weights, so the
target's default keeps the checkpoint's BF16; the DWDP context servers need the memory it frees. The ULL deployment
caps the context servers' max_tokens at 524288 because GPU 0 of each node also holds the CUDA contexts of that
node's other MPI ranks.

Layout of the target directory:

• Root: the model (modeling.py, weights.py, text_config.py, features.py) and one module per block: layers.py,
  mla.py, kda.py, moe.py, drafter.py, dflash.py, plus moe_runtime.py, attn_res.py, fp8_weight_read.py.
• ops/<category>/: the OpWrapper entries (attention, comm, embedding, gemm, moe, norm, spec, ssm); host glue used
  by a single entry sits next to it as a private module (_fp4_mla_decode.py, _mxfp8_gemm.py, _kda_prefill.py,
  _fla_index.py, _dspark_mla_rope.py).
• kernels/: the CuTe DSL / Triton kernels (kda/, mla/, moe/, spec/ and shared kernel files).
• modeling.py imports the block modules (MLA, drafter) eagerly, after features has selected the deployment, so
  an invalid TRTLLM_K3_ULL3_MLA_GATE_AT fails at import.

How the target follows the Modeling V2 rules:

• No torch custom-op registration. Every kernel the model launches goes through a stateless OpWrapper entry in
  ops/<category>/<entry>.py (attention, comm, embedding, gemm, moe, norm, spec, ssm) that calls the CuTe DSL /
  Triton launcher directly, with a plain-torch reference, is_valid, and CELLS (the shapes K3 runs, each with
  its deployment). Tests drive every cell under validating.
• Few private imports of Modeling V1 code: helpers the target needs are copied in with their license header
  (e.g. kernels/mla/mla_kv_fp8.py, the drafter's split-KV planner, the Helix post-process); the rest go through catalog
  entries.
• Split model file: modeling.py / weights.py instead of one 6k-line file; classes that shared a name with a
  Modeling V1 class are renamed (KimiK3DefaultForCausalLM, K3MoERuntime, K3MLADSparkForCausalLM, ...).

Known gaps (listed so review can decide):

• Entries without cells: k3_fp4_mla_decode and fp4_mla_kv_append (a reference needs the KV cache manager's
  ring / packed-V layouts), the three drafter all-reduce sites (k3_draft_ar_norm, k3_draft_embed,
  k3_draft_gemm_push; their 2-rank test compares them bitwise with the C++ fused all-reduce), and k3_ull_front,
  k3_moe_tail1, k3_shared_persist.
• Catalog bmm_out / cublas_mm declare ARCHS {SM_100, SM_103}; the target calls them on SM107.
• CODEOWNERS: the target falls under the existing kimi_k3_vl owner line.

Test Coverage

Unit tests (tests/unittest/_torch/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/, listed in
l0_gr100; multi-rank ones in l0_gb300_multi_gpus; CPU ones in l0_cpu):

• GB300 (SM103), this code (ocijhb, all Python from the branch): target + test_modeling_v2_contracts.py +
  test_modeling_v2_routing.py + test_disagg_utils.py 1512 passed / 260 skipped / 0 failed (1701850; skips are
  SM107-only kernels); the K3 l0_cpu entries 45 passed / 48 skipped / 0 failed (1702502).
• Rubin (SM107), hecate job 1020125 (same target code before the file consolidation, which only moves definitions): target 1269 passed / 8 skipped / 2 failed, contracts 19 / 0 / 0, routing
  346 passed / 70 skipped (other models' sm_103-only targets). The 8 skips are the context FMHA instrumentation
  checks, which need the internal CuTe DSL (cute.arch.nanosleep). The 2 failures were in the plain-torch
  reference of fp4_mla_gather_mxfp8 (a local reused the name of the context-offset input when the token count
  is not a multiple of 128), fixed in this head; on this head both cells and the rest of ops/attention pass
  (28 / 0, hecate 1020360).

Runtime layers 1–2 (GB300, built on the previous main base): the 14 test files they add or change pass (the cache
transceiver file 269 / 0 with UCX bound to the node's Ethernet NIC); the existing suites of the touched modules
(KV cache, speculative, disaggregated) show no new failures (3 test_kv_cache_v2_vocab_size qwen3_5_hybrid cases fail
identically on main); C++ mlaPreprocessTest 6 / 6 and kvCacheManagerV2CudaContextScopeTest 3 / 3. Their 54
SM107-only cases (test_fp4_mla.py, test_fp4_mla_kv_transfer.py, test_kimi_k3_mla.py,
test_dspark_cute_dsl_attention.py, test_kimi_kda_fused_verify_parity.py) on Rubin (hecate 1020178): 219 passed,
2 failed. The 2 are test_fp4_mla.py::test_fp4_mla_helix_scatter_uses_global_rope_and_local_cache_position
[active/inactive]; that test and the code it covers are unchanged from main (it compares the 128-row padded Q
buffer with the 4 active rows), so it fails on SM107 without this PR as well.

Bitwise identity with the pre-refactor model (same kernels, same launch order; GB300):

• Whole 3-layer K3 model (KDA, KDA, MLA; ctx / gen2 / gen1 phases) in the none / tp1 / dep16 / helix8 /
  helix16_moetp4ep4 profiles and at TP4 over MNNVL: OpWrapper conversion and module split each bit-identical to the
  pre-refactor code (jobs 1813764), MoE layer (M = 1..1024, tp1 / adp1 / TP4) and DSpark drafter (tp1, dep16, TP4,
  helix8, helix16) bit-identical (1814540).
• The final K3 code on this PR's runtime layers against the previous Rubin branch's code on its own runtime: whole
  model (tp1, dep16, helix8, helix16_moetp4ep4, TP4) and MoE layer + drafter bit-identical (ocijhb 1698165).
  Same build, final code against the pre-refactor code: whole model, MoE layer and drafter bit-identical in every
  profile and on every rank (ociaga 1815190).
• This head against the previous Rubin branch's code (flat build, all Python from this branch): whole model in every
  profile and at TP4, MoE layer and drafter bit-identical (ocijhb 1701855).

Accuracy (Rubin, GSM8K strict-match, disaggregated, the deployments in test_accuracy.py): 95.91 (helix8 bs20
DSpark5), 96.13 (dep16), 96.29 (ULL helix16 MoE TP4 x EP4) on the pre-refactor code; this head: TODO (hecate 1021419 / 1021420 / 1021421).

Performance (Rubin, same-allocation ABBA against the previous Rubin branch, three points c383 / c240 / c24):
TODO. The compared build has the same runtime as this PR; one known difference between the arms besides this code:
the previous branch also has a NIXL request-release gate (#21612, not part of this PR) that the compared build and
this PR do not; all three points use the NIXL Python transceiver.

PR Checklist

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