[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

tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py

@@ -131,6 +131,13 @@ class FastPath:
         "multi-request decode FMHA over an over-subscribed KV split, 72..512 rows",
         "construct",
     ),
+    FastPath(
+        "mixed_step_gen_rows",
+        "full-attention generation rows of mixed steps with 9..64 generation requests on "
+        "multi_request_split_kv, context rows on the thop op",
+        "construct",
+        ("multi_request_split_kv",),
+    ),
     # Dense projections and prefetch.
     FastPath("dense_gemms", "dense projections on one named kernel per shape", "construct"),
     FastPath(
@@ -305,6 +312,13 @@ class FastPath:
         ("persistent_experts", "moe_prologue"),
     ),
     FastPath("replicated_embedding", "token embedding gathered from a replicated table", "load"),
+    FastPath(
+        "model_boundary",
+        "decode input boundary (gather, embed norm, residual, layer 0 pre-QKV norm) in one launch; "
+        "last-layer join storing BF16",
+        "load",
+        ("replicated_embedding",),
+    ),
     # Drafter.
     FastPath("drafter_block_conv", "drafter block convolutions as one fused kernel", "draft"),
     FastPath(
@@ -358,6 +372,13 @@ class FastPath:
         "draft",
         ("drafter_block_decode",),
     ),
+    FastPath(
+        "wide_drafter_hops",
+        "the drafter all-reduce hops of 9..56-row (sm107: 9..64-row) block decodes on a wide "
+        "one-shot transport",
+        "draft",
+        ("drafter_hops",),
+    ),
     FastPath(
         "drafter_pdl_chain",
         "drafter and draft head launches releasing their GEMMs early",
@@ -383,6 +404,12 @@ class FastPath:
         "draft",
         ("split_sampler",),
     ),
+    FastPath(
+        "vocab_parallel_head",
+        "the verify head of 16..1024-row decode steps on each rank's own vocabulary shard (no logits all-gather)",
+        "draft",
+        ("split_sampler",),
+    ),
     # Executor.
     FastPath("deferred_input_prep", "a decode step's input writes as one launch", "executor"),
     FastPath("sampler_store", "the sampler's store update and snapshots as one launch", "executor"),

tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py

@@ -71,6 +71,8 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None:
         from .kernels.drafter_hops import hops as _drafter_hops
 
         _drafter_hops.install(getattr(shell.llm, "draft_model", None), shell._mapping)
+        if plan.enabled("wide_drafter_hops"):
+            _drafter_hops.install_wide(getattr(shell.llm, "draft_model", None), shell._mapping)
     if plan.enabled("drafter_pdl_chain"):
         # After the drafter_fixed_gemms GEMM methods and the drafter_hops hops:
         # the entry triggers rely on both.
@@ -82,7 +84,12 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None:
     if plan.enabled("split_sampler"):
         from .kernels.split_sampler import target_head as _target_head
 
-        _target_head.install(shell.llm, shell._mapping, plan.enabled("lm_head_inbox"))
+        _target_head.install(
+            shell.llm,
+            shell._mapping,
+            plan.enabled("lm_head_inbox"),
+            plan.enabled("vocab_parallel_head"),
+        )
         if plan.enabled("wide_split_sampler"):
             _target_head.widen(shell.llm)
     # Last: the arena compaction derives its worker class from whatever class every