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

CRITICAL FIX: FP32 RoPE cache + FP32 arithmetic for inverse RoPE round-trip

Browse Source 
BF16 cos/sin cache destroys cos²+sin²=1 identity (can be 0.996 in BF16).
This causes ~3% error per RoPE→inverse RoPE round-trip, accumulating
across 61 layers into garbage output. FP32 cache + FP32 arithmetic
gives exact round-trip (diff < 1e-7).

Also fixes: MoE expert loop indentation (was only running last expert).
This commit is contained in:
biondizzle
2026-05-31 09:14:59 +00:00
parent db74a887ab
commit 9d96c2fbbf
2 changed files with 37 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
single_shot_inference.py
View File

@@ -167,37 +167,44 @@ class mHCBlock:
def build_rope_cache(max_pos, rope_dim, device, theta=10000.0):
"""Build cos/sin caches for partial RoPE.
Returns: (cos_cache, sin_cache) each (max_pos, rope_dim//2) BF16
CRITICAL: FP32, not BF16! BF16 quantization destroys cos²+sin²=1
identity needed for inverse RoPE. BF16 cos²+sin² can be 0.996,
causing ~3% round-trip error that accumulates across 61 layers.
Returns: (cos_cache, sin_cache) each (max_pos, rope_dim//2) FP32
"""
half = rope_dim // 2
freqs = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
angles = torch.outer(torch.arange(max_pos, dtype=torch.float32), freqs)
return torch.cos(angles).bfloat16().to(device), torch.sin(angles).bfloat16().to(device)
return torch.cos(angles).to(device), torch.sin(angles).to(device)
def apply_rope_partial(x, positions, cos_cache, sin_cache, head_dim, rope_dim):
"""Apply partial GPT-J interleaved RoPE to the last rope_dim dims of each head."""
"""Apply partial GPT-J interleaved RoPE to the last rope_dim dims of each head.
Computes in FP32 for numerical stability (inverse RoPE requires cos²+sin²=1)."""
T, n_h, hd = x.shape
nope = hd - rope_dim
cos = cos_cache[positions].unsqueeze(1) # (T, 1, half) BF16
cos = cos_cache[positions].unsqueeze(1) # (T, 1, half) FP32
sin = sin_cache[positions].unsqueeze(1)
out = x.clone()
x_rope = x[:, :, nope:]
out[:, :, nope:][..., 0::2] = x_rope[..., 0::2] * cos - x_rope[..., 1::2] * sin
out[:, :, nope:][..., 1::2] = x_rope[..., 0::2] * sin + x_rope[..., 1::2] * cos
x_rope = x[:, :, nope:].float() # FP32 for accurate rotation
out[:, :, nope:] = (x_rope[..., 0::2] * cos - x_rope[..., 1::2] * sin).to(torch.bfloat16)
out[:, :, nope:][..., 1::2] = (x_rope[..., 0::2] * sin + x_rope[..., 1::2] * cos).to(torch.bfloat16)
return out
def apply_inverse_rope(o, positions, cos_cache, sin_cache, head_dim, rope_dim):
"""Apply inverse RoPE (conjugate rotation) to attention output."""
"""Apply inverse RoPE (conjugate rotation) to attention output.
Computes in FP32 for numerical stability."""
T, n_h, hd = o.shape
nope = hd - rope_dim
cos = cos_cache[positions].unsqueeze(1)
sin = sin_cache[positions].unsqueeze(1)
out = o.clone()
o_rope = o[:, :, nope:]
out[:, :, nope:][..., 0::2] = o_rope[..., 0::2] * cos + o_rope[..., 1::2] * sin
out[:, :, nope:][..., 1::2] = -o_rope[..., 0::2] * sin + o_rope[..., 1::2] * cos
o_rope = o[:, :, nope:].float()
out[:, :, nope:] = (o_rope[..., 0::2] * cos + o_rope[..., 1::2] * sin).to(torch.bfloat16)
out[:, :, nope:][..., 1::2] = (-o_rope[..., 0::2] * sin + o_rope[..., 1::2] * cos).to(torch.bfloat16)
return out
class SimpleKVCache:

27
tests/test_minimal_e2e.py
View File

@@ -54,31 +54,41 @@ class RMSNorm:
return (x_f * rms * self.weight).to(torch.bfloat16)
def build_rope_cache(max_pos, rope_dim, device, theta=10000.0):
"""Build FP32 cos/sin caches for RoPE.
CRITICAL: Must be FP32, not BF16! BF16 quantization destroys
cos²+sin²=1 identity needed for inverse RoPE round-trip.
BF16 cos²+sin² can be as low as 0.996, causing ~3% error.
"""
half = rope_dim // 2
freqs = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
angles = torch.outer(torch.arange(max_pos, dtype=torch.float32), freqs)
return torch.cos(angles).bfloat16().to(device), torch.sin(angles).bfloat16().to(device)
return torch.cos(angles).to(device), torch.sin(angles).to(device)
def apply_rope_partial(x, positions, cos_cache, sin_cache, head_dim, rope_dim):
"""Apply partial GPT-J interleaved RoPE. Computes in FP32 for accuracy."""
T, n_h, hd = x.shape
nope = hd - rope_dim
cos = cos_cache[positions].unsqueeze(1)
cos = cos_cache[positions].unsqueeze(1) # (T, 1, half) FP32
sin = sin_cache[positions].unsqueeze(1)
out = x.clone()
x_rope = x[:, :, nope:]
out[:, :, nope:][..., 0::2] = x_rope[..., 0::2] * cos - x_rope[..., 1::2] * sin
out[:, :, nope:][..., 1::2] = x_rope[..., 0::2] * sin + x_rope[..., 1::2] * cos
# Compute in FP32 for numerical stability
x_rope = x[:, :, nope:].float()
out[:, :, nope:] = (x_rope[..., 0::2] * cos - x_rope[..., 1::2] * sin).to(torch.bfloat16)
# Second pass for odd elements (need original even values)
out[:, :, nope:][..., 1::2] = (x_rope[..., 0::2] * sin + x_rope[..., 1::2] * cos).to(torch.bfloat16)
return out
def apply_inverse_rope(o, positions, cos_cache, sin_cache, head_dim, rope_dim):
"""Apply inverse RoPE (conjugate rotation). Computes in FP32 for accuracy."""
T, n_h, hd = o.shape
nope = hd - rope_dim
cos = cos_cache[positions].unsqueeze(1)
sin = sin_cache[positions].unsqueeze(1)
out = o.clone()
o_rope = o[:, :, nope:]
out[:, :, nope:][..., 0::2] = o_rope[..., 0::2] * cos + o_rope[..., 1::2] * sin
out[:, :, nope:][..., 1::2] = -o_rope[..., 0::2] * sin + o_rope[..., 1::2] * cos
o_rope = o[:, :, nope:].float()
out[:, :, nope:] = (o_rope[..., 0::2] * cos + o_rope[..., 1::2] * sin).to(torch.bfloat16)
out[:, :, nope:][..., 1::2] = (-o_rope[..., 0::2] * sin + o_rope[..., 1::2] * cos).to(torch.bfloat16)
return out
def load_weights_to_cpu(checkpoint_dir):
@@ -182,6 +192,7 @@ def test_layer0():
# Build mHC + norms for layer 0
li = 0
from single_shot_inference import mHCBlock
attn_mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=device)
attn_mhc.load_from_checkpoint(
all_weights[f"model.layers.{li}.attn_hc.fn"],