Merge pull request #5 from aymuos15/fix-blackwell-target

Drop a-suffix from Blackwell bw example arch targets

Author: patrick-toulme

examples/blackwell/layer_norm.py

@@ -1,7 +1,7 @@
 """Blackwell LayerNorm example using the maintained pyptx kernel path.
 
 Run ``python examples/blackwell/layer_norm.py`` to execute both a ``jax.jit``
-path and a PyTorch eager path on ``sm_100a``.
+path and a PyTorch eager path on ``sm_100``.
 """
 from __future__ import annotations
 
@@ -32,7 +32,7 @@ def _pick_rows_per_cta(B: int) -> int:
 def build_layer_norm(B: int, N: int, *, eps: float = 1e-5, rows_per_cta: int | None = None):
     if rows_per_cta is None:
         rows_per_cta = _pick_rows_per_cta(B)
-    return _build_layer_norm(B, N, eps=eps, rows_per_cta=rows_per_cta, arch="sm_100a")
+    return _build_layer_norm(B, N, eps=eps, rows_per_cta=rows_per_cta, arch="sm_100")
 
 
 def _run_jax_case(B: int, N: int) -> None:

examples/blackwell/rms_norm.py

@@ -1,7 +1,7 @@
 """Blackwell RMSNorm example using the maintained pyptx kernel path.
 
 Run ``python examples/blackwell/rms_norm.py`` to execute both a ``jax.jit``
-path and a PyTorch eager path on ``sm_100a``.
+path and a PyTorch eager path on ``sm_100``.
 """
 from __future__ import annotations
 
@@ -24,7 +24,7 @@
 
 
 def build_rms_norm(B: int, N: int, *, eps: float = 1e-6):
-    return _build_rms_norm(B, N, eps=eps, arch="sm_100a")
+    return _build_rms_norm(B, N, eps=eps, arch="sm_100")
 
 
 def _run_jax_case(B: int, N: int) -> None:

examples/blackwell/swiglu.py

@@ -1,7 +1,7 @@
 """Blackwell SwiGLU example using the maintained pyptx kernel path.
 
 Run ``python examples/blackwell/swiglu.py`` to execute both a ``jax.jit``
-path and a PyTorch eager path on ``sm_100a``.
+path and a PyTorch eager path on ``sm_100``.
 """
 from __future__ import annotations
 
@@ -32,7 +32,7 @@ def _pick_rows_per_cta(M: int) -> int:
 def build_fused_silu_mul(M: int, F: int, *, rows_per_cta: int | None = None):
     if rows_per_cta is None:
         rows_per_cta = _pick_rows_per_cta(M)
-    return _build_fused_silu_mul(M, F, rows_per_cta=rows_per_cta, arch="sm_100a")
+    return _build_fused_silu_mul(M, F, rows_per_cta=rows_per_cta, arch="sm_100")
 
 
 def _run_jax_case(M: int, F: int) -> None: