diff --git a/pkg/cpuid/cpuid_amd64.go b/pkg/cpuid/cpuid_amd64.go
index 16c55dff4..f5fd4ae34 100644
--- a/pkg/cpuid/cpuid_amd64.go
+++ b/pkg/cpuid/cpuid_amd64.go
@@ -384,17 +384,13 @@ const (
 // Extended state includes floating point registers, and other cpu state that's
 // not associated with the normal task context.
 //
-// We do not support enabling AMX within gVisor, therefore always exclude
-// AMXExtendedStateSize.
-// TODO(gvisor.dev/issues/9896): Implement AMX Support.
-//
 // Note: the return value matches the size of signal FP state frames.
 // Look at check_xstate_in_sigframe() in the kernel sources for more details.
 //
 //go:nosplit
 func (fs FeatureSet) ExtendedStateSize() (size, align uint) {
 	if fs.UseXsave() {
-		return uint(xsaveSize) - fs.AMXExtendedStateSize(), 64
+		return uint(xsaveSize), 64
 	}
 
 	// If we don't support xsave, we fall back to fxsave, which requires
@@ -423,7 +419,7 @@ func (fs FeatureSet) ValidXCR0Mask() uint64 {
 		return 0
 	}
 	ax, _, _, dx := fs.query(xSaveInfo)
-	return (uint64(dx)<<32 | uint64(ax)) ^ XCR0AMXMask
+	return (uint64(dx)<<32 | uint64(ax)) &^ XCR0AMXMask
 }
 
 // UseXsave returns the choice of fp state saving instruction.
diff --git a/pkg/ring0/entry_amd64.s b/pkg/ring0/entry_amd64.s
index ad49618b1..218eb3fea 100644
--- a/pkg/ring0/entry_amd64.s
+++ b/pkg/ring0/entry_amd64.s
@@ -233,9 +233,7 @@ TEXT ·doSwitchToUser(SB),NOSPLIT,$16-48
 	MOVB ·hasXSAVE(SB), BX
 	TESTB BX, BX
 	JZ no_xrstor
-	// Use xrstor to restore all available fp state. For now, we restore
-	// everything unconditionally by setting the implicit operand edx:eax
-	// (the "requested feature bitmap") to all 1's.
+	// Use xrstor to restore all available fp state.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x2f // XRSTOR64 0(DI)
@@ -281,7 +279,6 @@ done_sysret_or_iret:
 	TESTB BX, BX
 	JZ no_xsave
 	// Use xsave/xsaveopt to save all extended state.
-	// We save everything unconditionally by setting RFBM to all 1's.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	TESTB CX, CX
@@ -519,7 +516,6 @@ kernel:
 	TESTB BX, BX
 	JZ no_xsave
 	// Use xsave/xsaveopt to save all extended state.
-	// We save everything unconditionally by setting RFBM to all 1's.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	TESTB CX, CX
@@ -651,7 +647,6 @@ kernel:
 	TESTB BX, BX
 	JZ no_xsave
 	// Use xsave/xsaveopt to save all extended state.
-	// We save everything unconditionally by setting RFBM to all 1's.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	TESTB CX, CX
diff --git a/pkg/ring0/lib_amd64.go b/pkg/ring0/lib_amd64.go
index 585a5fa2a..73f774667 100644
--- a/pkg/ring0/lib_amd64.go
+++ b/pkg/ring0/lib_amd64.go
@@ -119,7 +119,8 @@ func Init(fs cpuid.FeatureSet) {
 	hasFSGSBASE = fs.HasFeature(cpuid.X86FeatureFSGSBase)
 	validXCR0Mask = uintptr(fs.ValidXCR0Mask())
 	if hasXSAVE {
-		localXCR0 = xgetbv(0)
+		XCR0AMXMask := uintptr((1 << 17) | (1 << 18))
+		localXCR0 = xgetbv(0) &^ XCR0AMXMask
 	}
 }
 
diff --git a/pkg/sentry/arch/fpu/fpu_amd64.go b/pkg/sentry/arch/fpu/fpu_amd64.go
index 6d4bfdfc5..a7e4bf3c3 100644
--- a/pkg/sentry/arch/fpu/fpu_amd64.go
+++ b/pkg/sentry/arch/fpu/fpu_amd64.go
@@ -70,8 +70,13 @@ const (
 func initX86FPState(data *byte, useXsave bool)
 
 func newX86FPStateSlice() State {
-	size, align := cpuid.HostFeatureSet().ExtendedStateSize()
-	capacity := size + FP_XSTATE_MAGIC2_SIZE
+	maxsize, align := cpuid.HostFeatureSet().ExtendedStateSize()
+	// We need capacity to be large enough to hold AMX bytes because of
+	// ptrace. PTRACE_SETREGSET/GETREGSET assume that AMX portions should
+	// always be used.
+	// TODO(gvisor.dev/issues/9896): Implement AMX Support.
+	capacity := maxsize + FP_XSTATE_MAGIC2_SIZE
+	size := maxsize - cpuid.HostFeatureSet().AMXExtendedStateSize()
 	// Always use at least 4096 bytes.
 	//
 	// For the KVM platform, this state is a fixed 4096 bytes, so make sure
@@ -129,6 +134,8 @@ func InitHostState() {
 		hostXCR0Mask = featureSet.ValidXCR0Mask()
 		hostUseXsave = featureSet.UseXsave()
 		hostFPSize, _ = featureSet.ExtendedStateSize()
+		// TODO(gvisor.dev/issues/9896): Implement AMX Support.
+		hostFPSize = hostFPSize - featureSet.AMXExtendedStateSize()
 	})
 }
 
diff --git a/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go b/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go
index 8d14eeeda..4732391ed 100644
--- a/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go
+++ b/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go
@@ -52,7 +52,8 @@ func (s *ArchState) Init() {
 	fs := cpuid.HostFeatureSet()
 
 	fpLenUint, _ := fs.ExtendedStateSize()
-	s.fpLen = uint32(fpLenUint)
+	// TODO(gvisor.dev/issues/9896): Implement AMX Support.
+	s.fpLen = uint32(fpLenUint - fs.AMXExtendedStateSize())
 	if fs.UseXsaveopt() {
 		s.xsaveMode = xsaveopt
 	} else if fs.UseXsave() {