Add warmup-based activation global scale computation in finalize_weights

The checkpoint input_scale is a calibration value that produces wrong gs
at runtime (too small → block scales saturate → garbage output → EOS).

Now calls compute_activation_global_scales() with sample data during weight
finalization, before cudagraph capture. This observes actual activation
magnitudes and computes correct L1 and L2 gs values.
This commit is contained in:
2026-05-17 10:48:24 +00:00
parent 4445882ba7
commit 04245b664b

View File

@@ -529,6 +529,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
l2_igs = w2_igs[:, 0]
else:
l2_igs = w2_igs
# Use checkpoint input_scale as initial guess, then warmup will override
self._cutedsl_runner._l1_activation_global_scale = l1_igs.mean().item()
self._cutedsl_runner._l2_activation_global_scale = l2_igs.mean().item()
@@ -544,6 +545,41 @@ class DeepseekV4MegaMoEExperts(nn.Module):
self.w2_weight_scale_2 = None
self.w2_input_scale = None
# Warmup: compute actual activation global scales from sample data.
# The checkpoint input_scale is a calibration value that doesn't match
# runtime activation magnitudes. We run a small forward pass to observe
# the actual amax and compute correct gs values.
self._warmup_activation_global_scales()
def _warmup_activation_global_scales(self) -> None:
"""Run a warmup forward pass to compute correct activation global scales.
Called once during finalize_weights, before cudagraph capture.
Uses quantize_to_nvfp4 (which calls .max()) to get the exact gs
from real activation magnitudes, then stores them for use by
quantize_activation_nvfp4 (no .max(), cudagraph-safe).
"""
import torch
runner = self._cutedsl_runner
device = runner.device
num_tokens = min(8, runner.max_num_tokens)
top_k = runner.top_k
with torch.no_grad():
# Sample hidden states: typical BF16 activations have amax ~1-10
hidden_states = torch.randn(num_tokens, runner.hidden_size,
dtype=torch.bfloat16, device=device)
# Assign all tokens to the first few experts evenly
topk_ids = torch.zeros(num_tokens, top_k, dtype=torch.int64, device=device)
for i in range(num_tokens):
for j in range(top_k):
topk_ids[i, j] = (runner.experts_start_idx + j) % (runner.experts_start_idx + runner.num_experts)
topk_weights = torch.ones(num_tokens, top_k, dtype=torch.float32, device=device) / top_k
runner.compute_activation_global_scales(
hidden_states, topk_weights, topk_ids
)
# Note: No explicit CuTeDSL warmup here. With FULL_AND_PIECEWISE
# CUDA graph mode, the kernel compiles during graph capture (startup).
# In eager mode, the first inference triggers JIT compilation.