Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models

Thank you for recognizing the “useful insights,” “strong results,” and “practical significance” of our work, as well as for your constructive comments! We have addressed all of your questions below:

1. Novelty of our techniques

We respectfully emphasize that the key contribution of this work lies in identifying critical factors affecting SLM real-device latency and offering generalizable insights and methodologies for latency-aware SLM design and training.

Specifically, while depth–width ratios and operator choices are known design factors, their impact on real-device latency and how to make effective trade-offs remain unclear, which are addressed by this work. Additionally, our weight normalization eliminates the training and inference overhead of nGPT (citation [28]) while achieving comparable or better accuracy, as analyzed in Section 2.3 of the main paper and Section 6.1 of the supplementary material.

As Reviewer ictF also noted, our work is “a decent attempt towards latency-optimal SLM” and has the potential to “produce significant community impact.”

2. Variance across runs

Since we adopt ≥100B token training for all experiments that report accuracy, we find that the results are stable across runs. For example, we rerun the experiments in Table 4 of our manuscript three times and show all results below: the PPL variation across runs is within 0.1, and the accuracy difference is within 0.16%, leading to consistent conclusions. Therefore, we didn’t include error bars in the manuscript.

Note: Acc is averaged over 8 commonsense reasoning tasks.

3. Clarifications on the search cost and the search focus

For the search cost, we clarify that each candidate architecture is trained for only 10B tokens, which takes <4 hours using 16 A100 GPUs, as mentioned in Section 5.2 of our supplementary material. This follows the “realistic training budgets” you suggested. (The 100B token setting you referred to applies only to the architecture ablation in Tables 1 and 3 of our manuscript, not to the search process.)

Regarding the search focus, our goal is to use evolutionary search to extend manual exploration of hybrid operators into an automated pipeline for obtaining more insights on LM operator combinations, rather than to build a comprehensive NAS framework. While our framework can support finer-grained searches with minimal changes, we deliberately avoid irregular architectures (e.g., with layerwise-varying #heads) as they may complicate deployment. Following your suggestion, we will explore larger design spaces in future work.

4. Benchmarks on more diverse tasks: Math/Coding/Long Context

Thank you for the question! We have added more benchmarks:

(a) Math / Coding

Since the paper submission, we have continued pretraining our models on an additional 2T tokens and benchmarked against baselines across 14 tasks, including Commonsense Reasoning (MMLU, Lambda, PIQA, ARCC, ARCE, HellaSwag, Winogrande, OBQA), Math (GSM8k, MathQA), and Coding (HumanEval, HumanEval-Plus, MBPP, MBPP-Plus).

As shown below, our Fast-SLM-2.7B achieves SOTA average accuracy among all models. Specifically, Fast-SLM-2.7B outperforms LLaMA3.2-3B/Qwen2.5-3B by 9.24%/0.90% in average accuracy with 1.46x/1.67x speed-ups, respectively. We will include these updated results in the final version.

(b) Long-context

To evaluate long-context tasks, we finetune Fast-SLM-2.7B (pretrained with a 4k sequence length) on 15B tokens using a 24k sequence length. We evaluate on the Needle-In-A-Haystack (NIAH) task from RULER, using 4k~32k input lengths. For reference, we also include results of h2o-danube3-4b (pretrained with an 8k sequence length). Both models use the NTK-based RoPE extension.

As shown above, we observe that:

(1) Fast-SLM-2.7B achieves strong NIAH scores, including reasonably good performance at the 32k input length;

(2) Compared to the pure Transformer model h2o-danube3-4b, Fast-SLM achieves better scores on inputs ≥16k in length. These results highlight the promise of Fast-SLMs in handling long-context inputs. A fair comparison with more models requires aligning long-context training settings and we will include evaluations on more long-context benchmarks in the final version, given more time.

We also refer to recent large hybrid LMs, e.g., Nemotron-H, MiniMax-M1, and Jamba 1.6, that replace attention with efficient linear alternatives, significantly reducing cache usage while maintaining strong long-context performance. This underscores the potential of hybrid models for long-context tasks.

5. Ablation study on deployment settings

We first clarify that CUDA Graph has become a standard component in nearly all popular LM deployment pipelines (e.g., TensorRT-LLM, vLLM, SGLang). This is because LMs are significantly memory-bound during decoding, and without CUDA Graph, CPU overhead can dominate the latency. CUDA Graph is indispensable for eliminating CPU overhead and ensuring reliable profiling of LM decoding.

Additionally, we integrated TensorRT-LLM’s attention kernel to enable CUDA Graph support for attention operators, as the vanilla attention from HuggingFace Transformers cannot be captured by CUDA Graph. The reason we did not directly adopt TensorRT-LLM is that it does not support linear attention like DeltaNet. Therefore, we combined these components to build a unified decoding pipeline capable of running hybrid models with arbitrary operators using CUDA Graph. This infrastructure is also a minor contribution of our work to the community.

Furthermore, following your suggestion, we have added an ablation study for each component:

The table shows that:

(1) The use of CUDA Graph can lead to 5x~7x gaps in decoding latency due to CPU overhead. This indicates that decoding with dominant CPU overhead, especially when it can be eliminated, is not a reliable profiling setting for GPU platforms;

(2) Without CUDA Graph, the latency reduction of Fast-SLM is smaller, as linear attention incurs more extensive CPU workloads than standard attention (according to github issues in the official Mamba repo), further amplifying CPU overhead.

(3) Further increasing the generation sequence length or batch size can lead to even greater speed-ups for our models.

Regarding non-GPU platforms, they fall outside the scope of this work. However, readers may refer to the analysis in our work to profile their target platforms and derive device-specific SLMs accordingly.

6. Discussions on safety issues with SLMs

Thank you for the suggestion! First, prior to the public release, we will evaluate our model on safety benchmarks, including the Aegis AI Content Safety Dataset and Garak, to ensure its safety. Second, and more broadly, we agree that on-device SLMs lower the barrier to misuse and require more safety-aware tuning to ensure reliable deployment. We will include this discussion in the final version.

7. Tool release for the depth-width scaling law

We will release the codebase for deriving the depth–width scaling law, along with our search pipeline and deployment workflow. We hope these tools will help the community develop SLMs tailored to their specific needs.

8. Benefits of Fast-SLM under small budgets

The benefits of Fast-SLM hold even under smaller training budgets. As shown in Tables 3/5 of our manuscript and in the ablation study provided to Reviewer eYGu, Fast-SLM and its individual components improve accuracy or efficiency with ≤100B tokens.

Moreover, Fast-SLM’s use of efficient operators enables faster training and fine-tuning, making it well-suited for dedicated downstream tasks.

9. Other potential risks in cache reduction

As noted in Bullet 4 of our response, Fast-SLM achieves higher accuracy across diverse tasks (e.g., commonsense reasoning/coding) compared to baselines, and maintains reasonably good long-context performance. This suggests that, for common downstream tasks, cache reduction does not hinder accuracy. Similar trends are seen in recent large hybrid LMs, e.g., Nemotron-H, MiniMax-M1, and Jamba 1.6.

Dear reviewer, please read the rebuttal and participate in discussions with authors. Thank you

Thank you for describing our work as “a decent attempt towards latency-optimal SLM” and recognizing its potential to “produce significant community impact,” as well as for your constructive comments! We will open-source all of our search/training recipes and final models upon acceptance. We have addressed all of your questions below:

1. Search depth-width ratios on Llama and transfer to hybrid models

(a) The reason for decoupling depth–width ratio search from operator search

Thank you for the question! First, we agree that jointly searching both operator combinations and depth–width ratios has greater potential to yield a strong architecture, and our codebase can support this without any modification. That said, we currently decouple depth–width ratio search from operator search for two main reasons: (1) We aim to highlight the insights from each design factor independently, with the other held fixed, so that the resulting findings can better benefit the community and inspire future work; (2) As you also noted, a large design space can make the search more difficult and time-consuming to converge.

Although this work focuses on providing insights into both design factors rather than proposing a NAS framework, we agree that a joint search has greater potential and will try it for future work given more time after the rebuttal period.

(b) The scaling law derived from Llama provides good starting points for selecting the depth–width ratio of hybrid models

Additionally, we would like to mention that the scaling law derived from Llama serves as a useful guide for selecting the depth–width ratio of hybrid models. In practice, we first build the scaling law using Llama models and then use it to predict promising candidates for our hybrid models, selecting the best one based on profiling results.

Specifically, the Llama-derived scaling law suggests that our Fast-SLM-960M should use 12 blocks, where each block includes one (linear) attention and one FFN. We benchmark Fast-SLM-960M-12b (our final adopted configuration) against its 8-block and 16-block variants under an iso-latency setting, i.e., all models are scaled to a comparable latency by adjusting the width.

As shown in the table below, Fast-SLM 12-block achieves the best PPL and accuracy across all three settings, suggesting that the depth–width ratio predicted by the LLaMA scaling law provides a strong starting point. Note that our goal with the scaling law is not to predict a single best setting, but rather to (1) emphasize the general trend of this efficiency-critical design factor, and (2) quickly obtain a near-optimal depth–width ratio under a target latency as a good starting point.

Note: Latency is measured on an A100 with an output length of 8k and batch size of 1. CR accuracy is averaged over 8 commonsense reasoning tasks.

2. Data points in Table 2 and Figure 5

Thank you for pointing out this potential confusion! As mentioned in Lines 206–207 of our submitted manuscript, our search metric is the model with the lowest PPL under a target latency. Accordingly, the final model we selected corresponds to the one with lowest PPL under approximately 18.3s latency in Figure 5 (which is annotated by a green dot).

The latency values in Figure 5 are obtained by summing the profiled latency of each operator from the LUT. The 17.71s shown in Table 2 refers to the same model, but with latency measured in an end-to-end manner, which results in a slightly lower absolute value. Although there is a small discrepancy between LUT-based and end-to-end latency measurements, we find that summing the profiled latencies from the LUT preserves the relative latency rankings across different models, particularly when the depth is fixed. This LUT-based approach is also widely adopted in prior NAS works such as FBNet [1]. We will clarify this point in the final version.

3. The importance of the aspect ratio to SLMs/LLMs

This is a great question! We have conducted additional analysis on this and plan to emphasize it in the final version. We completely agree that the impact of aspect ratio on both latency and accuracy strongly depends on model size, as elaborated below.

(1) In terms of latency, based on our fine-grained profiling on GPUs, the reason why aspect ratios matter for SLMs is that when the model is relatively small, GPU utilization tends to be low with relatively large overheads for kernel launch, making small and large kernels exhibit similar latency. In such scenarios, reducing the number of kernels is key to improving latency, and shallow-wide models can result in fewer but larger kernels. However, for LLMs with much larger kernels, GPU utilization is higher and kernel launch overhead becomes less significant, making aspect ratios less critical.

To demonstrate this, we compare two LLaMA models with 12 and 24 blocks, respectively. We vary their widths to equalize their latency and record the number of parameters, as more parameters may indicate potentially better accuracy. As shown in the table below, with a lower target latency and smaller model sizes, shallow-wide models can yield up to 10x more parameters under the same latency. For larger models, however, the additional parameters enabled by shallower models diminish, potentially resulting in smaller accuracy gains (or even worse accuracy). This indicates that aspect ratio has a greater impact on SLMs.

(2) In terms of accuracy, as observed in Figure 2 of our submitted manuscript, as model size increases, the accuracy gaps among different depth–width ratios gradually narrow. This also indicates that aspect ratio matters more for SLMs than for LLMs.

In summary, the above analysis suggests that our techniques and insights regarding depth–width ratios are particularly well-suited for SLM design. We believe this analysis strengthens our paper by identifying when these design factors matter most and when our techniques should be adopted. We will follow your suggestion and add further discussion in the final version.

4. How to construct the LUT

We followed common practice [1] to construct the LUT by profiling the latency of each operator, where the operator type and its width (i.e., dimensionality) define an operator, under the target deployment setting. Specifically, we profile Attention, Mamba2, DeltaNet, and FFN with various widths under the target setting and record their latencies in a table. For a model architecture with an arbitrary combination of these operators, we approximate the overall model latency by summing the latency of each operator. We find that this approach preserves the relative latency rankings across different models, particularly when the depth is fixed. We will clarify this in the final version.

[1] “FBNet: Hardware-Aware Efficient ConvNet Design via Differentiable Neural Architecture Search”, B. Wu et al., CVPR 2019.

Thank you for recognizing the insights our work offers to the community, as well as for your constructive comments! We have addressed all of your questions below:

1. Fast-SLM failed to reach SOTA

We respectfully clarify that accuracies on certain tasks are heavily influenced by the pretraining data. For example, the Qwen series achieves notably high scores on factual knowledge tasks like MMLU due to their proprietary pretraining corpora. Therefore, in the absence of access to the pretraining data of these models, we believe the average accuracy across diverse tasks offers a more meaningful comparison than single-task performance.

Additionally, since the paper submission, we have continued pretraining our models on an additional 2T tokens (totaling 4T tokens) and benchmarked them against baselines across 14 tasks, including Commonsense Reasoning (MMLU, Lambda, PIQA, ARCC, ARCE, HellaSwag, Winogrande, OBQA), Math (GSM8k, MathQA), and Coding (HumanEval, HumanEval-Plus, MBPP, MBPP-Plus).

As shown in the table below, our Fast-SLM-2.7B achieves SOTA average accuracy among all models. Specifically, Fast-SLM-960M surpasses Qwen3-0.6B by 3.77% in average accuracy while being 1.75× faster. Fast-SLM-2.7B outperforms LLaMA3.2-3B and Qwen2.5-3B by 9.24% and 0.90% in average accuracy, respectively, with 1.46x and 1.67x speed-ups. We will include these updated results in the final version and open-source our model.

2. Ablation experiments on Fast-SLM’s components

Thank you for the suggestion! We believe this analysis can strengthen our paper. Following your suggestion, we conducted the following ablation study: Starting from a 500M LLaMA model with 24 attention blocks, we progressively reach our final Fast-SLM through the following steps: (1) tune the depth/width ratio, (2) enhance it into Fast-SLM’s hybrid architecture, (3) add meta tokens, and (4) add weight normalization.

As shown in the table below, we observe that (1) tuning the depth/width ratio (into a shallow-wide model) reduces latency even though the parameter count is almost doubled, leading to better accuracy; (2) enhancing it into Fast-SLM’s hybrid architecture boosts task accuracy while significantly improving efficiency, e.g., a 3.2x latency reduction for (bs=32, output len=8k) generation; (3) adding meta tokens and weight normalization further improves task performance, yielding a +2.27% accuracy gain in total.

We will include this analysis in the final version of our paper.

Note: Acc is averaged over 8 commonsense reasoning tasks.

3. More analysis of operator combinations

As a recap of the context, for the 500M architecture ablation in Table 1 of our submitted manuscript, we find that adding Mamba2 with other operators can result in a 0.43~1.12 PPL reduction. For example, combining DeltaNet and Mamba2 leads to a 0.5/0.99 PPL reduction and a +0.2%/+0.3% accuracy improvement over pure DeltaNet and pure Mamba2, respectively. We humbly clarify that this improvement is non-trivial, as the reported gains of emerging operators over prior baselines are often of a similar scale. For example, DeltaNet [1] improves over its Mamba baseline by -0.15 PPL / +0.3% accuracy (see their Table 1 at 340M scale).

In addition, we have extensively experimented with different operator combinations, and in general, these hybrid architectures do not consistently outperform the pure models, often resulting in comparable PPL. The benefit of adding Mamba2 is a particularly interesting observation that we have consistently seen, and we highlight this aspect to provide insights to the community.

It is also worth noting that the goal of this combination study is to identify promising operators that are efficient and complementary to one another to serve as a search space for evolutionary search, rather than to find the single best hybrid architecture. Following your suggestion, we will include the comprehensive exploration of different operator combinations we did in the appendix of the final version.

[1] “Parallelizing Linear Transformers with the Delta Rule over Sequence Length”, S. Yang et al., NeurIPS 2024.

4. Reasons for choosing these two key factors

Thank you for the question! We also plan to emphasize this in the final version.

In general, there are two critical efficiency metrics for SLMs (corresponding to two use cases): (1) inference latency under batch size 1, and (2) generation throughput with the maximum batch size without out-of-memory.

For case (1), GPU utilization is often low with large overheads for kernel launch, making small and large kernels exhibit similar latency. In such scenarios, reducing the number of kernels is key to improving latency. Therefore, tuning the depth–width ratio, where shallow-wide models result in fewer but larger kernels, is crucial for optimizing batch size = 1 latency.

For case (2), GPUs are better saturated and become compute-bound, so adopting compute-efficient operators is critical. Building hybrid operators with efficient linear attention is an effective approach. (These hybrid operators can also help in case (1) by reducing latency per kernel.)

By optimizing these two macro-level design factors, our Fast-SLM-2.7B achieves efficiency across a wide range of deployment scenarios. For example, it is 1.67x / 3.34x faster than Qwen2.5-3B in case (1) and case (2), respectively. There may be other finer-grained or minor design factors, but we believe the above insights can also help guide those decisions. We will highlight this discussion in the final version.

Dear reviewer, please read the rebuttal and participate in discussions with authors. Thank you

Thank you for your constructive comments! We have addressed all of your questions below:

1. The generalizability of the architectural findings on other hardware platforms

Thank you for the suggestions! Following your suggestion, in addition to the data-center GPU A100 reported in our submitted manuscript, we have also verified our findings on (1) a consumer-level GPU, the RTX 3090, and (2) a workstation GPU, the RTX A5000. We will include results on other edge GPUs given more time, after obtaining access to them post-rebuttal.

a. Depth-width trade-offs

To demonstrate the importance of depth–width trade-offs and the advantages of shallow-wide models, we fix the model size of a LLaMA model to 500M / 3B and vary its depth and width. We then report the latency on the RTX 3090 and RTX A5000 (batch size = 1, input length = 128, output length = 8000).

As shown in the table below, shallow-wide models consistently exhibit lower latency across both platforms and model sizes, aligning with the observations in Section 2.1 of our submitted manuscript.

b. Efficiency of different operators

We benchmark the efficiency of different operators on the RTX 3090 and RTX A5000 (batch sizes = 1 and 16, input length = 128, output length = 8000), following the setting in Figure 4 and Section 2.2 of our submitted manuscript. All models are 500M.

As shown in the table below, Mamba2 and DeltaNet generally stand out as the most efficient operators across settings and platforms, echoing the observations in Section 2.2 of our submitted manuscript. Considering the necessity of attention for recall-intensive tasks, we adopt Attention, Mamba2, and DeltaNet as candidates in our architecture search in Section 2.2.

c. Benchmark the final model

We further provide benchmarks of the final delivered models on the RTX 3090 and RTX A5000 (batch sizes = 1 and 16, input length = 128, output length = 8000). As shown in the table below, Fast-SLM-2.7B continues to stand out with the lowest latency compared to other models of similar size. Specifically, Fast-SLM-960M and Fast-SLM-2.7B achieve 3.3x and 1.9x speed-ups over Qwen3-0.6B and Qwen2.5-3B, respectively, on the RTX 3090 with batch size = 16.

In summary, the key conclusion is that all findings remain valid across different hardware platforms. In general, the low GPU utilization during language model decoding highlights the need for fewer kernel launches, and the quadratic cost of attention motivates the adoption of more efficient operators with linear complexity. These two factors are broadly applicable across GPU platforms, reinforcing the importance of tuning depth–width ratios (as shallow-wide models result in fewer but larger kernels) and leveraging efficient/hybrid operators.

2. Clarifications on the depth-width scaling law

We first respectfully clarify that the depth–width scaling law is hardware-agnostic, meaning the fitted exponents are determined solely by the model. Specifically, as shown in Equation 1 of our submitted manuscript, the scaling law predicts the model loss based on depth, width, and data. In practice, given a target latency budget and deployment setting, the optimal depth–width ratio can be identified by profiling a range of configurations and selecting the one that satisfies the latency constraint while minimizing the loss predicted by the scaling law. For models other than Llama, given more time after the rebuttal period, we can follow your suggestion to train a series of models such as MoE and fit the scaling law to observe how the exponents vary.

In addition, we would like to mention that the scaling law derived from Llama can serve as a useful guide for selecting the depth–width ratio of other models. For example, the scaling law on Llama suggests that our Fast-SLM-960M should use 12 blocks, where one block includes one (linear) attention and one FFN. We benchmark Fast-SLM-960M-12b (the adopted configuration) against its 8-block and 16-block variants under an iso-latency setting, i.e., all models are scaled to a comparable latency by adjusting the width.

As shown in the table below, Fast-SLM 12-block achieves the best PPL and accuracy across all three settings, suggesting that the depth–width ratio predicted by the LLaMA scaling law provides a strong starting point. Note that our goal with the scaling law is not to predict a single best setting, but rather to (1) emphasize the general trend of this efficiency-critical design factor, and (2) quickly obtain a near-optimal depth–width ratio under a target latency as a good starting point.

Note: Latency is measured on an A100 with an output length of 8k and batch size of 1. CR accuracy is averaged over 8 commonsense reasoning tasks.