From beacc315694ab9a02c377194d1262b9810f47d53 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 15 May 2026 06:38:20 +0000 Subject: [PATCH] is paris in the top n? --- .../cutlass_nvfp4_gemm/kernel.py | 2 +- src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py | 12 ++++++++---- vllm/patches/deepseek_v4.py | 16 +++++++++++++++- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py index f1262d77..5c8b3ab2 100644 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py @@ -76,7 +76,7 @@ def cutlass_grouped_nvfp4_gemm( M_expert = token_indices.shape[0] # DEBUG: verify data going into GEMM - if e < 3 and M_expert > 0: + if MEGA_MOE_DEBUG and e < 3 and M_expert > 0: print(f"[GEMM-IN] expert={e} M={M_expert} N={N} K={K} " f"w shape={expert_w.shape} w_sf shape={expert_w_sf.shape} " f"w absmax={expert_w.view(torch.int8).abs().max().item()} " diff --git a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index 9724048e..093af80a 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -115,7 +115,8 @@ def nvfp4_mega_moe_l1( f"experts={num_experts_per_rank} native=1") # DEBUG: verify weight shapes after transpose - print(f"[L1-WT] l1_w shape={l1_weights.shape} l1_sf shape={l1_scales.shape} w_sf dtype={l1_scales.dtype}") + if MEGA_MOE_DEBUG: + print(f"[L1-WT] l1_w shape={l1_weights.shape} l1_sf shape={l1_scales.shape} w_sf dtype={l1_scales.dtype}") # Unpack uint32 packed UE4M3 scales to float8_e4m3fn x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf @@ -127,7 +128,8 @@ def nvfp4_mega_moe_l1( topk_ids, topk_weights, alpha=alpha, ) - print(f"[L1-GEMM-OUT] amax={output.abs().max().item():.4e} mean={output.float().mean().item():.4e} nonzero_frac={(output != 0).float().mean().item():.4f}") + if MEGA_MOE_DEBUG: + print(f"[L1-GEMM-OUT] amax={output.abs().max().item():.4e} mean={output.float().mean().item():.4e} nonzero_frac={(output != 0).float().mean().item():.4f}") return output # (num_tokens, 6144) bfloat16 @@ -315,7 +317,8 @@ def nvfp4_mega_moe_full( # ALWAYS-ON debug: alpha and scale ranges _x_sf_f32 = x_sf.to(torch.float32) _igs = l1_global_scale if isinstance(l1_global_scale, float) else l1_global_scale.item() if hasattr(l1_global_scale, 'item') else float(l1_global_scale) - print(f"[ALPHA L1] alpha={_igs:.4e} x_sf range [{_x_sf_f32.min().item():.4e}, {_x_sf_f32.max().item():.4e}] x_fp4_absmax={x_fp4.view(torch.int8).abs().max().item()}") + if MEGA_MOE_DEBUG: + print(f"[ALPHA L1] alpha={_igs:.4e} x_sf range [{_x_sf_f32.min().item():.4e}, {_x_sf_f32.max().item():.4e}] x_fp4_absmax={x_fp4.view(torch.int8).abs().max().item()}") # Convert global expert IDs to local expert IDs. # vLLM's symm_buffer stores global IDs (0..383) but our weight tensors @@ -368,7 +371,8 @@ def nvfp4_mega_moe_full( # ALWAYS-ON debug: L2 alpha and scale ranges _l1sf_f32 = l1_sf_out.to(torch.float32) _l2gs = l2_global_scale if isinstance(l2_global_scale, float) else l2_global_scale.item() if hasattr(l2_global_scale, 'item') else float(l2_global_scale) - print(f"[ALPHA L2] alpha={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}] activated amax={activated.abs().max().item():.4e}") + if MEGA_MOE_DEBUG: + print(f"[ALPHA L2] alpha={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}] activated amax={activated.abs().max().item():.4e}") # Step 5: L2 GEMM (native NVFP4 block-scaled MMA) l2_output = nvfp4_mega_moe_l2( diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index e8a16fe2..ad906734 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -373,7 +373,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): return False # DEBUG: log weight loads for expert params (weight only, not scales) - if shard_id in ("w1", "w3") and local_expert_id < 2 and loaded_weight.dtype in (torch.uint8, torch.int8): + if int(os.environ.get('MEGA_MOE_DEBUG', '0')) and shard_id in ("w1", "w3") and local_expert_id < 2 and loaded_weight.dtype in (torch.uint8, torch.int8): print(f"[WT-LOAD] {weight_name} expert={expert_id}→local={local_expert_id} " f"shard={shard_id} loaded_shape={tuple(loaded_weight.shape)} " f"param_shape={tuple(param.data[local_expert_id].shape)} " @@ -2202,6 +2202,20 @@ class DeepseekV4ForCausalLM(nn.Module): print(f"[LOGITS] top10 ids: {top_idx.tolist()}") print(f"[LOGITS] top10 vals: {[f'{v:.3f}' for v in top_vals.tolist()]}") print(f"[LOGITS] gap top1-top10: {(top_vals[0] - top_vals[-1]).item():.3f}") + + # Probe for "Paris" specifically + try: + from transformers import AutoTokenizer + if not hasattr(self, '_tok'): + self._tok = AutoTokenizer.from_pretrained('/model', trust_remote_code=True) + for v in [' Paris', 'Paris', ' paris', 'paris']: + tid = self._tok.encode(v, add_special_tokens=False)[0] + logit_val = last[tid].item() + rank = (last > logit_val).sum().item() + print(f"[LOGITS-PROBE] {repr(v)}→id={tid} logit={logit_val:.2f} rank={rank}") + except Exception as e: + print(f"[LOGITS-PROBE] failed: {e}") + return logits def forward(