[None][perf] Kimi-K3 KDA prefill grid sizing and MTP-verify kernels; DWDP cross-step expert prefetch with configurable slots

docs/source/developer-guide/moe-perf-guide.md

@@ -402,7 +402,7 @@ Usage: [`examples/wide_ep/ep_load_balancer/README.md`](https://github.com/NVIDIA
 
 [Distributed Weight Data Parallelism (DWDP)](https://doi.org/10.48550/arXiv.2604.01621) targets layer-wise synchronization and token-exchange overhead in the prefill phase, where differences in sequence lengths and KV-cache hits can leave workers waiting for the slowest rank. DWDP keeps a resident subset of experts on each worker and asynchronously prefetches the remaining expert weights from its NVLink peers. Each worker computes on its own tokens, replacing synchronized token dispatch/combine with weight transfers that overlap with computation.
 
-DWDP can be applied to the prefill phase of a serving system when the context batch is large enough for computation to hide the weight prefetch. Decode typically has too little computation to hide these transfers, so generation keeps the normal EP dispatch. Experimental cost of DWDP is a TTFT regression, from a lower service rate on the context stage, plus extra prefetch memory: each rank stores only its local experts, but adds a double-buffered region of about two layers of remote experts. DWDP cannot run on the same MoE path as EPLB. A disaggregated deployment can use DWDP on context servers and EPLB on generation servers.
+DWDP can be applied to the prefill phase of a serving system when the context batch is large enough for computation to hide the weight prefetch. Decode typically has too little computation to hide these transfers, so generation keeps the normal EP dispatch. Experimental cost of DWDP is a TTFT regression, from a lower service rate on the context stage, plus extra prefetch memory: each rank stores only its local experts, but adds `num_prefetch_slots` prefetch slots (two by default) of about one layer of remote experts each; more slots let the prefetch run further ahead of the compute. DWDP cannot run on the same MoE path as EPLB. A disaggregated deployment can use DWDP on context servers and EPLB on generation servers.
 
 Usage: [`examples/dwdp/README.md`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/dwdp/README.md) and [blog19](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog19_DWDP_Distributed_Weight_Data_Parallelism_for_High_Performance_LLM_Inference_on_NVL72.md).
 

docs/source/developer-guide/telemetry.md

@@ -85,6 +85,7 @@ unset or when the safety sanitizer rejects the runtime value.
 | `dwdp_config.num_experts_per_worker` | `int` | `value` |  |
 | `dwdp_config.num_groups` | `int` | `value` |  |
 | `dwdp_config.num_prefetch_experts` | `int` | `value` |  |
+| `dwdp_config.num_prefetch_slots` | `int` | `value` |  |
 | `enable_attention_dp` | `bool` | `value` |  |
 | `enable_autotuner` | `bool` | `value` |  |
 | `enable_chunked_prefill` | `bool` | `value` |  |

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -737,7 +737,7 @@ def _get_full_device_max_active_clusters(device_id: int, cluster_size: int,
                                              reserved_sms: int) -> int:
         """Return the cached full-device occupancy for a cluster shape.
 
-        ``reserved_sms`` (``TRTLLM_CUTEDSL_SM_RESERVE``, applied inside
+        ``reserved_sms`` (``cute_dsl_utils.sm_reserve()``, applied inside
         ``HardwareInfo.get_max_active_clusters``) is part of the key so a
         reserve set after the first query is not hidden by the cache.
         """

tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py

@@ -53,6 +53,7 @@ def kda_prefill(
     max_seqlen: int = 0,
     seq_order: torch.Tensor | None = None,
     cu_chunks: torch.Tensor | None = None,
+    sm_reserve: int = 0,
 ) -> torch.Tensor:
     """Update unique selected V-first state rows using packed raw Q/K/V.
 
@@ -70,6 +71,11 @@ def kda_prefill(
     upper bound for all replays, or zero to use the total token capacity.
     Small-BH and BT16 decomposition consume aligned, head-contiguous beta
     views directly; the engine requires compact beta storage.
+
+    ``sm_reserve`` SMs, held by a kernel resident for the whole step, are left
+    out of the Small-BH persistent grid; blocks beyond the free SMs would wait
+    behind the whole first wave. The routing above stays on the device SM count
+    so it matches BT16's own router.
     """
     device = q.device
     if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3), (10, 7)):
@@ -147,6 +153,8 @@ def kda_prefill(
     if max_seqlen < 0 or max_seqlen > q.shape[1]:
         raise ValueError("max_seqlen must be zero or a sequence-length bound within total tokens")
     sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    if not 0 <= sm_reserve < sm_count:
+        raise ValueError(f"sm_reserve={sm_reserve} must leave at least one of {sm_count} SMs")
     small_bh_eligible = state_pool.dtype == torch.bfloat16 and q.shape[1] > 1
     if mode == "small-bh" and not small_bh_eligible:
         raise ValueError("Small-BH requires BF16 state and more than one packed token")
@@ -193,7 +201,7 @@ def kda_prefill(
                 state_input=state_pool,
                 state_output=state_pool,
                 state_indices=state_indices,
-                max_active_clusters=sm_count,
+                max_active_clusters=sm_count - sm_reserve,
             )
         return output
     # Torch allocations belong to the CUDA graph's private memory pool during
@@ -252,6 +260,7 @@ def _kda_prefill_fake(
     max_seqlen=0,
     seq_order=None,
     cu_chunks=None,
+    sm_reserve=0,
 ):
     return torch.empty_like(q)