hsam007 commited on
Commit
ec0f162
·
verified ·
1 Parent(s): 34263ef

Remove sm_90 tile config (falls back to GB10 BM=32)

Browse files
torch-ext/dg_w4_kernels/fused_moe_w4_v2.py CHANGED
@@ -73,18 +73,10 @@ def moe_align_cuda(topk_ids, block_m, E):
73
  cap, nblk, E, BLOCK=1024, num_warps=4)
74
  return a_row, expert_ids, cap, inv
75
 
76
- # Per-arch launch configs (BM included: BM>=64 selects the wgmma lowering on sm_90).
77
  # GB10 (sm_121) keeps the shipped originals byte-for-byte — records unaffected.
78
- # sm_90 tier swept 2026-07-23 on H100 SXM5 (w4_h100_sweep.py --wgmma): gates = numerics
79
- # (rel_mean<5e-3) + run-to-run SELF-DETERMINISM per T. Winner is rel~1.4e-05 vs the BM=32
80
- # path and self-det=True at every T; 0.88-1.39 ms/layer (297-469 GB/s), 1.76-2.7x over the
81
- # mma tier, ~4.6x over the untuned GB10 config on H100.
82
  _CC = torch.cuda.get_device_capability() if torch.cuda.is_available() else (12, 1)
83
- if _CC[0] == 9: # Hopper (H100/H200): wgmma tier
84
- BM_V2 = 64
85
- G1 = dict(BN=64, num_warps=4, num_stages=2)
86
- G2 = dict(BN=128, num_warps=4, num_stages=4)
87
- elif _CC[0] == 12 and _CC[1] == 0: # sm_120: RTX 5090 / Blackwell consumer
88
  # Swept + e2e-gated on a 5090 (torch 2.13/triton 3.7): BM=16 is 1.08x on the isolated layer and
89
  # +3.4%/+3.9% end-to-end on gsm8k/humaneval vs the BM=32 fallback, accuracy unchanged (48/50 both).
90
  # Smaller BM wins because M/expert is ~16 at canvas 256, so BM=32 tiles are half padding.
 
73
  cap, nblk, E, BLOCK=1024, num_warps=4)
74
  return a_row, expert_ids, cap, inv
75
 
76
+ # Per-arch launch configs.
77
  # GB10 (sm_121) keeps the shipped originals byte-for-byte — records unaffected.
 
 
 
 
78
  _CC = torch.cuda.get_device_capability() if torch.cuda.is_available() else (12, 1)
79
+ if _CC[0] == 12 and _CC[1] == 0: # sm_120: RTX 5090 / Blackwell consumer
 
 
 
 
80
  # Swept + e2e-gated on a 5090 (torch 2.13/triton 3.7): BM=16 is 1.08x on the isolated layer and
81
  # +3.4%/+3.9% end-to-end on gsm8k/humaneval vs the BM=32 fallback, accuracy unchanged (48/50 both).
82
  # Smaller BM wins because M/expert is ~16 at canvas 256, so BM=32 tiles are half padding.