biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

is paris in the top n?

Browse Source
This commit is contained in:
biondizzle
2026-05-15 06:38:20 +00:00
parent 311b28bd9f
commit beacc31569
3 changed files with 24 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py
View File

@@ -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()} "

12
src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py
View File

@@ -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(

16
vllm/patches/deepseek_v4.py
View File

@@ -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(