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

Remove all inline diagnostics — incompatible with torch.compile

Browse Source 
Data-dependent expressions (amax().item(), isnan().any().item())
cause Dynamo guard failures even when gated by os.environ.
cudagraph_mode=NONE still uses torch.compile, so these break.
Will need enforce-eager for diagnostics going forward.
This commit is contained in:
biondizzle
2026-05-18 15:22:53 +00:00
parent a83d364d45
commit 301015b037
2 changed files with 0 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
vllm/patches/deepseek_v4.py
View File

@@ -1197,11 +1197,7 @@ class DeepseekV4DecoderLayer(nn.Module):
x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
)
x = self.attn_norm(x)
if os.environ.get('CLAWMINE_DEBUG', '0') == '1':
_print_tensor("pre_attn", x, self._layer_idx if hasattr(self, '_layer_idx') else -1)
x = self.attn(positions, x, None)
if os.environ.get('CLAWMINE_DEBUG', '0') == '1':
_print_tensor("post_attn", x, self._layer_idx if hasattr(self, '_layer_idx') else -1)
x = self.hc_post(x, residual, post, comb)
residual = x
@@ -1214,26 +1210,6 @@ class DeepseekV4DecoderLayer(nn.Module):
return x
def _print_tensor(label, t, layer_idx):
with torch.no_grad():
print(f"[CLAWMINE] L{layer_idx} {label}: amax={t.amax().item():.4f} NaN={torch.isnan(t).any().item()} shape={t.shape}")
def _diag_hidden_stats(hidden_states: torch.Tensor, layer_idx: int):
"""Print hidden state stats after each layer. Disabled unless
CLAWMINE_DEBUG=1. os.environ is evaluated at trace time, so
the data-dependent path is dead code when disabled."""
if os.environ.get('CLAWMINE_DEBUG', '0') != '1':
return
# Only reached when CLAWMINE_DEBUG=1 (must run with --enforce-eager)
with torch.no_grad():
amax = hidden_states.amax().item()
mean = hidden_states.float().mean().item()
has_nan = torch.isnan(hidden_states).any().item()
has_inf = torch.isinf(hidden_states).any().item()
print(f"[CLAWMINE] Layer {layer_idx}: amax={amax:.4f} mean={mean:.6f} NaN={has_nan} Inf={has_inf}")
@support_torch_compile
class DeepseekV4Model(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -1332,19 +1308,14 @@ class DeepseekV4Model(nn.Module):
) -> torch.Tensor | IntermediateTensors:
hidden_states = self.embed_input_ids(input_ids)
hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
if os.environ.get('CLAWMINE_DEBUG', '0') == '1':
_print_tensor("embed", hidden_states, -1)
if self.use_mega_moe:
input_ids = input_ids.to(torch.int64)
for layer_idx, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)):
layer._layer_idx = layer_idx
hidden_states = layer(
hidden_states,
positions,
input_ids,
)
# Diagnostic: print amax/mean every layer (eager-mode only, no Dynamo)
_diag_hidden_stats(hidden_states, layer_idx)
# Stash pre-hc_head residual for the MTP draft (captured copy_).

9
vllm/patches/deepseek_v4_attention.py
View File

@@ -3,7 +3,6 @@
"""
DeepseekV4 MLA Attention Layer
"""
import os
from collections.abc import Callable
from dataclasses import dataclass
@@ -304,10 +303,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
)
o = o_padded[:, : self.n_local_heads, :]
if os.environ.get('CLAWMINE_DEBUG', '0') == '1':
with torch.no_grad():
print(f"[CLAWMINE] attn L{self.layer_name}: o amax={o.amax().item():.4f} NaN={torch.isnan(o).any().item()}")
# O projection: inverse RoPE + FP8 quant + einsum + wo_b
o_fp8, o_scale = fused_inv_rope_fp8_quant(
o,
@@ -320,10 +315,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
tma_aligned_scales=self._tma_aligned_scales,
)
if os.environ.get('CLAWMINE_DEBUG', '0') == '1':
with torch.no_grad():
print(f"[CLAWMINE] attn L{self.layer_name}: o_fp8 amax={o_fp8.float().amax().item():.4f} o_scale amax={o_scale.float().amax().item():.6f}")
wo_a_fp8 = self.wo_a.weight
wo_a_scale = self.wo_a.weight_scale_inv