@@ -4,7 +4,7 @@ description: "Fine-tune DeepSeek-V4.1-Flash with NeMo AutoModel using checked-in
 slug: model-coverage/large-language-models/deepseek-ai/DeepSeek-V4.1-Flash
 ---
 
-[DeepSeek-V4.1 Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash)
+[DeepSeek-V4.1-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash)
 uses a 40-layer causal encoder/decoder Mixture-of-Experts (MoE) backbone with
 384 routed experts, one shared expert, and six selected experts per token. Its CSA2 attention shares
 compressed KV and index state between Full, Reuse, and Reindex layers.
@@ -41,7 +41,7 @@ The configuration requests 100 updates with validation every 50 updates,
 activation checkpointing, and reshard-after-forward. It uses no gradient
 accumulation or additional expert FSDP shard axis, and requests a one-hour
 allocation. W&B logging is disabled by default; set `wandb.enable: true` to
-enable online logging. TE FusedAdam retains FP32 master-weight precision
+enable online logging. Transformer Engine (TE) FusedAdam retains FP32 master-weight precision
 through int16 remainders for BF16 parameters.
 
 <Note>
@@ -57,6 +57,18 @@ The model and recipe use `attn: tilelang`, `linear: torch`,
 Set `model.backend.rms_norm: te` to use Transformer Engine RMSNorm for the
 text backbone, attention projections, compressor, and indexer. This preserves
 parameter names and dtypes; the default retains eager FP32 normalization.
+Set `model.backend.compile_hc: true` to run the single-pass mHC
+coefficient projection, stream collapse and stream expand through
+`torch.compile` (one fused kernel each instead of four to five FP32 passes;
+the eager `expand` materializes a `[batch, sequence, streams, streams, hidden]`
+FP32 temporary), and `model.backend.compile_norm: true` to compile the FP32
+RMSNorm chain the same way. Both are opt-in: compiled numerics are allclose to
+eager, not bitwise-identical (the 12-layer mini benchmark reproduced the eager
+loss to four decimals). The cores are specialized per shape: a new sequence
+length (or grad mode) compiles them again, so keep them off for runs with many
+distinct packed lengths.
+The Sinkhorn normalization keeps its TileKernels or torch backend. Both knobs are
+backend fields: they select kernels and are not written to `config.json`.
 TileLang and TileKernels are the optional dependencies also used by DeepSeek V4.
 HybridEP provides both token dispatch and expert-output combine. Router
 correction biases remain fixed during fine-tuning through
@@ -105,7 +117,8 @@ gradient-sized norm temporaries for contiguous CUDA owner shards. Set
 The GB200 recipes place 748B trainable parameters on 64 GPUs, about 117 GB per
 GPU for BF16 weights, BF16 gradients, BF16 moments, and int16 master
 remainders. On 80 GB H100 parts the same recipes need twice the ranks: with
-`ep_size: 64` on 128 or more GPUs the mesh derives an additional `ep_shard`
+remainders. On 80 GB H100 parts, the same recipes need twice the ranks: with
+`ep_size: 64` on 128 or more GPUs, the mesh derives an additional `ep_shard`
 FSDP axis over the routed experts, which halves the per-rank expert state.
 The variants below were validated on 8-GPU H100 nodes with InfiniBand, 100
 updates each, using the NGC `nemo-automodel:26.08` container plus this
@@ -130,7 +143,7 @@ Three points apply to every H100 variant:
 - **Set `NCCL_PXN_DISABLE=1`, or `NCCL_NET_PLUGIN=none NCCL_NET=IB`,** on
   multi-node InfiniBand systems. With the container's default NCCL network
   plugin and PXN both active, the Engram row-owner all-to-all received
-  byte-shifted token IDs from the first node's ranks and the recipe stopped in
+  byte-shifted token IDs from the first node's ranks, and the recipe stopped in
   `_validate_received_ids` during the first forward, with NCCL 2.30.5 and
   2.30.7 alike, and with both HPC-X 2.50 plugins the image ships (the default
   Spectrum-X `spcx` and IBext). Either switch alone avoids it: disabling PXN
@@ -144,7 +157,7 @@ Three points apply to every H100 variant:
   true`, torch 2.13 hits an activation-checkpoint recompute mismatch on the
   expert all-gather copy-out ([issue #3958](https://github.com/NVIDIA-NeMo/Automodel/issues/3958)).
   `moe.reshard_after_forward: false` avoids it at about 8.5 GB per rank and
-  becomes unnecessary once [PR #3961](https://github.com/NVIDIA-NeMo/Automodel/pull/3961)
+  becomes unnecessary after [PR #3961](https://github.com/NVIDIA-NeMo/Automodel/pull/3961)
   is merged.
 - **Point the JIT caches at node-local storage** (`HYBRID_EP_CACHE_DIR`,
   `TILELANG_CACHE_DIR`, `TRITON_CACHE_DIR`). HybridEP resolves its cache from
@@ -185,7 +198,7 @@ preserves released parameter names, required FP32 weights, and logical Engram
 rows while removing the original quantization metadata. Quantized re-export
 is not supported.
 
-Recipe initialization uses the NeMo AutoModel wrapper's internal Checkpointer and DCP
+Recipe initialization uses the NeMo AutoModel wrapper's internal Checkpointer and PyTorch Distributed Checkpoint (DCP)
 path to load the base checkpoint into the distributed model.
 
 Text inputs use unpacked `input_ids [batch, sequence]`, right padding, and
@@ -196,7 +209,7 @@ straight-through gradients; there is no quantization toggle.
 
 The released vision tower is enabled by the default model configuration and
 disabled by the HellaSwag recipe. `DeepseekV41Processor` supports ordinary
-text and image conversations and save/reload. Image inputs require unpacked
+text and image conversations and saving and reloading. Image inputs require unpacked
 `input_ids [batch, sequence]`, `pixel_values [all_patches, 3, patch_size, patch_size]`,
 `image_grid_hws [images, 2]`, and `vision_token_types [batch, sequence]`.
 The image encoder and projector reuse DeepSeek V4 modules; image delimiters
@@ -210,7 +223,7 @@ covered all 40 layers, 4,096 positions, and all 129,280 vocabulary entries
 against the pinned checkpoint's unchanged official `inference/` implementation.
 Official TP1 was compared with native EP4 and 4 Engram owners using
 `torch_mm`, native HybridEP dispatch and combine, and TileLang. Dense and expert
-reference weights were decoded for BF16 compute, with the original cache/index quantization retained;
+reference weights were decoded for BF16 compute, with the original cache and index quantization retained;
 this was not a default FP8/FP4 GEMM comparison.
 
 The native execution streamed one production block at a time while carrying
 
@@ -485,9 +485,14 @@ class BackendConfig:
             mixing chain is compiled as well. Compiled numerics are allclose to eager but
             not bitwise-identical.
         compile_norm: torch.compile the fp32 RMSNorm chain of models that opt in
-            (currently Kimi K3), fusing cast/pow/mean/rsqrt/mul into one kernel.
+            (currently Kimi K3 and DeepSeek V4.1), fusing cast/pow/mean/rsqrt/mul into one kernel.
             Same lazy once-per-process pattern as ``compile_situ``; numerics are
             allclose to eager but not bitwise-identical.
+        compile_hc: torch.compile the mHC (hyper-connection) cores of models that opt in
+            (currently DeepSeek V4.1): coefficient projection, stream collapse and stream expand,
+            one fused kernel each, static shapes (a new sequence length compiles them again; past
+            dynamo's recompile limit a core runs eager). Same lazy once-per-process pattern as
+            ``compile_situ``; numerics are allclose to eager but not bitwise-identical.
         shared_expert_overlap: run the shared experts of opted-in MoE models (currently Kimi K3)
             on a side CUDA stream so their GEMMs overlap the expert-parallel dispatch / combine
             communication of the routed path; numerics unchanged. Default False.
@@ -552,10 +557,14 @@ class BackendConfig:
     # Kimi K3): fused whole-tensor weighted-SiTU forward/backward, the dense SituAndMul core
     # and the attn-res chain. Numerics are allclose to eager, not bitwise-identical. Default False.
     compile_situ: bool = False
-    # When True, torch.compile the fp32 RMSNorm chain of opted-in models (currently Kimi K3),
-    # same lazy once-per-process pattern as compile_situ. Numerics are allclose to eager,
+    # When True, torch.compile the fp32 RMSNorm chain of opted-in models (currently Kimi K3 and
+    # DeepSeek V4.1), same lazy once-per-process pattern as compile_situ. Numerics are allclose to eager,
     # not bitwise-identical. Default False.
     compile_norm: bool = False
+    # When True, torch.compile the mHC projection / collapse / expand cores of opted-in models
+    # (currently DeepSeek V4.1), static shapes, same lazy once-per-process pattern as compile_situ.
+    # Numerics are allclose to eager, not bitwise-identical. Default False.
+    compile_hc: bool = False
     # When True, models that opt in (currently Kimi K3) run their shared experts on a side CUDA
     # stream, launched before the routed-expert path and joined after it, so the shared-expert
     # GEMMs overlap the expert-parallel dispatch / combine communication (Megatron-Core's
 
@@ -183,6 +183,9 @@ def __init__(
         self.hc_mult = hc_mult
         self.hc_sinkhorn_iters = hc_sinkhorn_iters
         self.hc_eps = hc_eps
+        # "torch": eager fp32 mHC math (released reference); "compile": the projection, collapse and expand
+        # cores wrapped once per process with torch.compile (layers.compile_hc_cores) — allclose, not bitwise,
+        # static shapes. Persisted by save_pretrained like every config field.
         self.engram_layer_ids = [1, 14] if engram_layer_ids is None else list(engram_layer_ids)
         self.engram_num_embeddings = (
             ([384006168, 384016682] if self.engram_layer_ids else [])
 
@@ -30,6 +30,100 @@
 from nemo_automodel.components.models.deepseek_v41.config import DeepseekV41TextConfig
 
 
+# ----------------------------------------------------------------------------------------- fp32 cores
+# The mHC coefficient projection, the stream collapse/expand and the fp32 RMSNorm each run several full fp32 passes
+# in eager mode, and ``expand`` materialises a ``[batch, sequence, streams, streams, hidden]`` fp32 temporary before
+# its stream sum. The arithmetic lives in the plain functions below so ``torch.compile`` can fuse each into one or
+# two kernels: ``BackendConfig.compile_hc`` wraps the three mHC cores, ``BackendConfig.compile_norm``
+# the RMSNorm core (same lazy once-per-process pattern as Kimi K3's ``compile_norm``). The modules call the
+# module-level dispatchers, so instances built before the hook ran still pick up the compiled core. Compiled
+# numerics are allclose to eager, not bitwise-identical. The mHC cores are compiled for static shapes only: with
+# dynamic shapes inductor keeps the broadcast temporary of ``expand`` and the fusion is lost, so every new sequence
+# length compiles them again. The RMSNorm core is compiled per shape as well: the dynamic-shape variant lost 15 % per
+# call and shifted the mini benchmark's first-step loss by 9e-4, the static kernels reproduce the eager loss.
+def _rms_norm_core(hidden_states: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    value = hidden_states.float()
+    value = value * torch.rsqrt(value.square().mean(-1, keepdim=True) + eps)
+    return (weight.float() * value).to(hidden_states.dtype)
+
+
+def _hc_project_core(
+    hidden_states: torch.Tensor,
+    fn: torch.Tensor,
+    base: torch.Tensor,
+    scale: torch.Tensor,
+    streams: int,
+    norm_eps: float,
+    eps: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Coefficient projection up to the Sinkhorn input: pre, post and the [streams, streams] combination logits."""
+    flat = hidden_states.flatten(2).float()
+    mixes = F.linear(flat, fn.float()) * torch.rsqrt(flat.square().mean(-1, keepdim=True) + norm_eps)
+    scales = torch.cat((scale[0].expand(streams), scale[1].expand(streams), scale[2].expand(streams * streams)))
+    # The released kernel fuses the affine multiply/add before sigmoid and
+    # Sinkhorn. Separate operations round its FP32 coefficients differently;
+    # these differences can survive the subsequent BF16 stream collapse. This
+    # holds on the eager path; the compiled core decomposes addcmul and is
+    # allclose only (see the module note above).
+    logits = torch.addcmul(base, mixes, scales)
+    pre = torch.sigmoid(logits[..., :streams]) + eps
+    post = 2 * torch.sigmoid(logits[..., streams : 2 * streams])
+    return pre, post, logits[..., 2 * streams :].unflatten(-1, (streams, streams))
+
+
+def _hc_collapse_core(hidden_states: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor:
+    return (pre_mix.unsqueeze(-1) * hidden_states.float()).sum(2).to(hidden_states.dtype)
+
+
+def _hc_expand_core(
+    output: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor
+) -> torch.Tensor:
+    update = post.unsqueeze(-1) * output.unsqueeze(-2)
+    mixed = (comb.unsqueeze(-1) * residual.unsqueeze(-2)).sum(2)
+    return (update + mixed).to(output.dtype)
+
+
+_rms_norm = _rms_norm_core
+_hc_project = _hc_project_core
+_hc_collapse = _hc_collapse_core
+_hc_expand = _hc_expand_core
+_NORM_COMPILED = False
+_HC_COMPILED = False
+
+
+def compile_norm_core() -> None:
+    """Wrap the fp32 RMSNorm core with ``torch.compile`` (``BackendConfig.compile_norm``); runs once per process.
+
+    Static shapes: one specialization per norm site (block, attention, compressor and indexer norms) and grad mode.
+    The dynamic-shape variant (Kimi K3's setting) was measured 15 % slower per call and moved the first-step loss of
+    the 12-layer mini benchmark by 9e-4, while the static kernels reproduced the eager loss to four decimals; no
+    ``fullgraph``, so a shape past dynamo's recompile limit falls back to eager instead of raising.
+    """
+    global _rms_norm, _NORM_COMPILED
+    if _NORM_COMPILED:
+        return
+    _rms_norm = torch.compile(_rms_norm_core, dynamic=False)
+    _NORM_COMPILED = True
+
+
+def compile_hc_cores() -> None:
+    """Wrap the mHC projection, collapse and expand cores with ``torch.compile`` (``BackendConfig.compile_hc``).
+
+    Runs once per process, static shapes (``dynamic=False``: the fusion of ``expand`` needs concrete shapes; a new
+    sequence length compiles the three cores again). No ``fullgraph``: past dynamo's recompile limit (many distinct
+    sequence lengths) a core runs eager instead of raising ``FailOnRecompileLimitHit``; the unit tests assert the
+    cores compile without graph breaks. The Sinkhorn normalisation keeps its own backend (TileKernels or torch):
+    compiling the 20-iteration torch loop is slower than the fused TileKernels kernel.
+    """
+    global _hc_project, _hc_collapse, _hc_expand, _HC_COMPILED
+    if _HC_COMPILED:
+        return
+    _hc_project = torch.compile(_hc_project_core, dynamic=False)
+    _hc_collapse = torch.compile(_hc_collapse_core, dynamic=False)
+    _hc_expand = torch.compile(_hc_expand_core, dynamic=False)
+    _HC_COMPILED = True
+
+
 class DeepseekV41RMSNorm(nn.Module):
     """Normalize in FP32 and multiply the scale before casting the result."""
 
@@ -47,9 +141,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         Returns:
             Tensor of shape [..., hidden], with the input dtype.
         """
-        value = hidden_states.float()
-        value = value * torch.rsqrt(value.square().mean(-1, keepdim=True) + self.eps)
-        return (self.weight.float() * value).to(hidden_states.dtype)
+        return _rms_norm(hidden_states, self.weight, self.eps)
 
 
 class DeepseekV41Mix(NamedTuple):
@@ -97,24 +189,10 @@ def forward(self, hidden_states: torch.Tensor) -> DeepseekV41Mix:
             Coefficients with pre/post tensors of shape [batch, sequence, streams]
             and comb of shape [batch, sequence, streams, streams]. All use FP32.
         """
-        flat = hidden_states.flatten(2).float()
-        mixes = F.linear(flat, self.fn.float()) * torch.rsqrt(flat.square().mean(-1, keepdim=True) + self.norm_eps)
-        streams = self.streams
-        scales = torch.cat(
-            (self.scale[0].expand(streams), self.scale[1].expand(streams), self.scale[2].expand(streams * streams))
-        )
-        # The released kernel fuses the affine multiply/add before sigmoid and
-        # Sinkhorn. Separate operations round its FP32 coefficients differently;
-        # these differences can survive the subsequent BF16 stream collapse.
-        logits = torch.addcmul(self.base, mixes, scales)
-        pre = torch.sigmoid(logits[..., :streams]) + self.eps
-        post = 2 * torch.sigmoid(logits[..., streams : 2 * streams])
-        comb = dsv4_sinkhorn_normalize(
-            logits[..., 2 * streams :].unflatten(-1, (streams, streams)),
-            backend=self.sinkhorn_backend,
-            repeat=self.iterations,
-            eps=self.eps,
+        pre, post, comb_logits = _hc_project(
+            hidden_states, self.fn, self.base, self.scale, self.streams, self.norm_eps, self.eps
         )
+        comb = dsv4_sinkhorn_normalize(comb_logits, backend=self.sinkhorn_backend, repeat=self.iterations, eps=self.eps)
         return DeepseekV41Mix(pre, post, comb)
 
     @staticmethod
@@ -128,7 +206,7 @@ def collapse(hidden_states: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor
         Returns:
             Tensor of shape [batch, sequence, hidden], with the input dtype.
         """
-        return (pre_mix.unsqueeze(-1) * hidden_states.float()).sum(2).to(hidden_states.dtype)
+        return _hc_collapse(hidden_states, pre_mix)
 
     @staticmethod
     def expand(output: torch.Tensor, residual: torch.Tensor, mix: DeepseekV41Mix) -> torch.Tensor:
@@ -143,6 +221,4 @@ def expand(output: torch.Tensor, residual: torch.Tensor, mix: DeepseekV41Mix) ->
         Returns:
             Tensor of shape [batch, sequence, streams, hidden], with output's dtype.
         """
-        update = mix.post.unsqueeze(-1) * output.unsqueeze(-2)
-        mixed = (mix.comb.unsqueeze(-1) * residual.unsqueeze(-2)).sum(2)
-        return (update + mixed).to(output.dtype)
+        return _hc_expand(output, residual, mix.post, mix.comb)
@@ -79,6 +79,8 @@
 from nemo_automodel.components.models.deepseek_v41.layers import (
     DeepseekV41HyperConnection,
     DeepseekV41RMSNorm,
+    compile_hc_cores,
+    compile_norm_core,
 )
 from nemo_automodel.components.models.deepseek_v41.packing import packed_layout
 from nemo_automodel.components.models.deepseek_v41.processing import (
@@ -237,6 +239,10 @@ def __init__(
         config = config.text_config
         self.config = config
         self.moe_config = moe_config
+        if backend.compile_hc:
+            compile_hc_cores()
+        if backend.compile_norm and backend.rms_norm != "te":
+            compile_norm_core()
         dtype = dtype_from_str(config.dtype, torch.bfloat16)
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, dtype=dtype)
         active_engram = any(i < config.num_hidden_layers for i in config.engram_layer_ids)