AI 基礎設施
Megatron-LM Expands Prefix-Aware Routing and Dynamic Batching, Enabling Hybrid Model Inference to Share KV Cache and Mamba States
NVIDIA summarized Megatron-LM changes from the first week of September, advancing the data-parallel inference frontend, prefix caching, and variable-length sequence processing toward a more complete serving path. The new features improve state reuse for hybrid Transformer/SSM models, but the company has not yet published throughput, time-to-first-token, or cross-node scaling data.
On September 7, NVIDIA published its [monthly Megatron-LM progress report](https://github.com/NVIDIA/Megatron-LM/issues/7100), summarizing 58 substantive merges completed through September 6. The most notable inference-side change is that the HTTP frontend can now scale with data-parallel replicas, while CPU work such as tokenization no longer blocks the coordinator’s event loop. The router now considers both replica load and prefix-cache hits, making requests with shared system prompts or long document prefixes more likely to be sent to GPUs that already hold the relevant state.
At the lower level, `DynamicInferenceContext` manages in-flight batches and a block-based KV cache, allowing requests of different lengths to be added, paused, or removed at each step. The new version also prevents idle dummy forward passes from clearing the prefix index and adds state slots for hybrid architectures such as Mamba and GatedDeltaNet. In addition to attention-layer KV states, repeated prefixes can reuse SSM and convolution states. NVIDIA’s [API documentation](https://docs.nvidia.com/megatron-core/developer-guide/latest/apidocs/core/core.inference.contexts.dynamic_context.html) indicates that these caches remain constrained by a preconfigured GPU memory budget and block boundaries.
On the training side, support has been expanded for variable-length and packed sequences, as well as runtime context parallelism, allowing each microbatch to use a different context-parallel subgroup. MFSDP also improves communication ordering, prefetching, and persistent parameter and gradient storage. This work is important for long-context and hybrid-attention models because fixed topologies can lead to idle resources or wasted memory when sequence lengths vary substantially.
However, this is a report on main-branch activity rather than a stable release. Some DeepSeek V4, disaggregated prefill/decode, and partial CUDA Graph features are still listed as works in progress. Deployers should verify the merge status of each item individually and measure cache-hit rates, TTFT, memory reclamation after request cancellation, and load skew under different routing coefficients. At present, the number of new features alone cannot be used to infer real-world performance gains.