diff --git a/single_shot_inference.py b/single_shot_inference.py index e7c3e64d..f75e0696 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -48,6 +48,7 @@ CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" MAX_NEW_TOKENS = 10 PROMPT = "The capital of France is" NUM_GPUS = 8 +SKIP_ROUTED_MOE = False # If True, only use shared expert (debug) # ===================================================================== # NVFP4 dequantization — matches checkpoint format exactly @@ -377,9 +378,8 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, k_expanded = k_full.expand(n_h, -1, -1).contiguous() v_expanded = v_full.expand(n_h, -1, -1).contiguous() # Attention sink (paper D5c) - # DISABLED for now to check impact sink_key = f"{pre}.sinks" - if False and sink_key in w and seq_len > 0: + if sink_key in w and seq_len > 0: sinks = w[sink_key].to(device=device) # (n_h,) BF16 sink_k = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=device) sink_v = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=device) @@ -521,8 +521,9 @@ def moe_forward(x, w, li, cfg, token_id, device): # ---- Run selected experts ---- T = x.shape[0] expert_outputs = [] - for i, eid in enumerate(expert_ids): - eid_int = eid.item() + if not SKIP_ROUTED_MOE: + for i, eid in enumerate(expert_ids): + eid_int = eid.item() epre = f"model.layers.{li}.mlp.experts.{eid_int}" gate = nvfp4_linear(x,