Skip to content

Fix FP4 init range for nvfp4_group_gemm#96

Merged
msaroufim merged 1 commit intogpu-mode:mainfrom
gau-nernst:fix_init
Jan 26, 2026
Merged

Fix FP4 init range for nvfp4_group_gemm#96
msaroufim merged 1 commit intogpu-mode:mainfrom
gau-nernst:fix_init

Conversation

@gau-nernst
Copy link
Contributor

@gau-nernst gau-nernst commented Jan 24, 2026

Key change: make the values of FP4 symmetric -> the result has mean of 0

Debug script using Modal (run with modal run debug.py)

import modal

from pathlib import Path

image = (
    modal.Image.from_registry(
        "nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04", add_python="3.12"
    )
    .entrypoint([])  # remove verbose logging by base image on entry
    .uv_pip_install("torch==2.9.1", index_url="https://download.pytorch.org/whl/cu130")
    .uv_pip_install("numpy")
    .add_local_python_source("reference", "task", "utils")
)
app = modal.App("debug", image=image)


@app.function(gpu="B200")
def run(task_config: dict):
    import torch
    from reference import generate_input, ref_kernel

    for cfg in task_config["benchmarks"]:
        data = generate_input(**cfg)
        out_list = ref_kernel(data)
        out = torch.cat([x.view(-1) for x in out_list], dim=0)

        print(cfg)
        print(
            f"  mean={out.mean().item():.2f}, std={out.std().item():.2f}, inf_any={out.isinf().any().item()}"
        )


@app.local_entrypoint()
def main():
    import yaml

    task_yaml = Path(__file__).parent / "task.yml"
    task_config = yaml.safe_load(open(task_yaml))
    run.remote(task_config)

Before

{'m': [80, 176, 128, 72, 64, 248, 96, 160], 'n': [4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096], 'k': [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168], 'g': 8, 'seed': 1111}
  mean=inf, std=nan, inf_any=True
{'m': [40, 76, 168, 72, 164, 148, 196, 160], 'n': [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168], 'k': [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048], 'g': 8, 'seed': 1111}
  mean=16944.00, std=1964.00, inf_any=False
{'m': [192, 320], 'n': [3072, 3072], 'k': [4096, 4096], 'g': 2, 'seed': 1111}
  mean=34048.00, std=2820.00, inf_any=False
{'m': [128, 384], 'n': [4096, 4096], 'k': [1536, 1536], 'g': 2, 'seed': 1111}
  mean=12816.00, std=1731.00, inf_any=False

After

{'m': [80, 176, 128, 72, 64, 248, 96, 160], 'n': [4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096], 'k': [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168], 'g': 8, 'seed': 1111}
  mean=0.65, std=184.12, inf_any=False
{'m': [40, 76, 168, 72, 164, 148, 196, 160], 'n': [7168, 7168, 7168, 7168, 7168, 7168, 7168, 7168], 'k': [2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048], 'g': 8, 'seed': 1111}
  mean=0.26, std=98.31, inf_any=False
{'m': [192, 320], 'n': [3072, 3072], 'k': [4096, 4096], 'g': 2, 'seed': 1111}
  mean=0.43, std=139.12, inf_any=False
{'m': [128, 384], 'n': [4096, 4096], 'k': [1536, 1536], 'g': 2, 'seed': 1111}
  mean=0.19, std=85.62, inf_any=False

@msaroufim
Copy link
Member

msaroufim commented Jan 24, 2026

@vickiw973 for review

vickiw973
vickiw973 approved these changes Jan 25, 2026
View reviewed changes
Copy link
Contributor

@vickiw973 vickiw973 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for the fix.

@msaroufim msaroufim self-requested a review January 26, 2026 00:47
@msaroufim msaroufim merged commit 07f0321 into gpu-mode:main Jan 26, 2026
@mobicham
Copy link

mobicham commented Jan 27, 2026

I didn't know it was easy to run evals like this on modal, super useful! Thank you 🙏

@gau-nernst gau-nernst deleted the fix_init branch January 28, 2026 05:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@msaroufim msaroufim msaroufim approved these changes

+1 more reviewer

@vickiw973 vickiw973 vickiw973 approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@gau-nernst @msaroufim @mobicham @vickiw973