diff --git a/sm4/cipher_asm.go b/sm4/cipher_asm.go
index 345f193..9cfc967 100644
--- a/sm4/cipher_asm.go
+++ b/sm4/cipher_asm.go
@@ -8,6 +8,9 @@ import (
 	"golang.org/x/sys/cpu"
 )
 
+var supportsAES = cpu.X86.HasAES
+var supportsGFMUL = cpu.X86.HasPCLMULQDQ
+
 //go:noescape
 func encryptBlocksAsm(xk *uint32, dst, src *byte)
 
@@ -21,9 +24,6 @@ type sm4CipherAsm struct {
 	sm4Cipher
 }
 
-var supportsAES = cpu.X86.HasAES
-var supportsGFMUL = cpu.X86.HasPCLMULQDQ
-
 func newCipher(key []byte) (cipher.Block, error) {
 	if !supportsAES {
 		return newCipherGeneric(key)
@@ -65,3 +65,13 @@ func (c *sm4CipherAsm) Decrypt(dst, src []byte) {
 	}
 	encryptBlockAsm(&c.dec[0], &dst[0], &src[0])
 }
+
+// expandKey is used by BenchmarkExpand to ensure that the asm implementation
+// of key expansion is used for the benchmark when it is available.
+func expandKey(key []byte, enc, dec []uint32) {
+	if supportsAES {
+		expandKeyAsm(&key[0], &ck[0], &enc[0], &dec[0])
+	} else {
+		expandKeyGo(key, enc, dec)
+	}
+}
diff --git a/sm4/cipher_generic.go b/sm4/cipher_generic.go
index f547cd1..cdbb6d2 100644
--- a/sm4/cipher_generic.go
+++ b/sm4/cipher_generic.go
@@ -12,3 +12,9 @@ import "crypto/cipher"
 func newCipher(key []byte) (cipher.Block, error) {
 	return newCipherGeneric(key)
 }
+
+// expandKey is used by BenchmarkExpand and should
+// call an assembly implementation if one is available.
+func expandKey(key []byte, enc, dec []uint32) {
+	expandKeyGo(key, enc, dec)
+}
\ No newline at end of file
diff --git a/sm4/cipher_test.go b/sm4/cipher_test.go
index 22ebfec..cad130b 100644
--- a/sm4/cipher_test.go
+++ b/sm4/cipher_test.go
@@ -89,6 +89,6 @@ func BenchmarkExpand(b *testing.B) {
 	c := &sm4Cipher{make([]uint32, rounds), make([]uint32, rounds)}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		expandKeyGo(tt.key, c.enc, c.dec)
+		expandKey(tt.key, c.enc, c.dec)
 	}
 }