Skip to content

optimze-elect-one-compilation-hint #504

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yanqinz2 merged 1 commit into from
Aug 7, 2026
Merged

optimze-elect-one-compilation-hint #504

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 17 additions & 9 deletions benchmark/gemm/frost/benchmark_block_scale_matmul.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,8 +145,11 @@ def _mkdata(batch: int, M: int, N: int, K: int, combo: str):
sfb_log = _rand_e8m0((batch, N, sf_k), dev)

c = torch.empty(batch, M, N, dtype=torch.bfloat16, device=dev)
sfa = torch.cat([_to_blocked(sfa_log[i]) for i in range(batch)]).view(batch, M, sf_k)
sfb = torch.cat([_to_blocked(sfb_log[i]) for i in range(batch)]).view(batch, N, sf_k)
# The blocked blob carries the PADDED dims (whole 128-row x 4-SF-K atoms) —
# that is what the kernel reads. The graph keeps the logical [batch, M, sf_k].
sf_k_pad = _ceil_div(sf_k, 4) * 4
sfa = torch.cat([_to_blocked(sfa_log[i]) for i in range(batch)]).view(batch, _ceil_div(M, 128) * 128, sf_k_pad)
sfb = torch.cat([_to_blocked(sfb_log[i]) for i in range(batch)]).view(batch, _ceil_div(N, 128) * 128, sf_k_pad)
return a, b, c, sfa, sfb


Expand Down Expand Up @@ -200,7 +203,7 @@ def _rotating(fn_of_set: Callable, pool: list) -> Callable:
return lambda i: fn_of_set(pool[i % n])


def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str):
def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str, verbose: bool = True):
"""(label, call) for cuBLAS's OWN block-scaled GEMM of this combo, via
torch.nn.functional.scaled_mm (cuBLASLt) — the apples-to-apples reference.
Returns None if the env can't run it (some cuBLASLt builds return no
Expand All @@ -217,7 +220,6 @@ def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str):
is_fp4, bs, _, _ = _COMBOS[combo]
ru = lambda x, m: ((x + m - 1) // m) * m
cd = lambda a, b: (a + b - 1) // b
sw = [SwizzleType.SWIZZLE_32_4_4]

if is_fp4:
a = torch.randint(0, 256, (M, K // 2), dtype=torch.uint8, device=dev).view(torch.float4_e2m1fn_x2)
Expand All @@ -235,11 +237,14 @@ def _scaled_mm_ref(batch: int, M: int, N: int, K: int, combo: str):
gb = torch.ones(1, device=dev)
recipe = [ScalingType.BlockWise1x16, ScalingType.TensorWise]
scA, scB = [sa, ga], [sb, gb]
# scaled_mm wants one swizzle PER scale in the recipe.
sw = [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE]
else: # mxfp4 / mxfp8 — no global scale
sa = torch.randn(na, device=dev).to(torch.float8_e8m0fnu)
sb = torch.randn(nb, device=dev).to(torch.float8_e8m0fnu)
recipe = [ScalingType.BlockWise1x32]
scA, scB = [sa], [sb]
sw = [SwizzleType.SWIZZLE_32_4_4]

# scaled_mm has no batched form, so a batched ref is B back-to-back single-GEMM
# calls (same operands); FLOPS counts all B so throughput stays comparable.
Expand All @@ -252,7 +257,10 @@ def call():
try:
call()
torch.cuda.synchronize()
except Exception:
except Exception as e:
# Print it — an unrunnable env and an API-shape bug look identical otherwise.
if verbose:
print(f" [scaled_mm '{combo}' reference unavailable: {type(e).__name__}: {e}]")
return None
suffix = f" ×{batch}" if batch > 1 else ""
return f"cuBLAS {combo} scaled_mm{suffix}", call
Expand All @@ -265,15 +273,15 @@ def _make_reference(batch: int, M: int, N: int, K: int, combo: str, ref_mode: st
/ bf16 (dense BF16) / auto (scaled_mm, falling back to BF16)."""
dev = "cuda"
if ref_mode in ("auto", "scaled_mm"):
ref = _scaled_mm_ref(batch, M, N, K, combo)
ref = _scaled_mm_ref(batch, M, N, K, combo, verbose=verbose)
if ref is not None:
return ref
if ref_mode == "scaled_mm":
sys.exit(
f"--ref scaled_mm: cuBLAS block-scaled GEMM for '{combo}' is not "
f"runnable here — torch.nn.functional.scaled_mm raised at the "
f"cuBLASLt heuristic (no algorithm for this driver / cuBLASLt "
f"build). Use --ref bf16, or --ref auto to fall back automatically."
f"runnable here — torch.nn.functional.scaled_mm raised (see the "
f"message above). Use --ref bf16, or --ref auto to fall back "
f"automatically."
)
if verbose:
print(" [note: scaled_mm reference unavailable in this env — " "falling back to dense BF16 cuBLAS]")
Expand Down
3 changes: 2 additions & 1 deletion benchmark/gemm/frost/benchmark_block_scale_matmul_swiglu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,8 @@ def _mk(rows):

def _sf(rows):
log = torch.randint(1, 4, (B, rows, sf_k), device=dev).to(torch.float8_e4m3fn)
blk = torch.stack([_to_blocked(log[b]) for b in range(B)]).view(B, rows, sf_k)
# _to_blocked pads to whole 128-row x 4-SF-K atoms — view the PADDED dims.
blk = torch.stack([_to_blocked(log[b]) for b in range(B)]).view(B, _ceil_div(rows, 128) * 128, _ceil_div(sf_k, 4) * 4)
return log, blk

sfa_log, sfa_b = _sf(M)
Expand Down
59 changes: 30 additions & 29 deletions python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,7 @@ def _kernel(

warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
elect_one = nvvm.elect_sync()

tidx = cute.arch.thread_idx()[0]
bidx = cute.arch.block_idx()[0]
Expand Down Expand Up @@ -252,7 +253,7 @@ def _kernel(
num_consumer_warps_per_cta = 7
clc_empty_count = num_consumer_warps_per_cta * cluster_size
if warp_idx == 0:
if nvvm.elect_sync():
if elect_one:
for i in range(ab_stages):
nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1)
nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count)
Expand Down Expand Up @@ -305,11 +306,11 @@ def _kernel(
while not nvvm.mbarrier_try_wait_parity(clc_empty_mbar_ptr.subview(stage), clc_empty_phase, time_limit=10_000_000):
pass

if nvvm.elect_sync():
if elect_one:
nvvm.mbarrier_arrive_expect_tx(clc_full_mbar_ptr.subview(stage), 16)

if is_cluster_leader_cta:
if nvvm.elect_sync():
if elect_one:
cute_clc.issue_clc_query(
clc_full_mbar_cute_base + stage,
clc_response_ptr_base + stage,
Expand All @@ -324,7 +325,7 @@ def _kernel(
is_valid_sched = vld

nvvm.bar_warp_sync(0xFFFFFFFF)
if nvvm.elect_sync():
if elect_one:
empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(stage), 0)
nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True)

Expand All @@ -347,7 +348,7 @@ def _kernel(
if warp_idx == tma_warp_id:
nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE)
if cutlass.const_expr(USE_PDL):
if nvvm.elect_sync():
if elect_one:
nvvm.griddepcontrol("wait")
ab_empty_phase_bit = cutlass.Int32(1)
ab_iter = cutlass.Int32(0)
Expand Down Expand Up @@ -379,7 +380,7 @@ def _kernel(

coord_k = k_tile_idx * cta_tile_mnk[2]
coord_sf_k = k_tile_idx * sf_tma_box_k
if nvvm.elect_sync():
if elect_one:
nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes)

for _ai in cutlass.range_constexpr(num_a_operands):
Expand All @@ -390,7 +391,7 @@ def _kernel(
sfa_m_block = coord_m_per_cta // 128
if cutlass.const_expr(multicast_a):
if n_rank == 0:
if nvvm.elect_sync():
if elect_one:
if cutlass.const_expr(a_is_m_major):
for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems):
nvvm.cp_async_bulk_tensor_shared_cluster_global(
Expand All @@ -417,7 +418,7 @@ def _kernel(
group=nvvm.CTAGroup.CTA_1,
)
else:
if nvvm.elect_sync():
if elect_one:
if cutlass.const_expr(a_is_m_major):
for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems):
nvvm.cp_async_bulk_tensor_shared_cluster_global(
Expand Down Expand Up @@ -445,7 +446,7 @@ def _kernel(
)
if cutlass.const_expr(multicast_a):
if n_rank == 0:
if nvvm.elect_sync():
if elect_one:
nvvm.cp_async_bulk_tensor_shared_cluster_global(
sSFA_stage,
tma_sfa_desc.get_ptr(),
Expand All @@ -456,7 +457,7 @@ def _kernel(
group=nvvm.CTAGroup.CTA_1,
)
else:
if nvvm.elect_sync():
if elect_one:
nvvm.cp_async_bulk_tensor_shared_cluster_global(
sSFA_stage,
tma_sfa_desc.get_ptr(),
Expand All @@ -475,7 +476,7 @@ def _kernel(
sfb_n_block = coord_n_per_cta // 128
if cutlass.const_expr(multicast_b):
if m_rank == 0:
if nvvm.elect_sync():
if elect_one:
if cutlass.const_expr(b_is_n_major):
for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems):
nvvm.cp_async_bulk_tensor_shared_cluster_global(
Expand All @@ -502,7 +503,7 @@ def _kernel(
group=nvvm.CTAGroup.CTA_1,
)
else:
if nvvm.elect_sync():
if elect_one:
if cutlass.const_expr(b_is_n_major):
for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems):
nvvm.cp_async_bulk_tensor_shared_cluster_global(
Expand Down Expand Up @@ -530,7 +531,7 @@ def _kernel(
)
if cutlass.const_expr(multicast_b):
if m_rank == 0:
if nvvm.elect_sync():
if elect_one:
nvvm.cp_async_bulk_tensor_shared_cluster_global(
sSFB_stage,
tma_sfb_desc.get_ptr(),
Expand All @@ -541,7 +542,7 @@ def _kernel(
group=nvvm.CTAGroup.CTA_1,
)
else:
if nvvm.elect_sync():
if elect_one:
nvvm.cp_async_bulk_tensor_shared_cluster_global(
sSFB_stage,
tma_sfb_desc.get_ptr(),
Expand Down Expand Up @@ -575,7 +576,7 @@ def _kernel(
)
tile_l = l_idx
nvvm.bar_warp_sync(0xFFFFFFFF)
if nvvm.elect_sync():
if elect_one:
empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0)
nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True)
tile_iter += 1
Expand All @@ -589,7 +590,7 @@ def _kernel(
if tail_stage == ab_stages:
tail_stage = cutlass.Int32(0)
tail_phase = tail_phase ^ 1
if nvvm.elect_sync():
if elect_one:
while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000):
pass

Expand Down Expand Up @@ -706,7 +707,7 @@ def _kernel(
for atom_r in cutlass.range(num_sf_atoms, unroll_full=True):
for _ai in cutlass.range_constexpr(num_a_operands):
for _m in cutlass.range_constexpr(num_blocks_m):
if nvvm.elect_sync():
if elect_one:
nvvm.tcgen05_cp(
s2t_shape,
sfa_dst_ptrs[_ai][_m],
Expand All @@ -716,7 +717,7 @@ def _kernel(
)
for _bj in cutlass.range_constexpr(num_b_operands):
for _m in cutlass.range_constexpr(num_blocks_n):
if nvvm.elect_sync():
if elect_one:
nvvm.tcgen05_cp(
s2t_shape,
sfb_dst_ptrs[_bj][_m],
Expand All @@ -732,7 +733,7 @@ def _kernel(
_bj = gemm_b_idx[g]
desc_a = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx)
desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx)
if nvvm.elect_sync():
if elect_one:
nvvm.tcgen05_mma_block_scale(
mma_block_scale_kind,
nvvm.CTAGroup.CTA_1,
Expand All @@ -747,15 +748,15 @@ def _kernel(
)
scale_d = cutlass.Boolean(True)

if nvvm.elect_sync():
if elect_one:
nvvm.tcgen05_commit(
ab_empty_mbar_ptr.subview(stage),
multicast_mask=ab_empty_arrive_mask,
group=nvvm.CTAGroup.CTA_1,
)
ab_iter += 1

if nvvm.elect_sync():
if elect_one:
nvvm.tcgen05_commit(
acc_full_mbar_ptr.subview(acc_stage),
group=nvvm.CTAGroup.CTA_1,
Expand All @@ -774,18 +775,18 @@ def _kernel(
cute.arch.fence_proxy("async.shared", space="cta")
is_valid = vld
nvvm.bar_warp_sync(0xFFFFFFFF)
if nvvm.elect_sync():
if elect_one:
empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0)
nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True)
tile_iter += 1

if cutlass.const_expr(USE_PDL):
if nvvm.elect_sync():
if elect_one:
nvvm.griddepcontrol("launch_dependents")
nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1)
tail_stage = acc_stage
tail_phase = acc_empty_phase_bit
if nvvm.elect_sync():
if elect_one:
for _ in range(acc_stages):
tail_stage = tail_stage + 1
if tail_stage == acc_stages:
Expand Down Expand Up @@ -878,7 +879,7 @@ def _kernel(

if use_acc_overlap and (not cd_out_is_m_major) and subtile_idx == acc_overlap_subtiles - 1:
nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC)
if nvvm.elect_sync():
if elect_one:
nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage))

col = coord_n + subtile_col_offset
Expand Down Expand Up @@ -934,7 +935,7 @@ def _kernel(
)

if warp_idx == 0:
if nvvm.elect_sync():
if elect_one:
if cutlass.const_expr(cd_out_is_m_major):
for _mb in cutlass.range_constexpr(cta_tile_mnk[0] // cd_mmajor_atom_m):
nvvm.cp_async_bulk_tensor_global_shared_cta(
Expand Down Expand Up @@ -971,7 +972,7 @@ def _kernel(

if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major):
nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC)
if nvvm.elect_sync():
if elect_one:
nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage))

consumer_stage = tile_iter % CLC_SCHED_STAGES
Expand All @@ -995,15 +996,15 @@ def _kernel(
)
tile_l = l_idx
nvvm.bar_warp_sync(0xFFFFFFFF)
if nvvm.elect_sync():
if elect_one:
empty_remote = nvvm.mapa(clc_empty_mbar_ptr.subview(consumer_stage), 0)
nvvm.mbarrier_arrive(empty_remote, scope=nvvm.MemScope.CLUSTER, relaxed=True)

tile_iter += 1

if cutlass.const_expr(use_acc_overlap):
nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC)
if nvvm.elect_sync():
if elect_one:
nvvm.mbarrier_arrive(tmem_dealloc_mbar_ptr)

# @@TMA_STORE_ONLY:BEGIN@@
Expand Down
Loading