faster inferneece

inference.py

@@ -17,6 +17,7 @@ Usage:
 """
 import argparse
 import inspect
+import time
 import torch
 import random
 import numpy as np
@@ -77,6 +78,7 @@ inputs = tokenizer(prompt, return_tensors="pt").to(
     next(model.context_tower.parameters()).device
 )
 
+t0 = time.perf_counter()
 if args.mode == "ar":
     outputs = model.generate(**inputs, max_new_tokens=args.max_new_tokens, do_sample=False)
 elif args.mode == "mock_ar":
@@ -120,5 +122,37 @@ else:
         generate_kwargs["step_callback"] = step_callback
     outputs = model.generate_mask_diffusion(inputs["input_ids"], **generate_kwargs)
 
-text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+if torch.cuda.is_available():
+    torch.cuda.synchronize()
+elapsed = max(time.perf_counter() - t0, 1e-9)
+
+prompt_len = inputs["input_ids"].shape[1]
+gen_ids = outputs[0][prompt_len:]
+n_new = int(gen_ids.shape[0])
+text = tokenizer.decode(gen_ids, skip_special_tokens=True)
+nfe = getattr(model, "_last_nfe", None)
+
+print("\n" + "=" * 70)
+print("--- Request 1/1 ---")
+print(f"Prompt: {prompt}")
+_nfe_str = f"{nfe} NFE, " if (args.mode == "mask_diffusion" and nfe is not None) else ""
+print(f"Generated ({_nfe_str}{n_new} tokens, {elapsed:.2f}s, {n_new / elapsed:.1f} tok/s):")
 print(text)
+print("=" * 70)
+if args.mode == "mask_diffusion":
+    print("Two-Tower mask-diffusion generation complete")
+    print("=" * 70)
+    print(f"  mode:                 {args.mode}")
+    print(f"  block_size:           {args.block_size}")
+    print(f"  steps_per_block:      {args.steps_per_block}")
+    print(f"  max_new_tokens:       {args.max_new_tokens}")
+    print(f"  num_blocks:           {args.max_new_tokens // args.block_size}")
+    print(f"  temperature:          {args.temperature}")
+    print(f"  top_k:                {args.top_k}")
+    print(f"  confidence_threshold: {args.confidence_threshold}")
+    print(f"  mask_token_id:        {args.mask_token_id}")
+    print(f"  NFE:                  {nfe}")
+    print(f"  wall_clock:           {elapsed:.2f}s")
+    print(f"  throughput:           {n_new / elapsed:.1f} tokens/s")
+    print(f"  model:                {args.model}")
+    print("=" * 70)

modeling_nemotron_h.py

@@ -853,8 +853,12 @@ class NemotronHMOE(nn.Module):
                 expert_output = expert(expert_input)
                 weighted_output = expert_output * expert_weights.unsqueeze(-1)
                 final_hidden_states.index_add_(0, token_indices, weighted_output)
-            else:
-                # Local empty expert: no-op compute that still marks params as used.
+            elif self.training:
+                # Training only: no-op forward on a zero token so DDP/grad hooks
+                # mark every expert's params as "used". It adds exactly 0 (no
+                # biases: relu2(0)=0, down_proj(0)=0), so it's numerically inert.
+                # Skipped at inference, where it would otherwise cost ~100+
+                # pointless expert MLP calls per MoE layer per step.
                 expert_dtype = expert.down_proj.weight.dtype
                 dummy_out = expert(torch.zeros_like(hidden_states[0]).unsqueeze(0).to(expert_dtype))
                 final_hidden_states = final_hidden_states + dummy_out

modeling_nemotron_twotower.py

@@ -564,7 +564,7 @@ class NemotronHTwoTowerForCausalLM(NemotronHPreTrainedModel, GenerationMixin):
             new_conv = comb[:, -d_conv:, :].transpose(1, 2).contiguous()
         return out, new_conv, new_ssm
 
-    def _run_denoiser_step_diffusion(self, block_ids, cache_state, t=None):
+    def _run_denoiser_step_diffusion(self, block_ids, cache_state, t=None, den_cache=None):
         """Diffusion denoiser forward over the FULL block (B, L) in one pass.
 
         Parity with mcore `_run_denoiser_step`:
@@ -601,8 +601,13 @@ class NemotronHTwoTowerForCausalLM(NemotronHPreTrainedModel, GenerationMixin):
             t_repr = self.t_embedder(t_dev)
             t_emb = self.t_block(t_repr)
 
-        # Fresh denoiser cache seeded from context: Mamba S-1 state + full context KV.
-        den_cache = self._build_denoiser_cache_diffusion(cache_state, den_device)
+        # Denoiser cache (context Mamba S-1 state + full context KV). It is
+        # READ-ONLY here and identical for every step within a block, so the
+        # caller should build it once per block and pass it in (avoids cloning +
+        # cuda:0->cuda:1 copying the whole context cache on every NFE). Fall back
+        # to building it if not provided.
+        if den_cache is None:
+            den_cache = self._build_denoiser_cache_diffusion(cache_state, den_device)
 
         hidden = tower.embeddings(den_input)
 
@@ -758,8 +763,14 @@ class NemotronHTwoTowerForCausalLM(NemotronHPreTrainedModel, GenerationMixin):
 
         cache_state = self._build_context_cache(input_ids)
         context_ids = input_ids.clone()
+        nfe = 0  # number of denoiser forward passes (network function evaluations)
 
+        den_device = next(self.denoiser_tower.parameters()).device
         for block_idx in range(num_blocks):
+            # Build the denoiser cache ONCE per block (context is fixed within a
+            # block); reused by every denoising step to avoid per-NFE clone+copy.
+            den_cache = self._build_denoiser_cache_diffusion(cache_state, den_device)
+
             # Initialize fully masked block
             xt = torch.full((B, block_size), mask_token_id, dtype=torch.long,
                             device=device)
@@ -777,7 +788,8 @@ class NemotronHTwoTowerForCausalLM(NemotronHPreTrainedModel, GenerationMixin):
                 t_vec = t_model.expand(B).to(device)
 
                 # Denoiser forward (logits come back on denoiser device, move to xt's device)
-                logits = self._run_denoiser_step_diffusion(xt, cache_state, t=t_vec)
+                logits = self._run_denoiser_step_diffusion(xt, cache_state, t=t_vec, den_cache=den_cache)
+                nfe += 1
                 logits = logits.to(device)
 
                 # p(x0|xt) with constraints
@@ -848,6 +860,8 @@ class NemotronHTwoTowerForCausalLM(NemotronHPreTrainedModel, GenerationMixin):
             if eos_token_id is not None and (xt == eos_token_id).any():
                 break
 
+        # Expose NFE (denoiser forward passes) for reporting, e.g. inference.py.
+        self._last_nfe = nfe
         return context_ids
 
     # ------------------------------------------------------------------