Add swiglu_limit=10.0 activation clamping (was missing)

DeepSeek-V4 uses SiluAndMulWithClamp(10.0) which clamps:
- silu(gate) to max 10.0
- up to [-10.0, 10.0]

Our runner was doing plain F.silu(gate) * up without clamping.
Large gate values could produce unbounded SiLU output, causing
numerical issues in the L2 GEMM. This is likely contributing to
garbage model output.
This commit is contained in:
2026-05-17 17:52:16 +00:00
parent 3f2f4e1882
commit a10c582cf4
2 changed files with 16 additions and 2 deletions

View File

@@ -46,6 +46,7 @@ class CuTeDSLMoERunner:
self.top_k = top_k
self.device = device
self.experts_start_idx = experts_start_idx
self._swiglu_limit = None # Set via set_swiglu_limit()
# Weight storage (set before _ensure_stacked)
self.l1_fp4 = None
@@ -92,6 +93,10 @@ class CuTeDSLMoERunner:
f"pid={os.getpid()}", flush=True)
self._buffers_allocated = False
def set_swiglu_limit(self, limit: float | None):
"""Set the swiglu_limit for activation clamping."""
self._swiglu_limit = limit
def _fill_token_indices(self):
"""Fill _token_indices with [0,0,..0, 1,1,..1, ...] (each token repeated top_k times).
@@ -438,10 +443,15 @@ class CuTeDSLMoERunner:
# Extract real token outputs from padded GEMM output
l1_out_real = l1_out[padded_dst]
# === SiLU(gate) * up ===
# === SiLU(gate) * up (with swiglu_limit clamp) ===
gate = l1_out_real[:, :self.intermediate_size]
up = l1_out_real[:, self.intermediate_size:]
activated = torch.nn.functional.silu(gate) * up
gate_silu = torch.nn.functional.silu(gate)
# Apply DeepSeek-V4 swiglu_limit: clamp both silu(gate) and up
if self._swiglu_limit is not None:
gate_silu = gate_silu.clamp(max=self._swiglu_limit)
up = up.clamp(min=-self._swiglu_limit, max=self._swiglu_limit)
activated = gate_silu * up
# === L2: down ===
padded_activated = self._shared_bufs['activated']

View File

@@ -551,6 +551,10 @@ class DeepseekV4MegaMoEExperts(nn.Module):
# the actual amax and compute correct gs values.
self._warmup_activation_global_scales()
# Set swiglu_limit for activation clamping in the runner
if self.swiglu_limit is not None:
self._cutedsl_runner.set_swiglu_limit(float(self.swiglu_limit))
def _warmup_activation_global_scales(self) -> None:
"""Run a warmup forward pass to compute correct activation global scales.