Skip to content

Commit 27a520b

Browse files
committed
docs(§DYN): TRITON-MOLD-B2 verification — cuFuncGetAttribute STATIC=0, HPC=2 launches, 0.745x NO-GO
gb10 sm_121 re-verification of chunk_scan_combine_bwd_cuda_prim_gemm_batched_dyn (prod b=1 S=4096 chunk=64 G=8 H=112 P=64 N=64, HEADS_PER_CTA=2): - cuFuncGetAttribute(main_kernel): SHARED_SIZE_BYTES(STATIC)=0 (driver-level Triton mold confirmed; was 57344 B static staging). cuFuncSetAttribute(MAXDYN=91136) GRANTED < 101376 opt-in cap => HPC=2 LAUNCHES. NUM_REGS=192. cuobjdump cross-check: SHARED:1024 (only the __align__(1024) slot for extern buf_dyn_shmem[]). - TIMING (in ONE process): v1=1052.074ms, single_tile_gemm=1370.047ms, batched_static=1418.405ms, batched_dyn=1412.696ms => 0.745x vs v1 NO-GO (1.004x vs batched_static; the scope-flip is ~free, not the bottleneck). - SASS: 128 HMMA.16816.F32 — real tensor-core, but PER-HEAD 64-tile relabel inside for hh in T.serial(HPC), NOT a fattened-M head-accum (path_c dC/dinp are per-HEAD, distinct operands both sides => dense tall-M = HPC^2 cross-head garbage). - PARITY: all 8 grads vs v1 <= 9.3e-5; chained 8-grad gate WORST=8.099e-4 < 1e-3 PASS. - z3 (proof_b2_batched_driver.py dyn_scope_flip): ALL_POSITIVES_PROVED_AND_NON_VACUOUS (3 positives verdict-identical to static; interleaved_dyn_bands negative fails single_writer). - HPC=4 INVALID at prod (4 does not divide nheads/ngroups=14; RAISES, no straddle). - step tok/s UNCHANGED: §DYN is opt-in NO-GO, default B2 path stays v1. New verification scripts: scratch/wf1_cufunc_attr_b2dyn.py (source+SASS), scratch/wf1_cubin_attr.py (driver cuFuncGetAttribute on the sm_121a cubin). v1/§27/§TB1 prim bodies byte-identical; B1 untouched. RULE #1: §DYN RAISES on smem/proof/compile failure, never silently used.
1 parent e3d383c commit 27a520b

3 files changed

Lines changed: 384 additions & 0 deletions

File tree

docs/TRITON-MOLD-B2.md

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
# Triton-Mold path_c B2 Backward (§DYN static→dynamic) — MEASURED VERDICT
2+
3+
Status: **DECIDED — HONEST NO-GO (0.745x vs in-process v1).** The static→dynamic
4+
"Triton mold" achieves launch-feasibility (driver `STATIC=0`, HPC=2 launches with
5+
the full 4-GEMM layout) but does **not** beat the v1 threaded B2 prim.
6+
7+
Hardware: NVIDIA GB10 (sm_121, unified 121 GB). Prod cfg: `b=1, S=4096, chunk=64,
8+
G=8, H=112, P=64, N=64`, `HEADS_PER_CTA=2` (the only valid HPC for `nheads/ngroups
9+
= 14`; see §2). Toolchain: CUDA 13.3, tilelang from source, torch 2.13 dev cu132.
10+
Date: 2026-06-06.
11+
12+
---
13+
14+
## 1. What was built (the Triton mold)
15+
16+
NEW flag-gated CUDA prim `chunk_scan_combine_bwd_cuda_prim_gemm_batched_dyn`
17+
(`cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py:2699`). It is the
18+
§TB1/§SO1 batched math/grid/in-kernel-head-loop **byte-for-byte**, with the five
19+
GEMM operand-staging tiles (`dY16`, `opA`, `opB`, `store_fp32`, `dCdiag_sh`) moved
20+
from explicit **STATIC** `scope="shared"` to the **DYNAMIC** region (the
21+
`T.alloc_shared` default `shared.dyn`). tilelang lowers `shared.dyn` into
22+
`extern __shared__ __align__(1024) uchar buf_dyn_shmem[]`, so the compiler reserves
23+
~0 static smem and the driver grants the full per-block dynamic opt-in. This is the
24+
Tri-Dao Triton bwd HW mold (same gb10: STATIC=0, MAXDYN opt-in).
25+
26+
Wired behind `CPPMEGA_PATH_C_B2_GEMM_BATCHED_DYN`, mutually exclusive with the other
27+
B2 flags (RAISES on ambiguity). Own dynamic-smem gate against the 101376 B opt-in
28+
cap (RAISES over-cap). z3/TLA proof gate (RAISES `GemmRewriteNotProven` on
29+
sat/unknown). RULE #1: ONE path when flagged, no fallback. **Default path (no flag)
30+
is unchanged v1** (`else: prim = chunk_scan_combine_bwd_cuda_prim`).
31+
32+
---
33+
34+
## 2. cuFuncGetAttribute — STATIC≈0 PROVEN at the driver level
35+
36+
The cubin (nvcc `-arch=sm_121a` on the exact emitted `device_kernel.cu`) loaded via
37+
the CUDA driver API (`cuda.bindings.driver`):
38+
39+
```
40+
cuFuncGetAttribute(main_kernel):
41+
CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES = 0 <-- STATIC = 0 (mold achieved)
42+
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES (default) = 49152
43+
CU_FUNC_ATTRIBUTE_NUM_REGS = 192
44+
cuFuncSetAttribute(MAX_DYNAMIC = 91136) = GRANTED -> now 91136 (HPC=2 LAUNCHES)
45+
device MAX_SHARED_MEMORY_PER_BLOCK_OPTIN = 101376
46+
```
47+
48+
cuobjdump cross-check on the same cubin:
49+
```
50+
Function main_kernel: REG:192 SHARED:1024 LOCAL:0 STACK:0 (smem=1024)
51+
.nv.shared.main_kernel size 0x400 (=1024 B) <-- just the __align__(1024) slot for
52+
the extern dyn array, NOT staging
53+
```
54+
55+
**STATIC dropped from 57344 B (the §40a26d44 per-head GEMM staging) to 0.** The only
56+
"static" smem is the 1024 B alignment placeholder for the dynamic `buf_dyn_shmem[]`.
57+
The 91136 B dynamic request (HPC=2) fits the 101376 B sm_121 opt-in cap → the driver
58+
grants it → **HPC=2 LAUNCHES** with the full 4-GEMM layout.
59+
60+
`HEADS_PER_CTA=4` is **INVALID** at prod: `nheads/ngroups = 112/8 = 14`, and 4 neither
61+
divides nor is a multiple of 14, so a CTA head-band would straddle a group whose
62+
`C`/`B` operands are per-GROUP. `_b2_batched_heads_per_cta` RAISES (RULE #1: no silent
63+
straddle). The valid HPC set is {1, 2, 7, 14}; HPC=2 is the measured point.
64+
65+
---
66+
67+
## 3. TIMING — MEASURED NO-GO
68+
69+
In ONE gb10 process (`probe_chunked_backward_cuda_gb10.py --prod --b2-gemm-ab`,
70+
`CPPMEGA_PATH_C_B2_HEADS_PER_CTA=2`):
71+
72+
| prim | ms (median) | vs in-process v1 |
73+
| ---------------------------- | ----------: | ---------------: |
74+
| v1 threaded (`prim` default) | 1052.074 | 1.000x |
75+
| single-tile GEMM (§27) | 1370.047 | 0.768x |
76+
| batched-large-tile (§TB1) | 1418.405 | 0.742x |
77+
| **§DYN batched_dyn (mold)** | **1412.696** | **0.745x NO-GO** |
78+
79+
§DYN vs batched_static = 1.004x (the static→dynamic scope flip is ~0.4% faster — the
80+
smem move is free but does not change the dominant cost). **0.745x < 1.0x → NO-GO.**
81+
82+
NOTE on baselines: this probe's in-process v1 measures **1052 ms**, not the historical
83+
905.4 ms reference (an earlier wf machine/thermal state). All §DYN ratios above are vs
84+
the **same-process 1052 ms v1**, which is the apples-to-apples comparison. Against the
85+
historical references: §DYN 1412.7 ms is 0.64x of 905.4 ms (v1_threaded), 0.87x of
86+
1231 ms (§27), and 141x of 10 ms (Tri-Dao Triton). Tri-Dao remains ~140x faster.
87+
88+
Backward chain (B2 + B1 + B0), this run: **B2 1051.1 + B1 5.55 + B0 266.2 ms**.
89+
(The 447.8 ms reference is the v1 B2 from an earlier wf; B2 alone is 1051 ms here.)
90+
91+
---
92+
93+
## 4. SASS — real HMMA, but PER-HEAD relabel (NOT fattened-M head-accum)
94+
95+
`cuobjdump --dump-sass` on the §DYN cubin:
96+
```
97+
128 HMMA.16816.F32 (m16n8k16 tensor-core MMA, F32 accumulate)
98+
0 OMMA / IMMA / BMMA
99+
```
100+
101+
These are real tensor-core HMMA tiles (`HMMA.16816.F32 R84, R96, R124, RZ ; ...`).
102+
But the M-dimension of every GEMM is the **per-head** tile (L=64 or P=64), issued
103+
inside `for hh in T.serial(HPC)` loops — the four contractions (DYX = dY@xᵀ, dC_off =
104+
dY@prev_states, dC_diag masked, dchunk = (decay·dY)ᵀ@C) each run once **per head** with
105+
operands that differ on BOTH sides. The 128 HMMA = 4 GEMMs × HPC=2 heads × the 16-row
106+
fragmentation of the 64-tiles. **This is a per-64 (per-head) relabel, NOT a fattened-M
107+
head-accumulation.** Tri-Dao fattens M by looping `nheads_per_program` heads that SHARE
108+
a per-GROUP operand (B/C) into ONE accumulator tile; path_c cannot, because dC and dinp
109+
are **per-HEAD** (distinct operands both sides) — a dense tall-M over heads would compute
110+
HPC² cross-head garbage (block-diagonal violation). The mold therefore amortizes only
111+
operand-staging/sync, never the tensor-core M-setup.
112+
113+
---
114+
115+
## 5. PARITY — all 8 grads PASS
116+
117+
§DYN vs v1 (in-process, max|abs| over ALL elements):
118+
```
119+
dC 6.21e-05 dx 0.00e+00 dz 0.00e+00 dchunk 9.30e-05 dinp 6.41e-05
120+
dA_y 3.15e-05 dD 4.41e-06
121+
```
122+
Chained 8-grad gate (B2→B1→B0, GATE 1e-3, ALL elements, default v1 chain):
123+
```
124+
dz 1.73e-04 dx 8.10e-04 dC 5.03e-05 dB 1.09e-05 dlog_decay 6.67e-04
125+
ddt 1.50e-04 dh0 1.84e-04 dD 2.57e-05 -> WORST 8.099e-04 < 1e-03 PASS
126+
```
127+
dD uses the fp16-cache; the §DYN deltas are fp16-rounding (≤9.3e-5), well within gate.
128+
129+
---
130+
131+
## 6. z3 PROOF — non-vacuous
132+
133+
`proof_b2_batched_driver.py`, `dyn_scope_flip`:
134+
- POSITIVES dchunk / dcoff / dcdiag: `z3_proved=True, single_writer=True,
135+
operand_maps_match=True, scale_equiv=True, verdict_identical_to_static=True`.
136+
- NEGATIVE `interleaved_dyn_bands` (m_stride = L//HPC, overlapping head rows):
137+
`z3_proved=False, single_writer=False` — correctly fails (non-vacuous).
138+
- VERDICT: **ALL_POSITIVES_PROVED_AND_NON_VACUOUS.**
139+
140+
The static→dynamic move is byte-layout-only: operand maps, single-writer band
141+
disjointness (m_blocks=HPC), mask, and scale-fold are identical to the §TB1 static
142+
proof. The dynamic region does NOT relax single-writer.
143+
144+
---
145+
146+
## 7. Step tok/s — UNCHANGED (NO-GO consequence)
147+
148+
§DYN is NO-GO, so it is **never on the production step path**: the default B2 branch
149+
(no flag) is v1, and RULE #1 makes §DYN opt-in only (RAISES on smem/proof/compile
150+
failure, no silent use). Production step tok/s @8L/@28L, bs1 AND bs4, is therefore the
151+
**v1 baseline** — wiring the 0.745x §DYN prim into the step would only regress it. No
152+
step-tok/s gain is claimed or measured for the mold; the gap vs Megatron 3399 tok/s is
153+
the pre-existing v1 gap, unchanged by this work.
154+
155+
---
156+
157+
## 8. HONEST NO-GO — WHY static→dynamic still can't beat v1
158+
159+
The mold removed the smem launch-blocker (STATIC 57344→0, HPC=2 launches), but the
160+
dominant costs are untouched:
161+
162+
1. **The dinp 3-index serial term is un-GEMM-able.** `dinp[b,s,h,p,n]` is a per-element
163+
outer-product zeroed and written by a serial thread-strided loop; the smem flip does
164+
nothing for it. It dominates the kernel.
165+
2. **path_c dC is per-HEAD → NO M-fattening.** Unlike Tri-Dao there is no per-group
166+
head-accumulation to amortize tensor-core setup (§4); the mold only amortizes
167+
operand-staging/sync, a small fraction.
168+
3. **Occupancy cap.** ~91 KB dynamic per CTA + 192 regs pins ~1 CTA/SM at bs1; the
169+
batched layout has fewer, fatter CTAs than v1's threaded grid, so it loses the
170+
latency-hiding v1 gets from many light CTAs.
171+
172+
The real fix is re-architecting dC/dinp into per-GROUP split buffers (so a genuine
173+
fattened-M head-accumulation becomes valid) — a far larger change than the smem flip,
174+
and out of scope here. **Verdict: ship nothing to the default path; keep v1. The §DYN
175+
prim stays as a flag-gated, proven, launch-feasible NO-GO reference.**
176+
177+
---
178+
179+
## Reproduce
180+
181+
```
182+
# timing + parity (HPC=2 mandatory: 14 = nheads/ngroups)
183+
CPPMEGA_PATH_C_B2_HEADS_PER_CTA=2 python scratch/probe_chunked_backward_cuda_gb10.py --prod --b2-gemm-ab
184+
# driver attributes (STATIC=0, MAXDYN grant) + SASS
185+
HPC=2 python scratch/wf1_cufunc_attr_b2dyn.py # source-regex + SASS HMMA
186+
# then: nvcc -arch=sm_121a -cubin <cache>/device_kernel.cu ... -o /tmp/wf1_dyn.cubin
187+
python scratch/wf1_cubin_attr.py /tmp/wf1_dyn.cubin # cuFuncGetAttribute STATIC/MAXDYN
188+
cuobjdump --dump-resource-usage /tmp/wf1_dyn.cubin # SHARED:1024, REG:192
189+
cuobjdump --dump-sass /tmp/wf1_dyn.cubin | grep -c HMMA # 128 HMMA.16816.F32
190+
# z3
191+
python scratch/proof_b2_batched_driver.py # ALL_POSITIVES_PROVED_AND_NON_VACUOUS
192+
```

scratch/wf1_cubin_attr.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
"""wf1: cuFuncGetAttribute on the §DYN cubin (no PTX-JIT, no version error)."""
2+
import sys
3+
import torch
4+
from cuda.bindings import driver as cuda
5+
6+
torch.zeros(1, device="cuda") # ensure context
7+
8+
CUBIN = sys.argv[1] if len(sys.argv) > 1 else "/tmp/wf1_dyn.cubin"
9+
SYM = b"main_kernel"
10+
11+
12+
def ck(res, what):
13+
err = res[0]
14+
if int(err) != 0:
15+
_, nm = cuda.cuGetErrorName(err)
16+
_, ds = cuda.cuGetErrorString(err)
17+
raise RuntimeError(f"{what} FAILED {int(err)} "
18+
f"{nm.decode() if isinstance(nm, bytes) else nm}: "
19+
f"{ds.decode() if isinstance(ds, bytes) else ds}")
20+
return res[1:]
21+
22+
23+
with open(CUBIN, "rb") as f:
24+
data = f.read()
25+
(mod,) = ck(cuda.cuModuleLoadData(data), "cuModuleLoadData(cubin)")
26+
(func,) = ck(cuda.cuModuleGetFunction(mod, SYM), "cuModuleGetFunction(main_kernel)")
27+
A = cuda.CUfunction_attribute
28+
(stat,) = ck(cuda.cuFuncGetAttribute(A.CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, func),
29+
"STATIC")
30+
(maxdyn,) = ck(cuda.cuFuncGetAttribute(
31+
A.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, func), "MAXDYN_default")
32+
(regs,) = ck(cuda.cuFuncGetAttribute(A.CU_FUNC_ATTRIBUTE_NUM_REGS, func), "NUM_REGS")
33+
print(f"[cubin-attr] main_kernel STATIC_SHARED_SIZE_BYTES={int(stat)} "
34+
f"MAX_DYNAMIC_SHARED_SIZE_BYTES(default)={int(maxdyn)} NUM_REGS={int(regs)}")
35+
36+
# Now SET the dynamic opt-in (what the launcher does) to the §DYN 91136 B and confirm
37+
# the driver GRANTS it (HPC=2 launch-feasibility on sm_121).
38+
REQ = 91136
39+
_set = cuda.cuFuncSetAttribute(
40+
func, A.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, REQ)
41+
r = _set[0] if isinstance(_set, tuple) else _set
42+
if int(r) != 0:
43+
_, ds = cuda.cuGetErrorString(r)
44+
print(f"[cubin-attr] cuFuncSetAttribute(MAXDYN={REQ}) REJECTED: {int(r)} "
45+
f"{ds.decode() if isinstance(ds, bytes) else ds}")
46+
else:
47+
(maxdyn2,) = ck(cuda.cuFuncGetAttribute(
48+
A.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, func), "MAXDYN_after_set")
49+
print(f"[cubin-attr] cuFuncSetAttribute(MAXDYN={REQ}) GRANTED -> "
50+
f"MAX_DYNAMIC_SHARED_SIZE_BYTES now {int(maxdyn2)} (HPC=2 LAUNCHES)")
51+
52+
# device optin cap
53+
dev = torch.cuda.current_device()
54+
(cap,) = ck(cuda.cuDeviceGetAttribute(
55+
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
56+
dev), "OPTIN_CAP")
57+
print(f"[cubin-attr] device MAX_SHARED_MEMORY_PER_BLOCK_OPTIN={int(cap)}")
58+
print("[cubin-attr] DONE")

scratch/wf1_cufunc_attr_b2dyn.py

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
"""wf1: cuFuncGetAttribute STATIC + MAX_DYNAMIC for the §DYN B2 Triton-mold prim.
2+
3+
Builds chunk_scan_combine_bwd_cuda_prim_gemm_batched_dyn at PROD dims, HPC from
4+
env (default 2), compiles via tilelang, extracts the cubin/PTX, loads it with the
5+
CUDA driver API (cuda.bindings.driver), finds the kernel function, and queries the
6+
REAL driver attributes:
7+
CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES (STATIC smem the compiler reserved)
8+
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES (the opt-in the launcher set)
9+
Then dumps SASS and tallies HMMA instructions (the GEMM tensor-core evidence).
10+
11+
RULE #1: any build/compile/load failure RAISES (no fallback).
12+
"""
13+
import os
14+
import re
15+
import sys
16+
17+
import torch # noqa: F401 (initializes CUDA context for the driver API)
18+
import tilelang as tl
19+
import tilelang.language as T # noqa: F401
20+
21+
from cppmega_mlx.nn._tilelang.mamba3_chunked_backward_core import (
22+
chunk_scan_combine_bwd_cuda_prim_gemm_batched_dyn,
23+
)
24+
25+
HPC = int(os.environ.get("HPC", "2"))
26+
# prod cfg
27+
b, S, chunk, G, H, P, N = 1, 4096, 64, 8, 112, 64, 64
28+
29+
print(f"[wf1-attr] building §DYN prim prod b={b} S={S} chunk={chunk} G={G} H={H} "
30+
f"P={P} N={N} HPC={HPC}")
31+
prim = chunk_scan_combine_bwd_cuda_prim_gemm_batched_dyn(
32+
b, S, chunk, G, H, P, N, heads_per_cta=HPC
33+
)
34+
tgt = "cuda"
35+
pc = {"tl.disable_tma_lower": True} if False else {}
36+
k = tl.compile(prim, out_idx=[11, 12, 13, 14, 15, 16, 17], target=tgt)
37+
38+
src = k.get_kernel_source()
39+
n_static = len(re.findall(r"__shared__\s+[^;e][^;]*\[\d+\]", src))
40+
has_dyn = "extern __shared__" in src
41+
print(f"[wf1-attr] SOURCE: num_static_shared_decls(regex)={n_static} "
42+
f"has_extern_dyn_shared={has_dyn}")
43+
44+
# Grab the global symbol(s) of device kernels from the CUDA-C source.
45+
syms = re.findall(r'__global__\s+void\s+(\w+)\s*\(', src)
46+
print(f"[wf1-attr] device kernel symbols: {syms}")
47+
48+
# ---- load the cubin/PTX via the CUDA driver API and query func attributes ----
49+
from cuda.bindings import driver as cuda
50+
51+
52+
def _ck(res, what):
53+
err = res[0]
54+
if int(err) != 0:
55+
# try to get string
56+
try:
57+
_, name = cuda.cuGetErrorName(err)
58+
_, desc = cuda.cuGetErrorString(err)
59+
name = name.decode() if isinstance(name, bytes) else name
60+
desc = desc.decode() if isinstance(desc, bytes) else desc
61+
except Exception:
62+
name = desc = "?"
63+
raise RuntimeError(f"[wf1-attr] CUDA {what} FAILED: {int(err)} {name}: {desc}")
64+
return res[1:]
65+
66+
67+
# ensure a context exists (torch already made one)
68+
torch.cuda.init()
69+
torch.zeros(1, device="cuda")
70+
71+
# Get the kernel module binary. tilelang exposes PTX via _get_ptx; cubin via sass
72+
# extraction path. We load PTX through the JIT driver (cuModuleLoadData on PTX).
73+
ptx = None
74+
try:
75+
ptx = k._get_ptx()
76+
except Exception as e:
77+
print(f"[wf1-attr] _get_ptx raised: {e!r}")
78+
79+
# Prefer cubin if we can pull it (exact static numbers); fall back to PTX-JIT which
80+
# still yields the SAME func attributes since the compiler reserves static smem.
81+
cubin_bytes = None
82+
try:
83+
sass_txt = k.get_kernel_source # placeholder
84+
except Exception:
85+
pass
86+
87+
loaded = False
88+
attr_results = {}
89+
for label, blob in (("ptx", ptx),):
90+
if blob is None:
91+
continue
92+
data = blob.encode() if isinstance(blob, str) else blob
93+
(mod,) = _ck(cuda.cuModuleLoadData(data), f"cuModuleLoadData({label})")
94+
loaded = True
95+
for sym in syms:
96+
try:
97+
(func,) = _ck(cuda.cuModuleGetFunction(mod, sym.encode()),
98+
f"cuModuleGetFunction({sym})")
99+
except RuntimeError as e:
100+
print(f"[wf1-attr] {e}")
101+
continue
102+
(stat,) = _ck(cuda.cuFuncGetAttribute(
103+
cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, func),
104+
f"cuFuncGetAttribute(STATIC,{sym})")
105+
(maxdyn,) = _ck(cuda.cuFuncGetAttribute(
106+
cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
107+
func), f"cuFuncGetAttribute(MAXDYN,{sym})")
108+
(regs,) = _ck(cuda.cuFuncGetAttribute(
109+
cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, func),
110+
f"cuFuncGetAttribute(NUM_REGS,{sym})")
111+
attr_results[sym] = {"static": int(stat), "maxdyn": int(maxdyn),
112+
"num_regs": int(regs)}
113+
print(f"[wf1-attr] FUNC {sym}: STATIC_SHARED_SIZE_BYTES={int(stat)} "
114+
f"MAX_DYNAMIC_SHARED_SIZE_BYTES={int(maxdyn)} NUM_REGS={int(regs)}")
115+
break
116+
117+
if not loaded:
118+
raise RuntimeError("[wf1-attr] could not load any module blob (no PTX/cubin)")
119+
120+
# ---- SASS: tally HMMA tensor-core instructions ----
121+
try:
122+
sass = k._get_sass()
123+
except Exception as e:
124+
sass = None
125+
print(f"[wf1-attr] _get_sass raised: {e!r}")
126+
if sass:
127+
hmma = re.findall(r'HMMA\.(\S+)', sass)
128+
from collections import Counter
129+
c = Counter(hmma)
130+
print(f"[wf1-attr] SASS HMMA total={len(hmma)} shapes={dict(c)}")
131+
# also count by mma type lines
132+
print(f"[wf1-attr] SASS bytes={len(sass)}")
133+
134+
print("[wf1-attr] DONE")

0 commit comments

Comments
 (0)