fix(qwen): preserve checkpoint exports and bound QSA memory

src/mcore_bridge/model/gpts/qwen4_exp.py

@@ -176,7 +176,7 @@ def _qsa_select(self, hidden_states, attn_kwargs, position_ids=None):
                                'Use --padding_free false with context_parallel_size 1 to take the bool-mask path, '
                                f'or set {QSA_SPARSE_KERNEL_ENV}=0 to fall back to full attention.')
         if cp_size > 1 and getattr(self.config, 'cp_comm_type', None) != 'all_gather':
-            raise RuntimeError(f"QSA sparse selection with context_parallel_size={cp_size} requires "
+            raise RuntimeError(f'QSA sparse selection with context_parallel_size={cp_size} requires '
                                f"cp_comm_type='all_gather' (got {getattr(self.config, 'cp_comm_type', None)!r}): the "
                                'selection has to see every key before attention runs, which ring/p2p cannot provide.')
         rotary_pos_emb = attn_kwargs.get('rotary_pos_emb')
@@ -287,6 +287,16 @@ def __init__(self, *args, **kwargs):
 class Qwen4ExpBridge(Qwen3NextBridge):
     hf_mixer_prefix = 'model.'
 
+    def _save_missing_weights(self, saver, saved_keys, source_model_dir=None) -> None:
+        # PLE export emits every shard. If it omits the scale, these are already
+        # dequantized parameters; copying the source FP8 scale would corrupt them.
+        accounted_keys = set(saved_keys)
+        suffix = 'ngram_embedding.shard_0.weight'
+        for key in saved_keys:
+            if key.endswith(f'ple.ple_embedding.{suffix}'):
+                accounted_keys.add(key[:-len(suffix)] + 'ngram_embedding.weight_scale')
+        super()._save_missing_weights(saver, accounted_keys, source_model_dir)
+
     def _get_hf_experts_attr(self, is_mtp: bool = False):
         # The checkpoint stores experts as packed per-layer tensors
         # (`mlp.experts.gate_up_proj` / `mlp.experts.down_proj`).

src/mcore_bridge/model/modules/kernels/ple_kernels.py

@@ -477,13 +477,19 @@ def backward(ctx, dout):
                     DIL=dilation,
                     BLOCK_W=BW)
 
+            # The convolution backward has consumed the recomputed norm output.
+            del normed
+
             # norm_conv backward: dwc on host, dx via kernel (fp32).
             x_hat = (gated.view(T, n, C) * rstdc.unsqueeze(-1)).view(T, W)
             dwc = (dnormed * x_hat).sum(dim=0).to(wc.dtype)
+            del x_hat
             dgated_norm = torch.empty(T, W, dtype=torch.float32, device=dev)
             if T > 0:
                 _ple_norm_bwd_kernel[(T * n, )](gated, wc, rstdc, dnormed, dgated_norm, T, N=n, C=C, BLOCK_C=block_c)
             dgated += dgated_norm
+            # Release token-sized FP32 temporaries before gate gradient buffers.
+            del gated, dnormed, dgated_norm
 
             dkey = torch.empty_like(key)
             dquery = torch.empty_like(hc_state)

src/mcore_bridge/model/modules/ple.py

@@ -265,8 +265,10 @@ def iter_export_table_to_hf(self, prefix=''):
             # scalar `weight_scale`, so divide by the scale stashed during loading
             # and cast back to fp8. Without a known scale the values cannot be
             # represented as fp8 + scale; keep the current dtype and warn.
-            scale = getattr(self, '_ngram_weight_scale', None)
-            if scale is None:
+            # Updated BF16 parameters cannot be losslessly requantized with the
+            # original checkpoint scale. Only immutable host tables reuse it.
+            scale = getattr(self, '_ngram_weight_scale', None) if self.cpu_offload else None
+            if scale is None and self.cpu_offload:
                 get_logger().warning(f'`{self._NGRAM_SCALE_KEY}` was not seen during loading; exporting the PLE ngram '
                                      'embedding without re-quantizing to fp8.')
             # Reduce on GPU: NCCL has no CPU backend, and the host table is pinned

src/mcore_bridge/model/modules/qsa_indexer.py

@@ -4,7 +4,7 @@
 from megatron.core.extensions.transformer_engine import TELinear
 from torch import nn
 
-# Byte budget for the transient score tile in select_token_indices_thd: the (token, block)
+# Byte budget for the transient score tile in packed and unpacked selection: the (token, block)
 # scoring is chunked over queries to avoid OOM from the full [T, n_heads, NB] fp32 tensor.
 # Tests monkeypatch this to compare one chunk (un-chunked reference) vs many chunks.
 _QSA_INDEX_SCORE_CHUNK_BYTES = 1024 * 1024 * 1024
@@ -180,18 +180,23 @@ def apply_rope(t, cos_, sin_):
         starts = torch.arange(max_blocks, device=device) * R
         block_keys = apply_rope(pooled, cos[:, starts], sin[:, starts])  # [b, nb, d]
 
-        # ---- score all (query, block) pairs ----
-        scores = torch.einsum('bqhd,bkd->bqhk', q.float(), block_keys.float())
-        scores = torch.relu(scores).sum(dim=2) / math.sqrt(self.index_head_dim)  # [b, s, nb]
-
-        # ---- restrict to blocks fully inside the causal prefix ----
-        n_blocks = (torch.arange(s, device=device) + 1) // R  # [s]
+        # Bound score workspace by queries, preserving all candidate blocks.
+        chunk_size = max(1, min(s, _QSA_INDEX_SCORE_CHUNK_BYTES // max(1, b * self.index_n_heads * max_blocks * 4)))
+        n_blocks = (torch.arange(s, device=device) + 1) // R
         block_ids = torch.arange(max_blocks, device=device)
-        scores = scores.masked_fill((block_ids[None, :] >= n_blocks[:, None])[None], float('-inf'))
-
         k = min(self.block_topk, max_blocks)
-        top_blocks = scores.topk(k, dim=-1).indices  # [b, s, k]
-        keep = top_blocks < n_blocks[None, :, None]  # drop the -inf padding slots
+        top_blocks = torch.empty((b, s, k), dtype=torch.long, device=device)
+        keep = torch.empty((b, s, k), dtype=torch.bool, device=device)
+        keys_float = block_keys.float()
+        for start in range(0, s, chunk_size):
+            end = min(start + chunk_size, s)
+            scores = torch.einsum('bqhd,bkd->bqhk', q[:, start:end].float(), keys_float)
+            scores = torch.relu(scores).sum(dim=2) / math.sqrt(self.index_head_dim)
+            scores = scores.masked_fill((block_ids[None, :] >= n_blocks[start:end, None])[None], float('-inf'))
+            selected = scores.topk(k, dim=-1).indices
+            top_blocks[:, start:end] = selected
+            keep[:, start:end] = selected < n_blocks[None, start:end, None]
+            del scores, selected
         return top_blocks, keep, n_blocks
 
     @torch.no_grad()