File size: 3,444 Bytes
e488ce7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | # SPDX-License-Identifier: GPL-3.0-or-later
"""Gate (a) for vram_lab: streamed block == stock DiTBlock on a random tiny block.
Run from the ComfyUI root with the ComfyUI venv:
CUDA_VISIBLE_DEVICES=1 python custom_nodes/ComfyUI-MAINodes/tests/test_vram_lab_block.py
Uses comfy.ops.disable_weight_init (bf16 weights, no quantization) so this
proves the schedule; the int8 row-exactness is a separate kernel property
measured 2026-08-18. Prints max abs diff per configuration; PASS threshold is
a handful of bf16 ulps at the residual's magnitude (matmul order differs
between chunked and whole GEMMs on cuBLAS, so bit-equality is not expected on
this path).
"""
import os
import sys
sys.path.insert(0, os.getcwd())
import torch
import comfy.options
comfy.options.enable_args_parsing()
import comfy.model_management # noqa: E402
import comfy.ops # noqa: E402
from comfy.ldm.minimax.model import DiTBlock, rope_rotation_table # noqa: E402
sys.path.insert(0, os.path.join(os.getcwd(), "custom_nodes"))
from importlib import import_module # noqa: E402
vram_lab = import_module("ComfyUI-MAINodes.vram_lab")
torch.manual_seed(0)
dev = "cuda"
dt = torch.bfloat16
hidden, heads, hd, ffn, t_dim = 512, 8, 64, 1024, 256
S = 6000
ops = comfy.ops.disable_weight_init
block = DiTBlock(hidden, heads, hd, ffn, t_dim, 1e-6, 1e-6, dtype=dt, device=dev, operations=ops)
with torch.no_grad():
for p in block.parameters():
p.copy_(torch.randn_like(p) * 0.02)
p.requires_grad_(False)
torch.set_grad_enabled(False)
x0 = torch.randn(S, hidden, device=dev, dtype=dt)
t_emb = torch.randn(3, t_dim, device=dev, dtype=dt)
# three segments with different modulation rows, like text | audio | video
segments = [(0, 300, 1), (300, 900, 2), (900, S, 0)]
rot_dim = 32
angles = torch.randn(S, rot_dim, device=dev) * 0.5
angles[:, rot_dim // 2:] = angles[:, :rot_dim // 2]
rope = rope_rotation_table(angles, dt) # [1, S, 1, rot/2, 2, 2]
with torch.no_grad():
ref = block(x0.clone(), t_emb, segments, rope, transformer_options={})
scale = ref.float().abs().max().item()
ulp = 2.0 ** (torch.floor(torch.log2(torch.tensor(scale))).item() - 7) # bf16 ulp at max magnitude
print(f"ref max |x| = {scale:.4f}, bf16 ulp there ~ {ulp:.2e}")
ok = True
for cfg in [dict(q_chunk=S, kv_chunk=S, mlp_chunk=S),
dict(q_chunk=1024, kv_chunk=1024, mlp_chunk=1024),
dict(q_chunk=700, kv_chunk=1500, mlp_chunk=0),
dict(q_chunk=1024, kv_chunk=1024, mlp_chunk=1024, kv_block=1024),
dict(q_chunk=512, kv_chunk=4096, mlp_chunk=2048, kv_block=700)]:
out = vram_lab.streamed_block_forward(block, x0.clone(), t_emb, segments, rope, {}, **cfg)
d = (out.float() - ref.float()).abs().max().item()
verdict = "PASS" if d <= 8 * ulp else "FAIL"
ok &= verdict == "PASS"
print(f"{verdict} {cfg}: max abs diff {d:.3e} ({d / ulp:.1f} ulp)")
# segment-modulation sanity: a wrong segment clip would show as a large diff at boundaries
out = vram_lab.streamed_block_forward(block, x0.clone(), t_emb, segments, rope, {}, q_chunk=250, kv_chunk=250, mlp_chunk=250)
d = (out.float() - ref.float()).abs()
print(f"chunk 250 (cuts every segment): max {d.max().item():.3e}, at rows 299/300 {d[299].max().item():.2e}/{d[300].max().item():.2e}")
ok &= d.max().item() <= 8 * ulp
print("ALL PASS" if ok else "SOME FAIL")
sys.exit(0 if ok else 1)
|