Other

本页暂无中文版。以下为英文原文。

6 ops, 27 workloads — FFT 1 · mHC 2 · Engram 3.

One table per op, one row per workload. Ratio is the fastest other implementation's device time divided by ours, so green is faster than it, plain is level with it, red is slower. Times are in ms. How these numbers are taken.

FFT

FFTC2CFwd

input: [*B, n]

• latency-4kinput: [4096]dtype=c64
• throughput-4kinput: [64, 4096]dtype=c64, c128
• longconv-16kinput: [64, 16384]dtype=c64
• longconv-64kinput: [8, 65536]dtype=c64
• spectrogram-1kinput: [256, 1024]dtype=c64
• fnet-512input: [4096, 512]dtype=c64
• shortconv-64input: [65536, 64]dtype=c64
• scientific-1minput: [4, 1048576]dtype=c128
• spectrogram-1k-b300input: [300, 1024]dtype=c64

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
latency-4k	c64	1.54×	0.00358	torch-compile torch-cufft	0.00554 0.00557	0.0686	0%	lat
throughput-4k	c64	1.28×	0.00445	torch-cufft torch-compile	0.0057 0.0057	3.54	21%	lat
	c128	1.18×	0.00726	torch-cufft torch-compile	0.00858 0.00858	2.17	26%	lat
longconv-16k	c64	1.07×	0.0154	torch-cufft torch-compile	0.0164 0.0164	4.77	24%	mem
longconv-64k	c64	1.24×	0.00893	torch-cufft torch-compile	0.0110 0.0110	4.7	21%	lat
spectrogram-1k	c64	1.15×	0.00352	torch-cufft torch-compile	0.00403 0.00406	3.72	26%	lat
fnet-512	c64	1.26×	0.0109	torch-compile torch-cufft	0.0138 0.0138	8.62	68%	mem
shortconv-64	c64	1.09×	0.0203	torch-compile torch-cufft	0.0222 0.0222	6.2	73%	mem
scientific-1m	c128	1.38×	0.1049	torch-compile torch-cufft	0.1448 0.1449	4	28%	mem
spectrogram-1k-b300	c64	1.14×	0.004	torch-cufft torch-compile	0.00458 0.00458	3.84	27%	lat

mHC

MHCPreFwd

phi: [C, Q]x: [B, C], bf16b: [Q]

alpha_pre=1alpha_post=1alpha_res=1sinkhorn_repeat=1dtype=f32

• hidden-1280phi: [5120, 24]x: [1, 5120], bf16b: [24]
• hidden-1920phi: [7680, 24]x: [2, 7680], bf16b: [24]
• hidden-2560phi: [10240, 24]x: [4, 10240], bf16b: [24]

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
hidden-1280	f32	4.40×	0.00472	torch-compile torch-ref	0.0208 0.0658	0.0629	2%	lat
hidden-1920	f32	6.43×	0.00509	torch-compile torch-ref	0.0327 0.0856	0.175	4%	lat
hidden-2560	f32	10.52×	0.0055	torch-compile torch-ref	0.0579 0.1246	0.432	5%	lat

MHCPostFwd

x_layer_out: [B, c_x]h_post: [B, n], f32x_res: [B, n * c_x]

dtype=bf16

• hidden-1280x_layer_out: [1, 1280]h_post: [1, 4], f32x_res: [1, 5120]
• hidden-1920x_layer_out: [2, 1920]h_post: [2, 4], f32x_res: [2, 7680]
• hidden-2560x_layer_out: [4, 2560]h_post: [4, 4], f32x_res: [4, 10240]

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
hidden-1280	bf16	0.98×	0.00144	torch-compile torch-ref	0.00141 0.0107	0.00711	0%	lat
hidden-1920	bf16	1.00×	0.0015	torch-compile torch-ref	0.0015 0.0111	0.0204	1%	lat
hidden-2560	bf16	0.98×	0.00166	torch-compile torch-ref	0.00163 0.0115	0.0492	2%	lat

Engram

EngramGateConvFwd

H, k, v: [M, seq_len, d]rms_w_h, rms_w_v: [d]conv_w: [4, d]

• train-s32H, k, v: [1, 32, 256]rms_w_h, rms_w_v: [256]conv_w: [4, 256]dtype=f16
• train-s64H, k, v: [2, 64, 512]rms_w_h, rms_w_v: [512]conv_w: [4, 512]dtype=f16
• train-s128H, k, v: [1, 128, 256]rms_w_h, rms_w_v: [256]conv_w: [4, 256]dtype=bf16
• train-s96H, k, v: [3, 96, 768]rms_w_h, rms_w_v: [768]conv_w: [4, 768]dtype=bf16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
train-s32	f16	2.54×	0.00442	torch-compile torch-ref	0.0112 0.0748	0.0532	0%	lat
train-s64	f16	2.14×	0.00512	torch-compile torch-ref	0.0110 0.0859	0.369	3%	lat
train-s128	bf16	2.81×	0.00454	torch-compile torch-ref	0.0128 0.0804	0.209	2%	lat
train-s96	bf16	2.15×	0.0072	torch-compile torch-ref	0.0155 0.1004	0.888	7%	lat

EngramGateConvBwd

dY, H, k, v, vhat: [M, seq_len, d]rms_w_h, rms_w_v: [d]conv_w: [4, d]alpha, rrms_h, rrms_k, rrms_v: [M, seq_len], f32

• train-s32dY, H, k, v, vhat: [1, 32, 256]rms_w_h, rms_w_v: [256]conv_w: [4, 256]alpha, rrms_h, rrms_k, rrms_v: [1, 32], f32dtype=f16
• train-s64dY, H, k, v, vhat: [2, 64, 512]rms_w_h, rms_w_v: [512]conv_w: [4, 512]alpha, rrms_h, rrms_k, rrms_v: [2, 64], f32dtype=f16
• train-s128dY, H, k, v, vhat: [1, 128, 256]rms_w_h, rms_w_v: [256]conv_w: [4, 256]alpha, rrms_h, rrms_k, rrms_v: [1, 128], f32dtype=bf16
• train-s96dY, H, k, v, vhat: [3, 96, 768]rms_w_h, rms_w_v: [768]conv_w: [4, 768]alpha, rrms_h, rrms_k, rrms_v: [3, 96], f32dtype=bf16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
train-s32	f16	18.60×	0.00906	torch	0.1684	0.0614	0%	lat
train-s64	f16	10.10×	0.0199	torch	0.2011	0.225	1%	lat
train-s128	bf16	11.91×	0.0155	torch	0.1845	0.145	1%	lat
train-s96	bf16	2.36×	0.1051	torch	0.2477	0.145	1%	mem

EngramDecodeFwd

e_t: [batch, d_mem]h_t: [batch, d]conv_state: [batch, L, d]W_K, W_V: [d_mem, d]rms_w_h, rms_w_v: [d]conv_w: [conv_kernel_size, d]

• serving-b1e_t: [1, 512]h_t: [1, 256]conv_state: [1, 12, 256]W_K, W_V: [512, 256]rms_w_h, rms_w_v: [256]conv_w: [4, 256]max_conv_len=12dilation=3dtype=f16
• serving-b4e_t: [4, 1024]h_t, conv_w: [4, 512]conv_state: [4, 20, 512]W_K, W_V: [1024, 512]rms_w_h, rms_w_v: [512]max_conv_len=20dilation=5dtype=f16
• serving-b8e_t: [8, 512]h_t: [8, 256]conv_state: [8, 18, 256]W_K, W_V: [512, 256]rms_w_h, rms_w_v: [256]conv_w: [4, 256]max_conv_len=18dilation=3dtype=bf16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
serving-b1	f16	1.88×	0.0071	torch-compile torch-ref	0.0133 0.0949	0.0747	2%	lat
serving-b4	f16	2.68×	0.0096	torch-compile torch-ref	0.0257 0.1221	0.879	5%	lat
serving-b8	bf16	3.29×	0.00653	torch-compile torch-ref	0.0215 0.1120	0.65	2%	lat