Add TIRx inference kernels and B200 denoising timings

dceb455 verified 5 days ago

8.04 kB

import functools
import os
import re

import torch

_ENABLED = os.getenv("MAGI_TIRX_DOWN", "0") == "1"
_VALIDATE = os.getenv("MAGI_TIRX_DOWN_VALIDATE", "0") == "1"
_checked = set()
_reported = False

KP = 13696
N = 5120


@functools.lru_cache(maxsize=2)
def executable(mp):
    import tirx_kernels.gemm.fp16_bf16_gemm as gemm
    from tirx_kernels.runner import compile_kernel

    cfg = dict(gemm.GEMM_CONFIGS[1024])
    cfg.update(
        cta_n=64,
        cta_k=128,
        pipe_depth=5,
        wb_pipe_depth=2,
        l2_group_size=2,
        overlap_epilogue=True,
    )

    original_cfg = gemm._cfg_for
    try:
        gemm._cfg_for = lambda n: dict(cfg)
        kernel = gemm.make_kernel("bf16", mp, N, KP)
        result = compile_kernel(kernel)
    finally:
        gemm._cfg_for = original_cfg

    print(f"TIRX_DOWN_COMPILED M={mp} N={N} K={KP}", flush=True)
    return result


def supported(module, z, weight):
    return (
        _ENABLED
        and not torch.is_grad_enabled()
        and not module.training
        and z.is_cuda
        and weight.device == z.device
        and z.dtype == torch.bfloat16
        and weight.dtype == torch.bfloat16
        and z.ndim == 3
        and weight.ndim == 3
        and z.shape[0] == 3
        and 0 < z.shape[1] <= 4096
        and z.shape[2] == 13652
        and tuple(weight.shape) == (3, 5120, 13652)
        and z.is_contiguous()
        and weight.is_contiguous()
        and not z.requires_grad
        and not weight.requires_grad
    )


def down(module, z, weight):
    global _reported

    m = z.shape[1]
    mp = 2048 if m <= 2048 else 4096

    # 保存原权重对象，避免地址复用导致误用缓存。
    # 权重被原地修改后，通过 version 重新建立缓存。
    # inference tensor没有版本计数；固定权重推理按对象身份缓存。
    version = None if torch.is_inference(weight) else weight._version
    cached = getattr(module, "_tirx_down_weight_cache", None)
    if (
        cached is None
        or cached[0] is not weight
        or cached[1] != version
    ):
        padded_weight = torch.zeros(
            3, N, KP, dtype=weight.dtype, device=weight.device
        )
        padded_weight[:, :, :13652].copy_(weight)
        module._tirx_down_weight_cache = (
            weight, version, padded_weight
        )
    else:
        padded_weight = cached[2]

    # 全部清零，确保 K 尾部和 M 尾部没有未初始化值。
    padded_input = torch.zeros(
        3, mp, KP, dtype=z.dtype, device=z.device
    )
    padded_input[:, :m, :13652].copy_(z)

    padded_output = torch.empty(
        3, mp, N, dtype=z.dtype, device=z.device
    )
    ex = executable(mp)
    for e in range(3):
        ex(padded_input[e], padded_weight[e], padded_output[e])

    output = padded_output[:, :m, :].contiguous()

    if not _reported:
        print(
            f"TIRX_DOWN_ACTIVE input={tuple(z.shape)} "
            f"padded_M={mp} padded_K={KP}",
            flush=True,
        )
        _reported = True

    # 检查前三种真实形状；这不是全模型质量验证。
    key = tuple(z.shape)
    if _VALIDATE and key not in _checked and len(_checked) < 3:
        expected = torch.bmm(z, weight.transpose(-2, -1))
        finite = bool(torch.isfinite(output).all().item())

        # 分块统计，控制验证时的临时显存。
        a = output.reshape(-1)
        b = expected.reshape(-1)
        max_abs = 0.0
        error_sq = 0.0
        reference_sq = 0.0
        for start in range(0, a.numel(), 1 << 20):
            af = a[start:start + (1 << 20)].float()
            bf = b[start:start + (1 << 20)].float()
            delta = af - bf
            max_abs = max(max_abs, delta.abs().max().item())
            error_sq += delta.square().sum(dtype=torch.float64).item()
            reference_sq += bf.square().sum(dtype=torch.float64).item()

        relative_l2 = (error_sq / max(reference_sq, 1e-30)) ** 0.5
        print(
            f"TIRX_DOWN_CHECK input={tuple(z.shape)} "
            f"finite={finite} max_abs={max_abs:.8g} "
            f"relative_l2={relative_l2:.8g}",
            flush=True,
        )

        # 仅用于拦截明显异常，不能代替质量评价。
        if not finite or relative_l2 > 1e-3:
            raise RuntimeError("TIRx down真实输入检查未通过")

        _checked.add(key)

    return output


def transform_source(source):
    pattern = re.compile(
        r"^(?P<indent>[ \t]*)"
        r"y\s*=\s*torch\.bmm\("
        r"z,\s*w_dn\.transpose\(-2,\s*-1\)"
        r"\)\.to\(dtype\)[ \t]*$",
        re.MULTILINE,
    )
    matches = list(pattern.finditer(source))
    if len(matches) != 1:
        raise RuntimeError(
            f"down BMM入口匹配到{len(matches)}处，未安装补丁"
        )

    match = matches[0]
    indent = match.group("indent")
    replacement = (
        indent + "if _tirx_down_supported(self, z, w_dn):\n"
        + indent + "    y = _tirx_down(self, z, w_dn)\n"
        + indent + "else:\n"
        + indent
        + "    y = torch.bmm(z, w_dn.transpose(-2, -1)).to(dtype)"
    )
    return source[:match.start()] + replacement + source[match.end():]


@torch.inference_mode()
def prewarm_model(model):
    from inference.model.dit.dit_module import SparseMoEMLP

    modules = [
        m for m in model.modules()
        if isinstance(m, SparseMoEMLP)
    ]
    if not modules:
        raise RuntimeError("down预热未找到SparseMoEMLP")

    print(
        f"TIRX_PREWARM_DOWN: 开始 modules={len(modules)} "
        f"enabled={_ENABLED}",
        flush=True,
    )

    first = True
    for index, module in enumerate(modules):
        if module.training:
            raise RuntimeError("down预热要求eval模式")

        experts = module.experts
        device = experts[0].down_proj.weight.device
        if device.type != "cuda":
            raise RuntimeError("down预热要求模型已在GPU，当前不支持CPU offload")

        if (
            len(experts) != 3
            or any(
                tuple(e.down_proj.weight.shape) != (5120, 13652)
                for e in experts
            )
        ):
            raise RuntimeError("down预热只适用于当前三专家、全宽13652模型")

        # 与原始 _cache_key("bmm", None) 完全一致。
        key = ("bmm", torch.bfloat16, device.type, device.index, -1)
        cache = module._grouped_weight_cache

        if key not in cache:
            up = torch.stack(
                [e.up_gate_proj.weight.to(torch.bfloat16) for e in experts],
                dim=0,
            )
            dn = torch.stack(
                [e.down_proj.weight.to(torch.bfloat16) for e in experts],
                dim=0,
            )
            cache[key] = (up, dn)

        if _ENABLED:
            weight = cache[key][1]
            version = (
                None if torch.is_inference(weight) else weight._version
            )
            padded = torch.zeros(
                3, N, KP, dtype=torch.bfloat16, device=device
            )
            padded[:, :, :13652].copy_(weight)
            module._tirx_down_weight_cache = (weight, version, padded)

            # 两档内核各启动一次；不调用down()，避免消耗真实输入检查名额。
            if first:
                for mp in (2048, 4096):
                    x = torch.zeros(
                        mp, KP, dtype=torch.bfloat16, device=device
                    )
                    y = torch.empty(
                        mp, N, dtype=torch.bfloat16, device=device
                    )
                    ex = executable(mp)
                    for _ in range(3):
                        ex(x, padded[0], y)
                    torch.cuda.synchronize(device)
                    del x, y
                first = False

        print(
            f"TIRX_PREWARM_DOWN: 权重准备 {index + 1}/{len(modules)}",
            flush=True,
        )

    torch.cuda.synchronize()
    print("TIRX_PREWARM_DOWN: 完成", flush=True)

	import functools
	import os
	import re

	import torch

	_ENABLED = os.getenv("MAGI_TIRX_DOWN", "0") == "1"
	_VALIDATE = os.getenv("MAGI_TIRX_DOWN_VALIDATE", "0") == "1"
	_checked = set()
	_reported = False

	KP = 13696
	N = 5120


	@functools.lru_cache(maxsize=2)
	def executable(mp):
	import tirx_kernels.gemm.fp16_bf16_gemm as gemm
	from tirx_kernels.runner import compile_kernel

	cfg = dict(gemm.GEMM_CONFIGS[1024])
	cfg.update(
	cta_n=64,
	cta_k=128,
	pipe_depth=5,
	wb_pipe_depth=2,
	l2_group_size=2,
	overlap_epilogue=True,
	)

	original_cfg = gemm._cfg_for
	try:
	gemm._cfg_for = lambda n: dict(cfg)
	kernel = gemm.make_kernel("bf16", mp, N, KP)
	result = compile_kernel(kernel)
	finally:
	gemm._cfg_for = original_cfg

	print(f"TIRX_DOWN_COMPILED M={mp} N={N} K={KP}", flush=True)
	return result


	def supported(module, z, weight):
	return (
	_ENABLED
	and not torch.is_grad_enabled()
	and not module.training
	and z.is_cuda
	and weight.device == z.device
	and z.dtype == torch.bfloat16
	and weight.dtype == torch.bfloat16
	and z.ndim == 3
	and weight.ndim == 3
	and z.shape[0] == 3
	and 0 < z.shape[1] <= 4096
	and z.shape[2] == 13652
	and tuple(weight.shape) == (3, 5120, 13652)
	and z.is_contiguous()
	and weight.is_contiguous()
	and not z.requires_grad
	and not weight.requires_grad
	)


	def down(module, z, weight):
	global _reported

	m = z.shape[1]
	mp = 2048 if m <= 2048 else 4096

	# 保存原权重对象，避免地址复用导致误用缓存。
	# 权重被原地修改后，通过 version 重新建立缓存。
	# inference tensor没有版本计数；固定权重推理按对象身份缓存。
	version = None if torch.is_inference(weight) else weight._version
	cached = getattr(module, "_tirx_down_weight_cache", None)
	if (
	cached is None
	or cached[0] is not weight
	or cached[1] != version
	):
	padded_weight = torch.zeros(
	3, N, KP, dtype=weight.dtype, device=weight.device
	)
	padded_weight[:, :, :13652].copy_(weight)
	module._tirx_down_weight_cache = (
	weight, version, padded_weight
	)
	else:
	padded_weight = cached[2]

	# 全部清零，确保 K 尾部和 M 尾部没有未初始化值。
	padded_input = torch.zeros(
	3, mp, KP, dtype=z.dtype, device=z.device
	)
	padded_input[:, :m, :13652].copy_(z)

	padded_output = torch.empty(
	3, mp, N, dtype=z.dtype, device=z.device
	)
	ex = executable(mp)
	for e in range(3):
	ex(padded_input[e], padded_weight[e], padded_output[e])

	output = padded_output[:, :m, :].contiguous()

	if not _reported:
	print(
	f"TIRX_DOWN_ACTIVE input={tuple(z.shape)} "
	f"padded_M={mp} padded_K={KP}",
	flush=True,
	)
	_reported = True

	# 检查前三种真实形状；这不是全模型质量验证。
	key = tuple(z.shape)
	if _VALIDATE and key not in _checked and len(_checked) < 3:
	expected = torch.bmm(z, weight.transpose(-2, -1))
	finite = bool(torch.isfinite(output).all().item())

	# 分块统计，控制验证时的临时显存。
	a = output.reshape(-1)
	b = expected.reshape(-1)
	max_abs = 0.0
	error_sq = 0.0
	reference_sq = 0.0
	for start in range(0, a.numel(), 1 << 20):
	af = a[start:start + (1 << 20)].float()
	bf = b[start:start + (1 << 20)].float()
	delta = af - bf
	max_abs = max(max_abs, delta.abs().max().item())
	error_sq += delta.square().sum(dtype=torch.float64).item()
	reference_sq += bf.square().sum(dtype=torch.float64).item()

	relative_l2 = (error_sq / max(reference_sq, 1e-30)) ** 0.5
	print(
	f"TIRX_DOWN_CHECK input={tuple(z.shape)} "
	f"finite={finite} max_abs={max_abs:.8g} "
	f"relative_l2={relative_l2:.8g}",
	flush=True,
	)

	# 仅用于拦截明显异常，不能代替质量评价。
	if not finite or relative_l2 > 1e-3:
	raise RuntimeError("TIRx down真实输入检查未通过")

	_checked.add(key)

	return output


	def transform_source(source):
	pattern = re.compile(
	r"^(?P<indent>[ \t]*)"
	r"y\s=\storch\.bmm\("
	r"z,\sw_dn\.transpose\(-2,\s-1\)"
	r"\)\.to\(dtype\)[ \t]*$",
	re.MULTILINE,
	)
	matches = list(pattern.finditer(source))
	if len(matches) != 1:
	raise RuntimeError(
	f"down BMM入口匹配到{len(matches)}处，未安装补丁"
	)

	match = matches[0]
	indent = match.group("indent")
	replacement = (
	indent + "if _tirx_down_supported(self, z, w_dn):\n"
	+ indent + " y = _tirx_down(self, z, w_dn)\n"
	+ indent + "else:\n"
	+ indent
	+ " y = torch.bmm(z, w_dn.transpose(-2, -1)).to(dtype)"
	)
	return source[:match.start()] + replacement + source[match.end():]


	@torch.inference_mode()
	def prewarm_model(model):
	from inference.model.dit.dit_module import SparseMoEMLP

	modules = [
	m for m in model.modules()
	if isinstance(m, SparseMoEMLP)
	]
	if not modules:
	raise RuntimeError("down预热未找到SparseMoEMLP")

	print(
	f"TIRX_PREWARM_DOWN: 开始 modules={len(modules)} "
	f"enabled={_ENABLED}",
	flush=True,
	)

	first = True
	for index, module in enumerate(modules):
	if module.training:
	raise RuntimeError("down预热要求eval模式")

	experts = module.experts
	device = experts[0].down_proj.weight.device
	if device.type != "cuda":
	raise RuntimeError("down预热要求模型已在GPU，当前不支持CPU offload")

	if (
	len(experts) != 3
	or any(
	tuple(e.down_proj.weight.shape) != (5120, 13652)
	for e in experts
	)
	):
	raise RuntimeError("down预热只适用于当前三专家、全宽13652模型")

	# 与原始 _cache_key("bmm", None) 完全一致。
	key = ("bmm", torch.bfloat16, device.type, device.index, -1)
	cache = module._grouped_weight_cache

	if key not in cache:
	up = torch.stack(
	[e.up_gate_proj.weight.to(torch.bfloat16) for e in experts],
	dim=0,
	)
	dn = torch.stack(
	[e.down_proj.weight.to(torch.bfloat16) for e in experts],
	dim=0,
	)
	cache[key] = (up, dn)

	if _ENABLED:
	weight = cache[key][1]
	version = (
	None if torch.is_inference(weight) else weight._version
	)
	padded = torch.zeros(
	3, N, KP, dtype=torch.bfloat16, device=device
	)
	padded[:, :, :13652].copy_(weight)
	module._tirx_down_weight_cache = (weight, version, padded)

	# 两档内核各启动一次；不调用down()，避免消耗真实输入检查名额。
	if first:
	for mp in (2048, 4096):
	x = torch.zeros(
	mp, KP, dtype=torch.bfloat16, device=device
	)
	y = torch.empty(
	mp, N, dtype=torch.bfloat16, device=device
	)
	ex = executable(mp)
	for _ in range(3):
	ex(x, padded[0], y)
	torch.cuda.synchronize(device)
	del x, y
	first = False

	print(
	f"TIRX_PREWARM_DOWN: 权重准备 {index + 1}/{len(modules)}",
	flush=True,
	)

	torch.cuda.synchronize()
	print("TIRX_PREWARM_DOWN: 完成", flush=True)