import torch

torch.manual_seed(0)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16

T, TOP_K, H = 8, 2, 4
S = T * TOP_K   # 16 routing slots


def run_mode(label, *, sentinel_makes_nan, prezero_before_mul):
    sentinel_mask = torch.zeros(S, dtype=torch.bool, device=DEVICE)
    sentinel_mask[S // 2:] = True

    sample_weights = torch.randn(S, device=DEVICE, dtype=torch.float32)
    sample_weights[sentinel_mask] = 0.0
    sample_weights = sample_weights.to(DTYPE).clone().requires_grad_(True)

    proj_out = torch.randn(S, H, device=DEVICE, dtype=DTYPE)
    if sentinel_makes_nan:
        proj_out[sentinel_mask] = float("nan")     # simulate uninitialized rows
    proj_out = proj_out.clone().requires_grad_(True)

    proj_out_used = (
        proj_out.masked_fill(sentinel_mask.unsqueeze(-1), 0.0)
        if prezero_before_mul else proj_out
    )

    weighted = proj_out_used * sample_weights.unsqueeze(-1)
    weighted_zero = weighted.masked_fill(sentinel_mask.unsqueeze(-1), 0.0)
    out = weighted_zero.view(T, TOP_K, H).sum(dim=1)
    out.backward(torch.randn_like(out))

    sw_nan = torch.isnan(sample_weights.grad).sum().item()
    print(f"[{label}]")
    print(f"  forward out finite     : {torch.isfinite(out).all().item()}")
    print(f"  sample_weights.grad nan: {sw_nan} of {S}")
    print(f"  sw.grad sentinel slice : {sample_weights.grad[sentinel_mask].float().tolist()}")


run_mode("A: proj_out fully init (baseline)",
         sentinel_makes_nan=False, prezero_before_mul=False)
run_mode("B: NaN at sentinels, NO pre-zero  (this PR's pattern)",
         sentinel_makes_nan=True,  prezero_before_mul=False)
run_mode("C: NaN at sentinels, pre-zero before mul  (proposed fix)",
         sentinel_makes_nan=True,  prezero_before_mul=True)

selected_hidden_states_g = hidden_states[perm // num_top_k]   # (S, H)
proj_out = _grouped_linear(selected_hidden_states_g, gate_up_proj, offsets)   # ←── leak 2
proj_out = proj_out.masked_fill(sentinel_mask, 0.0)
proj_out = self._apply_gate(proj_out)
proj_out = _grouped_linear(proj_out, down_proj, offsets)                      # ←── leak 3
proj_out.masked_fill_(sentinel_mask, 0.0)
weighted_out = proj_out * sample_weights_g.unsqueeze(-1)                      # ←── leak 1

# 1. Build the per-(token, slot) input tensor for the kernel by gathering rows from hidden_states:
selected_hidden_states_g = hidden_states[perm // num_top_k]   # shape (S, H)

# 2. Run the kernel. It only does work for rows [0, offsets[-1]) sentinels are excluded.
proj_out = _grouped_mm(selected_hidden_states_g, gate_up_proj, offsets)