vjindal/testing

setup.py

@@ -13,8 +13,8 @@ def get_default_dependencies():
 
     if platform in ["cuda", "cpu"]:
         return [
-            "torch>=2.1.2",
-            "triton>=2.3.1",
+            "torch==2.9.1.7+cu130",
+            "triton==3.4.0",
         ]
     elif platform == "rocm":
         return [
@@ -44,7 +44,8 @@ def get_optional_dependencies():
             "datasets>=2.19.2",
             "seaborn",
             "mkdocs-material",
-            "torchvision>=0.20",
+            "torchvision==0.24.1.6+cu130",
+            "numpy==1.26.4",
             "prek>=0.2.28",
         ]
     }

temp.py

@@ -0,0 +1,62 @@
+import os
+
+from transformers import AutoModelForCausalLM
+
+MODEL_DIR = "/shared/public/elr-models/Qwen/Qwen3-8B/9c925d64d72725edaf899c6cb9c377fd0709d9c5"
+
+os.environ["TRANSFORMERS_OFFLINE"] = "1"
+os.environ["HF_HUB_OFFLINE"] = "1"
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_DIR,
+    local_files_only=True,
+    trust_remote_code=True,
+    device_map="cpu",
+    torch_dtype="auto",
+    low_cpu_mem_usage=True,
+)
+model.eval()
+
+print(model)
+
+
+from liger_kernel.transformers import AutoLigerKernelForCausalLM
+
+model2 = AutoLigerKernelForCausalLM.from_pretrained(
+    MODEL_DIR,
+    local_files_only=True,
+    trust_remote_code=True,
+    device_map="cpu",
+    torch_dtype="auto",
+    low_cpu_mem_usage=True,
+)
+model2.eval()
+
+print(model2)
+
+# # Print "layers" (transformer blocks)
+# layers = None
+# for attr_path in [
+#     ("model", "layers"),  # Qwen/LLaMA-like
+#     ("model", "decoder", "layers"),
+#     ("transformer", "h"),  # GPT-2-like
+#     ("gpt_neox", "layers"),
+# ]:
+#     obj = model
+#     ok = True
+#     for a in attr_path:
+#         if not hasattr(obj, a):
+#             ok = False
+#             break
+#         obj = getattr(obj, a)
+#     if ok:
+#         layers = obj
+#         container_name = ".".join(attr_path)
+#         break
+
+# if layers is None:
+#     raise RuntimeError("Couldn't find a layers container on this model.")
+
+# print(f"{container_name} count = {len(layers)}")
+# for i, layer in enumerate(layers):
+#     print(f"{container_name}[{i}] -> {layer.__class__.__name__}")

temp2.py

@@ -0,0 +1,33 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from liger_kernel.transformers import AutoLigerKernelForCausalLM
+from liger_kernel.transformers.swiglu import LigerSwiGLUMLP
+
+
+MODEL_DIR = "/shared/public/elr-models/Qwen/Qwen3-8B/9c925d64d72725edaf899c6cb9c377fd0709d9c5"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
+model = AutoLigerKernelForCausalLM.from_pretrained(MODEL_DIR, torch_dtype="bfloat16").cuda()
+
+def hook(mod, inp, out):
+    x = inp[0]
+    print("MLP x:", tuple(x.shape), x.dtype, x.device, "intermediate:", mod.intermediate_size)
+
+# print(model)
+
+for m in model.modules():
+    print("module:", m)
+    print("module type:", type(m))
+    print("module name:", m.__class__.__name__)
+    if m.__class__.__name__ == "LigerSwiGLUMLP":
+        print("LigerSwiGLUMLP found")
+        m.register_forward_hook(hook)
+        break
+    print("--------------------------------")
+    if isinstance(m, LigerSwiGLUMLP):
+        m.register_forward_hook(hook)
+        break
+
+inputs = tokenizer("hello", return_tensors="pt").to("cuda")
+with torch.inference_mode():
+    _ = model(**inputs)          # triggers hook
+    # or: _ = model.generate(**inputs, max_new_tokens=1)

temp3.py

@@ -0,0 +1,73 @@
+import torch
+import triton
+import triton.testing
+from liger_kernel.ops.swiglu import _swiglu_forward_kernel
+
+@torch.no_grad()
+def bench_one(n_rows, n_cols=12288, dtype=torch.bfloat16, block_size=16384, warps=16, stages=4, iters=200):
+    a = torch.randn((n_rows, n_cols), device="cuda", dtype=dtype)
+    b = torch.randn_like(a)
+    c = torch.empty_like(a)
+    stride = c.stride(0)
+
+    def fn():
+        _swiglu_forward_kernel[(n_rows,)](
+            a, b, c, stride,
+            n_cols=n_cols,
+            BLOCK_SIZE=block_size,
+            num_warps=warps,
+            num_stages=stages,
+        )
+
+    # warmup
+    for _ in range(20):
+        fn()
+    torch.cuda.synchronize()
+
+    ms = triton.testing.do_bench(fn, rep=iters)
+    return ms
+
+def sweep():
+    n_cols = 12288
+    rows_list = [30000, 100000]  # represent different (batch*seq) regimes
+    configs = [
+        ("default", 16384, 4, 2),
+        ("default", 16384, 8, 2),
+        ("default", 16384, 16, 2),
+        ("default", 16384, 32, 2),
+        ("default", 16384, 4, 3),
+        ("default", 16384, 8, 3),
+        ("default", 16384, 16, 3),
+        ("default", 16384, 32, 3),
+        ("default", 16384, 4, 4),
+        ("default", 16384, 8, 4),
+        ("default", 16384, 16, 4),
+        ("default", 16384, 32, 4),
+        ("default", 16384, 4, 5),
+        ("default", 16384, 8, 5),
+        ("default", 16384, 16, 5),
+        ("default", 16384, 32, 5),
+        ("default", 16384, 4, 6),
+        ("default", 16384, 8, 6),
+        ("default", 16384, 16, 6),
+        ("default", 16384, 32, 6),
+        # ("ncols-8w", 12288, 8),
+        # ("ncols-16w", 12288, 16),
+        # ("ncols-32w", 12288, 32),
+        # ("big-32w", 32768, 32),
+    ]
+    for n_rows in rows_list:
+        print("\nrows:", n_rows)
+        best_ms = float('inf')
+        best_config = None
+        for name, bs, w, s in configs:
+            ms = bench_one(n_rows, n_cols=n_cols, block_size=bs, warps=w, stages=s)
+            # print(f"{name:10s}  BLOCK={bs:5d} warps={w:2d} stages={s:2d}  {ms:.4f} ms")
+            if ms < best_ms:
+                best_ms = ms
+                best_config = (name, bs, w, s)
+        if best_config is not None:
+            name, bs, w, s = best_config
+            print(f"BEST for rows={n_rows}: {name:10s}  BLOCK={bs:5d} warps={w:2d} stages={s:2d}  {best_ms:.4f} ms")
+
+sweep()